feat(net): add PollFd for waiting readiness

Author: Berrysoft

compio-net/src/lib.rs

@@ -5,13 +5,15 @@
 #![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
 #![warn(missing_docs)]
 
+mod poll_fd;
 mod resolve;
 mod socket;
 pub(crate) mod split;
 mod tcp;
 mod udp;
 mod unix;
 
+pub use poll_fd::*;
 pub use resolve::ToSocketAddrsAsync;
 pub(crate) use resolve::{each_addr, first_addr_buf};
 pub(crate) use socket::*;

compio-net/src/poll_fd.rs

@@ -0,0 +1,117 @@
+#[cfg(unix)]
+use std::os::fd::FromRawFd;
+#[cfg(windows)]
+use std::os::windows::io::{AsRawSocket, FromRawSocket, RawSocket};
+use std::{io, ops::Deref};
+
+use compio_buf::{BufResult, IntoInner};
+#[cfg(unix)]
+use compio_driver::op::{Interest, PollOnce};
+use compio_driver::{AsRawFd, RawFd, SharedFd, ToSharedFd};
+
+/// A wrapper for socket, providing functionalities to wait for readiness.
+#[derive(Debug)]
+pub struct PollFd<T: AsRawFd> {
+    inner: SharedFd<T>,
+}
+
+impl<T: AsRawFd> PollFd<T> {
+    /// Create [`PollFd`] without attaching the source. Ready-based sources need
+    /// not to be attached.
+    pub fn new(source: T) -> Self {
+        Self {
+            inner: SharedFd::new(source),
+        }
+    }
+
+    pub(crate) fn from_shared_fd(inner: SharedFd<T>) -> Self {
+        Self { inner }
+    }
+}
+
+#[cfg(unix)]
+impl<T: AsRawFd + 'static> PollFd<T> {
+    /// Wait for accept readiness, before calling `accept`, or after `accept`
+    /// returns `WouldBlock`.
+    pub async fn accept_ready(&self) -> io::Result<()> {
+        self.read_ready().await
+    }
+
+    /// Wait for connect readiness.
+    pub async fn connect_ready(&self) -> io::Result<()> {
+        self.write_ready().await
+    }
+
+    /// Wait for read readiness.
+    pub async fn read_ready(&self) -> io::Result<()> {
+        let op = PollOnce::new(self.to_shared_fd(), Interest::Readable);
+        let BufResult(res, _) = compio_runtime::submit(op).await;
+        res?;
+        Ok(())
+    }
+
+    /// Wait for write readiness.
+    pub async fn write_ready(&self) -> io::Result<()> {
+        let op = PollOnce::new(self.to_shared_fd(), Interest::Writable);
+        let BufResult(res, _) = compio_runtime::submit(op).await;
+        res?;
+        Ok(())
+    }
+}
+
+impl<T: AsRawFd> IntoInner for PollFd<T> {
+    type Inner = SharedFd<T>;
+
+    fn into_inner(self) -> Self::Inner {
+        self.inner
+    }
+}
+
+impl<T: AsRawFd> ToSharedFd<T> for PollFd<T> {
+    fn to_shared_fd(&self) -> SharedFd<T> {
+        self.inner.clone()
+    }
+}
+
+impl<T: AsRawFd> Clone for PollFd<T> {
+    fn clone(&self) -> Self {
+        Self {
+            inner: self.inner.clone(),
+        }
+    }
+}
+
+impl<T: AsRawFd> AsRawFd for PollFd<T> {
+    fn as_raw_fd(&self) -> RawFd {
+        self.inner.as_raw_fd()
+    }
+}
+
+#[cfg(windows)]
+impl<T: AsRawFd + AsRawSocket> AsRawSocket for PollFd<T> {
+    fn as_raw_socket(&self) -> RawSocket {
+        self.inner.as_raw_socket()
+    }
+}
+
+#[cfg(unix)]
+impl<T: AsRawFd + FromRawFd> FromRawFd for PollFd<T> {
+    unsafe fn from_raw_fd(fd: RawFd) -> Self {
+        Self::new(FromRawFd::from_raw_fd(fd))
+    }
+}
+
+#[cfg(windows)]
+impl<T: AsRawFd + FromRawSocket> FromRawSocket for PollFd<T> {
+    unsafe fn from_raw_socket(sock: RawSocket) -> Self {
+        Self::new(FromRawSocket::from_raw_socket(sock))
+    }
+}
+
+impl<T: AsRawFd> Deref for PollFd<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}

compio-net/src/socket.rs

