support IOCP accept

Author: loongs-zhang

.github/workflows/ci-preemptive.sh

@@ -34,3 +34,13 @@ if [ "${TARGET}" = "x86_64-unknown-linux-gnu" ]; then
     "${CARGO}" test --target "${TARGET}" --no-default-features --features io_uring,preemptive
     "${CARGO}" test --target "${TARGET}" --no-default-features --features io_uring,preemptive --release
 fi
+
+# test IOCP
+if [ "${OS}" = "windows-latest" ]; then
+    cd "${PROJECT_DIR}"/core
+    "${CARGO}" test --target "${TARGET}" --no-default-features --features iocp,preemptive
+    "${CARGO}" test --target "${TARGET}" --no-default-features --features iocp,preemptive --release
+    cd "${PROJECT_DIR}"/open-coroutine
+    "${CARGO}" test --target "${TARGET}" --no-default-features --features iocp,preemptive
+    "${CARGO}" test --target "${TARGET}" --no-default-features --features iocp,preemptive --release
+fi

.github/workflows/ci.sh

@@ -34,3 +34,13 @@ if [ "${TARGET}" = "x86_64-unknown-linux-gnu" ]; then
     "${CARGO}" test --target "${TARGET}" --no-default-features --features io_uring
     "${CARGO}" test --target "${TARGET}" --no-default-features --features io_uring --release
 fi
+
+# test IOCP
+if [ "${OS}" = "windows-latest" ]; then
+    cd "${PROJECT_DIR}"/core
+    "${CARGO}" test --target "${TARGET}" --no-default-features --features iocp
+    "${CARGO}" test --target "${TARGET}" --no-default-features --features iocp --release
+    cd "${PROJECT_DIR}"/open-coroutine
+    "${CARGO}" test --target "${TARGET}" --no-default-features --features iocp
+    "${CARGO}" test --target "${TARGET}" --no-default-features --features iocp --release
+fi

Cargo.toml

@@ -28,6 +28,7 @@ mio = { version = "1.0", default-features = false }
 
 cfg-if = "1.0.0"
 polling = "2.8.0"
+educe = "0.6.0"
 
 libc = "0.2"
 rand = "0.8"
@@ -50,7 +51,6 @@ once_cell = "1"
 dashmap = "6"
 num_cpus = "1"
 uuid = "1"
-derivative = "2"
 tempfile = "3"
 cc = "1"
 syn = "2"

core/Cargo.toml

@@ -29,7 +29,7 @@ uuid = { workspace = true, features = [
     "v4",
     "fast-rng",
 ], optional = true }
-derivative = { workspace = true, optional = true }
+educe = { workspace = true, optional = true }
 core_affinity = { workspace = true, optional = true }
 crossbeam-utils = { workspace = true, optional = true }
 psm.workspace = true
@@ -77,7 +77,7 @@ backtrace.workspace = true
 log = ["tracing", "tracing-subscriber", "time"]
 
 # low-level raw coroutine
-korosensei = ["corosensei", "uuid", "nix/pthread", "derivative"]
+korosensei = ["corosensei", "uuid", "nix/pthread", "educe"]
 
 # Provide preemptive scheduling implementation.
 # Enable for default.
@@ -89,6 +89,12 @@ net = ["korosensei", "polling", "mio", "crossbeam-utils", "core_affinity"]
 # Provide io_uring adaptation, this feature only works in linux.
 io_uring = ["net", "io-uring"]
 
+# Provide IOCP adaptation, this feature only works in windows.
+iocp = ["net"]
+
+# Provide completion IOCP adaptation
+completion_io = ["io_uring", "iocp"]
+
 # Provide syscall implementation.
 syscall = ["net"]
 

core/src/co_pool/task.rs

@@ -1,16 +1,15 @@
 use crate::catch;
-use derivative::Derivative;
 
 /// 做C兼容时会用到
 pub type UserTaskFunc = extern "C" fn(usize) -> usize;
 
 /// The task impls.
 #[repr(C)]
