diff --git a/tokio/src/runtime/driver/op.rs b/tokio/src/runtime/driver/op.rs index 0d7ca9455..40a135d74 100644 --- a/tokio/src/runtime/driver/op.rs +++ b/tokio/src/runtime/driver/op.rs @@ -106,7 +106,7 @@ pub(crate) trait Completable { /// Extracts the `CancelData` needed to safely cancel an in-flight io_uring operation. pub(crate) trait Cancellable { - fn cancell(self) -> CancelData; + fn cancel(self) -> CancelData; } impl Unpin for Op {} diff --git a/tokio/src/runtime/io/driver.rs b/tokio/src/runtime/io/driver.rs index fb496f140..1a49b4742 100644 --- a/tokio/src/runtime/io/driver.rs +++ b/tokio/src/runtime/io/driver.rs @@ -5,6 +5,7 @@ cfg_signal_internal_and_unix! { cfg_tokio_uring! { mod uring; use uring::UringContext; + use crate::loom::sync::atomic::AtomicUsize; } use crate::io::interest::Interest; @@ -52,6 +53,9 @@ pub(crate) struct Handle { #[cfg(all(tokio_uring, feature = "rt", feature = "fs", target_os = "linux",))] pub(crate) uring_context: Mutex, + + #[cfg(all(tokio_uring, feature = "rt", feature = "fs", target_os = "linux",))] + pub(crate) uring_state: AtomicUsize, } #[derive(Debug)] @@ -121,13 +125,10 @@ impl Driver { metrics: IoDriverMetrics::default(), #[cfg(all(tokio_uring, feature = "rt", feature = "fs", target_os = "linux",))] uring_context: Mutex::new(UringContext::new()), + #[cfg(all(tokio_uring, feature = "rt", feature = "fs", target_os = "linux",))] + uring_state: AtomicUsize::new(0), }; - #[cfg(all(tokio_uring, feature = "rt", feature = "fs", target_os = "linux",))] - { - handle.add_uring_source(Interest::READABLE)?; - } - Ok((driver, handle)) } diff --git a/tokio/src/runtime/io/driver/uring.rs b/tokio/src/runtime/io/driver/uring.rs index 8a129bc49..f828df314 100644 --- a/tokio/src/runtime/io/driver/uring.rs +++ b/tokio/src/runtime/io/driver/uring.rs @@ -2,18 +2,42 @@ use io_uring::{squeue::Entry, IoUring}; 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}; use super::{Handle, TOKEN_WAKEUP}; -use std::os::fd::AsRawFd; +use std::os::fd::{AsRawFd, RawFd}; use std::{io, mem, task::Waker}; const DEFAULT_RING_SIZE: u32 = 256; +#[repr(usize)] +#[derive(Debug, PartialEq, Eq)] +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: io_uring::IoUring, + pub(crate) uring: Option, pub(crate) ops: slab::Slab, } @@ -21,14 +45,42 @@ impl UringContext { pub(crate) fn new() -> Self { Self { ops: Slab::new(), - // TODO: make configurable - uring: IoUring::new(DEFAULT_RING_SIZE).unwrap(), + uring: None, } } + pub(crate) fn ring(&self) -> &io_uring::IoUring { + self.uring.as_ref().expect("io_uring not initialized") + } + + pub(crate) fn ring_mut(&mut self) -> &mut io_uring::IoUring { + self.uring.as_mut().expect("io_uring not initialized") + } + + /// 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) -> io::Result { + if self.uring.is_some() { + // Already initialized. + return Ok(false); + } + + self.uring.replace(IoUring::new(DEFAULT_RING_SIZE)?); + + Ok(true) + } + pub(crate) fn dispatch_completions(&mut self) { let ops = &mut self.ops; - let cq = self.uring.completion(); + let Some(mut uring) = self.uring.take() else { + // Uring is not initialized yet. + return; + }; + + let cq = uring.completion(); for cqe in cq { let idx = cqe.user_data() as usize; @@ -52,13 +104,15 @@ impl UringContext { } } + self.uring.replace(uring); + // `cq`'s drop gets called here, updating the latest head pointer } pub(crate) fn submit(&mut self) -> io::Result<()> { loop { // Errors from io_uring_enter: https://man7.org/linux/man-pages/man2/io_uring_enter.2.html#ERRORS - match self.uring.submit() { + match self.ring().submit() { Ok(_) => { return Ok(()); } @@ -83,8 +137,13 @@ impl UringContext { /// Drop the driver, cancelling any in-progress ops and waiting for them to terminate. impl Drop for UringContext { fn drop(&mut self) { + if self.uring.is_none() { + // Uring is not initialized or not supported. + return; + } + // Make sure we flush the submission queue before dropping the driver. - while !self.uring.submission().is_empty() { + while !self.ring_mut().submission().is_empty() { self.submit().expect("Internal error when dropping driver"); } @@ -109,11 +168,11 @@ impl Drop for UringContext { while !cancel_ops.is_empty() { // Wait until at least one completion is available. - self.uring + self.ring_mut() .submit_and_wait(1) .expect("Internal error when dropping driver"); - for cqe in self.uring.completion() { + for cqe in self.ring_mut().completion() { let idx = cqe.user_data() as usize; cancel_ops.remove(idx); } @@ -123,23 +182,69 @@ impl Drop for UringContext { impl Handle { #[allow(dead_code)] - pub(crate) fn add_uring_source(&self, interest: Interest) -> io::Result<()> { - // setup for io_uring - let uringfd = self.get_uring().lock().uring.as_raw_fd(); + fn add_uring_source(&self, uringfd: RawFd) -> io::Result<()> { let mut source = SourceFd(&uringfd); self.registry - .register(&mut source, TOKEN_WAKEUP, interest.to_mio()) + .register(&mut source, TOKEN_WAKEUP, Interest::READABLE.to_mio()) } pub(crate) fn get_uring(&self) -> &Mutex { &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. + pub(crate) fn check_and_init(&self) -> io::Result { + match State::from_usize(self.uring_state.load(Ordering::Acquire)) { + State::Uninitialized => match self.try_init() { + Ok(()) => { + self.set_uring_state(State::Initialized); + Ok(true) + } + // 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) + } + // For other system errors, we just return it. + Err(e) => Err(e), + }, + State::Unsupported => Ok(false), + State::Initialized => Ok(true), + } + } + + /// Initialize the io_uring context if it hasn't been initialized yet. + fn try_init(&self) -> io::Result<()> { + let mut guard = self.get_uring().lock(); + if guard.try_init()? { + self.add_uring_source(guard.ring().as_raw_fd())?; + } + + Ok(()) + } + + /// Register an operation with the io_uring. + /// + /// If this is the first io_uring operation, it will also initialize the io_uring context. + /// If io_uring isn't supported, this function returns an `ENOSYS` error, so the caller can + /// perform custom handling, such as falling back to an alternative mechanism. + /// /// # Safety /// /// Callers must ensure that parameters of the entry (such as buffer) are valid and will /// be valid for the entire duration of the operation, otherwise it may cause memory problems. pub(crate) unsafe fn register_op(&self, entry: Entry, waker: Waker) -> io::Result { + // Note: Maybe this check can be removed if upstream callers consistently use `check_and_init`. + if !self.check_and_init()? { + return Err(io::Error::from_raw_os_error(libc::ENOSYS)); + } + + // Uring is initialized. + let mut guard = self.get_uring().lock(); let ctx = &mut *guard; let index = ctx.ops.insert(Lifecycle::Waiting(waker)); @@ -155,7 +260,7 @@ impl Handle { }; // SAFETY: entry is valid for the entire duration of the operation - while unsafe { ctx.uring.submission().push(&entry).is_err() } { + while unsafe { ctx.ring_mut().submission().push(&entry).is_err() } { // If the submission queue is full, flush it to the kernel submit_or_remove(ctx)?; } @@ -180,8 +285,8 @@ impl Handle { // This Op will be cancelled. Here, we don't remove the lifecycle from the slab to keep // uring data alive until the operation completes. - let cancell_data = data.expect("Data should be present").cancell(); - match mem::replace(lifecycle, Lifecycle::Cancelled(cancell_data)) { + let cancel_data = data.expect("Data should be present").cancel(); + match mem::replace(lifecycle, Lifecycle::Cancelled(cancel_data)) { Lifecycle::Submitted | Lifecycle::Waiting(_) => (), // The driver saw the completion, but it was never polled. Lifecycle::Completed(_) => (),