Merge pull request #292 from George-Miao/fix/driver-type

Author: George-Miao

compio-driver/src/driver_type.rs

@@ -0,0 +1,115 @@
+use std::sync::atomic::{AtomicU8, Ordering};
+
+const UNINIT: u8 = u8::MAX;
+const IO_URING: u8 = 0;
+const POLLING: u8 = 1;
+const IOCP: u8 = 2;
+
+static DRIVER_TYPE: AtomicU8 = AtomicU8::new(UNINIT);
+
+/// Representing underlying driver type the fusion driver is using
+#[repr(u8)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum DriverType {
+    /// Using `polling` driver
+    Poll    = POLLING,
+
+    /// Using `io-uring` driver
+    IoUring = IO_URING,
+
+    /// Using `iocp` driver
+    IOCP    = IOCP,
+}
+
+impl DriverType {
+    fn from_num(n: u8) -> Self {
+        match n {
+            IO_URING => Self::IoUring,
+            POLLING => Self::Poll,
+            IOCP => Self::IOCP,
+            _ => unreachable!("invalid driver type"),
+        }
+    }
+
+    /// Get the underlying driver type
+    fn get() -> DriverType {
+        cfg_if::cfg_if! {
+            if #[cfg(windows)] {
+                DriverType::IOCP
+            } else if #[cfg(all(target_os = "linux", feature = "polling", feature = "io-uring"))] {
+                use io_uring::opcode::*;
+
+                // Add more opcodes here if used
+                const USED_OP: &[u8] = &[
+                    Read::CODE,
+                    Readv::CODE,
+                    Write::CODE,
+                    Writev::CODE,
+                    Fsync::CODE,
+                    Accept::CODE,
+                    Connect::CODE,
+                    RecvMsg::CODE,
+                    SendMsg::CODE,
+                    AsyncCancel::CODE,
+                    OpenAt::CODE,
+                    Close::CODE,
+                    Shutdown::CODE,
+                    // Linux kernel 5.19
+                    #[cfg(any(
+                        feature = "io-uring-sqe128",
+                        feature = "io-uring-cqe32",
+                        feature = "io-uring-socket"
+                    ))]
+                    Socket::CODE,
+                ];
+
+                (|| {
+                    let uring = io_uring::IoUring::new(2)?;
+                    let mut probe = io_uring::Probe::new();
+                    uring.submitter().register_probe(&mut probe)?;
+                    if USED_OP.iter().all(|op| probe.is_supported(*op)) {
+                        std::io::Result::Ok(DriverType::IoUring)
+                    } else {
+                        Ok(DriverType::Poll)
+                    }
+                })()
+                .unwrap_or(DriverType::Poll) // Should we fail here?
+            } else if #[cfg(all(target_os = "linux", feature = "io-uring"))] {
+                DriverType::IoUring
+            } else if #[cfg(unix)] {
+                DriverType::Poll
+            } else {
+                compile_error!("unsupported platform");
+            }
+        }
+    }
+
+    /// Get the underlying driver type and cache it. Following calls will return
+    /// the cached value.
+    pub fn current() -> DriverType {
+        match DRIVER_TYPE.load(Ordering::Acquire) {
+            UNINIT => {}
+            x => return DriverType::from_num(x),
+        }
+        let dev_ty = Self::get();
+
+        DRIVER_TYPE.store(dev_ty as u8, Ordering::Release);
+
+        dev_ty
+    }
+
+    /// Check if the current driver is `polling`
+    pub fn is_polling() -> bool {
+        Self::current() == DriverType::Poll
+    }
+
+    /// Check if the current driver is `io-uring`
+    pub fn is_iouring() -> bool {
+        Self::current() == DriverType::IoUring
+    }
+
+    /// Check if the current driver is `iocp`
+    pub fn is_iocp() -> bool {
+        Self::current() == DriverType::IOCP
+    }
+}

compio-driver/src/fusion/mod.rs

@@ -10,99 +10,13 @@ pub(crate) mod op;
 pub use std::os::fd::{AsRawFd, OwnedFd, RawFd};
 use std::{io, task::Poll, time::Duration};
 
