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; }