-#[derive(Derivative)]
-#[derivative(Debug)]
+#[derive(educe::Educe)]
+#[educe(Debug)]
 pub struct Task<'t> {
     name: String,
-    #[derivative(Debug = "ignore")]
+    #[educe(Debug(ignore))]
     func: Box<dyn FnOnce(Option<usize>) -> Option<usize> + 't>,
     param: Option<usize>,
 }

core/src/common/timer.rs

@@ -1,5 +1,4 @@
 use crate::impl_display_by_debug;
-use derivative::Derivative;
 use std::collections::{BTreeMap, VecDeque};
 use std::ops::{Deref, DerefMut};
 use std::sync::atomic::{AtomicUsize, Ordering};
@@ -62,11 +61,11 @@ impl_display_by_debug!(TimerEntry<T>);
 
 /// A queue for managing multiple `TimerEntry`.
 #[repr(C)]
-#[derive(Derivative)]
-#[derivative(Debug, Eq, PartialEq)]
+#[derive(educe::Educe)]
+#[educe(Debug, Eq, PartialEq)]
 pub struct TimerList<T> {
     inner: BTreeMap<u64, TimerEntry<T>>,
-    #[derivative(PartialEq = "ignore")]
+    #[educe(PartialEq(ignore))]
     total: AtomicUsize,
 }
 

core/src/coroutine/suspender.rs

@@ -5,6 +5,7 @@ use std::collections::VecDeque;
 use std::time::Duration;
 
 thread_local! {
+    #[allow(clippy::missing_const_for_thread_local)]
     static TIMESTAMP: RefCell<VecDeque<u64>> = const { RefCell::new(VecDeque::new()) };
 }
 
@@ -52,14 +53,13 @@ pub use korosensei::Suspender;
 #[cfg(feature = "korosensei")]
 mod korosensei {
     use corosensei::Yielder;
-    use derivative::Derivative;
 
     /// Ths suspender implemented for coroutine.
     #[repr(C)]
-    #[derive(Derivative)]
-    #[derivative(Debug)]
+    #[derive(educe::Educe)]
+    #[educe(Debug)]
     pub struct Suspender<'s, Param, Yield> {
-        #[derivative(Debug = "ignore")]
+        #[educe(Debug(ignore))]
         inner: &'s Yielder<Param, Yield>,
     }
 

core/src/net/event_loop.rs

@@ -6,6 +6,8 @@ use crate::scheduler::SchedulableCoroutine;
 use crate::{error, impl_current_for, impl_display_by_debug, info};
 use crossbeam_utils::atomic::AtomicCell;
 use dashmap::DashSet;
+#[cfg(all(target_os = "linux", feature = "io_uring"))]
+use libc::{epoll_event, iovec, msghdr, off_t, size_t, sockaddr, socklen_t, ssize_t};
 use once_cell::sync::Lazy;
 use rand::Rng;
 use std::ffi::{c_char, c_int, c_void, CStr, CString};
@@ -16,11 +18,15 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, Condvar, Mutex};
 use std::thread::JoinHandle;
 use std::time::Duration;
+#[cfg(all(windows, feature = "iocp"))]
+use windows_sys::Win32::Networking::WinSock::{
+    setsockopt, SOCKADDR, SOCKET, SOL_SOCKET, SO_UPDATE_ACCEPT_CONTEXT,
+};
 
 cfg_if::cfg_if! {
-    if #[cfg(all(target_os = "linux", feature = "io_uring"))] {
-        use libc::{epoll_event, iovec, msghdr, off_t, size_t, sockaddr, socklen_t, ssize_t};
+    if #[cfg(any(all(target_os = "linux", feature = "io_uring"), all(windows, feature = "iocp")))] {
         use dashmap::DashMap;
+        use std::ffi::c_longlong;
     }
 }
 
