steal other workers only before going to sleep

src/pool/spawn.rs

@@ -251,8 +251,8 @@ impl<T: TaskCell + Send> Local<T> {
         &self.core
     }
 
-    pub(crate) fn pop(&mut self) -> Option<Pop<T>> {
-        self.local_queue.pop()
+    pub(crate) fn pop(&mut self, steal_workers: bool) -> Option<Pop<T>> {
+        self.local_queue.pop(steal_workers)
     }
 
     /// Pops a task from the queue.
@@ -271,7 +271,8 @@ impl<T: TaskCell + Send> Local<T> {
                     if !self.core.mark_sleep() {
                         return false;
                     }
-                    task = self.local_queue.pop();
+                    // Steal other workers so we don't sleep while leaving other workers busy.
+                    task = self.local_queue.pop(true);
                     task.is_none()
                 },
                 || {},

src/pool/worker.rs

@@ -25,7 +25,9 @@ where
         // Wait some time before going to sleep, which is more expensive.
         let mut spin = SpinWait::new();
         loop {
-            if let Some(t) = self.local.pop() {
+            // By default we don't steal other workers to reduce the cost of pop.
+            // But we will always try to steal other workers before going to sleep.
+            if let Some(t) = self.local.pop(false) {
                 return Some(t);
             }
             if !spin.spin() {
@@ -50,7 +52,7 @@ where
         self.runner.end(&mut self.local);
 
         // Drain all futures in the queue
-        while self.local.pop().is_some() {}
+        while self.local.pop(true).is_some() {}
     }
 }
 

src/queue.rs

@@ -93,10 +93,10 @@ impl<T: TaskCell + Send> LocalQueue<T> {
 
     /// Gets a task cell from the queue. Returns `None` if there is no task cell
     /// available.
-    pub fn pop(&mut self) -> Option<Pop<T>> {
+    pub fn pop(&mut self, steal_workers: bool) -> Option<Pop<T>> {
         match &mut self.0 {
-            LocalQueueInner::SingleLevel(q) => q.pop(),
-            LocalQueueInner::Multilevel(q) => q.pop(),
+            LocalQueueInner::SingleLevel(q) => q.pop(steal_workers),
+            LocalQueueInner::Multilevel(q) => q.pop(steal_workers),
         }
     }
 

src/queue/multilevel.rs

@@ -97,7 +97,7 @@ where
         self.local_queue.push(task_cell);
     }
 
-    pub(super) fn pop(&mut self) -> Option<Pop<T>> {
+    pub(super) fn pop(&mut self, steal_workers: bool) -> Option<Pop<T>> {
         fn into_pop<T>(mut t: T, from_local: bool) -> Pop<T>
         where
             T: TaskCell,
@@ -130,7 +130,7 @@ where
                 Steal::Retry => need_retry = true,
                 _ => {}
             }
-            if !self.stealers.is_empty() {
+            if !self.stealers.is_empty() && steal_workers {
                 let mut found = None;
                 for (idx, stealer) in self.stealers.iter().enumerate() {
                     match stealer.steal_batch_and_pop(&self.local_queue) {
@@ -732,7 +732,7 @@ mod tests {
         let (injector, mut locals) = builder.build(1);
         injector.push(MockTask::new(0, Extras::multilevel_default()));
         thread::sleep(SLEEP_DUR);
-        let schedule_time = locals[0].pop().unwrap().schedule_time;
+        let schedule_time = locals[0].pop(true).unwrap().schedule_time;
         assert!(schedule_time.elapsed() >= SLEEP_DUR);
     }
 
@@ -818,14 +818,49 @@ mod tests {
             injector.push(MockTask::new(i, Extras::multilevel_default()));
         }
         let sum: u64 = (0..100)
-            .map(|_| locals[2].pop().unwrap().task_cell.sleep_ms)
+            .map(|_| locals[2].pop(false).unwrap().task_cell.sleep_ms)
             .sum();
         assert_eq!(sum, (0..100).sum());
-        assert!(locals.iter_mut().all(|c| c.pop().is_none()));
+        assert!(locals.iter_mut().all(|c| c.pop(true).is_none()));
     }
 
     #[test]
-    fn test_pop_by_steal_others() {
+    fn test_pop_without_stealing_workers() {
+        let builder = Builder::new(Config::default());
+        let (injector, mut locals) = builder.build_raw(3);
+        for i in 0..50 {
+            injector.push(MockTask::new(i, Extras::multilevel_default()));
+        }
+        assert!(injector.level_injectors[0]
+            .steal_batch(&locals[0].local_queue)
+            .is_success());
+        for i in 50..100 {
+            injector.push(MockTask::new(i, Extras::multilevel_default()));
+        }
+        assert!(injector.level_injectors[0]
+            .steal_batch(&locals[1].local_queue)
+            .is_success());
+
+        let mut sum = 0;
+        while let Some(task) = locals[2].pop(false) {
+            sum += task.task_cell.sleep_ms;
+        }
+        assert_ne!(
+            sum,
+            (0..100).sum(),
+            "locals[2] shall not pop all tasks without stealing others"
+        );
+
+        for &i in &[0, 1] {
+            while let Some(task) = locals[i].pop(false) {
+                sum += task.task_cell.sleep_ms;
+            }
+        }
+        assert_eq!(sum, (0..100).sum());
+    }
+
+    #[test]
+    fn test_pop_by_steal_workers() {
         let builder = Builder::new(Config::default());
         let (injector, mut locals) = builder.build_raw(3);
         for i in 0..50 {
@@ -841,10 +876,10 @@ mod tests {
             .steal_batch(&locals[1].local_queue)
             .is_success());
         let sum: u64 = (0..100)
-            .map(|_| locals[2].pop().unwrap().task_cell.sleep_ms)
+            .map(|_| locals[2].pop(true).unwrap().task_cell.sleep_ms)
             .sum();
         assert_eq!(sum, (0..100).sum());
-        assert!(locals.iter_mut().all(|c| c.pop().is_none()));
+        assert!(locals.iter_mut().all(|c| c.pop(true).is_none()));
     }
 
     #[test]
@@ -860,7 +895,7 @@ mod tests {
             .map(|mut consumer| {
                 let sum = sum.clone();
                 thread::spawn(move || {
-                    while let Some(pop) = consumer.pop() {
+                    while let Some(pop) = consumer.pop(true) {
                         sum.fetch_add(pop.task_cell.sleep_ms, SeqCst);
                     }
                 })
@@ -881,7 +916,7 @@ mod tests {
         let mut runner = runner_builder.build();
 
         remote.spawn(MockTask::new(100, Extras::new_multilevel(1, None)));
-        if let Some(Pop { task_cell, .. }) = locals[0].pop() {
+        if let Some(Pop { task_cell, .. }) = locals[0].pop(true) {
             assert!(runner.handle(&mut locals[0], task_cell));
         }
         assert!(

src/queue/single_level.rs

@@ -57,7 +57,7 @@ where
         self.local_queue.push(task_cell);
     }
 
-    pub fn pop(&mut self) -> Option<Pop<T>> {
+    pub fn pop(&mut self, steal_workers: bool) -> Option<Pop<T>> {
         fn into_pop<T>(mut t: T, from_local: bool) -> Pop<T>
         where
             T: TaskCell,
@@ -81,7 +81,7 @@ where
                 Steal::Retry => need_retry = true,
                 _ => {}
             }
-            if !self.stealers.is_empty() {
+            if !self.stealers.is_empty() && steal_workers {
                 let mut found = None;
                 for (idx, stealer) in self.stealers.iter().enumerate() {
                     match stealer.steal_batch_and_pop(&self.local_queue) {
@@ -184,7 +184,7 @@ mod tests {
         let (injector, mut locals) = super::create(1);
         injector.push(MockCell::new(0));
         thread::sleep(SLEEP_DUR);
-        let schedule_time = locals[0].pop().unwrap().schedule_time;
+        let schedule_time = locals[0].pop(true).unwrap().schedule_time;
         assert!(schedule_time.elapsed() >= SLEEP_DUR);
     }
 
@@ -195,14 +195,44 @@ mod tests {
             injector.push(MockCell::new(i));
         }
         let sum: i32 = (0..100)
-            .map(|_| locals[2].pop().unwrap().task_cell.value)
+            .map(|_| locals[2].pop(false).unwrap().task_cell.value)
             .sum();
         assert_eq!(sum, (0..100).sum());
-        assert!(locals.iter_mut().all(|c| c.pop().is_none()));
+        assert!(locals.iter_mut().all(|c| c.pop(true).is_none()));
     }
 
     #[test]
-    fn test_pop_by_steal_others() {
+    fn test_pop_without_stealing_workers() {
+        let (injector, mut locals) = super::create(3);
+        for i in 0..50 {
+            injector.push(MockCell::new(i));
+        }
+        assert!(injector.0.steal_batch(&locals[0].local_queue).is_success());
+        for i in 50..100 {
+            injector.push(MockCell::new(i));
+        }
+        assert!(injector.0.steal_batch(&locals[1].local_queue).is_success());
+
+        let mut sum = 0;
+        while let Some(task) = locals[2].pop(false) {
+            sum += task.task_cell.value;
+        }
+        assert_ne!(
+            sum,
+            (0..100).sum(),
+            "locals[2] shall not pop all tasks without stealing others"
+        );
+
+        for &i in &[0, 1] {
+            while let Some(task) = locals[i].pop(false) {
+                sum += task.task_cell.value;
+            }
+        }
+        assert_eq!(sum, (0..100).sum());
+    }
+
+    #[test]
+    fn test_pop_by_steal_workers() {
         let (injector, mut locals) = super::create(3);
         for i in 0..50 {
             injector.push(MockCell::new(i));
@@ -213,10 +243,10 @@ mod tests {
         }
         assert!(injector.0.steal_batch(&locals[1].local_queue).is_success());
         let sum: i32 = (0..100)
-            .map(|_| locals[2].pop().unwrap().task_cell.value)
+            .map(|_| locals[2].pop(true).unwrap().task_cell.value)
             .sum();
         assert_eq!(sum, (0..100).sum());
-        assert!(locals.iter_mut().all(|c| c.pop().is_none()));
+        assert!(locals.iter_mut().all(|c| c.pop(true).is_none()));
     }
 
     #[test]
@@ -231,7 +261,7 @@ mod tests {
             .map(|mut consumer| {
                 let sum = sum.clone();
                 thread::spawn(move || {
-                    while let Some(pop) = consumer.pop() {
+                    while let Some(pop) = consumer.pop(true) {
                         sum.fetch_add(pop.task_cell.value, Ordering::SeqCst);
                     }
                 })

src/task/callback.rs

@@ -204,7 +204,7 @@ mod tests {
         );
         assert_eq!(rx.recv().unwrap(), 42);
         assert_eq!(rx.recv().unwrap(), 42);
-        assert!(locals[0].pop().is_none());
+        assert!(locals[0].pop(true).is_none());
         assert!(rx.recv().is_err());
     }
 
@@ -230,7 +230,7 @@ mod tests {
         );
         assert_eq!(rx.recv().unwrap(), 42);
         assert_eq!(rx.recv().unwrap(), 42);
-        assert!(locals[0].pop().is_some());
+        assert!(locals[0].pop(true).is_some());
         assert!(rx.recv().is_err());
     }
 }

src/task/future.rs

@@ -476,7 +476,7 @@ mod tests {
 
         /// Run `Runner::handle` once.
         fn handle_once(&mut self) {
-            if let Some(t) = self.locals[0].pop() {
+            if let Some(t) = self.locals[0].pop(true) {
                 let runner = self.runner.clone();
                 runner.borrow_mut().handle(&mut self.locals[0], t.task_cell);
             }