@@ -14,6 +14,8 @@ use compio_driver::{
 use compio_runtime::Attacher;
 use socket2::{Domain, Protocol, SockAddr, Socket as Socket2, Type};
 
+use crate::PollFd;
+
 #[derive(Debug, Clone)]
 pub struct Socket {
     socket: Attacher<Socket2>,
@@ -34,6 +36,14 @@ impl Socket {
         self.socket.local_addr()
     }
 
+    pub fn to_poll_fd(&self) -> PollFd<Socket2> {
+        PollFd::from_shared_fd(self.to_shared_fd())
+    }
+
+    pub fn into_poll_fd(self) -> PollFd<Socket2> {
+        PollFd::from_shared_fd(self.socket.into_inner())
+    }
+
     #[cfg(windows)]
     pub async fn new(domain: Domain, ty: Type, protocol: Option<Protocol>) -> io::Result<Self> {
         use std::panic::resume_unwind;

compio-net/src/tcp.rs

@@ -3,9 +3,11 @@ use std::{future::Future, io, net::SocketAddr};
 use compio_buf::{BufResult, IoBuf, IoBufMut, IoVectoredBuf, IoVectoredBufMut};
 use compio_driver::impl_raw_fd;
 use compio_io::{AsyncRead, AsyncWrite};
-use socket2::{Protocol, SockAddr, Type};
+use socket2::{Protocol, SockAddr, Socket as Socket2, Type};
 
-use crate::{OwnedReadHalf, OwnedWriteHalf, ReadHalf, Socket, ToSocketAddrsAsync, WriteHalf};
+use crate::{
+    OwnedReadHalf, OwnedWriteHalf, PollFd, ReadHalf, Socket, ToSocketAddrsAsync, WriteHalf,
+};
 
 /// A TCP socket server, listening for connections.
 ///
@@ -203,6 +205,16 @@ impl TcpStream {
     pub fn into_split(self) -> (OwnedReadHalf<Self>, OwnedWriteHalf<Self>) {
         crate::into_split(self)
     }
+
+    /// Create [`PollFd`] from inner socket.
+    pub fn to_poll_fd(&self) -> PollFd<Socket2> {
+        self.inner.to_poll_fd()
+    }
+
+    /// Create [`PollFd`] from inner socket.
+    pub fn into_poll_fd(self) -> PollFd<Socket2> {
+        self.inner.into_poll_fd()
+    }
 }
 
 impl AsyncRead for TcpStream {

compio-net/src/unix.rs

@@ -3,9 +3,9 @@ use std::{future::Future, io, path::Path};
 use compio_buf::{BufResult, IoBuf, IoBufMut, IoVectoredBuf, IoVectoredBufMut};
 use compio_driver::impl_raw_fd;
 use compio_io::{AsyncRead, AsyncWrite};
-use socket2::{SockAddr, Type};
+use socket2::{SockAddr, Socket as Socket2, Type};
 
-use crate::{OwnedReadHalf, OwnedWriteHalf, ReadHalf, Socket, WriteHalf};
+use crate::{OwnedReadHalf, OwnedWriteHalf, PollFd, ReadHalf, Socket, WriteHalf};
 
 /// A Unix socket server, listening for connections.
 ///
@@ -187,6 +187,16 @@ impl UnixStream {
     pub fn into_split(self) -> (OwnedReadHalf<Self>, OwnedWriteHalf<Self>) {
         crate::into_split(self)
     }
+
+    /// Create [`PollFd`] from inner socket.
+    pub fn to_poll_fd(&self) -> PollFd<Socket2> {
+        self.inner.to_poll_fd()
+    }
+
+    /// Create [`PollFd`] from inner socket.
+    pub fn into_poll_fd(self) -> PollFd<Socket2> {
+        self.inner.into_poll_fd()
+    }
 }
 
 impl AsyncRead for UnixStream {

compio-net/tests/poll.rs

@@ -0,0 +1,88 @@
+use std::{
+    io,
+    net::{Ipv4Addr, SocketAddrV4},
+};
+
+use compio_net::PollFd;
+use socket2::{Domain, Protocol, SockAddr, Socket, Type};
+
+fn is_would_block(e: &io::Error) -> bool {
+    #[cfg(unix)]
+    {
+        e.kind() == io::ErrorKind::WouldBlock || e.raw_os_error() == Some(libc::EINPROGRESS)
+    }
+    #[cfg(not(unix))]
+    {
+        e.kind() == io::ErrorKind::WouldBlock
+    }
+}
+
+#[compio_macros::test]
+async fn poll_connect() {
+    let listener = Socket::new(Domain::IPV4, Type::STREAM, Some(Protocol::TCP)).unwrap();
+    listener.set_nonblocking(true).unwrap();
+    listener
+        .bind(&SockAddr::from(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)))
+        .unwrap();
+    listener.listen(4).unwrap();
+    let addr = listener.local_addr().unwrap();
+    let listener = PollFd::new(listener);
+    let accept_task = async {
+        loop {
+            listener.accept_ready().await.unwrap();
+            match listener.accept() {
+                Ok(res) => break res,
+                Err(e) if e.kind() == io::ErrorKind::WouldBlock => continue,
+                Err(e) => Err(e).unwrap(),
+            }
+        }
+    };
+
+    let client = Socket::new(Domain::IPV4, Type::STREAM, Some(Protocol::TCP)).unwrap();
+    client.set_nonblocking(true).unwrap();
+    let client = PollFd::new(client);
+    let res = client.connect(&addr);
+    let tx = if let Err(e) = res {
+        assert!(is_would_block(&e));
+        let (tx, _) = accept_task.await;
+        tx
+    } else {
+        let ((tx, _), res) = futures_util::join!(accept_task, client.connect_ready());
+        res.unwrap();
+        tx
+    };
+
+    tx.set_nonblocking(true).unwrap();
+    let tx = PollFd::new(tx);
+
+    let send_task = async {
+        loop {
+            match tx.send(b"Hello world!") {
+                Ok(res) => break res,
+                Err(e) if is_would_block(&e) => {}
+                Err(e) => Err(e).unwrap(),
+            }
+            tx.write_ready().await.unwrap();
+        }
+    };
+
+    let mut buffer = Vec::with_capacity(12);
+    let recv_task = async {
+        loop {
+            match client.recv(buffer.spare_capacity_mut()) {
+                Ok(res) => {
+                    unsafe { buffer.set_len(res) };
+                    break res;
+                }
+                Err(e) if is_would_block(&e) => {}
+                Err(e) => Err(e).unwrap(),
+            }
+            client.read_ready().await.unwrap();
+        }
+    };
+
+    let (write, read) = futures_util::join!(send_task, recv_task);
+    assert_eq!(write, 12);
+    assert_eq!(read, 12);
+    assert_eq!(buffer, b"Hello world!");
+}