Merge pull request #265 from Berrysoft/dev/poll-fd

Author: George-Miao

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/mod.rs

@@ -0,0 +1,90 @@
+cfg_if::cfg_if! {
+    if #[cfg(windows)] {
+        #[path = "windows.rs"]
+        mod sys;
+    } else if #[cfg(unix)] {
+        #[path = "unix.rs"]
+        mod sys;
+    }
+}
+
+#[cfg(windows)]
+use std::os::windows::io::{AsRawSocket, RawSocket};
+use std::{io, ops::Deref};
+
+use compio_buf::IntoInner;
+use compio_driver::{AsRawFd, RawFd, SharedFd, ToSharedFd};
+
+/// A wrapper for socket, providing functionalities to wait for readiness.
+#[derive(Debug)]
+pub struct PollFd<T: AsRawFd>(sys::PollFd<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) -> io::Result<Self> {
+        Self::from_shared_fd(SharedFd::new(source))
+    }
+
+    pub(crate) fn from_shared_fd(inner: SharedFd<T>) -> io::Result<Self> {
+        Ok(Self(sys::PollFd::new(inner)?))
+    }
+}
+
+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.0.accept_ready().await
+    }
+
+    /// Wait for connect readiness.
+    pub async fn connect_ready(&self) -> io::Result<()> {
+        self.0.connect_ready().await
+    }
+
+    /// Wait for read readiness.
+    pub async fn read_ready(&self) -> io::Result<()> {
+        self.0.read_ready().await
+    }
+
+    /// Wait for write readiness.
+    pub async fn write_ready(&self) -> io::Result<()> {
+        self.0.write_ready().await
+    }
+}
+
+impl<T: AsRawFd> IntoInner for PollFd<T> {
+    type Inner = SharedFd<T>;
+
+    fn into_inner(self) -> Self::Inner {
+        self.0.into_inner()
+    }
+}
+
+impl<T: AsRawFd> ToSharedFd<T> for PollFd<T> {
+    fn to_shared_fd(&self) -> SharedFd<T> {
+        self.0.to_shared_fd()
+    }
+}
+
+impl<T: AsRawFd> AsRawFd for PollFd<T> {
+    fn as_raw_fd(&self) -> RawFd {
+        self.0.as_raw_fd()
+    }
+}
+
+#[cfg(windows)]
+impl<T: AsRawFd + AsRawSocket> AsRawSocket for PollFd<T> {
+    fn as_raw_socket(&self) -> RawSocket {
+        self.0.as_raw_socket()
+    }
+}
+
+impl<T: AsRawFd> Deref for PollFd<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}

compio-net/src/poll_fd/unix.rs

@@ -0,0 +1,70 @@
+use std::{io, ops::Deref};
+
+use compio_buf::{BufResult, IntoInner};
+use compio_driver::{
+    op::{Interest, PollOnce},
+    AsRawFd, RawFd, SharedFd, ToSharedFd,
+};
+
+#[derive(Debug)]
+pub struct PollFd<T: AsRawFd> {
+    inner: SharedFd<T>,
+}
+
+impl<T: AsRawFd> PollFd<T> {
+    pub fn new(inner: SharedFd<T>) -> io::Result<Self> {
+        Ok(Self { inner })
+    }
+}
+
+impl<T: AsRawFd + 'static> PollFd<T> {
+    pub async fn accept_ready(&self) -> io::Result<()> {
+        self.read_ready().await
+    }
+
+    pub async fn connect_ready(&self) -> io::Result<()> {
+        self.write_ready().await
+    }
+
+    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(())
+    }
+
+    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> AsRawFd for PollFd<T> {
+    fn as_raw_fd(&self) -> RawFd {
+        self.inner.as_raw_fd()
+    }
+}
+
+impl<T: AsRawFd> Deref for PollFd<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}

compio-net/src/poll_fd/windows.rs

