feat(driver,iocp): add iocp-global feature

Author: Berrysoft

compio-driver/Cargo.toml

@@ -42,7 +42,7 @@ socket2 = { workspace = true }
 [target.'cfg(windows)'.dependencies]
 compio-buf = { workspace = true, features = ["arrayvec"] }
 aligned-array = "1.0.1"
-crossbeam-skiplist = { workspace = true }
+crossbeam-skiplist = { workspace = true, optional = true }
 once_cell = { workspace = true }
 widestring = { workspace = true }
 windows-sys = { workspace = true, features = [
@@ -71,6 +71,7 @@ polling = "3.3.0"
 os_pipe = { workspace = true }
 
 [target.'cfg(unix)'.dependencies]
+crossbeam-channel = { workspace = true }
 crossbeam-queue = { workspace = true }
 libc = { workspace = true }
 
@@ -84,6 +85,8 @@ polling = ["dep:polling", "dep:os_pipe"]
 io-uring-sqe128 = []
 io-uring-cqe32 = []
 
+iocp-global = ["dep:crossbeam-skiplist"]
+
 # Nightly features
 once_cell_try = []
 nightly = ["once_cell_try"]

compio-driver/src/iocp/cp/global.rs

@@ -0,0 +1,162 @@
+#[cfg(feature = "once_cell_try")]
+use std::sync::OnceLock;
+use std::{
+    io,
+    os::windows::io::{AsRawHandle, RawHandle},
+    sync::atomic::{AtomicUsize, Ordering},
+    time::Duration,
+};
+
+use crossbeam_channel::{unbounded, Receiver, RecvTimeoutError, Sender, TryRecvError};
+use crossbeam_skiplist::SkipMap;
+#[cfg(not(feature = "once_cell_try"))]
+use once_cell::sync::OnceCell as OnceLock;
+use windows_sys::Win32::Foundation::ERROR_TIMEOUT;
+
+use super::CompletionPort;
+use crate::{Entry, Overlapped, RawFd};
+
+struct GlobalPort {
+    port: CompletionPort,
+    drivers: SkipMap<usize, Sender<Entry>>,
+}
+
+impl GlobalPort {
+    pub fn new() -> io::Result<Self> {
+        Ok(Self {
+            port: CompletionPort::new()?,
+            drivers: SkipMap::new(),
+        })
+    }
+
+    pub fn register(&self, driver: usize) -> Receiver<Entry> {
+        let (sender, receiver) = unbounded();
+        self.drivers.insert(driver, sender);
+        receiver
+    }
+
+    pub fn attach(&self, fd: RawFd) -> io::Result<()> {
+        self.port.attach(fd)
+    }
+
+    pub fn post<T: ?Sized>(
+        &self,
+        res: io::Result<usize>,
+        optr: *mut Overlapped<T>,
+    ) -> io::Result<()> {
+        self.port.post(res, optr)
+    }
+
+    pub fn push(&self, driver: usize, entry: Entry) {
+        self.drivers
+            .get(&driver)
+            .expect("driver should register first")
+            .value()
+            .send(entry)
+            .ok(); // It's OK if the driver has been dropped.
+    }
+}
+
+impl AsRawHandle for GlobalPort {
+    fn as_raw_handle(&self) -> RawHandle {
+        self.port.as_raw_handle()
+    }
+}
+
+static IOCP_PORT: OnceLock<GlobalPort> = OnceLock::new();
+
+#[inline]
+fn iocp_port() -> io::Result<&'static GlobalPort> {
+    IOCP_PORT.get_or_try_init(GlobalPort::new)
+}
+
+fn iocp_start() -> io::Result<()> {
+    let port = iocp_port()?;
+    std::thread::spawn(move || {
+        loop {
+            for (driver, entry) in port.port.poll(None)? {
+                port.push(driver, entry);
+            }
+        }
+        #[allow(unreachable_code)]
+        io::Result::Ok(())
+    });
+    Ok(())
+}
+
+static DRIVER_COUNTER: AtomicUsize = AtomicUsize::new(0);
+static IOCP_INIT_ONCE: OnceLock<()> = OnceLock::new();
+
+pub struct Port {
+    id: usize,
+    port: &'static GlobalPort,
+    receiver: Receiver<Entry>,
+}
+
+impl Port {
+    pub fn new() -> io::Result<Self> {
+        IOCP_INIT_ONCE.get_or_try_init(iocp_start)?;
+
+        let id = DRIVER_COUNTER.fetch_add(1, Ordering::AcqRel);
+        let port = iocp_port()?;
+        let receiver = port.register(id);
+        Ok(Self { id, port, receiver })
+    }
+
+    pub fn id(&self) -> usize {
+        self.id
+    }
+
+    pub fn attach(&mut self, fd: RawFd) -> io::Result<()> {
+        self.port.attach(fd)
+    }
+
+    pub fn handle(&self) -> PortHandle {
+        PortHandle::new(self.port)
+    }
+
+    pub fn poll(&self, timeout: Option<Duration>) -> io::Result<impl Iterator<Item = Entry> + '_> {
+        let e = if let Some(timeout) = timeout {
+            match self.receiver.recv_timeout(timeout) {
+                Ok(e) => e,
+                Err(e) => match e {
+                    RecvTimeoutError::Timeout => {
+                        return Err(io::Error::from_raw_os_error(ERROR_TIMEOUT as _));
+                    }
+                    RecvTimeoutError::Disconnected => {
+                        unreachable!("IOCP thread should not exit")
+                    }
+                },
+            }
+        } else {
+            self.receiver.recv().expect("IOCP thread should not exit")
+        };
+        Ok(Some(e)
+            .into_iter()
+            .chain(std::iter::from_fn(|| match self.receiver.try_recv() {
+                Ok(e) => Some(e),
+                Err(e) => match e {
+                    TryRecvError::Empty => None,
+                    TryRecvError::Disconnected => unreachable!("IOCP thread should not exit"),
+                },
+            })))
+    }
+}
+
+pub struct PortHandle {
+    port: &'static GlobalPort,
+}
+
+impl PortHandle {
+    fn new(port: &'static GlobalPort) -> Self {
+        Self { port }
+    }
+
+    pub fn post<T: ?Sized>(
+        &self,
+        res: io::Result<usize>,
+        optr: *mut Overlapped<T>,
+    ) -> io::Result<()> {
+        self.port.post(res, optr)
+    }
+}

