feat(driver,iocp): use IOCP instead of crossbeam channel

Berrysoft · Berrysoft · commit e9b7e6f1461c · 2024-03-04T16:07:02.000+08:00
diff --git a/compio-driver/Cargo.toml b/compio-driver/Cargo.toml
@@ -42,7 +42,6 @@ socket2 = { workspace = true }
 [target.'cfg(windows)'.dependencies]
 compio-buf = { workspace = true, features = ["arrayvec"] }
 aligned-array = "1.0.1"
-crossbeam-skiplist = { workspace = true, optional = true }
 once_cell = { workspace = true }
 widestring = { workspace = true }
 windows-sys = { workspace = true, features = [
@@ -85,7 +84,7 @@ polling = ["dep:polling", "dep:os_pipe"]
 io-uring-sqe128 = []
 io-uring-cqe32 = []
 
-iocp-global = ["dep:crossbeam-skiplist"]
+iocp-global = []
 
 # Nightly features
 once_cell_try = []
diff --git a/compio-driver/src/iocp/cp/global.rs b/compio-driver/src/iocp/cp/global.rs
@@ -3,45 +3,27 @@ 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,
-    System::IO::{PostQueuedCompletionStatus, OVERLAPPED},
-};
+use windows_sys::Win32::System::IO::PostQueuedCompletionStatus;
 
 use super::CompletionPort;
 use crate::{syscall, 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 deregister(&self, driver: usize) {
-        self.drivers.remove(&driver);
-    }
-
     pub fn attach(&self, fd: RawFd) -> io::Result<()> {
         self.port.attach(fd)
     }
@@ -53,12 +35,6 @@ impl GlobalPort {
     ) -> io::Result<()> {
         self.port.post(res, optr)
     }
-
-    pub fn push(&self, driver: usize, entry: Entry) {
-        if let Some(e) = self.drivers.get(&driver) {
-            e.value().send(entry).ok(); // It's OK if the driver has been dropped.
-        }
-    }
 }
 
 impl AsRawHandle for GlobalPort {
@@ -78,8 +54,20 @@ fn iocp_start() -> io::Result<()> {
     let port = iocp_port()?;
     std::thread::spawn(move || {
         loop {
-            for (driver, entry) in port.port.poll(None, None)? {
-                port.push(driver.0, entry);
+            for entry in port.port.poll_raw(None)? {
+                // 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 };
+                syscall!(
+                    BOOL,
+                    PostQueuedCompletionStatus(
+                        overlapped.driver,
+                        entry.dwNumberOfBytesTransferred,
+                        entry.lpCompletionKey,
+                        entry.lpOverlapped,
+                    )
+                )
+                .ok();
             }
         }
         #[allow(unreachable_code)]
@@ -88,68 +76,39 @@ fn iocp_start() -> io::Result<()> {
     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>,
+    port: CompletionPort,
+    global_port: &'static GlobalPort,
 }
 
 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) -> PortId {
-        PortId(self.id)
+        Ok(Self {
+            port: CompletionPort::new()?,
+            global_port: iocp_port()?,
+        })
     }
 
     pub fn attach(&mut self, fd: RawFd) -> io::Result<()> {
-        self.port.attach(fd)
+        self.global_port.attach(fd)
     }
 
     pub fn handle(&self) -> PortHandle {
-        PortHandle::new(self.port)
+        PortHandle::new(self.global_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"),
-                },
-            })))
+        self.port.poll(timeout, None)
     }
 }
 
-impl Drop for Port {
-    fn drop(&mut self) {
-        self.port.deregister(self.id);
+impl AsRawHandle for Port {
+    fn as_raw_handle(&self) -> RawHandle {
+        self.port.as_raw_handle()
     }
 }
 
@@ -170,18 +129,3 @@ impl PortHandle {
         self.port.post(res, optr)
     }
 }
-
-/// The unique ID of IOCP driver.
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub struct PortId(usize);
-
-impl PortId {
-    /// Post raw entry to IOCP.
-    pub fn post_raw(&self, transferred: u32, key: usize, optr: *mut OVERLAPPED) -> io::Result<()> {
-        syscall!(
-            BOOL,
-            PostQueuedCompletionStatus(iocp_port()?.as_raw_handle() as _, transferred, key, optr)
-        )?;
-        Ok(())
-    }
-}
diff --git a/compio-driver/src/iocp/cp/mod.rs b/compio-driver/src/iocp/cp/mod.rs
@@ -9,7 +9,7 @@ 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,
+        ERROR_NO_DATA, FACILITY_NTWIN32, HANDLE, INVALID_HANDLE_VALUE, NTSTATUS, STATUS_PENDING,
         STATUS_SUCCESS,
     },
     Storage::FileSystem::SetFileCompletionNotificationModes,
@@ -86,13 +86,10 @@ impl CompletionPort {
         Ok(())
     }
 
