add workstealng tests

Author: jdidion

Cargo.toml

@@ -40,7 +40,7 @@ name = "perf"
 harness = false
 
 [features]
-default = ["affinity", "batching", "retry"]
+default = []
 affinity = ["dep:core_affinity"]
 batching = []
 retry = []

src/hive/inner/queue/channel.rs

@@ -10,7 +10,7 @@ use std::sync::Arc;
 use std::time::Duration;
 
 // time to wait when polling the global queue
-const RECV_TIMEOUT: Duration = Duration::from_secs(1);
+const RECV_TIMEOUT: Duration = Duration::from_millis(100);
 
 /// Type alias for the input task channel sender
 type TaskSender<W> = crossbeam_channel::Sender<Task<W>>;

src/hive/inner/queue/workstealing.rs

@@ -14,6 +14,11 @@ use parking_lot::RwLock;
 use rand::prelude::*;
 use std::ops::Deref;
 use std::sync::Arc;
+use std::thread;
+use std::time::Duration;
+
+/// Time to wait after trying to pop and finding all queues empty.
+const EMPTY_DELAY: Duration = Duration::from_millis(100);
 
 pub struct WorkstealingTaskQueues<W: Worker> {
     global: Arc<GlobalQueue<W>>,
@@ -40,7 +45,7 @@ impl<W: Worker> TaskQueues<W> for WorkstealingTaskQueues<W> {
 
     fn update_for_threads(&self, start_index: usize, end_index: usize, config: &Config) {
         let local_queues = self.local.read();
-        assert!(local_queues.len() > end_index);
+        assert!(local_queues.len() >= end_index);
         (start_index..end_index).for_each(|thread_index| local_queues[thread_index].update(config));
     }
 
@@ -103,14 +108,22 @@ impl<W: Worker> GlobalQueue<W> {
     }
 
     /// Tries to steal a task from a random worker using its `Stealer`.
-    fn try_steal(&self) -> Option<Task<W>> {
+    fn try_steal_from_worker(&self) -> Result<Task<W>, PopTaskError> {
         let stealers = self.stealers.read();
         let n = stealers.len();
         // randomize the stealing order, to prevent always stealing from the same thread
         std::iter::from_fn(|| Some(rand::rng().random_range(0..n)))
             .take(n)
             .filter_map(|i| stealers[i].steal().success())
             .next()
+            .ok_or_else(|| {
+                if self.is_closed() && self.queue.is_empty() {
+                    PopTaskError::Closed
+                } else {
+                    thread::park_timeout(EMPTY_DELAY);
+                    PopTaskError::Empty
+                }
+            })
     }
 
     /// Tries to steal a task from the global queue, otherwise tries to steal a task from another
@@ -119,7 +132,7 @@ impl<W: Worker> GlobalQueue<W> {
         if let Some(task) = self.queue.steal().success() {
             Ok(task)
         } else {
-            self.try_steal().ok_or(PopTaskError::Empty)
+            self.try_steal_from_worker()
         }
     }
 
@@ -137,9 +150,10 @@ impl<W: Worker> GlobalQueue<W> {
             .steal_batch_with_limit_and_pop(local_batch, limit + 1)
             .success()
         {
-            return Ok(task);
+            Ok(task)
+        } else {
+            self.try_steal_from_worker()
         }
-        self.try_steal().ok_or(PopTaskError::Empty)
     }
 
     fn is_closed(&self) -> bool {

src/hive/inner/shared.rs

@@ -717,8 +717,8 @@ mod retry {
 
 #[cfg(test)]
 mod tests {
-    use crate::bee::stock::ThunkWorker;
     use crate::bee::DefaultQueen;
+    use crate::bee::stock::ThunkWorker;
     use crate::hive::ChannelTaskQueues;
 
     type VoidThunkWorker = ThunkWorker<()>;

src/hive/mod.rs

@@ -2025,6 +2025,7 @@ mod batching_tests {
         hive: &Hive<DefaultQueen<ThunkWorker<ThreadId>>, T>,
         num_threads: usize,
         batch_limit: usize,
+        assert_exact: bool,
     ) {
         let tasks_per_thread = batch_limit + 2;
         let (tx, rx) = crate::hive::outcome_channel();
@@ -2038,52 +2039,83 @@ mod batching_tests {
         hive.join();
         let thread_counts = count_thread_ids(rx, task_ids);
         assert_eq!(thread_counts.len(), num_threads);
-        assert!(
-            thread_counts
-                .values()
-                .all(|&count| count == tasks_per_thread)
+        assert_eq!(
+            thread_counts.values().sum::<usize>(),
+            tasks_per_thread * num_threads
         );
+        if assert_exact {
+            assert!(
+                thread_counts
+                    .values()
+                    .all(|&count| count == tasks_per_thread)
+            );
+        } else {
+            assert!(thread_counts.values().all(|&count| count > 0));
+        }
     }
 
     #[rstest]
-    fn test_batching<B, F>(#[values(channel_builder, workstealing_builder)] builder_factory: F)
-    where
-        B: TaskQueuesBuilder,
-        F: Fn(bool) -> B,
-    {
+    fn test_batching_channel() {
         const NUM_THREADS: usize = 4;
         const BATCH_LIMIT: usize = 24;
-        let hive = builder_factory(false)
+        let hive = channel_builder(false)
             .with_worker_default()
             .num_threads(NUM_THREADS)
             .batch_limit(BATCH_LIMIT)
             .build();
-        run_test(&hive, NUM_THREADS, BATCH_LIMIT);
+        run_test(&hive, NUM_THREADS, BATCH_LIMIT, true);
     }
 
     #[rstest]
-    fn test_set_batch_limit<B, F>(
-        #[values(channel_builder, workstealing_builder)] builder_factory: F,
-    ) where
-        B: TaskQueuesBuilder,
-        F: Fn(bool) -> B,
-    {
+    fn test_batching_workstealing() {
+        const NUM_THREADS: usize = 4;
+        const BATCH_LIMIT: usize = 24;
+        let hive = workstealing_builder(false)
+            .with_worker_default()
+            .num_threads(NUM_THREADS)
+            .batch_limit(BATCH_LIMIT)
+            .build();
+        run_test(&hive, NUM_THREADS, BATCH_LIMIT, false);
+    }
+
+    #[rstest]
+    fn test_set_batch_limit_channel() {
         const NUM_THREADS: usize = 4;
         const BATCH_LIMIT_0: usize = 10;
-        const BATCH_LIMIT_1: usize = 20;
-        const BATCH_LIMIT_2: usize = 50;
-        let hive = builder_factory(false)
+        const BATCH_LIMIT_1: usize = 50;
+        const BATCH_LIMIT_2: usize = 20;
+        let hive = channel_builder(false)
             .with_worker_default()
             .num_threads(NUM_THREADS)
             .batch_limit(BATCH_LIMIT_0)
             .build();
-        run_test(&hive, NUM_THREADS, BATCH_LIMIT_0);
+        run_test(&hive, NUM_THREADS, BATCH_LIMIT_0, true);
         // increase batch size
-        hive.set_worker_batch_limit(BATCH_LIMIT_2);
-        run_test(&hive, NUM_THREADS, BATCH_LIMIT_2);
+        hive.set_worker_batch_limit(BATCH_LIMIT_1);
+        run_test(&hive, NUM_THREADS, BATCH_LIMIT_1, true);
         // decrease batch size
+        hive.set_worker_batch_limit(BATCH_LIMIT_2);
+        run_test(&hive, NUM_THREADS, BATCH_LIMIT_2, true);
+    }
+
+    #[rstest]
+    fn test_set_batch_limit_workstealing() {
+        const NUM_THREADS: usize = 4;
+        const BATCH_LIMIT_0: usize = 10;
+        const BATCH_LIMIT_1: usize = 50;
+        const BATCH_LIMIT_2: usize = 20;
+        let hive = workstealing_builder(false)
+            .with_worker_default()
+            .num_threads(NUM_THREADS)
+            .batch_limit(BATCH_LIMIT_0)
+            .build();
+        run_test(&hive, NUM_THREADS, BATCH_LIMIT_0, false);
+        // increase batch size
         hive.set_worker_batch_limit(BATCH_LIMIT_1);
-        run_test(&hive, NUM_THREADS, BATCH_LIMIT_1);
+        run_test(&hive, NUM_THREADS, BATCH_LIMIT_1, false);
+        // decrease batch size
+        hive.set_worker_batch_limit(BATCH_LIMIT_2);
+        run_test(&hive, NUM_THREADS, BATCH_LIMIT_2, false);
     }
 
     #[rstest]