compio_driver/
driver_type.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
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
    }
}