-    // If current_driver is specified, any entry that doesn't belong the driver will
-    // be reposted. The driver id will be used as IOCP handle.
-    pub fn poll(
+    pub fn poll_raw(
         &self,
         timeout: Option<Duration>,
-        current_driver: Option<PortId>,
-    ) -> io::Result<impl Iterator<Item = (PortId, Entry)>> {
+    ) -> io::Result<impl Iterator<Item = OVERLAPPED_ENTRY>> {
         const DEFAULT_CAPACITY: usize = 1024;
 
         let mut entries = ArrayVec::<OVERLAPPED_ENTRY, { DEFAULT_CAPACITY }>::new();
@@ -115,22 +112,33 @@ impl CompletionPort {
         trace!("recv_count: {recv_count}");
         unsafe { entries.set_len(recv_count as _) };
 
-        Ok(entries.into_iter().map(move |entry| {
-            let transferred = entry.dwNumberOfBytesTransferred;
-            trace!("entry transferred: {transferred}");
+        Ok(entries.into_iter())
+    }
+
+    // If current_driver is specified, any entry that doesn't belong the driver will
+    // be reposted. The driver id will be used as IOCP handle.
+    pub fn poll(
+        &self,
+        timeout: Option<Duration>,
+        current_driver: Option<HANDLE>,
+    ) -> io::Result<impl Iterator<Item = Entry>> {
+        Ok(self.poll_raw(timeout)?.map(move |entry| {
             // 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 };
             if let Some(current_driver) = current_driver {
                 if overlapped.driver != current_driver {
-                    overlapped
-                        .driver
-                        .post_raw(
+                    // Repose the entry to correct port.
+                    syscall!(
+                        BOOL,
+                        PostQueuedCompletionStatus(
+                            overlapped.driver,
                             entry.dwNumberOfBytesTransferred,
                             entry.lpCompletionKey,
                             entry.lpOverlapped,
                         )
-                        .ok();
+                    )
+                    .ok();
                 }
             }
             let res = if matches!(
@@ -140,7 +148,7 @@ impl CompletionPort {
                 if entry.lpCompletionKey != 0 {
                     Ok(entry.lpCompletionKey)
                 } else {
-                    Ok(transferred as _)
+                    Ok(entry.dwNumberOfBytesTransferred as _)
                 }
             } else {
                 let error = unsafe { RtlNtStatusToDosError(overlapped.base.Internal as _) };
@@ -149,7 +157,7 @@ impl CompletionPort {
                     _ => Err(io::Error::from_raw_os_error(error as _)),
                 }
             };
-            (overlapped.driver, Entry::new(overlapped.user_data, res))
+            Entry::new(overlapped.user_data, res)
         }))
     }
 }
diff --git a/compio-driver/src/iocp/cp/multi.rs b/compio-driver/src/iocp/cp/multi.rs
@@ -1,12 +1,12 @@
-use std::{io, os::windows::io::AsRawHandle, sync::Arc, time::Duration};
-
-use windows_sys::Win32::{
-    Foundation::HANDLE,
-    System::IO::{PostQueuedCompletionStatus, OVERLAPPED},
+use std::{
+    io,
+    os::windows::io::{AsRawHandle, RawHandle},
+    sync::Arc,
+    time::Duration,
 };
 
 use super::CompletionPort;
-use crate::{syscall, Entry, Overlapped, RawFd};
+use crate::{Entry, Overlapped, RawFd};
 
 pub struct Port {
     port: Arc<CompletionPort>,
@@ -19,10 +19,6 @@ impl Port {
         })
     }
 
-    pub fn id(&self) -> PortId {
-        PortId(self.port.as_raw_handle() as _)
-    }
-
     pub fn attach(&mut self, fd: RawFd) -> io::Result<()> {
         self.port.attach(fd)
     }
@@ -32,14 +28,14 @@ impl Port {
     }
 
     pub fn poll(&self, timeout: Option<Duration>) -> io::Result<impl Iterator<Item = Entry> + '_> {
-        let current_id = self.id();
-        self.port.poll(timeout, Some(current_id)).map(move |it| {
-            it.filter_map(
-                move |(id, entry)| {
-                    if id == current_id { Some(entry) } else { None }
-                },
-            )
-        })
+        let current_id = self.as_raw_handle() as _;
+        self.port.poll(timeout, Some(current_id))
+    }
+}
+
+impl AsRawHandle for Port {
+    fn as_raw_handle(&self) -> RawHandle {
+        self.port.as_raw_handle()
     }
 }
 
@@ -60,18 +56,3 @@ impl PortHandle {
         self.port.post(res, optr)
     }
 }
-
-/// The unique ID of IOCP driver.
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub struct PortId(HANDLE);
-
-impl PortId {
-    /// Post raw entry to IOCP.
-    pub fn post_raw(&self, transferred: u32, key: usize, optr: *mut OVERLAPPED) -> io::Result<()> {
-        syscall!(
-            BOOL,
-            PostQueuedCompletionStatus(self.0, transferred, key, optr)
-        )?;
-        Ok(())
-    }
-}
diff --git a/compio-driver/src/iocp/mod.rs b/compio-driver/src/iocp/mod.rs