compio-driver/src/iocp/cp/mod.rs

@@ -0,0 +1,155 @@
+use std::{
+    io,
+    os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle, RawHandle},
+    time::Duration,
+};
+
+use compio_buf::arrayvec::ArrayVec;
+use compio_log::*;
+use windows_sys::Win32::{
+    Foundation::{
+        RtlNtStatusToDosError, ERROR_BAD_COMMAND, ERROR_HANDLE_EOF, ERROR_IO_INCOMPLETE,
+        ERROR_NO_DATA, FACILITY_NTWIN32, INVALID_HANDLE_VALUE, NTSTATUS, STATUS_PENDING,
+        STATUS_SUCCESS,
+    },
+    Storage::FileSystem::SetFileCompletionNotificationModes,
+    System::{
+        SystemServices::ERROR_SEVERITY_ERROR,
+        Threading::INFINITE,
+        WindowsProgramming::{FILE_SKIP_COMPLETION_PORT_ON_SUCCESS, FILE_SKIP_SET_EVENT_ON_HANDLE},
+        IO::{
+            CreateIoCompletionPort, GetQueuedCompletionStatusEx, PostQueuedCompletionStatus,
+            OVERLAPPED_ENTRY,
+        },
+    },
+};
+
+use crate::{syscall, Entry, Overlapped, RawFd};
+
+cfg_if::cfg_if! {
+    if #[cfg(feature = "iocp-global")] {
+        mod global;
+        pub use global::*;
+    } else {
+        mod multi;
+        pub use multi::*;
+    }
+}
+
+struct CompletionPort {
+    port: OwnedHandle,
+}
+
+impl CompletionPort {
+    pub fn new() -> io::Result<Self> {
+        let port = syscall!(BOOL, CreateIoCompletionPort(INVALID_HANDLE_VALUE, 0, 0, 1))?;
+        trace!("new iocp handle: {port}");
+        let port = unsafe { OwnedHandle::from_raw_handle(port as _) };
+        Ok(Self { port })
+    }
+
+    pub fn attach(&self, fd: RawFd) -> io::Result<()> {
+        syscall!(
+            BOOL,
+            CreateIoCompletionPort(fd as _, self.port.as_raw_handle() as _, 0, 0)
+        )?;
+        syscall!(
+            BOOL,
+            SetFileCompletionNotificationModes(
+                fd as _,
+                (FILE_SKIP_COMPLETION_PORT_ON_SUCCESS | FILE_SKIP_SET_EVENT_ON_HANDLE) as _
+            )
+        )?;
+        Ok(())
+    }
+
+    pub fn post<T: ?Sized>(
+        &self,
+        res: io::Result<usize>,
+        optr: *mut Overlapped<T>,
+    ) -> io::Result<()> {
+        if let Err(e) = &res {
+            let code = e.raw_os_error().unwrap_or(ERROR_BAD_COMMAND as _);
+            unsafe { &mut *optr }.base.Internal = ntstatus_from_win32(code) as _;
+        }
+        // We have to use CompletionKey to transfer the result because it is large
+        // enough. It is OK because we set it to zero when attaching handles to IOCP.
+        syscall!(
+            BOOL,
+            PostQueuedCompletionStatus(
+                self.port.as_raw_handle() as _,
+                0,
+                res.unwrap_or_default(),
+                optr.cast()
+            )
+        )?;
+        Ok(())
+    }
+
+    pub fn poll(
+        &self,
+        timeout: Option<Duration>,
+    ) -> io::Result<impl Iterator<Item = (usize, Entry)>> {
+        const DEFAULT_CAPACITY: usize = 1024;
+
+        let mut entries = ArrayVec::<OVERLAPPED_ENTRY, { DEFAULT_CAPACITY }>::new();
+        let mut recv_count = 0;
+        let timeout = match timeout {
+            Some(timeout) => timeout.as_millis() as u32,
+            None => INFINITE,
+        };
+        syscall!(
+            BOOL,
+            GetQueuedCompletionStatusEx(
+                self.port.as_raw_handle() as _,
+                entries.as_mut_ptr(),
+                DEFAULT_CAPACITY as _,
+                &mut recv_count,
+                timeout,
+                0
+            )
+        )?;
+        trace!("recv_count: {recv_count}");
+        unsafe { entries.set_len(recv_count as _) };
+
+        Ok(entries.into_iter().map(|entry| {
+            let transferred = entry.dwNumberOfBytesTransferred;
+            trace!("entry transferred: {transferred}");
+            // Any thin pointer is OK because we don't use the type of opcode.
+            let overlapped_ptr: *mut Overlapped<()> = entry.lpOverlapped.cast();
+            let overlapped = unsafe { &*overlapped_ptr };
+            let res = if matches!(
+                overlapped.base.Internal as NTSTATUS,
+                STATUS_SUCCESS | STATUS_PENDING
+            ) {
+                if entry.lpCompletionKey != 0 {
+                    Ok(entry.lpCompletionKey)
+                } else {
+                    Ok(transferred as _)
+                }
+            } else {
+                let error = unsafe { RtlNtStatusToDosError(overlapped.base.Internal as _) };
+                match error {
+                    ERROR_IO_INCOMPLETE | ERROR_HANDLE_EOF | ERROR_NO_DATA => Ok(0),
+                    _ => Err(io::Error::from_raw_os_error(error as _)),
+                }
+            };
+            (overlapped.driver, Entry::new(overlapped.user_data, res))
+        }))
+    }
+}
+
+impl AsRawHandle for CompletionPort {
+    fn as_raw_handle(&self) -> RawHandle {
+        self.port.as_raw_handle()
+    }
+}
+
+#[inline]
+fn ntstatus_from_win32(x: i32) -> NTSTATUS {
+    if x <= 0 {
+        x
+    } else {
+        ((x) & 0x0000FFFF) | (FACILITY_NTWIN32 << 16) as NTSTATUS | ERROR_SEVERITY_ERROR as NTSTATUS
+    }
+}