@@ -0,0 +1,218 @@
+use std::{
+    io,
+    ops::Deref,
+    os::windows::io::{AsRawHandle, AsRawSocket, FromRawHandle, OwnedHandle, RawSocket},
+    pin::Pin,
+    ptr::null,
+    sync::atomic::{AtomicI32, AtomicUsize, Ordering},
+    task::Poll,
+};
+
+use compio_buf::{BufResult, IntoInner};
+use compio_driver::{syscall, AsRawFd, OpCode, OpType, RawFd, SharedFd, ToSharedFd};
+use windows_sys::Win32::{
+    Foundation::ERROR_IO_PENDING,
+    Networking::WinSock::{
+        WSAEnumNetworkEvents, WSAEventSelect, FD_ACCEPT, FD_CONNECT, FD_MAX_EVENTS, FD_READ,
+        FD_WRITE, WSANETWORKEVENTS,
+    },
+    System::{Threading::CreateEventW, IO::OVERLAPPED},
+};
+
+#[derive(Debug)]
+pub struct PollFd<T: AsRawFd> {
+    inner: SharedFd<T>,
+    event: WSAEvent,
+}
+
+impl<T: AsRawFd> PollFd<T> {
+    pub fn new(inner: SharedFd<T>) -> io::Result<Self> {
+        Ok(Self {
+            inner,
+            event: WSAEvent::new()?,
+        })
+    }
+}
+
+impl<T: AsRawFd + 'static> PollFd<T> {
+    pub async fn accept_ready(&self) -> io::Result<()> {
+        self.event.wait(self.to_shared_fd(), FD_ACCEPT).await
+    }
+
+    pub async fn connect_ready(&self) -> io::Result<()> {
+        self.event.wait(self.to_shared_fd(), FD_CONNECT).await
+    }
+
+    pub async fn read_ready(&self) -> io::Result<()> {
+        self.event.wait(self.to_shared_fd(), FD_READ).await
+    }
+
+    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>;
+
+    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> AsRawFd for PollFd<T> {
+    fn as_raw_fd(&self) -> RawFd {
+        self.inner.as_raw_fd()
+    }
+}
+
+impl<T: AsRawFd + AsRawSocket> AsRawSocket for PollFd<T> {
+    fn as_raw_socket(&self) -> RawSocket {
+        self.inner.as_raw_socket()
+    }
+}
+
+impl<T: AsRawFd> Deref for PollFd<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+#[derive(Debug)]
+pub struct WSAEvent {
+    ev_object: SharedFd<OwnedHandle>,
+    ev_record: [AtomicUsize; FD_MAX_EVENTS as usize],
+    events: AtomicI32,
+}
+
+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),
+        })
+    }
+
+    pub async fn wait<T: AsRawFd + 'static>(
+        &self,
+        mut 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() 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 mut ev_object = self.ev_object.clone();
+
+        let index = event.ilog2() 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,
+        };
+        loop {
+            let op = WaitWSAEvent::new(socket, ev_object, event);
+            let BufResult(res, op) = compio_runtime::submit(op).await;
+            WaitWSAEvent {
+                socket,
+                ev_object,
+                ..
+            } = op;
+            match res {
+                Ok(_) => break Ok(()),
+                Err(e) if e.kind() == io::ErrorKind::WouldBlock => {}
+                Err(e) => break Err(e),
+            }
+        }
+    }
+}
+
+struct WaitWSAEvent<T> {
+    socket: SharedFd<T>,
+    ev_object: SharedFd<OwnedHandle>,
+    event: i32,
+}
+
+impl<T> WaitWSAEvent<T> {
+    pub fn new(socket: SharedFd<T>, ev_object: SharedFd<OwnedHandle>, event: i32) -> Self {
+        Self {
+            socket,
+            ev_object,
+            event,
+        }
+    }
+}
+
+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() };
+        syscall!(
+            SOCKET,
+            WSAEnumNetworkEvents(
+                self.socket.as_raw_fd() as _,
+                self.ev_object.as_raw_handle() as _,
+                &mut events
+            )
+        )?;
+        let res = if (events.lNetworkEvents & self.event) != 0 {
+            events.iErrorCode[self.event.ilog2() as usize]
+        } else {
+            ERROR_IO_PENDING as _
+        };
+        if res == 0 {
+            Poll::Ready(Ok(0))
+        } else {
+            Poll::Ready(Err(io::Error::from_raw_os_error(res)))
+        }
+    }
+}