Skip to main content

aioduct/runtime/
compio_rt.rs

1use std::future::Future;
2use std::io;
3use std::net::SocketAddr;
4use std::pin::Pin;
5use std::task::{Context, Poll};
6use std::time::Duration;
7
8use hyper::rt::{self, Read, Write};
9use pin_project_lite::pin_project;
10
11use super::{ConnectorLocal, RuntimeCompletion, RuntimeLocal};
12
13/// Compio async runtime implementation using native io_uring/IOCP for TCP I/O.
14pub struct CompioRuntime;
15
16// ── New trait impls (v0.2) ──────────────────────────────────────────────────
17
18impl RuntimeCompletion for CompioRuntime {
19    type Sleep = CompioSleep;
20
21    fn sleep(duration: Duration) -> Self::Sleep {
22        CompioSleep::new(async_io::Timer::after(duration))
23    }
24
25    fn block_on<F: Future>(future: F) -> Result<F::Output, crate::error::Error> {
26        let rt = compio_runtime::Runtime::new().map_err(crate::error::Error::Io)?;
27        Ok(rt.block_on(future))
28    }
29}
30
31impl RuntimeLocal for CompioRuntime {
32    fn spawn_local<F: Future<Output = ()> + 'static>(future: F) {
33        compio_runtime::spawn(future).detach();
34    }
35}
36
37// CompioRuntime does NOT implement RuntimePoll — it's completion-based,
38// single-threaded, and cannot migrate futures between threads.
39
40// ── SocketConfig ──────────────────────────────────────────────────────────
41
42impl super::SocketConfig for CompioTcpStream {
43    fn set_keepalive(
44        &self,
45        time: Duration,
46        interval: Option<Duration>,
47        retries: Option<u32>,
48    ) -> io::Result<()> {
49        use socket2::SockRef;
50        let sock_ref = SockRef::from(&self.socket_handle);
51        let mut keepalive = socket2::TcpKeepalive::new().with_time(time);
52        if let Some(interval) = interval {
53            keepalive = keepalive.with_interval(interval);
54        }
55        #[cfg(any(
56            target_os = "linux",
57            target_os = "macos",
58            target_os = "ios",
59            target_os = "freebsd",
60            target_os = "netbsd",
61        ))]
62        if let Some(retries) = retries {
63            keepalive = keepalive.with_retries(retries);
64        }
65        #[cfg(not(any(
66            target_os = "linux",
67            target_os = "macos",
68            target_os = "ios",
69            target_os = "freebsd",
70            target_os = "netbsd",
71        )))]
72        let _ = retries;
73        sock_ref.set_tcp_keepalive(&keepalive)
74    }
75
76    #[cfg(target_os = "linux")]
77    fn set_fast_open(&self) -> io::Result<()> {
78        use std::os::unix::io::AsRawFd;
79
80        unsafe extern "C" {
81            fn setsockopt(
82                sockfd: std::ffi::c_int,
83                level: std::ffi::c_int,
84                optname: std::ffi::c_int,
85                optval: *const std::ffi::c_void,
86                optlen: u32,
87            ) -> std::ffi::c_int;
88        }
89
90        let fd = self.socket_handle.as_raw_fd();
91        const IPPROTO_TCP: std::ffi::c_int = 6;
92        const TCP_FASTOPEN_CONNECT: std::ffi::c_int = 30;
93        let optval: std::ffi::c_int = 1;
94        unsafe {
95            let ret = setsockopt(
96                fd,
97                IPPROTO_TCP,
98                TCP_FASTOPEN_CONNECT,
99                &optval as *const std::ffi::c_int as *const std::ffi::c_void,
100                std::mem::size_of::<std::ffi::c_int>() as u32,
101            );
102            if ret != 0 {
103                return Err(io::Error::last_os_error());
104            }
105        }
106        Ok(())
107    }
108
109    #[cfg(target_os = "linux")]
110    fn bind_device(&self, interface: &str) -> io::Result<()> {
111        use socket2::SockRef;
112        let sock_ref = SockRef::from(&self.socket_handle);
113        sock_ref.bind_device(Some(interface.as_bytes()))
114    }
115}
116
117// ── TcpConnector ──────────────────────────────────────────────────────────
118
119/// TCP connector for the Compio runtime.
120///
121/// Uses native `compio_net::TcpStream` for io_uring (Linux) / IOCP (Windows)
122/// I/O, bridged to futures-io via `compio_io::compat::AsyncStream`.
123#[derive(Clone, Copy, Default)]
124pub struct TcpConnector;
125
126impl ConnectorLocal for TcpConnector {
127    type Stream = CompioTcpStream;
128
129    async fn connect(&self, addr: SocketAddr) -> io::Result<Self::Stream> {
130        let stream = compio_net::TcpStream::connect(addr).await?;
131        stream.set_nodelay(true)?;
132        Ok(CompioTcpStream::new(stream))
133    }
134
135    async fn connect_bound(
136        &self,
137        addr: SocketAddr,
138        local: std::net::IpAddr,
139    ) -> io::Result<Self::Stream> {
140        use socket2::{Domain, Protocol, SockAddr, Socket, Type};
141
142        let std_stream = compio_runtime::spawn_blocking(move || {
143            let domain = if addr.is_ipv4() {
144                Domain::IPV4
145            } else {
146                Domain::IPV6
147            };
148            let socket = Socket::new(domain, Type::STREAM, Some(Protocol::TCP))?;
149            socket.bind(&SockAddr::from(std::net::SocketAddr::new(local, 0)))?;
150            socket.connect(&SockAddr::from(addr))?;
151            socket.set_tcp_nodelay(true)?;
152            Ok::<std::net::TcpStream, io::Error>(socket.into())
153        })
154        .await
155        .map_err(|e| io::Error::other(format!("{e:?}")))?;
156        let std_stream = std_stream?;
157        std_stream.set_nonblocking(true)?;
158        let compio_stream = compio_net::TcpStream::from_std(std_stream)?;
159        Ok(CompioTcpStream::new(compio_stream))
160    }
161
162    fn from_std_tcp(&self, stream: std::net::TcpStream) -> io::Result<Self::Stream> {
163        stream.set_nonblocking(true)?;
164        stream.set_nodelay(true)?;
165        let compio_stream = compio_net::TcpStream::from_std(stream)?;
166        Ok(CompioTcpStream::new(compio_stream))
167    }
168
169    fn into_std_tcp(&self, stream: Self::Stream) -> io::Result<std::net::TcpStream> {
170        use socket2::SockRef;
171        let sock = SockRef::from(&stream.socket_handle).try_clone()?;
172        drop(stream);
173        let std_stream: std::net::TcpStream = sock.into();
174        std_stream.set_nonblocking(false)?;
175        Ok(std_stream)
176    }
177}
178
179// ── CompioTcpStream ──────────────────────────────────────────────────────────
180
181pin_project! {
182    /// Compound stream type that keeps a `compio_net::TcpStream` handle alongside
183    /// the async I/O bridge for socket operations (keepalive, fast open, etc.).
184    ///
185    /// `compio_net::TcpStream` is clone-cheap (shared fd), so this is essentially
186    /// free. The `socket_handle` is used by `set_tcp_keepalive` etc. to access the
187    /// raw socket via `AsFd`.
188    ///
189    pub struct CompioTcpStream {
190        io: Pin<Box<CompioIo<compio_io::compat::AsyncStream<compio_net::TcpStream>>>>,
191        pub(crate) socket_handle: compio_net::TcpStream,
192    }
193}
194
195impl CompioTcpStream {
196    pub(crate) fn new(stream: compio_net::TcpStream) -> Self {
197        let socket_handle = stream.clone();
198        Self {
199            io: Box::pin(CompioIo::new(compio_io::compat::AsyncStream::new(stream))),
200            socket_handle,
201        }
202    }
203}
204
205impl Read for CompioTcpStream {
206    fn poll_read(
207        self: Pin<&mut Self>,
208        cx: &mut Context<'_>,
209        buf: rt::ReadBufCursor<'_>,
210    ) -> Poll<io::Result<()>> {
211        Read::poll_read(self.project().io.as_mut(), cx, buf)
212    }
213}
214
215impl Write for CompioTcpStream {
216    fn poll_write(
217        self: Pin<&mut Self>,
218        cx: &mut Context<'_>,
219        buf: &[u8],
220    ) -> Poll<io::Result<usize>> {
221        Write::poll_write(self.project().io.as_mut(), cx, buf)
222    }
223
224    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
225        Write::poll_flush(self.project().io.as_mut(), cx)
226    }
227
228    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
229        Write::poll_shutdown(self.project().io.as_mut(), cx)
230    }
231
232    fn poll_write_vectored(
233        self: Pin<&mut Self>,
234        cx: &mut Context<'_>,
235        bufs: &[io::IoSlice<'_>],
236    ) -> Poll<io::Result<usize>> {
237        Write::poll_write_vectored(self.project().io.as_mut(), cx, bufs)
238    }
239
240    fn is_write_vectored(&self) -> bool {
241        self.io.is_write_vectored()
242    }
243}
244
245pin_project! {
246    /// Compio-backed sleep future using async_io::Timer.
247    ///
248    /// Uses async_io's timer because compio_runtime's `TimerFuture` is `!Send`.
249    pub struct CompioSleep {
250        #[pin]
251        inner: async_io::Timer,
252    }
253}
254
255impl CompioSleep {
256    pub(crate) fn new(inner: async_io::Timer) -> Self {
257        Self { inner }
258    }
259}
260
261impl Future for CompioSleep {
262    type Output = ();
263
264    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
265        match self.project().inner.poll(cx) {
266            Poll::Ready(_instant) => Poll::Ready(()),
267            Poll::Pending => Poll::Pending,
268        }
269    }
270}
271
272pin_project! {
273    /// Adapter bridging futures-io `AsyncRead`/`AsyncWrite` to hyper's `Read`/`Write` for compio.
274    pub struct CompioIo<T> {
275        #[pin]
276        inner: T,
277    }
278}
279
280impl<T> CompioIo<T> {
281    /// Wrap an async-io type.
282    pub fn new(inner: T) -> Self {
283        Self { inner }
284    }
285
286    /// Get a reference to the inner I/O type.
287    pub fn inner(&self) -> &T {
288        &self.inner
289    }
290}
291
292impl<T> Read for CompioIo<T>
293where
294    T: futures_io::AsyncRead,
295{
296    fn poll_read(
297        self: Pin<&mut Self>,
298        cx: &mut Context<'_>,
299        mut buf: rt::ReadBufCursor<'_>,
300    ) -> Poll<io::Result<()>> {
301        let slice = unsafe {
302            let uninit = buf.as_mut();
303            std::ptr::write_bytes(uninit.as_mut_ptr(), 0, uninit.len());
304            std::slice::from_raw_parts_mut(uninit.as_mut_ptr() as *mut u8, uninit.len())
305        };
306        match futures_io::AsyncRead::poll_read(self.project().inner, cx, slice) {
307            Poll::Ready(Ok(n)) => {
308                unsafe { buf.advance(n) };
309                Poll::Ready(Ok(()))
310            }
311            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
312            Poll::Pending => Poll::Pending,
313        }
314    }
315}
316
317impl<T> Write for CompioIo<T>
318where
319    T: futures_io::AsyncWrite,
320{
321    fn poll_write(
322        self: Pin<&mut Self>,
323        cx: &mut Context<'_>,
324        buf: &[u8],
325    ) -> Poll<io::Result<usize>> {
326        futures_io::AsyncWrite::poll_write(self.project().inner, cx, buf)
327    }
328
329    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
330        futures_io::AsyncWrite::poll_flush(self.project().inner, cx)
331    }
332
333    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
334        futures_io::AsyncWrite::poll_close(self.project().inner, cx)
335    }
336
337    fn poll_write_vectored(
338        self: Pin<&mut Self>,
339        cx: &mut Context<'_>,
340        bufs: &[io::IoSlice<'_>],
341    ) -> Poll<io::Result<usize>> {
342        futures_io::AsyncWrite::poll_write_vectored(self.project().inner, cx, bufs)
343    }
344
345    fn is_write_vectored(&self) -> bool {
346        true
347    }
348}
349
350// ── DefaultResolver ───────────────────────────────────────────────────────
351
352/// Default DNS resolver for compio using blocking `getaddrinfo` via
353/// `compio_runtime::spawn_blocking`.
354pub struct DefaultResolver;
355
356impl super::Resolve for DefaultResolver {
357    fn resolve(
358        &self,
359        host: &str,
360        port: u16,
361    ) -> Pin<Box<dyn Future<Output = io::Result<SocketAddr>> + Send>> {
362        let addr = format!("{host}:{port}");
363        Box::pin(AssertSend(async move {
364            let addrs = compio_runtime::spawn_blocking(move || {
365                use std::net::ToSocketAddrs;
366                addr.to_socket_addrs().map(|iter| iter.collect::<Vec<_>>())
367            })
368            .await
369            .map_err(|e| io::Error::other(format!("{e:?}")))?;
370            let addrs = addrs?;
371            addrs.into_iter().next().ok_or_else(|| {
372                io::Error::new(io::ErrorKind::AddrNotAvailable, "no addresses found")
373            })
374        }))
375    }
376
377    fn resolve_all(
378        &self,
379        host: &str,
380        port: u16,
381    ) -> Pin<Box<dyn Future<Output = io::Result<Vec<SocketAddr>>> + Send>> {
382        let addr = format!("{host}:{port}");
383        Box::pin(AssertSend(async move {
384            let addrs = compio_runtime::spawn_blocking(move || {
385                use std::net::ToSocketAddrs;
386                addr.to_socket_addrs().map(|iter| iter.collect::<Vec<_>>())
387            })
388            .await
389            .map_err(|e| io::Error::other(format!("{e:?}")))?;
390            let addrs = addrs?;
391            if addrs.is_empty() {
392                return Err(io::Error::new(
393                    io::ErrorKind::AddrNotAvailable,
394                    "no addresses found",
395                ));
396            }
397            Ok(addrs)
398        }))
399    }
400}
401
402/// Wrapper that unsafely implements `Send` for a `!Send` future.
403///
404/// # Safety
405///
406/// Only safe in compio's thread-per-core model where futures never cross
407/// thread boundaries.
408struct AssertSend<F>(F);
409
410unsafe impl<F> Send for AssertSend<F> {}
411
412impl<F: Future> Future for AssertSend<F> {
413    type Output = F::Output;
414
415    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
416        let inner = unsafe { self.map_unchecked_mut(|s| &mut s.0) };
417        inner.poll(cx)
418    }
419}
420
421#[cfg(test)]
422#[allow(deprecated)]
423mod tests {
424    use super::*;
425
426    #[test]
427    fn connector_connect_works() {
428        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
429        let addr = listener.local_addr().unwrap();
430        compio_runtime::Runtime::new().unwrap().block_on(async {
431            let connector = TcpConnector;
432            let stream = connector.connect(addr).await.unwrap();
433            assert!(Write::is_write_vectored(&stream));
434        });
435    }
436
437    // ── New trait tests (v0.2) ──────────────────────────────────────────────
438
439    #[test]
440    fn runtime_completion_sleep() {
441        use crate::runtime::RuntimeCompletion;
442        compio_runtime::Runtime::new().unwrap().block_on(async {
443            let start = std::time::Instant::now();
444            <CompioRuntime as RuntimeCompletion>::sleep(Duration::from_millis(10)).await;
445            assert!(start.elapsed() >= Duration::from_millis(10));
446        });
447    }
448
449    #[test]
450    fn runtime_local_spawn_local() {
451        use crate::runtime::RuntimeLocal;
452        use std::sync::Arc;
453        use std::sync::atomic::{AtomicBool, Ordering};
454        compio_runtime::Runtime::new().unwrap().block_on(async {
455            let flag = Arc::new(AtomicBool::new(false));
456            let flag2 = flag.clone();
457            CompioRuntime::spawn_local(async move {
458                flag2.store(true, Ordering::SeqCst);
459            });
460            compio_runtime::time::sleep(Duration::from_millis(10)).await;
461            assert!(flag.load(Ordering::SeqCst));
462        });
463    }
464
465    #[test]
466    fn connector_connect_bound_works() {
467        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
468        let addr = listener.local_addr().unwrap();
469        compio_runtime::Runtime::new().unwrap().block_on(async {
470            let connector = TcpConnector;
471            let local: std::net::IpAddr = "127.0.0.1".parse().unwrap();
472            let stream =
473                crate::runtime::ConnectorLocal::connect_bound(&connector, addr, local).await;
474            assert!(stream.is_ok());
475        });
476    }
477
478    #[test]
479    fn connector_from_std_tcp_works() {
480        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
481        let addr = listener.local_addr().unwrap();
482        let std_stream = std::net::TcpStream::connect(addr).unwrap();
483        compio_runtime::Runtime::new().unwrap().block_on(async {
484            let connector = TcpConnector;
485            let result = crate::runtime::ConnectorLocal::from_std_tcp(&connector, std_stream);
486            assert!(result.is_ok());
487        });
488    }
489
490    // ── ConnectorLocal connect_bound IPv6 path ─────────────────────────────
491
492    #[test]
493    fn connector_connect_bound_ipv6() {
494        let listener = match std::net::TcpListener::bind("[::1]:0") {
495            Ok(l) => l,
496            Err(_) => return, // Skip if IPv6 not available
497        };
498        let addr = listener.local_addr().unwrap();
499        compio_runtime::Runtime::new().unwrap().block_on(async {
500            let connector = TcpConnector;
501            let local: std::net::IpAddr = "::1".parse().unwrap();
502            let stream =
503                crate::runtime::ConnectorLocal::connect_bound(&connector, addr, local).await;
504            assert!(stream.is_ok());
505        });
506    }
507
508    // ── DefaultResolver resolve_all error ─────────────────────────────
509
510    #[test]
511    fn default_resolver_resolve_all_invalid_host_errors() {
512        use crate::runtime::Resolve;
513        compio_runtime::Runtime::new().unwrap().block_on(async {
514            let resolver = DefaultResolver;
515            let result = resolver
516                .resolve_all("this.host.does.not.exist.invalid", 80)
517                .await;
518            assert!(result.is_err());
519        });
520    }
521
522    #[test]
523    fn default_resolver_resolve_single() {
524        use crate::runtime::Resolve;
525        compio_runtime::Runtime::new().unwrap().block_on(async {
526            let resolver = DefaultResolver;
527            let addr = resolver.resolve("localhost", 80).await.unwrap();
528            assert_eq!(addr.port(), 80);
529        });
530    }
531
532    #[test]
533    fn block_on_works() {
534        use crate::runtime::RuntimeCompletion;
535        let result = CompioRuntime::block_on(async { 42 }).unwrap();
536        assert_eq!(result, 42);
537    }
538
539    #[test]
540    fn compio_io_inner_accessor() {
541        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
542        let addr = listener.local_addr().unwrap();
543        let io = CompioIo::new(
544            async_io::Async::<std::net::TcpStream>::try_from(
545                std::net::TcpStream::connect(addr).unwrap(),
546            )
547            .unwrap(),
548        );
549        let _inner = io.inner();
550    }
551
552    // ── CompioIo tests ─────────────────────────────────────────────────
553
554    #[test]
555    fn compio_io_new_and_inner() {
556        let val = 42u32;
557        let io = CompioIo::new(val);
558        assert_eq!(*io.inner(), 42u32);
559    }
560
561    // ── DefaultResolver resolve_all ────────────────────────────────────
562
563    #[test]
564    fn default_resolver_resolve_all_multiple() {
565        use crate::runtime::Resolve;
566        compio_runtime::Runtime::new().unwrap().block_on(async {
567            let resolver = DefaultResolver;
568            let addrs = resolver.resolve_all("localhost", 80).await.unwrap();
569            assert!(!addrs.is_empty());
570            for addr in &addrs {
571                assert_eq!(addr.port(), 80);
572            }
573        });
574    }
575
576    #[test]
577    fn default_resolver_invalid_host_errors() {
578        use crate::runtime::Resolve;
579        compio_runtime::Runtime::new().unwrap().block_on(async {
580            let resolver = DefaultResolver;
581            let result = resolver
582                .resolve("this.host.does.not.exist.invalid", 80)
583                .await;
584            assert!(result.is_err());
585        });
586    }
587
588    // ── CompioSleep new() constructor ──────────────────────────────────
589
590    #[test]
591    fn compio_sleep_new_completes() {
592        compio_runtime::Runtime::new().unwrap().block_on(async {
593            let timer = async_io::Timer::after(Duration::from_millis(5));
594            let sleep = CompioSleep::new(timer);
595            let start = std::time::Instant::now();
596            sleep.await;
597            assert!(start.elapsed() >= Duration::from_millis(5));
598        });
599    }
600
601    // ── RuntimeLocal ───────────────────────────────────────────────────
602
603    #[test]
604    fn runtime_completion_block_on_nested() {
605        use crate::runtime::RuntimeCompletion;
606        // block_on should work for simple computations
607        let result = CompioRuntime::block_on(async { "hello".len() }).unwrap();
608        assert_eq!(result, 5);
609    }
610}