rt: move deferred task list to scheduler (#5741)

Previously, the deferred task list (list of tasks that yielded and are
waiting to be woken) was stored on the global runtime context. Because
the scheduler is responsible for waking these tasks, it took additional
TLS reads to perform the wake operation.

Instead, this commit moves the list of deferred tasks into the scheduler
context. This makes it easily accessible from the scheduler itself.

tokio/src/runtime/context.rs

@@ -10,10 +10,11 @@ cfg_rt! {
     mod scoped;
     use scoped::Scoped;
 
-    use crate::runtime::{scheduler, task::Id, Defer};
+    use crate::runtime::{scheduler, task::Id};
 
     use std::cell::RefCell;
     use std::marker::PhantomData;
+    use std::task::Waker;
     use std::time::Duration;
 
     cfg_taskdump! {
@@ -45,11 +46,6 @@ struct Context {
     #[cfg(feature = "rt")]
     runtime: Cell<EnterRuntime>,
 
-    /// Yielded task wakers are stored here and notified after resource drivers
-    /// are polled.
-    #[cfg(feature = "rt")]
-    defer: RefCell<Option<Defer>>,
-
     #[cfg(any(feature = "rt", feature = "macros"))]
     rng: FastRand,
 
@@ -93,9 +89,6 @@ tokio_thread_local! {
             #[cfg(feature = "rt")]
             runtime: Cell::new(EnterRuntime::NotEntered),
 
-            #[cfg(feature = "rt")]
-            defer: RefCell::new(None),
-
             #[cfg(any(feature = "rt", feature = "macros"))]
             rng: FastRand::new(RngSeed::new()),
 
@@ -170,12 +163,6 @@ cfg_rt! {
 
         #[allow(dead_code)] // Only tracking the guard.
         pub(crate) handle: SetCurrentGuard,
-
-        /// If true, then this is the root runtime guard. It is possible to nest
-        /// runtime guards by using `block_in_place` between the calls. We need
-        /// to track the root guard as this is the guard responsible for freeing
-        /// the deferred task queue.
-        is_root: bool,
     }
 
     /// Guard tracking that a caller has entered a blocking region.
@@ -240,20 +227,9 @@ cfg_rt! {
                 // Set the entered flag
                 c.runtime.set(EnterRuntime::Entered { allow_block_in_place });
 
-                // Initialize queue to track yielded tasks
-                let mut defer = c.defer.borrow_mut();
-
-                let is_root = if defer.is_none() {
-                    *defer = Some(Defer::new());
-                    true
-                } else {
-                    false
-                };
-
                 Some(EnterRuntimeGuard {
                     blocking: BlockingRegionGuard::new(),
                     handle: c.set_current(handle),
-                    is_root,
                 })
             }
         })
@@ -292,11 +268,17 @@ cfg_rt! {
         DisallowBlockInPlaceGuard(reset)
     }
 
-    pub(crate) fn with_defer<R>(f: impl FnOnce(&mut Defer) -> R) -> Option<R> {
-        CONTEXT.with(|c| {
-            let mut defer = c.defer.borrow_mut();
-            defer.as_mut().map(f)
-        })
+    #[track_caller]
+    pub(crate) fn defer(waker: &Waker) {
+        with_scheduler(|maybe_scheduler| {
+            if let Some(scheduler) = maybe_scheduler {
+                scheduler.defer(waker);
+            } else {
+                // Called from outside of the runtime, immediately wake the
+                // task.
+                waker.wake_by_ref();
+            }
+        });
     }
 
     pub(super) fn set_scheduler<R>(v: &scheduler::Context, f: impl FnOnce() -> R) -> R {
@@ -342,10 +324,6 @@ cfg_rt! {
             CONTEXT.with(|c| {
                 assert!(c.runtime.get().is_entered());
                 c.runtime.set(EnterRuntime::NotEntered);
-
-                if self.is_root {
-                    *c.defer.borrow_mut() = None;
-                }
             });
         }
     }
@@ -354,6 +332,7 @@ cfg_rt! {
         fn new() -> BlockingRegionGuard {
             BlockingRegionGuard { _p: PhantomData }
         }
+
         /// Blocks the thread on the specified future, returning the value with
         /// which that future completes.
         pub(crate) fn block_on<F>(&mut self, f: F) -> Result<F::Output, AccessError>
@@ -397,10 +376,6 @@ cfg_rt! {
                     return Err(());
                 }
 
-                // Wake any yielded tasks before parking in order to avoid
-                // blocking.
-                with_defer(|defer| defer.wake());
-
                 park.park_timeout(when - now);
             }
         }

tokio/src/runtime/defer.rs

@@ -1,38 +0,0 @@
-use std::task::Waker;
-
-pub(crate) struct Defer {
-    deferred: Vec<Waker>,
-}
-
-impl Defer {
-    pub(crate) fn new() -> Defer {
-        Defer {
-            deferred: Default::default(),
-        }
-    }
-
-    pub(crate) fn defer(&mut self, waker: &Waker) {
-        // If the same task adds itself a bunch of times, then only add it once.
-        if let Some(last) = self.deferred.last() {
-            if last.will_wake(waker) {
-                return;
-            }
-        }
-        self.deferred.push(waker.clone());
-    }
-
-    pub(crate) fn is_empty(&self) -> bool {
-        self.deferred.is_empty()
-    }
-
-    pub(crate) fn wake(&mut self) {
-        for waker in self.deferred.drain(..) {
-            waker.wake();
-        }
-    }
-
-    #[cfg(tokio_taskdump)]
-    pub(crate) fn take_deferred(&mut self) -> Vec<Waker> {
-        std::mem::take(&mut self.deferred)
-    }
-}

tokio/src/runtime/mod.rs

@@ -230,9 +230,6 @@ cfg_rt! {
         pub use crate::util::rand::RngSeed;
     }
 
-    mod defer;
-    pub(crate) use defer::Defer;
-
     cfg_taskdump! {
         pub mod dump;
         pub use dump::Dump;

tokio/src/runtime/park.rs

@@ -284,11 +284,6 @@ impl CachedParkThread {
                 return Ok(v);
             }
 
-            // Wake any yielded tasks before parking in order to avoid
-            // blocking.
-            #[cfg(feature = "rt")]
-            crate::runtime::context::with_defer(|defer| defer.wake());
-
             self.park();
         }
     }

tokio/src/runtime/scheduler/current_thread.rs

@@ -2,9 +2,9 @@ use crate::future::poll_fn;
 use crate::loom::sync::atomic::AtomicBool;
 use crate::loom::sync::Arc;
 use crate::runtime::driver::{self, Driver};
+use crate::runtime::scheduler::{self, Defer};
 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::runtime::{blocking, context, Config, MetricsBatch, SchedulerMetrics, WorkerMetrics};
 use crate::sync::notify::Notify;
 use crate::util::atomic_cell::AtomicCell;
 use crate::util::{waker_ref, RngSeedGenerator, Wake, WakerRef};
@@ -15,6 +15,7 @@ use std::fmt;
 use std::future::Future;
 use std::sync::atomic::Ordering::{AcqRel, Release};
 use std::task::Poll::{Pending, Ready};
+use std::task::Waker;
 use std::time::Duration;
 
 /// Executes tasks on the current thread
@@ -98,6 +99,9 @@ pub(crate) struct Context {
     /// Scheduler core, enabling the holder of `Context` to execute the
     /// scheduler.
     core: RefCell<Option<Box<Core>>>,
+
+    /// Deferred tasks, usually ones that called `task::yield_now()`.
+    pub(crate) defer: Defer,
 }
 
 type Notified = task::Notified<Arc<Handle>>;
@@ -201,6 +205,7 @@ impl CurrentThread {
             context: scheduler::Context::CurrentThread(Context {
                 handle: handle.clone(),
                 core: RefCell::new(Some(core)),
+                defer: Defer::new(),
             }),
             scheduler: self,
         })
@@ -320,21 +325,11 @@ impl Core {
     }
 }
 
