feat(net): use WSAEventSelect on Windows for PollFd

Author: Berrysoft

compio-net/src/poll_fd.rs

@@ -1,31 +1,38 @@
-#[cfg(unix)]
-use std::os::fd::FromRawFd;
 #[cfg(windows)]
-use std::os::windows::io::{AsRawSocket, FromRawSocket, RawSocket};
+use std::os::windows::io::{AsRawSocket, RawSocket};
 use std::{io, ops::Deref};
 
-use compio_buf::{BufResult, IntoInner};
-#[cfg(unix)]
-use compio_driver::op::{Interest, PollOnce};
+use compio_buf::IntoInner;
 use compio_driver::{AsRawFd, RawFd, SharedFd, ToSharedFd};
+#[cfg(windows)]
+use windows_sys::Win32::Networking::WinSock::{FD_ACCEPT, FD_CONNECT, FD_READ, FD_WRITE};
+#[cfg(unix)]
+use {
+    compio_buf::BufResult,
+    compio_driver::op::{Interest, PollOnce},
+};
 
 /// A wrapper for socket, providing functionalities to wait for readiness.
 #[derive(Debug)]
 pub struct PollFd<T: AsRawFd> {
     inner: SharedFd<T>,
+    #[cfg(windows)]
+    event: sys::WSAEvent,
 }
 
 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 fn new(source: T) -> io::Result<Self> {
+        Self::from_shared_fd(SharedFd::new(source))
     }
 
-    pub(crate) fn from_shared_fd(inner: SharedFd<T>) -> Self {
-        Self { inner }
+    pub(crate) fn from_shared_fd(inner: SharedFd<T>) -> io::Result<Self> {
+        Ok(Self {
+            inner,
+            #[cfg(windows)]
+            event: sys::WSAEvent::new()?,
+        })
     }
 }
 
@@ -59,6 +66,30 @@ impl<T: AsRawFd + 'static> PollFd<T> {
     }
 }
 
+#[cfg(windows)]
+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.event.wait(self.to_shared_fd(), FD_ACCEPT).await
+    }
+
+    /// Wait for connect readiness.
+    pub async fn connect_ready(&self) -> io::Result<()> {
+        self.event.wait(self.to_shared_fd(), FD_CONNECT).await
+    }
+
+    /// Wait for read readiness.
+    pub async fn read_ready(&self) -> io::Result<()> {
+        self.event.wait(self.to_shared_fd(), FD_READ).await
+    }
+
+    /// Wait for write readiness.
+    pub async fn write_ready(&self) -> io::Result<()> {
+        self.event.wait(self.to_shared_fd(), FD_WRITE).await
+    }
+}
+
 impl<T: AsRawFd> IntoInner for PollFd<T> {
     type Inner = SharedFd<T>;
 
@@ -73,14 +104,6 @@ impl<T: AsRawFd> ToSharedFd<T> for PollFd<T> {
     }
 }
 
-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()
@@ -94,24 +117,150 @@ impl<T: AsRawFd + AsRawSocket> AsRawSocket for PollFd<T> {
     }
 }
 
-#[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))
+impl<T: AsRawFd> Deref for PollFd<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
     }
 }
 
 #[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))
+mod sys {
+    use std::{
+        io,
+        os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle},
+        pin::Pin,
+        ptr::null,
+        sync::atomic::{AtomicI32, AtomicUsize, Ordering},
+        task::Poll,
+    };
+
+    use compio_buf::{BufResult, IntoInner};
+    use compio_driver::{syscall, AsRawFd, OpCode, OpType, SharedFd};
+    use windows_sys::Win32::{
+        Networking::WinSock::{WSAEnumNetworkEvents, WSAEventSelect, WSANETWORKEVENTS},
+        System::{Threading::CreateEventW, IO::OVERLAPPED},
+    };
+
+    const EVENT_COUNT: usize = 5;
+
+    #[derive(Debug)]
+    pub struct WSAEvent {
+        ev_object: SharedFd<OwnedHandle>,
+        ev_record: [AtomicUsize; EVENT_COUNT],
+        events: AtomicI32,
     }
-}
 
