rt: use task::Inject with current_thread scheduler (#5702)

Previously, the current_thread scheduler used its own injection queue
instead of sharing the same one as the multi-threaded scheduler. This
patch updates the current_thread scheduler to use the same injection
queue as the multi-threaded one (`task::Inject`).

`task::Inject` includes an optimization where it does not need to
acquire the mutex if the queue is empty.

benches/Cargo.toml

@@ -40,6 +40,11 @@ name = "sync_watch"
 path = "sync_watch.rs"
 harness = false
 
+[[bench]]
+name = "rt_current_thread"
+path = "rt_current_thread.rs"
+harness = false
+
 [[bench]]
 name = "rt_multi_threaded"
 path = "rt_multi_threaded.rs"

benches/rt_current_thread.rs

@@ -0,0 +1,83 @@
+//! Benchmark implementation details of the threaded scheduler. These benches are
+//! intended to be used as a form of regression testing and not as a general
+//! purpose benchmark demonstrating real-world performance.
+
+use tokio::runtime::{self, Runtime};
+
+use bencher::{benchmark_group, benchmark_main, Bencher};
+
+const NUM_SPAWN: usize = 1_000;
+
+fn spawn_many_local(b: &mut Bencher) {
+    let rt = rt();
+    let mut handles = Vec::with_capacity(NUM_SPAWN);
+
+    b.iter(|| {
+        rt.block_on(async {
+            for _ in 0..NUM_SPAWN {
+                handles.push(tokio::spawn(async move {}));
+            }
+
+            for handle in handles.drain(..) {
+                handle.await.unwrap();
+            }
+        });
+    });
+}
+
+fn spawn_many_remote_idle(b: &mut Bencher) {
+    let rt = rt();
+    let rt_handle = rt.handle();
+    let mut handles = Vec::with_capacity(NUM_SPAWN);
+
+    b.iter(|| {
+        for _ in 0..NUM_SPAWN {
+            handles.push(rt_handle.spawn(async {}));
+        }
+
+        rt.block_on(async {
+            for handle in handles.drain(..) {
+                handle.await.unwrap();
+            }
+        });
+    });
+}
+
+fn spawn_many_remote_busy(b: &mut Bencher) {
+    let rt = rt();
+    let rt_handle = rt.handle();
+    let mut handles = Vec::with_capacity(NUM_SPAWN);
+
+    rt.spawn(async {
+        fn iter() {
+            tokio::spawn(async { iter() });
+        }
+
+        iter()
+    });
+
+    b.iter(|| {
+        for _ in 0..NUM_SPAWN {
+            handles.push(rt_handle.spawn(async {}));
+        }
+
+        rt.block_on(async {
+            for handle in handles.drain(..) {
+                handle.await.unwrap();
+            }
+        });
+    });
+}
+
+fn rt() -> Runtime {
+    runtime::Builder::new_current_thread().build().unwrap()
+}
+
+benchmark_group!(
+    scheduler,
+    spawn_many_local,
+    spawn_many_remote_idle,
+    spawn_many_remote_busy
+);
+
+benchmark_main!(scheduler);

tokio/src/runtime/scheduler/current_thread.rs

@@ -1,8 +1,8 @@
 use crate::future::poll_fn;
 use crate::loom::sync::atomic::AtomicBool;
-use crate::loom::sync::{Arc, Mutex};
+use crate::loom::sync::Arc;
 use crate::runtime::driver::{self, Driver};
-use crate::runtime::task::{self, JoinHandle, OwnedTasks, Schedule, Task};
+use crate::runtime::task::{self, Inject, JoinHandle, OwnedTasks, Schedule, Task};
 use crate::runtime::{blocking, context, scheduler, Config};
 use crate::runtime::{MetricsBatch, SchedulerMetrics, WorkerMetrics};
 use crate::sync::notify::Notify;
@@ -66,8 +66,8 @@ struct Core {
 
 /// Scheduler state shared between threads.
 struct Shared {
-    /// Remote run queue. None if the `Runtime` has been dropped.
-    queue: Mutex<Option<VecDeque<Notified>>>,
+    /// Remote run queue
+    inject: Inject<Arc<Handle>>,
 
     /// Collection of all active tasks spawned onto this executor.
     owned: OwnedTasks<Arc<Handle>>,
@@ -115,7 +115,7 @@ impl CurrentThread {
 
         let handle = Arc::new(Handle {
             shared: Shared {
-                queue: Mutex::new(Some(VecDeque::with_capacity(INITIAL_CAPACITY))),
+                inject: Inject::new(),
                 owned: OwnedTasks::new(),
                 woken: AtomicBool::new(false),
                 config,
@@ -217,15 +217,12 @@ impl CurrentThread {
                 drop(task);
             }
 
-            // Drain remote queue and set it to None
-            let remote_queue = handle.shared.queue.lock().take();
+            // Close the injection queue
+            handle.shared.inject.close();
 
-            // Using `Option::take` to replace the shared queue with `None`.
-            // We already shut down every task, so we just need to drop the task.
-            if let Some(remote_queue) = remote_queue {
-                for task in remote_queue {
-                    drop(task);
-                }
+            // Drain remote queue
+            while let Some(task) = handle.shared.inject.pop() {
+                drop(task);
             }
 
             assert!(handle.shared.owned.is_empty());
@@ -259,9 +256,12 @@ impl Core {
 
     fn next_task(&mut self, handle: &Handle) -> Option<Notified> {
         if self.tick % handle.shared.config.global_queue_interval == 0 {
-            handle.pop().or_else(|| self.next_local_task(handle))
+            handle
+                .next_remote_task()
+                .or_else(|| self.next_local_task(handle))
         } else {
-            self.next_local_task(handle).or_else(|| handle.pop())
+            self.next_local_task(handle)
+                .or_else(|| handle.next_remote_task())
         }
     }
 
@@ -440,14 +440,11 @@ impl Handle {
             };
             let local = &mut core.tasks;
 
-            let mut injection = self.shared.queue.lock();
-            let injection = if let Some(injection) = injection.as_mut() {
-                injection
-            } else {
+            if self.shared.inject.is_closed() {
                 return;
-            };
+            }
 
-            traces = trace_current_thread(&self.shared.owned, local, injection)
+            traces = trace_current_thread(&self.shared.owned, local, &self.shared.inject)
                 .into_iter()
                 .map(dump::Task::new)
                 .collect();
@@ -461,11 +458,8 @@ impl Handle {
         dump::Dump::new(traces)
     }
 
-    fn pop(&self) -> Option<Notified> {
-        match self.shared.queue.lock().as_mut() {
-            Some(queue) => queue.pop_front(),
-            None => None,
-        }
+    fn next_remote_task(&self) -> Option<Notified> {
+        self.shared.inject.pop()
     }
 
     fn waker_ref(me: &Arc<Self>) -> WakerRef<'_> {
@@ -488,14 +482,7 @@ cfg_metrics! {
         }
 
         pub(crate) fn injection_queue_depth(&self) -> usize {
-            // TODO: avoid having to lock. The multi-threaded injection queue
-            // could probably be used here.
-            self.shared
-                .queue
-                .lock()
-                .as_ref()
-                .map(|queue| queue.len())
-                .unwrap_or(0)
+            self.shared.inject.len()
         }
 
         pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics {
@@ -549,14 +536,9 @@ impl Schedule for Arc<Handle> {
                 // Track that a task was scheduled from **outside** of the runtime.
                 self.shared.scheduler_metrics.inc_remote_schedule_count();
 
-                // If the queue is None, then the runtime has shut down. We
-                // don't need to do anything with the notification in that case.
-                let mut guard = self.shared.queue.lock();
-                if let Some(queue) = guard.as_mut() {
-                    queue.push_back(task);
-                    drop(guard);
-                    self.driver.unpark();
-                }
+                // Schedule the task
+                self.shared.inject.push(task);
+                self.driver.unpark();
             }
         });
     }

tokio/src/runtime/task/core.rs

@@ -278,15 +278,11 @@ impl<T: Future, S: Schedule> Core<T, S> {
     }
 }
 
-cfg_rt_multi_thread! {
-    impl Header {
-        pub(super) unsafe fn set_next(&self, next: Option<NonNull<Header>>) {
-            self.queue_next.with_mut(|ptr| *ptr = next);
-        }
-    }
-}
-
 impl Header {
+    pub(super) unsafe fn set_next(&self, next: Option<NonNull<Header>>) {
+        self.queue_next.with_mut(|ptr| *ptr = next);
+    }
+
     // safety: The caller must guarantee exclusive access to this field, and
     // must ensure that the id is either 0 or the id of the OwnedTasks
     // containing this task.

tokio/src/runtime/task/inject.rs

@@ -52,6 +52,12 @@ impl<T: 'static> Inject<T> {
         self.len() == 0
     }
 
+    // Kind of annoying to have to include the cfg here
+    #[cfg(any(tokio_taskdump, all(feature = "rt-multi-thread", not(tokio_wasi))))]
+    pub(crate) fn is_closed(&self) -> bool {
+        self.pointers.lock().is_closed
+    }
+
     /// Closes the injection queue, returns `true` if the queue is open when the
     /// transition is made.
     pub(crate) fn close(&self) -> bool {
@@ -65,10 +71,6 @@ impl<T: 'static> Inject<T> {
         true
     }
 
-    pub(crate) fn is_closed(&self) -> bool {
-        self.pointers.lock().is_closed
-    }
-
     pub(crate) fn len(&self) -> usize {
         self.len.load(Acquire)
     }
@@ -104,71 +106,6 @@ impl<T: 'static> Inject<T> {
         self.len.store(len + 1, Release);
     }
 
-    /// Pushes several values into the queue.
-    #[inline]
-    pub(crate) fn push_batch<I>(&self, mut iter: I)
-    where
-        I: Iterator<Item = task::Notified<T>>,
-    {
-        let first = match iter.next() {
-            Some(first) => first.into_raw(),
-            None => return,
-        };
-
-        // Link up all the tasks.
-        let mut prev = first;
-        let mut counter = 1;
-
-        // We are going to be called with an `std::iter::Chain`, and that
-        // iterator overrides `for_each` to something that is easier for the
-        // compiler to optimize than a loop.
-        iter.for_each(|next| {
-            let next = next.into_raw();
-
-            // safety: Holding the Notified for a task guarantees exclusive
-            // access to the `queue_next` field.
-            set_next(prev, Some(next));
-            prev = next;
-            counter += 1;
-        });
-
-        // Now that the tasks are linked together, insert them into the
-        // linked list.
-        self.push_batch_inner(first, prev, counter);
-    }
-
-    /// Inserts several tasks that have been linked together into the queue.
-    ///
-    /// The provided head and tail may be be the same task. In this case, a
-    /// single task is inserted.
-    #[inline]
-    fn push_batch_inner(
-        &self,
-        batch_head: NonNull<task::Header>,
-        batch_tail: NonNull<task::Header>,
-        num: usize,
-    ) {
-        debug_assert!(get_next(batch_tail).is_none());
-
-        let mut p = self.pointers.lock();
-
-        if let Some(tail) = p.tail {
-            set_next(tail, Some(batch_head));
-        } else {
-            p.head = Some(batch_head);
-        }
-
-        p.tail = Some(batch_tail);
-
-        // Increment the count.
-        //
-        // 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() };
-
-        self.len.store(len + num, Release);
-    }
-
     pub(crate) fn pop(&self) -> Option<task::Notified<T>> {
         // Fast path, if len == 0, then there are no values
         if self.is_empty() {
@@ -201,6 +138,75 @@ impl<T: 'static> Inject<T> {
     }
 }
 
+cfg_rt_multi_thread! {
+    impl<T: 'static> Inject<T> {
+        /// Pushes several values into the queue.
+        #[inline]
+        pub(crate) fn push_batch<I>(&self, mut iter: I)
+        where
+            I: Iterator<Item = task::Notified<T>>,
+        {
+            let first = match iter.next() {
+                Some(first) => first.into_raw(),
+                None => return,
+            };
+
+            // Link up all the tasks.
+            let mut prev = first;
+            let mut counter = 1;
+
+            // We are going to be called with an `std::iter::Chain`, and that
+            // iterator overrides `for_each` to something that is easier for the
+            // compiler to optimize than a loop.
+            iter.for_each(|next| {
+                let next = next.into_raw();
+
+                // safety: Holding the Notified for a task guarantees exclusive
+                // access to the `queue_next` field.
+                set_next(prev, Some(next));
+                prev = next;
+                counter += 1;
+            });
+
+            // Now that the tasks are linked together, insert them into the
+            // linked list.
+            self.push_batch_inner(first, prev, counter);
+        }
+
+        /// Inserts several tasks that have been linked together into the queue.
+        ///
+        /// The provided head and tail may be be the same task. In this case, a
+        /// single task is inserted.
+        #[inline]
+        fn push_batch_inner(
+            &self,
+            batch_head: NonNull<task::Header>,
+            batch_tail: NonNull<task::Header>,
+            num: usize,
+        ) {
+            debug_assert!(get_next(batch_tail).is_none());
+
+            let mut p = self.pointers.lock();
+
+            if let Some(tail) = p.tail {
+                set_next(tail, Some(batch_head));
+            } else {
+                p.head = Some(batch_head);
+            }
+
+            p.tail = Some(batch_tail);
+
+            // Increment the count.
+            //
+            // 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() };
+
+            self.len.store(len + num, Release);
+        }
+    }
+}
+
 impl<T: 'static> Drop for Inject<T> {
     fn drop(&mut self) {
         if !std::thread::panicking() {

tokio/src/runtime/task/mod.rs

@@ -182,10 +182,8 @@ mod id;
 #[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
 pub use id::{id, try_id, Id};
 
-cfg_rt_multi_thread! {
-    mod inject;
-    pub(super) use self::inject::Inject;
-}
+mod inject;
+pub(super) use self::inject::Inject;
 
 #[cfg(feature = "rt")]
 mod abort;
@@ -370,25 +368,23 @@ impl<S: 'static> Notified<S> {
     }
 }
 
-cfg_rt_multi_thread! {
-    impl<S: 'static> Notified<S> {
-        unsafe fn from_raw(ptr: NonNull<Header>) -> Notified<S> {
-            Notified(Task::from_raw(ptr))
-        }
+impl<S: 'static> Notified<S> {
+    unsafe fn from_raw(ptr: NonNull<Header>) -> Notified<S> {
+        Notified(Task::from_raw(ptr))
     }
+}
 
-    impl<S: 'static> Task<S> {
-        fn into_raw(self) -> NonNull<Header> {
-            let ret = self.raw.header_ptr();
-            mem::forget(self);
-            ret
-        }
+impl<S: 'static> Task<S> {
+    fn into_raw(self) -> NonNull<Header> {
+        let ret = self.raw.header_ptr();
+        mem::forget(self);
+        ret
     }
+}
 
-    impl<S: 'static> Notified<S> {
-        fn into_raw(self) -> NonNull<Header> {
-            self.0.into_raw()
-        }
+impl<S: 'static> Notified<S> {
+    fn into_raw(self) -> NonNull<Header> {
+        self.0.into_raw()
     }
 }
 

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

@@ -1,5 +1,6 @@
 use crate::loom::sync::Arc;
 use crate::runtime::scheduler::current_thread;
+use crate::runtime::task::Inject;
 use backtrace::BacktraceFrame;
 use std::cell::Cell;
 use std::collections::VecDeque;
@@ -236,11 +237,14 @@ impl<T: Future> Future for Root<T> {
 pub(in crate::runtime) fn trace_current_thread(
     owned: &OwnedTasks<Arc<current_thread::Handle>>,
     local: &mut VecDeque<Notified<Arc<current_thread::Handle>>>,
-    injection: &mut VecDeque<Notified<Arc<current_thread::Handle>>>,
+    injection: &Inject<Arc<current_thread::Handle>>,
 ) -> Vec<Trace> {
     // clear the local and injection queues
     local.clear();
-    injection.clear();
+
+    while let Some(task) = injection.pop() {
+        drop(task);
+    }
 
     // notify each task
     let mut tasks = vec![];

tokio/tests/rt_common.rs

@@ -1317,4 +1317,39 @@ rt_test! {
             }
         });
     }
+
+    #[test]
+    #[cfg(not(target_os="wasi"))]
+    fn shutdown_concurrent_spawn() {
+        const NUM_TASKS: usize = 10_000;
+        for _ in 0..5 {
+            let (tx, rx) = std::sync::mpsc::channel();
+            let rt = rt();
+
+            let mut txs = vec![];
+
+            for _ in 0..NUM_TASKS {
+                let (tx, rx) = tokio::sync::oneshot::channel();
+                txs.push(tx);
+                rt.spawn(async move {
+                    rx.await.unwrap();
+                });
+            }
+
+            // Prime the tasks
+            rt.block_on(async { tokio::task::yield_now().await });
+
+            let th = std::thread::spawn(move || {
+                tx.send(()).unwrap();
+                for tx in txs.drain(..) {
+                    let _ = tx.send(());
+                }
+            });
+
+            rx.recv().unwrap();
+            drop(rt);
+
+            th.join().unwrap();
+        }
+    }
 }