@@ -32,11 +38,17 @@ pub(crate) struct EventLoop<'e> {
     stop: Arc<(Mutex<bool>, Condvar)>,
     shared_stop: Arc<(Mutex<AtomicUsize>, Condvar)>,
     cpu: usize,
-    #[cfg(all(target_os = "linux", feature = "io_uring"))]
+    #[cfg(any(
+        all(target_os = "linux", feature = "io_uring"),
+        all(windows, feature = "iocp")
+    ))]
     operator: crate::net::operator::Operator<'e>,
     #[allow(clippy::type_complexity)]
-    #[cfg(all(target_os = "linux", feature = "io_uring"))]
-    syscall_wait_table: DashMap<usize, Arc<(Mutex<Option<ssize_t>>, Condvar)>>,
+    #[cfg(any(
+        all(target_os = "linux", feature = "io_uring"),
+        all(windows, feature = "iocp")
+    ))]
+    syscall_wait_table: DashMap<usize, Arc<(Mutex<Option<c_longlong>>, Condvar)>>,
     selector: Poller,
     pool: CoroutinePool<'e>,
     phantom_data: PhantomData<&'e EventLoop<'e>>,
@@ -90,9 +102,15 @@ impl<'e> EventLoop<'e> {
             stop: Arc::new((Mutex::new(false), Condvar::new())),
             shared_stop,
             cpu,
-            #[cfg(all(target_os = "linux", feature = "io_uring"))]
+            #[cfg(any(
+                all(target_os = "linux", feature = "io_uring"),
+                all(windows, feature = "iocp")
+            ))]
             operator: crate::net::operator::Operator::new(cpu)?,
-            #[cfg(all(target_os = "linux", feature = "io_uring"))]
+            #[cfg(any(
+                all(target_os = "linux", feature = "io_uring"),
+                all(windows, feature = "iocp")
+            ))]
             syscall_wait_table: DashMap::new(),
             selector: Poller::new()?,
             pool: CoroutinePool::new(name, stack_size, min_size, max_size, keep_alive_time),
@@ -223,7 +241,29 @@ impl<'e> EventLoop<'e> {
             }
         }
 
