diff --git a/tokio/Cargo.toml b/tokio/Cargo.toml index 35c49ffca..44de04d16 100644 --- a/tokio/Cargo.toml +++ b/tokio/Cargo.toml @@ -85,7 +85,7 @@ sync = [] test-util = ["rt", "sync", "time"] time = [] # Unstable feature. Requires `--cfg tokio_unstable` to enable. -io-uring = ["dep:io-uring", "libc", "mio/os-poll", "mio/os-ext", "dep:slab", "once_cell"] +io-uring = ["dep:io-uring", "libc", "mio/os-poll", "mio/os-ext", "dep:slab"] # Unstable feature. Requires `--cfg tokio_unstable` to enable. taskdump = ["dep:backtrace"] @@ -115,7 +115,6 @@ libc = { version = "0.2.168", optional = true } mio = { version = "1.0.1", default-features = false, features = ["os-poll", "os-ext"], optional = true } slab = { version = "0.4.9", optional = true } backtrace = { version = "0.3.58", optional = true } -once_cell = { version = "1.21.3", optional = true } [target.'cfg(unix)'.dependencies] libc = { version = "0.2.168", optional = true } diff --git a/tokio/src/runtime/io/driver.rs b/tokio/src/runtime/io/driver.rs index fd8b5ad16..04540cf2b 100644 --- a/tokio/src/runtime/io/driver.rs +++ b/tokio/src/runtime/io/driver.rs @@ -5,8 +5,7 @@ cfg_signal_internal_and_unix! { cfg_io_uring! { mod uring; use uring::UringContext; - // TODO: Replace with std after https://github.com/rust-lang/rust/issues/109737 - use once_cell::sync::OnceCell; + use crate::loom::sync::atomic::AtomicUsize; } use crate::io::interest::Interest; @@ -68,7 +67,7 @@ pub(crate) struct Handle { feature = "fs", target_os = "linux", ))] - pub(crate) uring_probe: OnceCell>, + pub(crate) uring_state: AtomicUsize, } #[derive(Debug)] @@ -151,7 +150,7 @@ impl Driver { feature = "fs", target_os = "linux", ))] - uring_probe: OnceCell::new(), + uring_state: AtomicUsize::new(0), }; Ok((driver, handle)) diff --git a/tokio/src/runtime/io/driver/uring.rs b/tokio/src/runtime/io/driver/uring.rs index a43e77978..5c61739b2 100644 --- a/tokio/src/runtime/io/driver/uring.rs +++ b/tokio/src/runtime/io/driver/uring.rs @@ -2,6 +2,7 @@ use io_uring::{squeue::Entry, IoUring, Probe}; use mio::unix::SourceFd; use slab::Slab; +use crate::loom::sync::atomic::Ordering; use crate::runtime::driver::op::{Cancellable, Lifecycle}; use crate::{io::Interest, loom::sync::Mutex}; @@ -12,8 +13,32 @@ use std::{io, mem, task::Waker}; const DEFAULT_RING_SIZE: u32 = 256; +#[repr(usize)] +#[derive(Debug, PartialEq, Eq, Copy, Clone)] +enum State { + Uninitialized = 0, + Initialized = 1, + Unsupported = 2, +} + +impl State { + fn as_usize(&self) -> usize { + *self as usize + } + + fn from_usize(value: usize) -> Self { + match value { + 0 => State::Uninitialized, + 1 => State::Initialized, + 2 => State::Unsupported, + _ => unreachable!("invalid Uring state: {}", value), + } + } +} + pub(crate) struct UringContext { pub(crate) uring: Option, + pub(crate) probe: io_uring::Probe, pub(crate) ops: slab::Slab, } @@ -22,6 +47,7 @@ impl UringContext { Self { ops: Slab::new(), uring: None, + probe: Probe::new(), } } @@ -33,12 +59,16 @@ impl UringContext { self.uring.as_mut().expect("io_uring not initialized") } + pub(crate) fn is_opcode_supported(&self, opcode: u8) -> bool { + self.probe.is_supported(opcode) + } + /// Perform `io_uring_setup` system call, and Returns true if this /// actually initialized the io_uring. /// /// If the machine doesn't support io_uring, then this will return an /// `ENOSYS` error. - pub(crate) fn try_init(&mut self, probe: &mut Probe) -> io::Result { + pub(crate) fn try_init(&mut self) -> io::Result { if self.uring.is_some() { // Already initialized. return Ok(false); @@ -46,7 +76,7 @@ impl UringContext { let uring = IoUring::new(DEFAULT_RING_SIZE)?; - match uring.submitter().register_probe(probe) { + match uring.submitter().register_probe(&mut self.probe) { Ok(_) => {} Err(e) if e.raw_os_error() == Some(libc::EINVAL) => { // The kernel does not support IORING_REGISTER_PROBE. @@ -164,6 +194,10 @@ impl Handle { &self.uring_context } + fn set_uring_state(&self, state: State) { + self.uring_state.store(state.as_usize(), Ordering::Release); + } + /// Check if the io_uring context is initialized. If not, it will try to initialize it. /// Then, check if the provided opcode is supported. /// @@ -173,34 +207,41 @@ impl Handle { /// this returns `Ok(false)`. /// An error is returned if an io_uring syscall returns an unexpected error value. pub(crate) fn check_and_init(&self, opcode: u8) -> io::Result { - let probe = self.uring_probe.get_or_try_init(|| { - let mut probe = Probe::new(); - match self.try_init(&mut probe) { - Ok(()) => Ok(Some(probe)), - // If the system doesn't support io_uring, we set the probe to `None`. - Err(e) if e.raw_os_error() == Some(libc::ENOSYS) => Ok(None), + match State::from_usize(self.uring_state.load(Ordering::Acquire)) { + State::Uninitialized => match self.try_init_and_check_opcode(opcode) { + Ok(opcode_supported) => { + self.set_uring_state(State::Initialized); + Ok(opcode_supported) + } + // If the system doesn't support io_uring, we set the state to Unsupported. + Err(e) if e.raw_os_error() == Some(libc::ENOSYS) => { + self.set_uring_state(State::Unsupported); + Ok(false) + } // If we get EPERM, io-uring syscalls may be blocked (for example, by seccomp). // In this case, we try to fall back to spawn_blocking for this and future operations. // See also: https://github.com/tokio-rs/tokio/issues/7691 - Err(e) if e.raw_os_error() == Some(libc::EPERM) => Ok(None), + Err(e) if e.raw_os_error() == Some(libc::EPERM) => { + self.set_uring_state(State::Unsupported); + Ok(false) + } // For other system errors, we just return it. Err(e) => Err(e), - } - })?; - - Ok(probe - .as_ref() - .is_some_and(|probe| probe.is_supported(opcode))) + }, + State::Unsupported => Ok(false), + State::Initialized => Ok(self.get_uring().lock().is_opcode_supported(opcode)), + } } /// Initialize the io_uring context if it hasn't been initialized yet. - fn try_init(&self, probe: &mut Probe) -> io::Result<()> { + /// Then, check whether the given opcode is supported. + fn try_init_and_check_opcode(&self, opcode: u8) -> io::Result { let mut guard = self.get_uring().lock(); - if guard.try_init(probe)? { + if guard.try_init()? { self.add_uring_source(guard.ring().as_raw_fd())?; } - Ok(()) + Ok(guard.is_opcode_supported(opcode)) } /// Register an operation with the io_uring.