runtime: optional eager I/O driver/timer handoff when polling tasks (#8010)

tokio/src/runtime/builder.rs

@@ -142,6 +142,10 @@ pub struct Builder {
     pub(super) unhandled_panic: UnhandledPanic,
 
     timer_flavor: TimerFlavor,
+
+    /// Whether or not to enable eager hand-off for the I/O and time drivers (in
+    /// `tokio_unstable`).
+    enable_eager_driver_handoff: bool,
 }
 
 cfg_unstable! {
@@ -334,6 +338,9 @@ impl Builder {
             disable_lifo_slot: false,
 
             timer_flavor: TimerFlavor::Traditional,
+
+            // Eager driver handoff is disabled by default.
+            enable_eager_driver_handoff: false,
         }
     }
 
@@ -414,6 +421,40 @@ impl Builder {
         self
     }
 
+    /// Enable eager hand-off of the I/O and time drivers for multi-threaded
+    /// runtimes, which is disabled by default.
+    ///
+    /// When this option is enabled, a worker thread which has parked on the I/O
+    /// or time driver will notify another worker thread once it is preparing to
+    /// begin polling a task from the run queue, so that the notified worker can
+    /// begin polling the I/O or time driver. This can reduce the latency with
+    /// which I/O and timer notifications are processed, especially when some
+    /// tasks have polls that take a long time to complete. In addition, it can
+    /// reduce the risk of a deadlock which may occur when a task blocks the
+    /// worker thread which is holding the I/O or time driver until some other
+    /// task, which is waiting for a notification from *that* driver, unblocks
+    /// it.
+    ///
+    /// This option is disabled by default, as enabling it may potentially
+    /// increase contention due to extra synchronization in cross-driver
+    /// wakeups.
+    ///
+    /// This option only applies to multi-threaded runtimes. Attempting to use
+    /// this option with any other runtime type will have no effect.
+    ///
+    /// **Note**: This is an [unstable API][unstable]. Eager driver hand-off is
+    /// an experimental feature whose behavior may be removed or changed in 1.x
+    /// releases. See [the documentation on unstable features][unstable] for
+    /// details.
+    ///
+    /// [unstable]: crate#unstable-features
+    #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))]
+    #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "rt-multi-thread"))))]
+    pub fn enable_eager_driver_handoff(&mut self) -> &mut Self {
+        self.enable_eager_driver_handoff = true;
+        self
+    }
+
     /// Sets the number of worker threads the `Runtime` will use.
     ///
     /// This can be any number above 0 though it is advised to keep this value
@@ -1661,6 +1702,10 @@ impl Builder {
                 #[cfg(tokio_unstable)]
                 unhandled_panic: self.unhandled_panic.clone(),
                 disable_lifo_slot: self.disable_lifo_slot,
+                // This setting never makes sense for a current thread runtime,
+                // as it only configures how the I/O driver is stolen across
+                // workers.
+                enable_eager_driver_handoff: false,
                 seed_generator: seed_generator_1,
                 metrics_poll_count_histogram: self.metrics_poll_count_histogram_builder(),
             },
@@ -1843,6 +1888,7 @@ cfg_rt_multi_thread! {
                     #[cfg(tokio_unstable)]
                     unhandled_panic: self.unhandled_panic.clone(),
                     disable_lifo_slot: self.disable_lifo_slot,
+                    enable_eager_driver_handoff: self.enable_eager_driver_handoff,
                     seed_generator: seed_generator_1,
                     metrics_poll_count_histogram: self.metrics_poll_count_histogram_builder(),
                 },
@@ -1880,7 +1926,11 @@ impl fmt::Debug for Builder {
             .field("after_start", &self.after_start.as_ref().map(|_| "..."))
             .field("before_stop", &self.before_stop.as_ref().map(|_| "..."))
             .field("before_park", &self.before_park.as_ref().map(|_| "..."))
-            .field("after_unpark", &self.after_unpark.as_ref().map(|_| "..."));
+            .field("after_unpark", &self.after_unpark.as_ref().map(|_| "..."))
+            .field(
+                "enable_eager_driver_handoff",
+                &self.enable_eager_driver_handoff,
+            );
 
         if self.name.is_none() {
             debug.finish_non_exhaustive()

tokio/src/runtime/config.rs

@@ -53,4 +53,9 @@ pub(crate) struct Config {
     #[cfg(tokio_unstable)]
     /// How to respond to unhandled task panics.
     pub(crate) unhandled_panic: crate::runtime::UnhandledPanic,
+
+    /// If `true`, an idle worker is woken whenever a worker thread transitions
+    /// from polling the I/O driver to polling its own tasks (requires
+    /// `tokio_unstable`).
+    pub(crate) enable_eager_driver_handoff: bool,
 }

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

@@ -21,6 +21,13 @@ pub(crate) struct Unparker {
     inner: Arc<Inner>,
 }
 
+/// Represents how a worker thread was parked
+#[derive(Copy, Clone, Eq, PartialEq)]
+pub(crate) enum HadDriver {
+    Yes,
+    No,
+}
+
 struct Inner {
     /// Avoids entering the park if possible
     state: AtomicUsize,
@@ -66,8 +73,8 @@ impl Parker {
         }
     }
 
-    pub(crate) fn park(&mut self, handle: &driver::Handle) {
+    pub(crate) fn park(&mut self, handle: &driver::Handle) -> HadDriver {
-        self.inner.park(handle);
+        self.inner.park(handle)
     }
 
     /// Parks the current thread for up to `duration`.
@@ -75,11 +82,16 @@ impl Parker {
     /// This function tries to acquire the driver lock. If it succeeds, it
     /// parks using the driver. Otherwise, it fails back to using a condvar,
     /// unless the duration is zero, in which case it returns immediately.
-    pub(crate) fn park_timeout(&mut self, handle: &driver::Handle, duration: Duration) {
+    pub(crate) fn park_timeout(
+        &mut self,
+        handle: &driver::Handle,
+        duration: Duration,
+    ) -> HadDriver {
         if let Some(mut driver) = self.inner.shared.driver.try_lock() {
-            self.inner.park_driver(&mut driver, handle, Some(duration));
+            self.inner.park_driver(&mut driver, handle, Some(duration))
         } else if !duration.is_zero() {
             self.inner.park_condvar(Some(duration));
+            HadDriver::No
         } else {
             // https://github.com/tokio-rs/tokio/issues/6536
             // Hacky, but it's just for loom tests. The counter gets incremented during
@@ -87,6 +99,7 @@ impl Parker {
             // lock.
             #[cfg(loom)]
             CURRENT_THREAD_PARK_COUNT.with(|count| count.fetch_add(1, SeqCst));
+            HadDriver::No
         }
     }
 
@@ -116,7 +129,7 @@ impl Unparker {
 
 impl Inner {
     /// Parks the current thread for at most `dur`.
-    fn park(&self, handle: &driver::Handle) {
+    fn park(&self, handle: &driver::Handle) -> HadDriver {
         // If we were previously notified then we consume this notification and
         // return quickly.
         if self
@@ -124,13 +137,14 @@ impl Inner {
             .compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst)
             .is_ok()
         {
-            return;
+            return HadDriver::No;
         }
 
         if let Some(mut driver) = self.shared.driver.try_lock() {
-            self.park_driver(&mut driver, handle, None);
+            self.park_driver(&mut driver, handle, None)
         } else {
             self.park_condvar(None);
+            HadDriver::No
         }
     }
 
@@ -216,12 +230,12 @@ impl Inner {
         driver: &mut Driver,
         handle: &driver::Handle,
         duration: Option<Duration>,
-    ) {
+    ) -> HadDriver {
         if duration.as_ref().is_some_and(Duration::is_zero) {
             // zero duration doesn't actually park the thread, it just
             // polls the I/O events, timers, etc.
             driver.park_timeout(handle, Duration::ZERO);
-            return;
+            return HadDriver::Yes;
         }
 
         match self
@@ -239,7 +253,7 @@ impl Inner {
                 let old = self.state.swap(EMPTY, SeqCst);
                 debug_assert_eq!(old, NOTIFIED, "park state changed unexpectedly");
 
-                return;
+                return HadDriver::No;
             }
             Err(actual) => panic!("inconsistent park state; actual = {actual}"),
         }
@@ -256,6 +270,8 @@ impl Inner {
             PARKED_DRIVER => {} // no notification, alas
             n => panic!("inconsistent park_timeout state: {n}"),
         }
+
+        HadDriver::Yes
     }
 
     fn unpark(&self, driver: &driver::Handle) {

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

@@ -59,7 +59,7 @@
 use crate::loom::sync::{Arc, Mutex};
 use crate::runtime;
 use crate::runtime::scheduler::multi_thread::{
-    idle, queue, Counters, Handle, Idle, Overflow, Parker, Stats, TraceStatus, Unparker,
+    idle, park, queue, Counters, Handle, Idle, Overflow, Parker, Stats, TraceStatus, Unparker,
 };
 use crate::runtime::scheduler::{inject, Defer, Lock};
 use crate::runtime::task::OwnedTasks;
@@ -132,6 +132,15 @@ struct Core {
     /// True if the scheduler is being traced
     is_traced: bool,
 
+    /// Whether or not the worker has just returned from a park in which we
+    /// parked on the I/O driver.
+    had_driver: park::HadDriver,
+
+    /// If `true`, the worker should eagerly notify another worker when polling
+    /// the first task after returning from a park in which it parked on the I/O
+    /// or time driver.
+    enable_eager_driver_handoff: bool,
+
     /// Parker
     ///
     /// Stored in an `Option` as the parker is added / removed to make the
@@ -280,6 +289,8 @@ pub(super) fn create(
             is_searching: false,
             is_shutdown: false,
             is_traced: false,
+            enable_eager_driver_handoff: config.enable_eager_driver_handoff,
+            had_driver: park::HadDriver::No,
             park: Some(park),
             global_queue_interval: stats.tuned_global_queue_interval(&config),
             stats,
@@ -616,7 +627,30 @@ impl Context {
 
         // Make sure the worker is not in the **searching** state. This enables
         // another idle worker to try to steal work.
-        core.transition_from_searching(&self.worker);
+        let notified_parked_worker = core.transition_from_searching(&self.worker);
+
+        // If the setting to wake eagerly when releasing the I/O driver is
+        // enabled, and this worker had the driver, wake a parked worker to come
+        // grab it from us.
+        //
+        // Note that this is only done when we are *actually* about to poll a
+        // task, rather than whenever the worker has unparked. When the worker
+        // has been unparked, it may not actually have any tasks to poll, and if
+        // it's still holding the I/O driver, it should just go back to polling
+        // the driver again, rather than trying to wake someone else spuriously.
+        //
+        // Note that this explicitly checks `cfg!(tokio_unstable)` in addition,
+        // as that should result in this whole expression being eliminated at
+        // compile-time when unstable features are disabled.
+        if cfg!(tokio_unstable)
+            && core.enable_eager_driver_handoff
+            && core.had_driver == park::HadDriver::Yes
+            && !notified_parked_worker
+        // don't do it a second time
+        {
+            core.had_driver = park::HadDriver::No;
+            self.worker.handle.notify_parked_local();
+        }
 
         self.assert_lifo_enabled_is_correct(&core);
 
@@ -825,11 +859,11 @@ impl Context {
         };
 
         // Park thread
-        if let Some(timeout) = duration {
+        let had_driver = if let Some(timeout) = duration {
-            park.park_timeout(&self.worker.handle.driver, timeout);
+            park.park_timeout(&self.worker.handle.driver, timeout)
         } else {
-            park.park(&self.worker.handle.driver);
+            park.park(&self.worker.handle.driver)
-        }
+        };
 
         self.defer.wake();
 
@@ -854,6 +888,8 @@ impl Context {
 
         // Place `park` back in `core`
         core.park = Some(park);
+        core.had_driver = had_driver;
+
         if core.should_notify_others() {
             self.worker.handle.notify_parked_local();
         }
@@ -1117,13 +1153,13 @@ impl Core {
         self.is_searching
     }
 
-    fn transition_from_searching(&mut self, worker: &Worker) {
+    fn transition_from_searching(&mut self, worker: &Worker) -> bool {
         if !self.is_searching {
-            return;
+            return false;
         }
 
         self.is_searching = false;
-        worker.handle.transition_worker_from_searching();
+        worker.handle.transition_worker_from_searching()
     }
 
     fn has_tasks(&self) -> bool {
@@ -1370,12 +1406,18 @@ impl Handle {
         }
     }
 
-    fn notify_parked_local(&self) {
+    /// Notify a parked worker.
+    ///
+    /// Returns `true` if a worker was notified, `false` otherwise.
+    fn notify_parked_local(&self) -> bool {
         super::counters::inc_num_inc_notify_local();
 
         if let Some(index) = self.shared.idle.worker_to_notify(&self.shared) {
             super::counters::inc_num_unparks_local();
             self.shared.remotes[index].unpark.unpark(&self.driver);
+            true
+        } else {
+            false
         }
     }
 
@@ -1404,11 +1446,14 @@ impl Handle {
         }
     }
 
-    fn transition_worker_from_searching(&self) {
+    /// Returns `true` if another parked worker was notified, `false` otherwise.
+    fn transition_worker_from_searching(&self) -> bool {
         if self.shared.idle.transition_worker_from_searching() {
             // We are the final searching worker. Because work was found, we
             // need to notify another worker.
-            self.notify_parked_local();
+            self.notify_parked_local()
+        } else {
+            false
         }
     }
 

tokio/tests/rt_unstable_eager_driver_handoff.rs

@@ -0,0 +1,149 @@
+// These tests only work on Unix platforms because they rely on Unix pipes
+// as a way of generating I/O events from within the same process.
+//
+// Also, Miri doesn't like it when you leak a thread, which will happen in
+// the "deadlock" case below, so we skip these tests on Miri.
+#![cfg(all(not(miri), unix, feature = "full"))]
+
+use std::sync::mpsc::RecvTimeoutError;
+use std::time::Duration;
+
+/// Test that, without `enable_eager_driver_handoff`, we can reliably reproduce
+/// a deadlock when a task blocks indefinitely. If this test fails, it means
+/// that the test `eager_driver_handoff_fixes_deadlock` is not actually testing
+/// a condition that can deadlock the runtime.
+#[test]
+fn deadlocks_consistently() {
+    let rt = rt_builder().build().unwrap();
+    assert_eq!(
+        do_test(rt),
+        Err(RecvTimeoutError::Timeout),
+        "runtime did not deadlock! the `eager_driver_handoff_fixes_deadlock` \
+         test may no longer reproduce the bug it is intended to test a fix \
+         for!",
+    );
+}
+
+/// This is the one that actually tests whether eager driver handoff works as
+/// expected: it runs the same reproducer as `deadlocks_consistently` a single
+/// timebut with `enable_eager_driver_handoff` enabled. If this test fails, it
+/// means that the eager driver handoff fix is not working as expected.
+#[test]
+#[cfg(tokio_unstable)]
+fn eager_driver_handoff_fixes_deadlock() {
+    let rt = rt_builder().enable_eager_driver_handoff().build().unwrap();
+    assert_eq!(
+        do_test(rt),
+        Ok(()),
+        "the runtime should not deadlock because the driver is \
+         eagerly handed off"
+    );
+}
+
+/// Reproduces a deadlock occurring when the worker thread holding the I/O or
+/// time driver runs a task that blocks that worker until another task, which is
+/// *waiting on the I/O or time driver*, performs an action that unblocks it.
+/// This general class of problem is described in:
+/// <https://github.com/tokio-rs/tokio/issues/4730>.
+///
+/// The deadlock occurs as follows:
+///
+/// 1. Both worker threads are idle. Worker A parks on the I/O driver (holding
+///    the driver lock), Worker B parks on a condvar.
+/// 2. An I/O event fires. Worker A (the driver holder) wakes up and processes the
+///    I/O  event, placing the woken task in its run queue.
+/// 3. Worker A begins polling the task. The task blocks the worker thread
+///    until it is woken by the other task, so Worker A never returns to
+///    poll the I/O driver.
+/// 4. Worker B remains parked on the condvar indefinitely — nobody calls
+///    `unpark` on it, so it never wakes to take over the I/O driver.
+/// 5. Subsequent I/O events are never processed because no worker is polling
+///    the time driver. The runtime is wedged.
+///
+/// To trigger this reliably, the test uses two tasks which send and receive
+/// on a pair of Unix pipes. We use pipes to reliably trigger I/O events from
+/// within the same process.
+///
+/// The "bad" task's pipe is written to immediately, while the "good" task is
+/// still waiting for a read on its pipe. Then, when it's woken, it writes to
+/// the other pipe, waking the "good" task, and then blocks on a notification
+/// from a blocking channel which is only written to by the "good" task.
+///
+/// Because this task first waits on a pipe read, this ensures that it is
+/// first polled from the worker thread that has last parked on the I/O
+/// driver, so when it blocks, it prevents the driver from making progress.
+///
+/// The specific scenario we've constructed here is, admittedly, somewhat
+/// contrived: most Tokio applications aren't going to spawn tasks that attempt
+/// to communicate using a `std::sync::mpsc` blocking channel. Instead, you can
+/// imagine the blocking channel as a stand-in for some other operation that
+/// blocks a thread until some event occurs, which is triggered by an operation
+/// performed by the other task, which is waiting on an asynchronous event
+/// before performing it. The task that blocks the runtime could, for example,
+/// be waiting on a blocking syscall that completes only when the other task
+/// does something.
+fn do_test(rt: tokio::runtime::Runtime) -> Result<(), RecvTimeoutError> {
+    use tokio::io::{AsyncReadExt, AsyncWriteExt};
+    use tokio::net::unix::pipe::pipe;
+
+    // A blocking MPSC from the standard library is used to wait for the test to
+    // complete within a reasonable amount of time, so that we can determine
+    // whether or not the runtime has deadlocked.
+    let (done_tx, done_rx) = std::sync::mpsc::channel();
+
+    std::thread::spawn(move || {
+        rt.block_on(async {
+            let (mut pipe1_tx, mut pipe1_rx) = pipe().expect("ceci n'est pas une pipe");
+            let (mut pipe2_tx, mut pipe2_rx) = pipe().expect("ceci n'est pas une pipe");
+
+            // Note that we use a *blocking* MPSC from the standard library
+            // here, since the entire purpose of the channel is to temporarily
+            // block a worker thread when calling `recv`.
+            let (deadlock_tx, deadlock_rx) = std::sync::mpsc::channel();
+
+            let bad_task = tokio::spawn(async move {
+                // Wait on a pipe for a bit to ensure that we end up on the worker
+                // holding the time driver.
+                let mut buf = [0u8; 1];
+                pipe1_rx.read_exact(&mut buf).await.unwrap();
+                // Okay, we have now definitely woken up on the worker
+                // thread that polled the I/0 driver. Now, block this
+                // worker thread until woken by the *other* task. If the
+                // other task's I/O driver notification wakes it up,
+                // this task will be unblocked and the runtime will keep running.
+                //
+                // However, if *this* task blocking the thread prevents the
+                // I/O driver from ever running, the runtime will be wedged forever.
+                pipe2_tx.write_all(&[2]).await.unwrap();
+                deadlock_rx.recv().unwrap();
+            });
+
+            let good_task = tokio::spawn(async move {
+                let mut buf = [0u8; 1];
+                pipe2_rx.read_exact(&mut buf).await.unwrap();
+                deadlock_tx.send(()).unwrap();
+            });
+
+            tokio::time::sleep(Duration::from_millis(100)).await;
+
+            pipe1_tx.write_all(&[1]).await.unwrap();
+
+            good_task.await.unwrap();
+            bad_task.await.unwrap();
+        });
+
+        done_tx.send(()).unwrap();
+    });
+
+    done_rx.recv_timeout(Duration::from_secs(10))
+}
+
+/// Base runtime builder for both tests in this module: two worker threads, time
+/// driver enabled. This returns a `Builder`, rather than a `Runtime`, so that
+/// the `eager_driver_handoff_fixes_deadlock` test can configure the runtime to
+/// enable eager driver handoff before building it.
+fn rt_builder() -> tokio::runtime::Builder {
+    let mut builder = tokio::runtime::Builder::new_multi_thread();
+    builder.enable_all().worker_threads(2);
+    builder
+}