From fccc28f1d0041fc5bdd798dc2645dc3b5fc71ef1 Mon Sep 17 00:00:00 2001 From: Eliza Weisman Date: Wed, 8 Apr 2026 13:46:27 -0700 Subject: [PATCH] runtime: optional eager I/O driver/timer handoff when polling tasks (#8010) --- tokio/src/runtime/builder.rs | 52 +++++- tokio/src/runtime/config.rs | 5 + .../runtime/scheduler/multi_thread/park.rs | 36 +++-- .../runtime/scheduler/multi_thread/worker.rs | 69 ++++++-- .../tests/rt_unstable_eager_driver_handoff.rs | 149 ++++++++++++++++++ 5 files changed, 288 insertions(+), 23 deletions(-) create mode 100644 tokio/tests/rt_unstable_eager_driver_handoff.rs diff --git a/tokio/src/runtime/builder.rs b/tokio/src/runtime/builder.rs index 62f183b5c..6032f2c60 100644 --- a/tokio/src/runtime/builder.rs +++ b/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() diff --git a/tokio/src/runtime/config.rs b/tokio/src/runtime/config.rs index ae972a9e7..ad25eb326 100644 --- a/tokio/src/runtime/config.rs +++ b/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, } diff --git a/tokio/src/runtime/scheduler/multi_thread/park.rs b/tokio/src/runtime/scheduler/multi_thread/park.rs index 79a3507f8..efe0ff067 100644 --- a/tokio/src/runtime/scheduler/multi_thread/park.rs +++ b/tokio/src/runtime/scheduler/multi_thread/park.rs @@ -21,6 +21,13 @@ pub(crate) struct Unparker { inner: Arc, } +/// 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) { - self.inner.park(handle); + pub(crate) fn park(&mut self, handle: &driver::Handle) -> HadDriver { + 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, - ) { + ) -> 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) { diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index beec19c08..86690b795 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/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 { - park.park_timeout(&self.worker.handle.driver, timeout); + let had_driver = if let Some(timeout) = duration { + 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 } } diff --git a/tokio/tests/rt_unstable_eager_driver_handoff.rs b/tokio/tests/rt_unstable_eager_driver_handoff.rs new file mode 100644 index 000000000..d41b14cac --- /dev/null +++ b/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: +/// . +/// +/// 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 +}