Skip to content

Commit

Permalink
feat(driver): make DriverType available on all platform
Browse files Browse the repository at this point in the history
  • Loading branch information
George-Miao committed Aug 29, 2024
1 parent e438822 commit 2eb98d5
Show file tree
Hide file tree
Showing 4 changed files with 121 additions and 87 deletions.
115 changes: 115 additions & 0 deletions compio-driver/src/driver_type.rs
Original file line number Diff line number Diff line change
@@ -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
}
}
89 changes: 2 additions & 87 deletions compio-driver/src/fusion/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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`
Expand Down Expand Up @@ -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"),
}
}

Expand Down
1 change: 1 addition & 0 deletions compio-driver/src/fusion/op.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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"),
}
}
}
Expand Down
3 changes: 3 additions & 0 deletions compio-driver/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down

0 comments on commit 2eb98d5

Please sign in to comment.