-fn did_defer_tasks() -> bool {
-    context::with_defer(|deferred| !deferred.is_empty()).unwrap()
-}
-
-fn wake_deferred_tasks() {
-    context::with_defer(|deferred| deferred.wake());
-}
-
 #[cfg(tokio_taskdump)]
-fn wake_deferred_tasks_and_free() {
-    let wakers = context::with_defer(|deferred| deferred.take_deferred());
-    if let Some(wakers) = wakers {
-        for waker in wakers {
-            waker.wake();
-        }
+fn wake_deferred_tasks_and_free(context: &Context) {
+    let wakers = context.defer.take_deferred();
+    for waker in wakers {
+        waker.wake();
     }
 }
 
@@ -372,7 +367,7 @@ impl Context {
 
             let (c, _) = self.enter(core, || {
                 driver.park(&handle.driver);
-                wake_deferred_tasks();
+                self.defer.wake();
             });
 
             core = c;
@@ -398,7 +393,7 @@ impl Context {
 
         let (mut core, _) = self.enter(core, || {
             driver.park_timeout(&handle.driver, Duration::from_millis(0));
-            wake_deferred_tasks();
+            self.defer.wake();
         });
 
         core.driver = Some(driver);
@@ -418,6 +413,10 @@ impl Context {
         let core = self.core.borrow_mut().take().expect("core missing");
         (core, ret)
     }
+
+    pub(crate) fn defer(&self, waker: &Waker) {
+        self.defer.defer(waker);
+    }
 }
 
 // ===== impl Handle =====
@@ -479,12 +478,15 @@ impl Handle {
                 .into_iter()
                 .map(dump::Task::new)
                 .collect();
-        });
 
-        // Taking a taskdump could wakes every task, but we probably don't want
-        // the `yield_now` vector to be that large under normal circumstances.
-        // Therefore, we free its allocation.
-        wake_deferred_tasks_and_free();
+            // Avoid double borrow panic
+            drop(maybe_core);
+
+            // Taking a taskdump could wakes every task, but we probably don't want
+            // the `yield_now` vector to be that large under normal circumstances.
+            // Therefore, we free its allocation.
+            wake_deferred_tasks_and_free(context);
+        });
 
         dump::Dump::new(traces)
     }
