compio_driver/
driver_type.rs

1use std::sync::atomic::{AtomicU8, Ordering};
2
3const UNINIT: u8 = u8::MAX;
4const IO_URING: u8 = 0;
5const POLLING: u8 = 1;
6const IOCP: u8 = 2;
7
8static DRIVER_TYPE: AtomicU8 = AtomicU8::new(UNINIT);
9
10/// Representing underlying driver type the fusion driver is using
11#[repr(u8)]
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
13pub enum DriverType {
14    /// Using `polling` driver
15    Poll    = POLLING,
16
17    /// Using `io-uring` driver
18    IoUring = IO_URING,
19
20    /// Using `iocp` driver
21    IOCP    = IOCP,
22}
23
24impl DriverType {
25    fn from_num(n: u8) -> Self {
26        match n {
27            IO_URING => Self::IoUring,
28            POLLING => Self::Poll,
29            IOCP => Self::IOCP,
30            _ => unreachable!("invalid driver type"),
31        }
32    }
33
34    /// Get the underlying driver type
35    fn get() -> DriverType {
36        cfg_if::cfg_if! {
37            if #[cfg(windows)] {
38                DriverType::IOCP
39            } else if #[cfg(all(target_os = "linux", feature = "polling", feature = "io-uring"))] {
40                use io_uring::opcode::*;
41
42                // Add more opcodes here if used
43                const USED_OP: &[u8] = &[
44                    Read::CODE,
45                    Readv::CODE,
46                    Write::CODE,
47                    Writev::CODE,
48                    Fsync::CODE,
49                    Accept::CODE,
50                    Connect::CODE,
51                    RecvMsg::CODE,
52                    SendMsg::CODE,
53                    AsyncCancel::CODE,
54                    OpenAt::CODE,
55                    Close::CODE,
56                    Shutdown::CODE,
57                    // Linux kernel 5.19
58                    #[cfg(any(
59                        feature = "io-uring-sqe128",
60                        feature = "io-uring-cqe32",
61                        feature = "io-uring-socket"
62                    ))]
63                    Socket::CODE,
64                ];
65
66                (|| {
67                    let uring = io_uring::IoUring::new(2)?;
68                    let mut probe = io_uring::Probe::new();
69                    uring.submitter().register_probe(&mut probe)?;
70                    if USED_OP.iter().all(|op| probe.is_supported(*op)) {
71                        std::io::Result::Ok(DriverType::IoUring)
72                    } else {
73                        Ok(DriverType::Poll)
74                    }
75                })()
76                .unwrap_or(DriverType::Poll) // Should we fail here?
77            } else if #[cfg(all(target_os = "linux", feature = "io-uring"))] {
78                DriverType::IoUring
79            } else if #[cfg(unix)] {
80                DriverType::Poll
81            } else {
82                compile_error!("unsupported platform");
83            }
84        }
85    }
86
87    /// Get the underlying driver type and cache it. Following calls will return
88    /// the cached value.
89    pub fn current() -> DriverType {
90        match DRIVER_TYPE.load(Ordering::Acquire) {
91            UNINIT => {}
92            x => return DriverType::from_num(x),
93        }
94        let dev_ty = Self::get();
95
96        DRIVER_TYPE.store(dev_ty as u8, Ordering::Release);
97
98        dev_ty
99    }
100
101    /// Check if the current driver is `polling`
102    pub fn is_polling() -> bool {
103        Self::current() == DriverType::Poll
104    }
105
106    /// Check if the current driver is `io-uring`
107    pub fn is_iouring() -> bool {
108        Self::current() == DriverType::IoUring
109    }
110
111    /// Check if the current driver is `iocp`
112    pub fn is_iocp() -> bool {
113        Self::current() == DriverType::IOCP
114    }
115}