-impl<T: AsRawFd> Deref for PollFd<T> {
-    type Target = T;
+    impl WSAEvent {
+        pub fn new() -> io::Result<Self> {
+            Ok(Self {
+                ev_object: SharedFd::new(unsafe {
+                    OwnedHandle::from_raw_handle(syscall!(
+                        HANDLE,
+                        CreateEventW(null(), 1, 0, null())
+                    )? as _)
+                }),
+                ev_record: Default::default(),
+                events: AtomicI32::new(0),
+            })
+        }
 
-    fn deref(&self) -> &Self::Target {
-        &self.inner
+        pub async fn wait<T: AsRawFd + 'static>(
+            &self,
+            socket: SharedFd<T>,
+            event: u32,
+        ) -> io::Result<()> {
+            struct EventGuard<'a> {
+                wsa_event: &'a WSAEvent,
+                event: i32,
+            }
+
+            impl Drop for EventGuard<'_> {
+                fn drop(&mut self) {
+                    let index = (self.event.ilog2() - 1) as usize;
+                    if self.wsa_event.ev_record[index].fetch_sub(1, Ordering::Relaxed) == 1 {
+                        self.wsa_event
+                            .events
+                            .fetch_add(!self.event, Ordering::Relaxed);
+                    }
+                }
+            }
+
+            let event = event as i32;
+            let ev_object = self.ev_object.clone();
+
+            let index = (event.ilog2() - 1) as usize;
+            let events = if self.ev_record[index].fetch_add(1, Ordering::Relaxed) == 0 {
+                self.events.fetch_or(event, Ordering::Relaxed) | event
+            } else {
+                self.events.load(Ordering::Relaxed)
+            };
+            syscall!(
+                SOCKET,
+                WSAEventSelect(
+                    socket.as_raw_fd() as _,
+                    ev_object.as_raw_handle() as _,
+                    events
+                )
+            )?;
+            let _guard = EventGuard {
+                wsa_event: self,
+                event,
+            };
+            let op = WaitWSAEvent::new(socket, ev_object, index + 1);
+            let BufResult(res, _) = compio_runtime::submit(op).await;
+            res?;
+            Ok(())
+        }
+    }
+
+    struct WaitWSAEvent<T> {
+        socket: SharedFd<T>,
+        ev_object: SharedFd<OwnedHandle>,
+        index: usize,
+    }
+
+    impl<T> WaitWSAEvent<T> {
+        pub fn new(socket: SharedFd<T>, ev_object: SharedFd<OwnedHandle>, index: usize) -> Self {
+            Self {
+                socket,
+                ev_object,
+                index,
+            }
+        }
+    }
+
+    impl<T> IntoInner for WaitWSAEvent<T> {
+        type Inner = SharedFd<OwnedHandle>;
+
+        fn into_inner(self) -> Self::Inner {
+            self.ev_object
+        }
+    }
+
+    impl<T: AsRawFd> OpCode for WaitWSAEvent<T> {
+        fn op_type(&self) -> OpType {
+            OpType::Event(self.ev_object.as_raw_fd())
+        }
+
+        unsafe fn operate(self: Pin<&mut Self>, _optr: *mut OVERLAPPED) -> Poll<io::Result<usize>> {
+            let mut events: WSANETWORKEVENTS = unsafe { std::mem::zeroed() };
+            events.lNetworkEvents = 10;
+            syscall!(
+                SOCKET,
+                WSAEnumNetworkEvents(
+                    self.socket.as_raw_fd() as _,
+                    self.ev_object.as_raw_handle() as _,
+                    &mut events
+                )
+            )?;
+            let res = events.iErrorCode[self.index + 1];
+            if res == 0 {
+                Poll::Ready(Ok(0))
+            } else {
+                Poll::Ready(Err(io::Error::from_raw_os_error(res)))
+            }
+        }
     }
 }

compio-net/src/socket.rs

@@ -36,11 +36,11 @@ impl Socket {
         self.socket.local_addr()
     }
 
-    pub fn to_poll_fd(&self) -> PollFd<Socket2> {
+    pub fn to_poll_fd(&self) -> io::Result<PollFd<Socket2>> {
         PollFd::from_shared_fd(self.to_shared_fd())
     }
 
-    pub fn into_poll_fd(self) -> PollFd<Socket2> {
+    pub fn into_poll_fd(self) -> io::Result<PollFd<Socket2>> {
         PollFd::from_shared_fd(self.socket.into_inner())
     }
 

compio-net/src/tcp.rs

@@ -207,12 +207,12 @@ impl TcpStream {
     }
 
     /// Create [`PollFd`] from inner socket.
-    pub fn to_poll_fd(&self) -> PollFd<Socket2> {
+    pub fn to_poll_fd(&self) -> io::Result<PollFd<Socket2>> {
         self.inner.to_poll_fd()
     }
 
     /// Create [`PollFd`] from inner socket.
-    pub fn into_poll_fd(self) -> PollFd<Socket2> {
+    pub fn into_poll_fd(self) -> io::Result<PollFd<Socket2>> {
         self.inner.into_poll_fd()
     }
 }

compio-net/src/unix.rs

@@ -189,12 +189,12 @@ impl UnixStream {
     }
 
     /// Create [`PollFd`] from inner socket.
-    pub fn to_poll_fd(&self) -> PollFd<Socket2> {
+    pub fn to_poll_fd(&self) -> io::Result<PollFd<Socket2>> {
         self.inner.to_poll_fd()
     }
 
     /// Create [`PollFd`] from inner socket.
-    pub fn into_poll_fd(self) -> PollFd<Socket2> {
+    pub fn into_poll_fd(self) -> io::Result<PollFd<Socket2>> {
         self.inner.into_poll_fd()
     }
 }

compio-net/tests/poll.rs

@@ -26,21 +26,21 @@ async fn poll_connect() {
         .unwrap();
     listener.listen(4).unwrap();
     let addr = listener.local_addr().unwrap();
-    let listener = PollFd::new(listener);
+    let listener = PollFd::new(listener).unwrap();
     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(),
+                Err(e) => panic!("{:?}", e),
             }
         }
     };
 
     let client = Socket::new(Domain::IPV4, Type::STREAM, Some(Protocol::TCP)).unwrap();
     client.set_nonblocking(true).unwrap();
-    let client = PollFd::new(client);
+    let client = PollFd::new(client).unwrap();
     let res = client.connect(&addr);
     let tx = if let Err(e) = res {
         assert!(is_would_block(&e));
@@ -53,14 +53,14 @@ async fn poll_connect() {
     };
 
     tx.set_nonblocking(true).unwrap();
-    let tx = PollFd::new(tx);
+    let tx = PollFd::new(tx).unwrap();
 
     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(),
+                Err(e) => panic!("{:?}", e),
             }
             tx.write_ready().await.unwrap();
         }
@@ -75,7 +75,7 @@ async fn poll_connect() {
                     break res;
                 }
                 Err(e) if is_would_block(&e) => {}
-                Err(e) => Err(e).unwrap(),
+                Err(e) => panic!("{:?}", e),
             }
             client.read_ready().await.unwrap();
         }