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

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

@@ -437,10 +437,15 @@ impl<T> Steal<T> {
     }
 
     /// 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<T>,
         dst_stats: &mut Stats,
+        steal_lifo: bool,
     ) -> Option<task::Notified<T>> {
         // Safety: the caller is the only thread that mutates `dst.tail` and
         // holds a mutable reference.
@@ -462,14 +467,18 @@ impl<T> Steal<T> {
         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);

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()
     }

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

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<NoopSchedule>) -> 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();
 

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