-        #[cfg(all(target_os = "linux", feature = "io_uring"))]
+        cfg_if::cfg_if! {
+            if #[cfg(all(target_os = "linux", feature = "io_uring"))] {
+                left_time = self.adapt_io_uring(left_time)?;
+            } else if #[cfg(all(windows, feature = "iocp"))] {
+                left_time = self.adapt_iocp(left_time)?;
+            }
+        }
+
+        // use epoll/kevent/iocp
+        let mut events = Events::with_capacity(1024);
+        self.selector.select(&mut events, left_time)?;
+        #[allow(clippy::explicit_iter_loop)]
+        for event in events.iter() {
+            let token = event.get_token();
+            if event.readable() || event.writable() {
+                unsafe { self.resume(token) };
+            }
+        }
+        Ok(())
+    }
+
+    #[cfg(all(target_os = "linux", feature = "io_uring"))]
+    fn adapt_io_uring(&self, mut left_time: Option<Duration>) -> std::io::Result<Option<Duration>> {
         if crate::net::operator::support_io_uring() {
             // use io_uring
             let (count, mut cq, left) = self.operator.select(left_time, 0)?;
@@ -239,7 +279,7 @@ impl<'e> EventLoop<'e> {
                     if let Some((_, pair)) = self.syscall_wait_table.remove(&token) {
                         let (lock, cvar) = &*pair;
                         let mut pending = lock.lock().expect("lock failed");
-                        *pending = Some(result);
+                        *pending = Some(result.try_into().expect("result overflow"));
                         cvar.notify_one();
                     }
                     unsafe { self.resume(token) };
@@ -249,18 +289,47 @@ impl<'e> EventLoop<'e> {
                 left_time = Some(left.unwrap_or(Duration::ZERO));
             }
         }
+        Ok(left_time)
+    }
 
-        // use epoll/kevent/iocp
-        let mut events = Events::with_capacity(1024);
-        self.selector.select(&mut events, left_time)?;
-        #[allow(clippy::explicit_iter_loop)]
-        for event in events.iter() {
-            let token = event.get_token();
-            if event.readable() || event.writable() {
+    #[cfg(all(windows, feature = "iocp"))]
+    fn adapt_iocp(&self, mut left_time: Option<Duration>) -> std::io::Result<Option<Duration>> {
+        // use IOCP
+        let (count, mut cq, left) = self.operator.select(left_time, 0)?;
+        if count > 0 {
+            for cqe in &mut cq {
+                let token = cqe.token;
+                // resolve completed read/write tasks
+                // todo refactor IOCP impl
+                let result = match cqe.syscall {
+                    Syscall::accept => {
+                        unsafe {
+                            _ = setsockopt(
+                                cqe.socket,
+                                SOL_SOCKET,
+                                SO_UPDATE_ACCEPT_CONTEXT,
+                                (&cqe.from_fd as *const SOCKET).cast(),
+                                size_of::<SOCKET>() as c_int,
+                            )
+                        };
+                        cqe.socket.try_into().expect("result overflow")
+                    }
+                    _ => panic!("unsupported"),
+                };
+                eprintln!("IOCP finish {token} {result}");
+                if let Some((_, pair)) = self.syscall_wait_table.remove(&token) {
+                    let (lock, cvar) = &*pair;
+                    let mut pending = lock.lock().expect("lock failed");
+                    *pending = Some(result);
+                    cvar.notify_one();
+                }
                 unsafe { self.resume(token) };
             }
         }
-        Ok(())
+        if left != left_time {
+            left_time = Some(left.unwrap_or(Duration::ZERO));
+        }
+        Ok(left_time)
     }
 
     #[allow(clippy::unused_self)]
@@ -404,7 +473,7 @@ macro_rules! impl_io_uring {
             pub(super) fn $syscall(
                 &self,
                 $($arg: $arg_type),*
-            ) -> std::io::Result<Arc<(Mutex<Option<ssize_t>>, Condvar)>> {
+            ) -> std::io::Result<Arc<(Mutex<Option<c_longlong>>, Condvar)>> {
                 let token = EventLoop::token(Syscall::$syscall);
                 self.operator.$syscall(token, $($arg, )*)?;
                 let arc = Arc::new((Mutex::new(None), Condvar::new()));
@@ -439,6 +508,29 @@ impl_io_uring!(writev(fd: c_int, iov: *const iovec, iovcnt: c_int) -> ssize_t);
 impl_io_uring!(pwritev(fd: c_int, iov: *const iovec, iovcnt: c_int, offset: off_t) -> ssize_t);
 impl_io_uring!(sendmsg(fd: c_int, msg: *const msghdr, flags: c_int) -> ssize_t);
 
+macro_rules! impl_iocp {
+    ( $syscall: ident($($arg: ident : $arg_type: ty),*) -> $result: ty ) => {
+        #[cfg(all(windows, feature = "iocp"))]
+        impl EventLoop<'_> {
+            pub(super) fn $syscall(
+                &self,
+                $($arg: $arg_type),*
+            ) -> std::io::Result<Arc<(Mutex<Option<c_longlong>>, Condvar)>> {
+                let token = EventLoop::token(Syscall::$syscall);
+                self.operator.$syscall(token, $($arg, )*)?;
+                let arc = Arc::new((Mutex::new(None), Condvar::new()));
+                assert!(
+                    self.syscall_wait_table.insert(token, arc.clone()).is_none(),
+                    "The previous token was not retrieved in a timely manner"
+                );
+                Ok(arc)
+            }
+        }
+    }
+}
+
+impl_iocp!(accept(fd: SOCKET, addr: *mut SOCKADDR, len: *mut c_int) -> c_int);
+
 #[cfg(all(test, not(all(unix, feature = "preemptive"))))]
 mod tests {
     use crate::net::event_loop::EventLoop;