-pub use driver_type::DriverType;
 pub(crate) use iour::{sockaddr_storage, socklen_t};
 pub use iour::{OpCode as IourOpCode, OpEntry};
 pub use poll::{Decision, OpCode as PollOpCode};
 
+pub use crate::driver_type::DriverType; // Re-export so current user won't be broken
 use crate::{Key, OutEntries, ProactorBuilder};
 
-mod driver_type {
-    use std::sync::atomic::{AtomicU8, Ordering};
-
-    const UNINIT: u8 = u8::MAX;
-    const IO_URING: u8 = 0;
-    const POLLING: u8 = 1;
-
-    static DRIVER_TYPE: AtomicU8 = AtomicU8::new(UNINIT);
-
-    /// Representing underlying driver type the fusion driver is using
-    #[repr(u8)]
-    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
-    pub enum DriverType {
-        /// Using `polling` driver
-        Poll    = POLLING,
-
-        /// Using `io-uring` driver
-        IoUring = IO_URING,
-    }
-
-    impl DriverType {
-        fn from_num(n: u8) -> Self {
-            match n {
-                IO_URING => Self::IoUring,
-                POLLING => Self::Poll,
-                _ => unreachable!("invalid driver type"),
-            }
-        }
-
-        /// Get the underlying driver type
-        pub fn current() -> DriverType {
-            match DRIVER_TYPE.load(Ordering::Acquire) {
-                UNINIT => {}
-                x => return DriverType::from_num(x),
-            }
-
-            let dev_ty = if uring_available() {
-                DriverType::IoUring
-            } else {
-                DriverType::Poll
-            };
-
-            DRIVER_TYPE.store(dev_ty as u8, Ordering::Release);
-
-            dev_ty
-        }
-    }
-
-    fn uring_available() -> bool {
-        use io_uring::opcode::*;
-
-        // Add more opcodes here if used
-        const USED_OP: &[u8] = &[
-            Read::CODE,
-            Readv::CODE,
-            Write::CODE,
-            Writev::CODE,
-            Fsync::CODE,
-            Accept::CODE,
-            Connect::CODE,
-            RecvMsg::CODE,
-            SendMsg::CODE,
-            AsyncCancel::CODE,
-            OpenAt::CODE,
-            Close::CODE,
-            Shutdown::CODE,
-            // Linux kernel 5.19
-            #[cfg(any(
-                feature = "io-uring-sqe128",
-                feature = "io-uring-cqe32",
-                feature = "io-uring-socket"
-            ))]
-            Socket::CODE,
-        ];
-
-        Ok(())
-            .and_then(|_| {
-                let uring = io_uring::IoUring::new(2)?;
-                let mut probe = io_uring::Probe::new();
-                uring.submitter().register_probe(&mut probe)?;
-                std::io::Result::Ok(USED_OP.iter().all(|op| probe.is_supported(*op)))
-            })
-            .unwrap_or(false)
-    }
-}
-
 /// Fused [`OpCode`]
 ///
 /// This trait encapsulates both operation for `io-uring` and `polling`
@@ -131,6 +45,7 @@ impl Driver {
             DriverType::IoUring => Ok(Self {
                 fuse: FuseDriver::IoUring(iour::Driver::new(builder)?),
             }),
+            _ => unreachable!("Fuse driver will only be enabled on linux"),
         }
     }
 

compio-driver/src/fusion/op.rs

@@ -61,6 +61,7 @@ macro_rules! op {
                         DriverType::IoUring => Self {
                             inner: [< $name Inner >]::IoUring(iour::$name::new($($arg),*)),
                         },
+                        _ => unreachable!("Fuse driver will only be enabled on linux"),
                     }
                 }
             }

compio-driver/src/lib.rs

@@ -37,6 +37,9 @@ pub use asyncify::*;
 mod fd;
 pub use fd::*;
 
+mod driver_type;
+pub use driver_type::*;
+
 cfg_if::cfg_if! {
     if #[cfg(windows)] {
         #[path = "iocp/mod.rs"]