rt: add sharding support for uring

Author: mox692

tokio/src/runtime/builder.rs

@@ -1457,7 +1457,7 @@ impl Builder {
         use crate::runtime::scheduler;
         use crate::runtime::Config;
 
-        let (driver, driver_handle) = driver::Driver::new(self.get_cfg())?;
+        let (driver, driver_handle) = driver::Driver::new(self.get_cfg(), 1)?;
 
         // Blocking pool
         let blocking_pool = blocking::create_blocking_pool(self, self.max_blocking_threads);
@@ -1612,7 +1612,7 @@ cfg_rt_multi_thread! {
 
             let worker_threads = self.worker_threads.unwrap_or_else(num_cpus);
 
-            let (driver, driver_handle) = driver::Driver::new(self.get_cfg())?;
+            let (driver, driver_handle) = driver::Driver::new(self.get_cfg(), worker_threads)?;
 
             // Create the blocking pool
             let blocking_pool =

tokio/src/runtime/driver.rs

@@ -43,8 +43,9 @@ pub(crate) struct Cfg {
 }
 
 impl Driver {
-    pub(crate) fn new(cfg: Cfg) -> io::Result<(Self, Handle)> {
-        let (io_stack, io_handle, signal_handle) = create_io_stack(cfg.enable_io, cfg.nevents)?;
+    pub(crate) fn new(cfg: Cfg, num_workers: usize) -> io::Result<(Self, Handle)> {
+        let (io_stack, io_handle, signal_handle) =
+            create_io_stack(cfg.enable_io, cfg.nevents, num_workers)?;
 
         let clock = create_clock(cfg.enable_pause_time, cfg.start_paused);
 
@@ -136,12 +137,12 @@ cfg_io_driver! {
         Disabled(UnparkThread),
     }
 
-    fn create_io_stack(enabled: bool, nevents: usize) -> io::Result<(IoStack, IoHandle, SignalHandle)> {
+    fn create_io_stack(enabled: bool, nevents: usize, num_workers: usize) -> io::Result<(IoStack, IoHandle, SignalHandle)> {
         #[cfg(loom)]
         assert!(!enabled);
 
         let ret = if enabled {
-            let (io_driver, io_handle) = crate::runtime::io::Driver::new(nevents)?;
+            let (io_driver, io_handle) = crate::runtime::io::Driver::new(nevents,num_workers)?;
 
             let (signal_driver, signal_handle) = create_signal_driver(io_driver, &io_handle)?;
             let process_driver = create_process_driver(signal_driver);
@@ -202,7 +203,7 @@ cfg_not_io_driver! {
     #[derive(Debug)]
     pub(crate) struct IoStack(ParkThread);
 
-    fn create_io_stack(_enabled: bool, _nevents: usize) -> io::Result<(IoStack, IoHandle, SignalHandle)> {
+    fn create_io_stack(_enabled: bool, _nevents: usize,  _num_worker: usize) -> io::Result<(IoStack, IoHandle, SignalHandle)> {
         let park_thread = ParkThread::new();
         let unpark_thread = park_thread.unpark();
         Ok((IoStack(park_thread), unpark_thread, Default::default()))

tokio/src/runtime/driver/op.rs

@@ -1,4 +1,5 @@
 use crate::runtime::Handle;
+use crate::runtime::OpId;
 use io_uring::cqueue;
 use io_uring::squeue::Entry;
 use std::future::Future;
@@ -35,6 +36,8 @@ pub(crate) enum State {
 }
 
 pub(crate) struct Op<T: Cancellable> {
+    // Uring shard id this operation belongs to
+    pub(crate) shard_id: usize,
     // Handle to the runtime
     handle: Handle,
     // State of this Op
@@ -51,7 +54,9 @@ impl<T: Cancellable> Op<T> {
     #[allow(dead_code)]
     pub(crate) unsafe fn new(entry: Entry, data: T) -> Self {
         let handle = Handle::current();
+        let shard_size = handle.inner.driver().io().uring_context.len();
         Self {
+            shard_id: OpId::next().as_u64() as usize % shard_size,
             handle,
             data: Some(data),
             state: State::Initialize(Some(entry)),
@@ -71,7 +76,11 @@ impl<T: Cancellable> Drop for Op<T> {
             State::Polled(index) => {
                 let data = self.take_data();
                 let handle = &mut self.handle;
-                handle.inner.driver().io().cancel_op(index, data);
+                handle
+                    .inner
+                    .driver()
+                    .io()
+                    .cancel_op(self.shard_id, index, data);
             }
             // This Op has not been polled yet.
             // We don't need to do anything here.
@@ -124,13 +133,13 @@ impl<T: Cancellable + Completable + Send> Future for Op<T> {
                 let entry = entry_opt.take().expect("Entry must be present");
                 let waker = cx.waker().clone();
                 // SAFETY: entry is valid for the entire duration of the operation
-                let idx = unsafe { driver.register_op(entry, waker)? };
+                let idx = unsafe { driver.register_op(this.shard_id, entry, waker)? };
                 this.state = State::Polled(idx);
                 Poll::Pending
             }
 
             State::Polled(idx) => {
-                let mut ctx = driver.get_uring().lock();
+                let mut ctx = driver.get_uring(this.shard_id).lock();
                 let lifecycle = ctx.ops.get_mut(*idx).expect("Lifecycle must be present");
 
                 match mem::replace(lifecycle, Lifecycle::Submitted) {

tokio/src/runtime/io/driver.rs

@@ -52,10 +52,10 @@ pub(crate) struct Handle {
     pub(crate) metrics: IoDriverMetrics,
 
     #[cfg(all(tokio_uring, feature = "rt", feature = "fs", target_os = "linux",))]
-    pub(crate) uring_context: Mutex<UringContext>,
+    pub(crate) uring_context: Box<[Mutex<UringContext>]>,
 
     #[cfg(all(tokio_uring, feature = "rt", feature = "fs", target_os = "linux",))]
-    pub(crate) uring_state: AtomicUsize,
+    pub(crate) uring_state: Box<[AtomicUsize]>,
 }
 
 #[derive(Debug)]
@@ -102,7 +102,10 @@ fn _assert_kinds() {
 impl Driver {
     /// Creates a new event loop, returning any error that happened during the
     /// creation.
-    pub(crate) fn new(nevents: usize) -> io::Result<(Driver, Handle)> {
+    pub(crate) fn new(
+        nevents: usize,
+        #[allow(unused)] num_workers: usize,
+    ) -> io::Result<(Driver, Handle)> {
         let poll = mio::Poll::new()?;
         #[cfg(not(target_os = "wasi"))]
         let waker = mio::Waker::new(poll.registry(), TOKEN_WAKEUP)?;
@@ -124,9 +127,15 @@ impl Driver {
             waker,
             metrics: IoDriverMetrics::default(),
             #[cfg(all(tokio_uring, feature = "rt", feature = "fs", target_os = "linux",))]
-            uring_context: Mutex::new(UringContext::new()),
+            uring_context: (0..num_workers)
+                .map(|_| Mutex::new(UringContext::new()))
+                .collect::<Vec<_>>()
+                .into_boxed_slice(),
             #[cfg(all(tokio_uring, feature = "rt", feature = "fs", target_os = "linux",))]
-            uring_state: AtomicUsize::new(0),
+            uring_state: (0..num_workers)
+                .map(|_| AtomicUsize::new(0))
+                .collect::<Vec<_>>()
+                .into_boxed_slice(),
         };
 
         Ok((driver, handle))
@@ -200,9 +209,13 @@ impl Driver {
 
         #[cfg(all(tokio_uring, feature = "rt", feature = "fs", target_os = "linux",))]
         {
-            let mut guard = handle.get_uring().lock();
-            let ctx = &mut *guard;
-            ctx.dispatch_completions();
+            let shard_size = handle.uring_context.len();
+
+            for shard_id in 0..shard_size {
+                let mut guard = handle.get_uring(shard_id).lock();
+                let ctx = &mut *guard;
+                ctx.dispatch_completions();
+            }
         }
 
         handle.metrics.incr_ready_count_by(ready_count);

tokio/src/runtime/io/driver/uring.rs

@@ -188,25 +188,26 @@ impl Handle {
             .register(&mut source, TOKEN_WAKEUP, Interest::READABLE.to_mio())
     }
 
-    pub(crate) fn get_uring(&self) -> &Mutex<UringContext> {
-        &self.uring_context
+    pub(crate) fn get_uring(&self, shard_id: usize) -> &Mutex<UringContext> {
+        let shard_id = shard_id % self.uring_context.len();
+        &self.uring_context[shard_id]
     }
 
-    fn set_uring_state(&self, state: State) {
-        self.uring_state.store(state.as_usize(), Ordering::Release);
+    fn set_uring_state(&self, shard_id: usize, state: State) {
+        self.uring_state[shard_id].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<bool> {
-        match State::from_usize(self.uring_state.load(Ordering::Acquire)) {
-            State::Uninitialized => match self.try_init() {
+    pub(crate) fn check_and_init(&self, shard_id: usize) -> io::Result<bool> {
+        match State::from_usize(self.uring_state[shard_id].load(Ordering::Acquire)) {
+            State::Uninitialized => match self.try_init(shard_id) {
                 Ok(()) => {
-                    self.set_uring_state(State::Initialized);
+                    self.set_uring_state(shard_id, 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);
+                    self.set_uring_state(shard_id, State::Unsupported);
                     Ok(false)
                 }
                 // For other system errors, we just return it.
@@ -218,8 +219,8 @@ impl Handle {
     }
 
     /// 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();
+    fn try_init(&self, shard_id: usize) -> io::Result<()> {
+        let mut guard = self.get_uring(shard_id).lock();
         if guard.try_init()? {
             self.add_uring_source(guard.ring().as_raw_fd())?;
         }
@@ -237,15 +238,20 @@ impl Handle {
     ///
     /// 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<usize> {
+    pub(crate) unsafe fn register_op(
+        &self,
+        shard_id: usize,
+        entry: Entry,
+        waker: Waker,
+    ) -> io::Result<usize> {
         // Note: Maybe this check can be removed if upstream callers consistently use `check_and_init`.
-        if !self.check_and_init()? {
+        if !self.check_and_init(shard_id)? {
             return Err(io::Error::from_raw_os_error(libc::ENOSYS));
         }
 
         // Uring is initialized.
 
-        let mut guard = self.get_uring().lock();
+        let mut guard = self.get_uring(shard_id).lock();
         let ctx = &mut *guard;
         let index = ctx.ops.insert(Lifecycle::Waiting(waker));
         let entry = entry.user_data(index as u64);
@@ -273,8 +279,8 @@ impl Handle {
 
     // TODO: Remove this annotation when operations are actually supported
     #[allow(unused_variables, unreachable_code)]
-    pub(crate) fn cancel_op<T: Cancellable>(&self, index: usize, data: Option<T>) {
-        let mut guard = self.get_uring().lock();
+    pub(crate) fn cancel_op<T: Cancellable>(&self, shard_id: usize, index: usize, data: Option<T>) {
+        let mut guard = self.get_uring(shard_id).lock();
         let ctx = &mut *guard;
         let ops = &mut ctx.ops;
         let Some(lifecycle) = ops.get_mut(index) else {

tokio/src/runtime/mod.rs

@@ -421,3 +421,8 @@ cfg_rt! {
     /// After thread starts / before thread stops
     type Callback = std::sync::Arc<dyn Fn() + Send + Sync>;
 }
+
+cfg_tokio_uring! {
+    pub(crate) mod op_id;
+    pub(crate) use op_id::OpId;
+}

tokio/src/runtime/op_id.rs

@@ -0,0 +1,32 @@
+// TODO: Put together with other similar id related utils.
+
+use std::num::NonZeroU64;
+
+#[derive(Eq, PartialEq, Clone, Copy, Hash, Debug)]
+pub(crate) struct OpId(NonZeroU64);
+
+impl OpId {
+    pub(crate) fn next() -> Self {
+        use crate::loom::sync::atomic::Ordering::Relaxed;
+        use crate::loom::sync::atomic::StaticAtomicU64;
+
+        #[cfg(all(test, loom))]
+        crate::loom::lazy_static! {
+            static ref NEXT_ID: StaticAtomicU64 = StaticAtomicU64::new(1);
+        }
+
+        #[cfg(not(all(test, loom)))]
+        static NEXT_ID: StaticAtomicU64 = StaticAtomicU64::new(1);
+
+        loop {
+            let id = NEXT_ID.fetch_add(1, Relaxed);
+            if let Some(id) = NonZeroU64::new(id) {
+                return Self(id);
+            }
+        }
+    }
+
+    pub(crate) fn as_u64(&self) -> u64 {
+        self.0.get()
+    }
+}