diff --git a/src/pool/spawn.rs b/src/pool/spawn.rs index e090cfc..aa65090 100644 --- a/src/pool/spawn.rs +++ b/src/pool/spawn.rs @@ -251,8 +251,8 @@ impl Local { &self.core } - pub(crate) fn pop(&mut self) -> Option> { - self.local_queue.pop() + pub(crate) fn pop(&mut self, steal_workers: bool) -> Option> { + self.local_queue.pop(steal_workers) } /// Pops a task from the queue. @@ -271,7 +271,8 @@ impl Local { 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() }, || {}, diff --git a/src/pool/worker.rs b/src/pool/worker.rs index 514fb09..a2b20d2 100644 --- a/src/pool/worker.rs +++ b/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() {} } } diff --git a/src/queue.rs b/src/queue.rs index 2109e2c..b5a6e54 100644 --- a/src/queue.rs +++ b/src/queue.rs @@ -93,10 +93,10 @@ impl LocalQueue { /// Gets a task cell from the queue. Returns `None` if there is no task cell /// available. - pub fn pop(&mut self) -> Option> { + pub fn pop(&mut self, steal_workers: bool) -> Option> { 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), } } diff --git a/src/queue/multilevel.rs b/src/queue/multilevel.rs index ae7760f..1319819 100644 --- a/src/queue/multilevel.rs +++ b/src/queue/multilevel.rs @@ -97,7 +97,7 @@ where self.local_queue.push(task_cell); } - pub(super) fn pop(&mut self) -> Option> { + pub(super) fn pop(&mut self, steal_workers: bool) -> Option> { fn into_pop(mut t: T, from_local: bool) -> Pop 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!( diff --git a/src/queue/single_level.rs b/src/queue/single_level.rs index 2de2881..3adbc7e 100644 --- a/src/queue/single_level.rs +++ b/src/queue/single_level.rs @@ -57,7 +57,7 @@ where self.local_queue.push(task_cell); } - pub fn pop(&mut self) -> Option> { + pub fn pop(&mut self, steal_workers: bool) -> Option> { fn into_pop(mut t: T, from_local: bool) -> Pop 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); } }) diff --git a/src/task/callback.rs b/src/task/callback.rs index 5186d91..58c1ef5 100644 --- a/src/task/callback.rs +++ b/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()); } } diff --git a/src/task/future.rs b/src/task/future.rs index 7596b80..3d29a89 100644 --- a/src/task/future.rs +++ b/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); }