rt: move the multi-thread inject queue to its own mutex (#8382)

This migrates the inject queue to its own lock, instead of sharing the
scheduler's `synced` mutex with the idle worker state. They originally
shared a lock as part of #5747 and #5754, but the unified critical
sections that was intended to enable were removed with the alternative
multi-threaded scheduler in #7275.

This is a first step towards #7973.

tokio/src/runtime/scheduler/inject.rs

@@ -36,7 +36,7 @@ impl<T: 'static> Inject<T> {
     }
 
     // Kind of annoying to have to include the cfg here
-    #[cfg(all(tokio_unstable, feature = "taskdump"))]
+    #[cfg(any(all(tokio_unstable, feature = "taskdump"), feature = "rt-multi-thread"))]
     pub(crate) fn is_closed(&self) -> bool {
         let synced = self.synced.lock();
         self.shared.is_closed(&synced)

tokio/src/runtime/scheduler/inject/rt_multi_thread.rs

@@ -1,34 +1,18 @@
-use super::{Shared, Synced};
+use super::{Inject, Pop};
 
-use crate::runtime::scheduler::Lock;
 use crate::runtime::task;
 
 use std::sync::atomic::Ordering::Release;
 
-impl<'a> Lock<Synced> for &'a mut Synced {
-    type Handle = &'a mut Synced;
-
-    fn lock(self) -> Self::Handle {
-        self
+impl<T: 'static> Inject<T> {
+    pub(crate) fn is_empty(&self) -> bool {
+        self.shared.is_empty()
     }
-}
 
-impl AsMut<Synced> for Synced {
-    fn as_mut(&mut self) -> &mut Synced {
-        self
-    }
-}
-
-impl<T: 'static> Shared<T> {
     /// Pushes several values into the queue.
-    ///
-    /// # Safety
-    ///
-    /// Must be called with the same `Synced` instance returned by `Inject::new`
     #[inline]
-    pub(crate) unsafe fn push_batch<L, I>(&self, shared: L, mut iter: I)
+    pub(crate) fn push_batch<I>(&self, mut iter: I)
     where
-        L: Lock<Synced>,
         I: Iterator<Item = task::Notified<T>>,
     {
         let first = match iter.next() {
@@ -56,10 +40,9 @@ impl<T: 'static> Shared<T> {
         // Now that the tasks are linked together, insert them into the
         // linked list.
         //
-        // Safety: exactly the same safety requirements as `push_batch` method.
-        unsafe {
-            self.push_batch_inner(shared, first, prev, counter);
-        }
+        // safety: the batch was linked just above from `Notified`s this
+        // function took ownership of, satisfying both obligations.
+        unsafe { self.push_batch_inner(first, prev, counter) };
     }
 
     /// Inserts several tasks that have been linked together into the queue.
@@ -69,28 +52,32 @@ impl<T: 'static> Shared<T> {
     ///
     /// # Safety
     ///
-    /// Must be called with the same `Synced` instance returned by `Inject::new`
+    /// The caller must own the `Notified` for each of the `num` tasks, and
+    /// the tasks must be linked from `batch_head` to `batch_tail` through
+    /// their `queue_next` fields, with `batch_tail`'s `queue_next` unset.
     #[inline]
-    unsafe fn push_batch_inner<L>(
+    unsafe fn push_batch_inner(
         &self,
-        shared: L,
         batch_head: task::RawTask,
         batch_tail: task::RawTask,
         num: usize,
-    ) where
-        L: Lock<Synced>,
-    {
+    ) {
         debug_assert!(unsafe { batch_tail.get_queue_next().is_none() });
 
-        let mut synced = shared.lock();
+        let mut synced = self.synced.lock();
 
-        if synced.as_mut().is_closed {
+        if synced.is_closed {
+            // Drop the lock before dropping the tasks: dropping a task can
+            // run arbitrary user `Drop` code, which may reentrantly acquire
+            // this lock by scheduling a task.
             drop(synced);
 
             let mut curr = Some(batch_head);
 
             while let Some(task) = curr {
-                // Safety: exactly the same safety requirements as `push_batch_inner`.
+                // safety: per this function's contract, the caller owns each
+                // task's `Notified` and linked the batch through `queue_next`;
+                // reconstituting the `Notified` here takes that ownership.
                 curr = unsafe { task.get_queue_next() };
 
                 let _ = unsafe { task::Notified::<T>::from_raw(task) };
@@ -99,8 +86,6 @@ impl<T: 'static> Shared<T> {
             return;
         }
 
-        let synced = synced.as_mut();
-
         if let Some(tail) = synced.tail {
             unsafe {
                 tail.set_queue_next(Some(batch_head));
@@ -115,8 +100,29 @@ impl<T: 'static> Shared<T> {
         //
         // safety: All updates to the len atomic are guarded by the mutex. As
         // such, a non-atomic load followed by a store is safe.
-        let len = unsafe { self.len.unsync_load() };
+        let len = unsafe { self.shared.len.unsync_load() };
 
-        self.len.store(len + num, Release);
+        self.shared.len.store(len + num, Release);
+    }
+
+    /// Pops up to `n` values from the queue, passing an iterator over them to
+    /// `f`. The queue lock is held while `f` runs, so any values `f` does not
+    /// consume are removed from the queue and dropped before the lock is
+    /// released.
+    pub(crate) fn pop_n<R>(&self, n: usize, f: impl FnOnce(Pop<'_, T>) -> R) -> R {
+        let mut synced = self.synced.lock();
+        // safety: passing correct `Synced`
+        f(unsafe { self.shared.pop_n(&mut synced, n) })
+    }
+
+    /// Pops every task from the queue into `dst`, holding the queue lock for
+    /// the entire drain so it is atomic with respect to concurrent pushes.
+    #[cfg(all(tokio_unstable, feature = "taskdump"))]
+    pub(crate) fn drain_into(&self, dst: &mut Vec<task::Notified<T>>) {
+        let mut synced = self.synced.lock();
+        // safety: passing correct `Synced`
+        while let Some(task) = unsafe { self.shared.pop(&mut synced) } {
+            dst.push(task);
+        }
     }
 }

tokio/src/runtime/scheduler/lock.rs

@@ -1,6 +0,0 @@
-/// A lock (mutex) yielding generic data.
-pub(crate) trait Lock<T> {
-    type Handle: AsMut<T>;
-
-    fn lock(self) -> Self::Handle;
-}

tokio/src/runtime/scheduler/mod.rs

@@ -17,9 +17,6 @@ cfg_rt_multi_thread! {
     mod block_in_place;
     pub(crate) use block_in_place::block_in_place;
 
-    mod lock;
-    use lock::Lock;
-
     pub(crate) mod multi_thread;
     pub(crate) use multi_thread::MultiThread;
 }

tokio/src/runtime/scheduler/multi_thread/idle.rs

@@ -151,6 +151,9 @@ impl Idle {
     }
 
     fn notify_should_wakeup(&self) -> bool {
+        // This must be a `SeqCst` RMW rather than a load: it is what makes
+        // the caller's preceding inject-queue push visible to a parking
+        // worker's subsequent unlocked queue-emptiness check.
         let state = State(self.state.fetch_add(0, SeqCst));
         state.num_searching() == 0 && state.num_unparked() < self.num_workers
     }

tokio/src/runtime/scheduler/multi_thread/mod.rs

@@ -26,8 +26,6 @@ pub(crate) use worker::{Context, Launch, Shared};
 cfg_taskdump! {
     mod trace;
     use trace::TraceStatus;
-
-    pub(crate) use worker::Synced;
 }
 
 cfg_not_taskdump! {

tokio/src/runtime/scheduler/multi_thread/worker.rs

@@ -61,7 +61,7 @@ use crate::runtime;
 use crate::runtime::scheduler::multi_thread::{
     idle, park, queue, Counters, Handle, Idle, Overflow, Parker, Stats, TraceStatus, Unparker,
 };
-use crate::runtime::scheduler::{inject, Defer, Lock};
+use crate::runtime::scheduler::{Defer, Inject};
 use crate::runtime::task::OwnedTasks;
 use crate::runtime::{
     blocking, driver, scheduler, task, Config, SchedulerMetrics, TimerFlavor, WorkerMetrics,
@@ -174,7 +174,7 @@ pub(crate) struct Shared {
     /// Global task queue used for:
     ///  1. Submit work to the scheduler while **not** currently on a worker thread.
     ///  2. Submit work to the scheduler when a worker run queue is saturated
-    pub(super) inject: inject::Shared<Arc<Handle>>,
+    pub(super) inject: Inject<Arc<Handle>>,
 
     /// Coordinates idle workers
     idle: Idle,
@@ -220,13 +220,16 @@ pub(crate) struct Synced {
     /// Synchronized state for `Idle`.
     pub(super) idle: idle::Synced,
 
-    /// Synchronized state for `Inject`.
-    pub(crate) inject: inject::Synced,
-
     #[cfg(all(tokio_unstable, feature = "time"))]
     /// Timers pending to be registered.
     /// This is used to register a timer but the [`Core`]
     /// is not available in the current thread.
+    ///
+    /// This must stay under the same mutex as `idle`: a parking worker drains
+    /// it (via `try_lock`) only after publishing its parked state under this
+    /// lock, and `notify_if_work_pending` does not check for pending timers,
+    /// so sharing the lock with the parking transition is what prevents a
+    /// timer push from being stranded while every worker sleeps.
     inject_timers: Vec<time_alt::EntryHandle>,
 }
 
@@ -316,7 +319,6 @@ pub(super) fn create(
     }
 
     let (idle, idle_synced) = Idle::new(size);
-    let (inject, inject_synced) = inject::Shared::new();
     let schedule_latency_start = config.track_task_schedule_latency.then(Instant::now);
 
     let remotes_len = remotes.len();
@@ -325,12 +327,11 @@ pub(super) fn create(
         task_hooks: TaskHooks::from_config(&config),
         shared: Shared {
             remotes: remotes.into_boxed_slice(),
-            inject,
+            inject: Inject::new(),
             idle,
             owned: OwnedTasks::new(size),
             synced: Mutex::new(Synced {
                 idle: idle_synced,
-                inject: inject_synced,
                 #[cfg(all(tokio_unstable, feature = "time"))]
                 inject_timers: Vec::new(),
             }),
@@ -1141,17 +1142,15 @@ impl Core {
             // and not pushed onto the local queue.
             let n = usize::max(1, n);
 
-            let mut synced = worker.handle.shared.synced.lock();
-            // safety: passing in the correct `inject::Synced`.
-            let mut tasks = unsafe { worker.inject().pop_n(&mut synced.inject, n) };
+            worker.inject().pop_n(n, |mut tasks| {
+                // Pop the first task to return immediately
+                let ret = tasks.next();
 
-            // Pop the first task to return immediately
-            let ret = tasks.next();
+                // Push the rest of the on the run queue
+                self.run_queue.push_back(tasks);
 
-            // Push the rest of the on the run queue
-            self.run_queue.push_back(tasks);
-
-            ret
+                ret
+            })
         }
     }
 
@@ -1291,8 +1290,7 @@ impl Core {
 
         if !self.is_shutdown {
             // Check if the scheduler has been shutdown
-            let synced = worker.handle.shared.synced.lock();
-            self.is_shutdown = worker.inject().is_closed(&synced.inject);
+            self.is_shutdown = worker.inject().is_closed();
         }
 
         if !self.is_traced {
@@ -1344,7 +1342,7 @@ impl Core {
 
 impl Worker {
     /// Returns a reference to the scheduler's injection queue.
-    fn inject(&self) -> &inject::Shared<Arc<Handle>> {
+    fn inject(&self) -> &Inject<Arc<Handle>> {
         &self.handle.shared.inject
     }
 }
@@ -1417,23 +1415,13 @@ impl Handle {
     }
 
     fn next_remote_task(&self) -> Option<Notified> {
-        if self.shared.inject.is_empty() {
-            return None;
-        }
-
-        let mut synced = self.shared.synced.lock();
-        // safety: passing in correct `idle::Synced`
-        unsafe { self.shared.inject.pop(&mut synced.inject) }
+        self.shared.inject.pop()
     }
 
     fn push_remote_task(&self, task: Notified) {
         self.shared.scheduler_metrics.inc_remote_schedule_count();
 
-        let mut synced = self.shared.synced.lock();
-        // safety: passing in correct `idle::Synced`
-        unsafe {
-            self.shared.inject.push(&mut synced.inject, task);
-        }
+        self.shared.inject.push(task);
     }
 
     #[cfg(all(tokio_unstable, feature = "time"))]
@@ -1458,11 +1446,7 @@ impl Handle {
     }
 
     pub(super) fn close(&self) {
-        if self
-            .shared
-            .inject
-            .close(&mut self.shared.synced.lock().inject)
-        {
+        if self.shared.inject.close() {
             self.notify_all();
         }
     }
@@ -1558,29 +1542,7 @@ impl Overflow<Arc<Handle>> for Handle {
     where
         I: Iterator<Item = task::Notified<Arc<Handle>>>,
     {
-        unsafe {
-            self.shared.inject.push_batch(self, iter);
-        }
-    }
-}
-
-pub(crate) struct InjectGuard<'a> {
-    lock: crate::loom::sync::MutexGuard<'a, Synced>,
-}
-
-impl<'a> AsMut<inject::Synced> for InjectGuard<'a> {
-    fn as_mut(&mut self) -> &mut inject::Synced {
-        &mut self.lock.inject
-    }
-}
-
-impl<'a> Lock<inject::Synced> for &'a Handle {
-    type Handle = InjectGuard<'a>;
-
-    fn lock(self) -> Self::Handle {
-        InjectGuard {
-            lock: self.shared.synced.lock(),
-        }
+        self.shared.inject.push_batch(iter);
     }
 }
 

tokio/src/runtime/scheduler/multi_thread/worker/taskdump.rs

@@ -35,12 +35,9 @@ impl Handle {
 
         let owned = &self.shared.owned;
         let mut local = self.shared.steal_all();
-        let synced = &self.shared.synced;
         let injection = &self.shared.inject;
 
-        // safety: `trace_multi_thread` is invoked with the same `synced` that `injection`
-        // was created with.
-        let traces = unsafe { trace_multi_thread(owned, &mut local, synced, injection) }
+        let traces = trace_multi_thread(owned, &mut local, injection)
             .into_iter()
             .map(|(id, trace)| dump::Task::new(id, trace))
             .collect();

tokio/src/runtime/task/trace/mod.rs

@@ -386,21 +386,13 @@ pub(in crate::runtime) fn trace_current_thread(
 }
 
 cfg_rt_multi_thread! {
-    use crate::loom::sync::Mutex;
     use crate::runtime::scheduler::multi_thread;
-    use crate::runtime::scheduler::multi_thread::Synced;
-    use crate::runtime::scheduler::inject::Shared;
 
-    /// Trace and poll all tasks of the `current_thread` runtime.
-    ///
-    /// ## Safety
-    ///
-    /// Must be called with the same `synced` that `injection` was created with.
-    pub(in crate::runtime) unsafe fn trace_multi_thread(
+    /// Trace and poll all tasks of the `multi_thread` runtime.
+    pub(in crate::runtime) fn trace_multi_thread(
         owned: &OwnedTasks<Arc<multi_thread::Handle>>,
         local: &mut multi_thread::queue::Local<Arc<multi_thread::Handle>>,
-        synced: &Mutex<Synced>,
-        injection: &Shared<Arc<multi_thread::Handle>>,
+        injection: &Inject<Arc<multi_thread::Handle>>,
     ) -> Vec<(Id, Trace)> {
         let mut dequeued = Vec::new();
 
@@ -410,13 +402,7 @@ cfg_rt_multi_thread! {
         }
 
         // clear the injection queue
-        let mut synced = synced.lock();
-        // Safety: exactly the same safety requirements as `trace_multi_thread` function.
-        while let Some(notified) = unsafe { injection.pop(&mut synced.inject) } {
-            dequeued.push(notified);
-        }
-
-        drop(synced);
+        injection.drain_into(&mut dequeued);
 
         // precondition: we have drained the tasks from the local and injection
         // queues.

tokio/src/runtime/tests/inject.rs

@@ -1,54 +1,44 @@
-use crate::runtime::scheduler::inject;
+use crate::runtime::scheduler::Inject;
 
 #[test]
 fn push_and_pop() {
     const N: usize = 2;
 
-    let (inject, mut synced) = inject::Shared::new();
+    let inject = Inject::new();
 
     for i in 0..N {
         assert_eq!(inject.len(), i);
         let (task, _) = super::unowned(async {});
-        unsafe { inject.push(&mut synced, task) };
+        inject.push(task);
     }
 
     for i in 0..N {
         assert_eq!(inject.len(), N - i);
-        assert!(unsafe { inject.pop(&mut synced) }.is_some());
+        assert!(inject.pop().is_some());
     }
 
     println!("--------------");
 
-    assert!(unsafe { inject.pop(&mut synced) }.is_none());
+    assert!(inject.pop().is_none());
 }
 
 #[test]
 fn push_batch_and_pop() {
-    let (inject, mut inject_synced) = inject::Shared::new();
+    let inject = Inject::new();
 
-    unsafe {
-        inject.push_batch(
-            &mut inject_synced,
-            (0..10).map(|_| super::unowned(async {}).0),
-        );
+    inject.push_batch((0..10).map(|_| super::unowned(async {}).0));
 
-        assert_eq!(5, inject.pop_n(&mut inject_synced, 5).count());
-        assert_eq!(5, inject.pop_n(&mut inject_synced, 5).count());
-        assert_eq!(0, inject.pop_n(&mut inject_synced, 5).count());
-    }
+    assert_eq!(5, inject.pop_n(5, |tasks| tasks.count()));
+    assert_eq!(5, inject.pop_n(5, |tasks| tasks.count()));
+    assert_eq!(0, inject.pop_n(5, |tasks| tasks.count()));
 }
 
 #[test]
 fn pop_n_drains_on_drop() {
-    let (inject, mut inject_synced) = inject::Shared::new();
+    let inject = Inject::new();
 
-    unsafe {
-        inject.push_batch(
-            &mut inject_synced,
-            (0..10).map(|_| super::unowned(async {}).0),
-        );
-        let _ = inject.pop_n(&mut inject_synced, 10);
+    inject.push_batch((0..10).map(|_| super::unowned(async {}).0));
+    inject.pop_n(10, |_| ());
 
-        assert_eq!(inject.len(), 0);
-    }
+    assert_eq!(inject.len(), 0);
 }