From d45fbde6f8e0c89674bb93a14adf74cbdf44a6d2 Mon Sep 17 00:00:00 2001 From: Eliza Weisman Date: Thu, 7 May 2026 14:00:40 -0700 Subject: [PATCH] rt: only steal tasks after trying every queue (#8120) As discussed in [this comment][1], this commit changes work-stealing behavior so that a searching worker will now only attempt to steal from LIFO slots if it has already attempted to steal tasks from *every* other worker's main queue and found nothing. This is accomplished by adding a flag in `Stealer::steal_into` to control whether or not we attempt to steal LIFO slot tasks, and changing the work-stealing loop to make multiple passes over the set of workers, and set the `steal_lifo` flag when we are on the final pass. This way, we still bias towards stealing from run queues even while trying to steal from LIFO slots, in case more tasks were notified since the previous time we checked. This behavior is loosely inspired by [what the Go runtime does][2], if you can make sense of their weird and bizarre single-letter naming scheme in which "G"s are goroutines (tasks), "P"s are worker threads, "runnext" is the LIFO slot, and "sudog"s are ... also involved in some way). [1]: https://github.com/tokio-rs/tokio/pull/8120#discussion_r3201495853 [2]: https://github.com/golang/go/blob/release-branch.go1.26/src/runtime/proc.go#L3828-L3895 --- .../runtime/scheduler/multi_thread/queue.rs | 23 ++++++--- .../runtime/scheduler/multi_thread/worker.rs | 48 +++++++++++++------ .../scheduler/multi_thread/worker/taskdump.rs | 4 +- .../runtime/tests/loom_multi_thread/queue.rs | 16 ++++--- tokio/src/runtime/tests/queue.rs | 6 +-- 5 files changed, 65 insertions(+), 32 deletions(-) diff --git a/tokio/src/runtime/scheduler/multi_thread/queue.rs b/tokio/src/runtime/scheduler/multi_thread/queue.rs index 6fef9e846..bfa33552b 100644 --- a/tokio/src/runtime/scheduler/multi_thread/queue.rs +++ b/tokio/src/runtime/scheduler/multi_thread/queue.rs @@ -437,10 +437,15 @@ impl Steal { } /// Steals half the tasks from self and place them into `dst`. + /// + /// If `steal_lifo` is set, this method will attempt to steal from the LIFO + /// slot if this worker's run queue is empty. Otherwise, it will only steal + /// from the main queue. pub(crate) fn steal_into( &self, dst: &mut Local, dst_stats: &mut Stats, + steal_lifo: bool, ) -> Option> { // Safety: the caller is the only thread that mutates `dst.tail` and // holds a mutable reference. @@ -462,14 +467,18 @@ impl Steal { let mut n = self.steal_into2(dst, dst_tail); if n == 0 { - // If no tasks were stolen, let's see if there's one in the LIFO - // slot. - let lifo = self.0.lifo.take(); - if lifo.is_some() { - dst_stats.incr_steal_count(1); - dst_stats.incr_steal_operations(); + if steal_lifo { + // If no tasks were stolen, let's see if there's one in the + // LIFO slot. + let lifo = self.0.lifo.take(); + if lifo.is_some() { + dst_stats.incr_steal_count(1); + dst_stats.incr_steal_operations(); + } + return lifo; + } else { + return None; } - return lifo; } dst_stats.incr_steal_count(n as u16); diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index c682b5206..e7712d3de 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -1135,27 +1135,45 @@ impl Core { return None; } + // How many passes over the set of worker threads shall we make before + // giving up? On the final pass we shall additionally attempt to steal + // tasks from each worker's LIFO slot should the run queue be empty. + // + // Two passes was chosen arbitrarily. The Go runtime makes four similar + // workstealing passes, but it also attempts to steal both timers and GC + // work as well as goroutines (tasks), so that's a bit different than + // our behavior. See: + // https://github.com/golang/go/blob/release-branch.go1.26/src/runtime/proc.go#L3828-L3895 + const PASSES: usize = 2; + let num = worker.handle.shared.remotes.len(); - // Start from a random worker - let start = self.rand.fastrand_n(num as u32) as usize; + for i in 1..=PASSES { + // If we are making our final pass over the other workers, we shall + // also attempt to steal a task from the LIFO slot if we were not + // able to steal from the main queue. + let steal_lifo = i == PASSES; - for i in 0..num { - let i = (start + i) % num; + // Start from a random worker + let start = self.rand.fastrand_n(num as u32) as usize; - // Don't steal from ourself! We know we don't have work. - if i == worker.index { - continue; - } + for i in 0..num { + let i = (start + i) % num; - let target = &worker.handle.shared.remotes[i]; - if let Some(task) = target - .steal - .steal_into(&mut self.run_queue, &mut self.stats) - { - return Some(task); + // Don't steal from ourself! We know we don't have work. + if i == worker.index { + continue; + } + + let target = &worker.handle.shared.remotes[i]; + if let Some(task) = + target + .steal + .steal_into(&mut self.run_queue, &mut self.stats, steal_lifo) + { + return Some(task); + } } } - // Fallback on checking the global queue worker.handle.next_remote_task() } diff --git a/tokio/src/runtime/scheduler/multi_thread/worker/taskdump.rs b/tokio/src/runtime/scheduler/multi_thread/worker/taskdump.rs index 312673034..02ba905d2 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker/taskdump.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker/taskdump.rs @@ -68,7 +68,9 @@ impl Shared { for remote in self.remotes.iter() { let steal = &remote.steal; while !steal.is_empty() { - if let Some(task) = steal.steal_into(&mut local, &mut stats) { + // Always pass the `steal_lifo` flag here to ensure we also + // steal from the LIFO slot if we have drained the main queue. + if let Some(task) = steal.steal_into(&mut local, &mut stats, true) { local.push_back([task].into_iter()); } } diff --git a/tokio/src/runtime/tests/loom_multi_thread/queue.rs b/tokio/src/runtime/tests/loom_multi_thread/queue.rs index 2098b548a..abf4fbfa9 100644 --- a/tokio/src/runtime/tests/loom_multi_thread/queue.rs +++ b/tokio/src/runtime/tests/loom_multi_thread/queue.rs @@ -21,7 +21,9 @@ fn basic() { let mut n = 0; for _ in 0..3 { - if steal.steal_into(&mut local, &mut stats).is_some() { + // steal_lifo = false here is fine since there are no LIFO slot + // tasks in this test. + if steal.steal_into(&mut local, &mut stats, false).is_some() { n += 1; } @@ -76,7 +78,7 @@ fn basic_lifo() { let mut n = 0; for _ in 0..3 { - if steal.steal_into(&mut local, &mut stats).is_some() { + if steal.steal_into(&mut local, &mut stats, true).is_some() { n += 1; } @@ -133,7 +135,7 @@ fn steal_overflow() { let (_, mut local) = queue::local(); let mut n = 0; - if steal.steal_into(&mut local, &mut stats).is_some() { + if steal.steal_into(&mut local, &mut stats, false).is_some() { n += 1; } @@ -256,7 +258,7 @@ fn steal_tasks(steal: queue::Steal) -> usize { let mut stats = new_stats(); let (_, mut local) = queue::local(); - if steal.steal_into(&mut local, &mut stats).is_none() { + if steal.steal_into(&mut local, &mut stats, true).is_none() { return 0; } @@ -290,7 +292,9 @@ fn chained_steal() { let th = thread::spawn(move || { let mut stats = new_stats(); let (_, mut local) = queue::local(); - s1.steal_into(&mut local, &mut stats); + // again, steal_lifo = false here is fine since there are no LIFO slot + // tasks in this test. + s1.steal_into(&mut local, &mut stats, false); while local.pop().is_some() {} }); @@ -298,7 +302,7 @@ fn chained_steal() { // Drain our tasks, then attempt to steal while l1.pop().is_some() {} - s2.steal_into(&mut l1, &mut stats); + s2.steal_into(&mut l1, &mut stats, false); th.join().unwrap(); diff --git a/tokio/src/runtime/tests/queue.rs b/tokio/src/runtime/tests/queue.rs index 9047f4ad7..75e538215 100644 --- a/tokio/src/runtime/tests/queue.rs +++ b/tokio/src/runtime/tests/queue.rs @@ -125,7 +125,7 @@ fn steal_batch() { local1.push_back_or_overflow(task, &inject, &mut stats); } - assert!(steal1.steal_into(&mut local2, &mut stats).is_some()); + assert!(steal1.steal_into(&mut local2, &mut stats, false).is_some()); cfg_unstable_metrics! { assert_metrics!(stats, steal_count == 2); @@ -172,7 +172,7 @@ fn stress1() { let mut n = 0; for _ in 0..NUM_STEAL { - if steal.steal_into(&mut local, &mut stats).is_some() { + if steal.steal_into(&mut local, &mut stats, false).is_some() { n += 1; } @@ -233,7 +233,7 @@ fn stress2() { let mut n = 0; for _ in 0..NUM_STEAL { - if steal.steal_into(&mut local, &mut stats).is_some() { + if steal.steal_into(&mut local, &mut stats, false).is_some() { n += 1; }