@@ -671,7 +673,7 @@ impl CoreGuard<'_> {
                         None => {
                             core.metrics.end_processing_scheduled_tasks();
 
-                            core = if did_defer_tasks() {
+                            core = if !context.defer.is_empty() {
                                 context.park_yield(core, handle)
                             } else {
                                 context.park(core, handle)

tokio/src/runtime/scheduler/defer.rs

@@ -0,0 +1,43 @@
+use std::cell::RefCell;
+use std::task::Waker;
+
+pub(crate) struct Defer {
+    deferred: RefCell<Vec<Waker>>,
+}
+
+impl Defer {
+    pub(crate) fn new() -> Defer {
+        Defer {
+            deferred: Default::default(),
+        }
+    }
+
+    pub(crate) fn defer(&self, waker: &Waker) {
+        let mut deferred = self.deferred.borrow_mut();
+
+        // If the same task adds itself a bunch of times, then only add it once.
+        if let Some(last) = deferred.last() {
+            if last.will_wake(waker) {
+                return;
+            }
+        }
+
+        deferred.push(waker.clone());
+    }
+
+    pub(crate) fn is_empty(&self) -> bool {
+        self.deferred.borrow().is_empty()
+    }
+
+    pub(crate) fn wake(&self) {
+        while let Some(waker) = self.deferred.borrow_mut().pop() {
+            waker.wake();
+        }
+    }
+
+    #[cfg(tokio_taskdump)]
+    pub(crate) fn take_deferred(&self) -> Vec<Waker> {
+        let mut deferred = self.deferred.borrow_mut();
+        std::mem::take(&mut *deferred)
+    }
+}

tokio/src/runtime/scheduler/mod.rs

@@ -1,6 +1,9 @@
 cfg_rt! {
     pub(crate) mod current_thread;
     pub(crate) use current_thread::CurrentThread;
+
+    mod defer;
+    use defer::Defer;
 }
 
 cfg_rt_multi_thread! {
@@ -56,6 +59,7 @@ cfg_rt! {
     use crate::runtime::context;
     use crate::task::JoinHandle;
     use crate::util::RngSeedGenerator;
+    use std::task::Waker;
 
     impl Handle {
         #[track_caller]
@@ -203,6 +207,14 @@ cfg_rt! {
             }
         }
 
+        pub(crate) fn defer(&self, waker: &Waker) {
+            match self {
+                Context::CurrentThread(context) => context.defer(waker),
+                #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
+                Context::MultiThread(context) => context.defer(waker),
+            }
+        }
+
         cfg_rt_multi_thread! {
             #[track_caller]
             pub(crate) fn expect_multi_thread(&self) -> &multi_thread::Context {

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

@@ -62,6 +62,7 @@ use crate::runtime::context;
 use crate::runtime::scheduler::multi_thread::{
     idle, queue, Counters, Handle, Idle, Parker, Stats, Unparker,
 };
+use crate::runtime::scheduler::Defer;
 use crate::runtime::task::{Inject, OwnedTasks};
 use crate::runtime::{
     blocking, coop, driver, scheduler, task, Config, SchedulerMetrics, WorkerMetrics,
@@ -70,6 +71,7 @@ use crate::util::atomic_cell::AtomicCell;
 use crate::util::rand::{FastRand, RngSeedGenerator};
 
 use std::cell::RefCell;
+use std::task::Waker;
 use std::time::Duration;
 
 /// A scheduler worker
@@ -189,6 +191,10 @@ pub(crate) struct Context {
 
     /// Core data
     core: RefCell<Option<Box<Core>>>,
+
+    /// Tasks to wake after resource drivers are polled. This is mostly to
+    /// handle yielded tasks.
+    pub(crate) defer: Defer,
 }
 
 /// Starts the workers
@@ -432,6 +438,7 @@ fn run(worker: Arc<Worker>) {
     let cx = scheduler::Context::MultiThread(Context {
         worker,
         core: RefCell::new(None),
+        defer: Defer::new(),
     });
 
     context::set_scheduler(&cx, || {
@@ -444,7 +451,7 @@ fn run(worker: Arc<Worker>) {
         // Check if there are any deferred tasks to notify. This can happen when
         // the worker core is lost due to `block_in_place()` being called from
         // within the task.
-        wake_deferred_tasks();
+        cx.defer.wake();
     });
 }
 
@@ -484,7 +491,7 @@ impl Context {
                 core = self.run_task(task, core)?;
             } else {
                 // Wait for work
-                core = if did_defer_tasks() {
+                core = if !self.defer.is_empty() {
                     self.park_timeout(core, Some(Duration::from_millis(0)))
                 } else {
                     self.park(core)
@@ -669,7 +676,7 @@ impl Context {
             park.park(&self.worker.handle.driver);
         }
 
-        wake_deferred_tasks();
+        self.defer.wake();
 
         // Remove `core` from context
         core = self.core.borrow_mut().take().expect("core missing");
@@ -685,6 +692,10 @@ impl Context {
 
         core
     }
+
+    pub(crate) fn defer(&self, waker: &Waker) {
+        self.defer.defer(waker);
+    }
 }
 
 impl Core {
@@ -1054,14 +1065,6 @@ impl Handle {
     }
 }
 
-fn did_defer_tasks() -> bool {
-    context::with_defer(|deferred| !deferred.is_empty()).unwrap()
-}
-
-fn wake_deferred_tasks() {
-    context::with_defer(|deferred| deferred.wake());
-}
-
 #[track_caller]
 fn with_current<R>(f: impl FnOnce(Option<&Context>) -> R) -> R {
     use scheduler::Context::MultiThread;

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

@@ -1,6 +1,8 @@
 use crate::loom::sync::Arc;
-use crate::runtime::scheduler::current_thread;
+use crate::runtime::context;
+use crate::runtime::scheduler::{self, current_thread};
 use crate::runtime::task::Inject;
+
 use backtrace::BacktraceFrame;
 use std::cell::Cell;
 use std::collections::VecDeque;
@@ -179,10 +181,15 @@ pub(crate) fn trace_leaf(cx: &mut task::Context<'_>) -> Poll<()> {
     if did_trace {
         // Use the same logic that `yield_now` uses to send out wakeups after
         // the task yields.
-        let defer = crate::runtime::context::with_defer(|rt| {
-            rt.defer(cx.waker());
+        context::with_scheduler(|scheduler| {
+            if let Some(scheduler) = scheduler {
+                match scheduler {
+                    scheduler::Context::CurrentThread(s) => s.defer.defer(cx.waker()),
+                    #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
+                    scheduler::Context::MultiThread(s) => s.defer.defer(cx.waker()),
+                }
+            }
         });
-        debug_assert!(defer.is_some());
 
         Poll::Pending
     } else {

tokio/src/task/yield_now.rs

@@ -54,15 +54,7 @@ pub async fn yield_now() {
 
             self.yielded = true;
 
-            let defer = context::with_defer(|rt| {
-                rt.defer(cx.waker());
-            });
-
-            if defer.is_none() {
-                //  Not currently in a runtime, just notify ourselves
-                //  immediately.
-                cx.waker().wake_by_ref();
-            }
+            context::defer(cx.waker());
 
             Poll::Pending
         }

tokio/tests/task_yield_now.rs

@@ -0,0 +1,16 @@
+#![cfg(all(feature = "full", tokio_unstable))]
+
+use tokio::task;
+use tokio_test::task::spawn;
+
+// `yield_now` is tested within the runtime in `rt_common`.
+#[test]
+fn yield_now_outside_of_runtime() {
+    let mut task = spawn(async {
+        task::yield_now().await;
+    });
+
+    assert!(task.poll().is_pending());
+    assert!(task.is_woken());
+    assert!(task.poll().is_ready());
+}