From dc759b9dd61bec0e951bd225b6529a33088e8ae2 Mon Sep 17 00:00:00 2001 From: Eliza Weisman Date: Sat, 18 Apr 2026 14:11:21 -0700 Subject: [PATCH] defer lifo unpark w timer --- .../runtime/scheduler/multi_thread/idle.rs | 55 +++++++++++++++---- .../runtime/scheduler/multi_thread/queue.rs | 4 ++ .../runtime/scheduler/multi_thread/worker.rs | 54 +++++++++++++++--- 3 files changed, 95 insertions(+), 18 deletions(-) diff --git a/tokio/src/runtime/scheduler/multi_thread/idle.rs b/tokio/src/runtime/scheduler/multi_thread/idle.rs index 77cfb157e..0000a5419 100644 --- a/tokio/src/runtime/scheduler/multi_thread/idle.rs +++ b/tokio/src/runtime/scheduler/multi_thread/idle.rs @@ -23,8 +23,14 @@ pub(super) struct Synced { sleepers: Vec, } -const UNPARK_SHIFT: usize = 16; -const UNPARK_MASK: usize = !SEARCH_MASK; +pub(super) struct TransitionToParked { + pub(super) is_last_searcher: bool, + pub(super) any_lifo: bool, +} + +const UNPARK_SHIFT: usize = (usize::BITS as usize / 2) - 2; +const ANY_LIFO: usize = 1 << (usize::BITS - 1); +const UNPARK_MASK: usize = !(SEARCH_MASK | ANY_LIFO); const SEARCH_MASK: usize = (1 << UNPARK_SHIFT) - 1; #[derive(Copy, Clone)] @@ -32,6 +38,10 @@ struct State(usize); impl Idle { pub(super) fn new(num_workers: usize) -> (Idle, Synced) { + assert!( + num_workers <= UNPARK_MASK, + "{num_workers} is too many workers (max is {UNPARK_MASK})" + ); let init = State::new(num_workers); let idle = Idle { @@ -88,7 +98,7 @@ impl Idle { shared: &Shared, worker: usize, is_searching: bool, - ) -> bool { + ) -> TransitionToParked { // Acquire the lock let mut lock = shared.synced.lock(); @@ -144,18 +154,25 @@ impl Idle { false } + pub(super) fn put_lifo(&self) -> bool { + State(self.state.fetch_or(ANY_LIFO, SeqCst)).any_lifo() + } + + pub(super) fn clear_lifo(&self) { + self.state.fetch_and(!ANY_LIFO, SeqCst); + } + + pub(super) fn should_attempt_lifo_steal(&self) -> bool { + let state = State(self.state.fetch_add(0, SeqCst)); + state.any_lifo() + } + /// Returns `true` if `worker_id` is contained in the sleep set. pub(super) fn is_parked(&self, shared: &Shared, worker_id: usize) -> bool { let lock = shared.synced.lock(); lock.idle.sleepers.contains(&worker_id) } - /// Returns `true` if all other workers are currently parked. - pub(super) fn all_parked(&self) -> bool { - let state = State(self.state.fetch_add(0, SeqCst)); - state.num_unparked() <= 1 - } - fn notify_should_wakeup(&self) -> bool { let state = State(self.state.fetch_add(0, SeqCst)); state.num_searching() == 0 && state.num_unparked() < self.num_workers @@ -192,7 +209,7 @@ impl State { /// Track a sleeping worker /// /// Returns `true` if this is the final searching worker. - fn dec_num_unparked(cell: &AtomicUsize, is_searching: bool) -> bool { + fn dec_num_unparked(cell: &AtomicUsize, is_searching: bool) -> TransitionToParked { let mut dec = 1 << UNPARK_SHIFT; if is_searching { @@ -200,7 +217,11 @@ impl State { } let prev = State(cell.fetch_sub(dec, SeqCst)); - is_searching && prev.num_searching() == 1 + let is_last_searcher = is_searching && prev.num_searching() == 1; + TransitionToParked { + is_last_searcher, + any_lifo: prev.any_lifo(), + } } /// Number of workers currently searching @@ -212,6 +233,10 @@ impl State { fn num_unparked(self) -> usize { (self.0 & UNPARK_MASK) >> UNPARK_SHIFT } + + fn any_lifo(self) -> bool { + self.0 & ANY_LIFO == ANY_LIFO + } } impl From for State { @@ -244,3 +269,11 @@ fn test_state() { assert_eq!(10, state.num_unparked()); assert_eq!(0, state.num_searching()); } + +#[test] +fn masks() { + println!("UNPARK_SHIFT = {UNPARK_SHIFT}"); + println!("UNPARK_MASK = {UNPARK_MASK:064b}"); + println!("SEARCH_MASK = {SEARCH_MASK:064b}"); + println!("ANY_LIFO = {ANY_LIFO:064b}"); +} diff --git a/tokio/src/runtime/scheduler/multi_thread/queue.rs b/tokio/src/runtime/scheduler/multi_thread/queue.rs index 7836183a4..7964e2ee3 100644 --- a/tokio/src/runtime/scheduler/multi_thread/queue.rs +++ b/tokio/src/runtime/scheduler/multi_thread/queue.rs @@ -440,6 +440,10 @@ impl Steal { self.len() == 0 } + pub(crate) fn has_lifo(&self) -> bool { + self.0.lifo.is_some() + } + /// Steals half the tasks from self and place them into `dst`. pub(crate) fn steal_into( &self, diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index 29a4a2349..80f10f05e 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -256,6 +256,15 @@ type Notified = task::Notified>; /// improvements. const MAX_LIFO_POLLS_PER_TICK: usize = 3; +/// Maximum time a parked worker will sleep before waking to check +/// for tasks stranded in other workers' LIFO slots. +const LIFO_EXCLUSIVITY_TIMEOUT: Duration = Duration::from_millis(100); + +enum TransitionToParked { + No, + Yes(Option), +} + #[allow(clippy::too_many_arguments)] pub(super) fn create( size: usize, @@ -804,13 +813,13 @@ impl Context { f(); } - if core.transition_to_parked(&self.worker) { + if let TransitionToParked::Yes(timeout) = core.transition_to_parked(&self.worker) { while !core.is_shutdown && !core.is_traced { core.stats.about_to_park(); core.stats .submit(&self.worker.handle.shared.worker_metrics[self.worker.index]); - core = self.park_internal(core, None); + core = self.park_internal(core, timeout); core.stats.unparked(); @@ -820,6 +829,10 @@ impl Context { if core.transition_from_parked(&self.worker) { break; } + + if self.steal_stranded_lifo(&mut core) { + break; + } } } @@ -829,6 +842,26 @@ impl Context { core } + fn steal_stranded_lifo(&self, _core: &mut Core) -> bool { + if !self.worker.handle.shared.idle.should_attempt_lifo_steal() { + return false; + } + + let remotes = self.worker.handle.shared.remotes.iter().enumerate(); + for (i, remote) in remotes { + if i != self.worker.index && remote.steal.has_lifo() { + self.worker + .handle + .shared + .idle + .unpark_worker_by_id(&self.worker.handle.shared, self.worker.index); + return true; + } + } + + false + } + fn park_yield(&self, core: Box) -> Box { self.park_internal(core, Some(Duration::from_millis(0))) } @@ -1193,16 +1226,19 @@ impl Core { /// Prepares the worker state for parking. /// /// Returns true if the transition happened, false if there is work to do first. - fn transition_to_parked(&mut self, worker: &Worker) -> bool { + fn transition_to_parked(&mut self, worker: &Worker) -> TransitionToParked { // Workers should not park if they have work to do if self.has_tasks() || self.is_traced { - return false; + return TransitionToParked::No; } // When the final worker transitions **out** of searching to parked, it // must check all the queues one last time in case work materialized // between the last work scan and transitioning out of searching. - let is_last_searcher = worker.handle.shared.idle.transition_worker_to_parked( + let idle::TransitionToParked { + is_last_searcher, + any_lifo, + } = worker.handle.shared.idle.transition_worker_to_parked( &worker.handle.shared, worker.index, self.is_searching, @@ -1216,7 +1252,11 @@ impl Core { worker.handle.notify_if_work_pending(); } - true + TransitionToParked::Yes(if any_lifo { + Some(LIFO_EXCLUSIVITY_TIMEOUT) + } else { + None + }) } /// Returns `true` if the transition happened. @@ -1362,7 +1402,7 @@ impl Handle { .push_back_or_overflow(prev, self, &mut core.stats); true } else { - self.shared.idle.all_parked() + !self.shared.idle.put_lifo() } };