feat: replace channel with semaphore for better performance

crates/noworkers/examples/parallel_processing.rs

@@ -62,7 +62,7 @@ async fn example_parallel_map(dataset: Vec<i32>) -> anyhow::Result<()> {
 
                 // Update progress
                 let count = processed.fetch_add(1, Ordering::SeqCst) + 1;
-                if count % 10 == 0 {
+                if count.is_multiple_of(10) {
                     println!("  Processed {}/{} items", count, total);
                 }
 
@@ -127,9 +127,7 @@ async fn example_pipeline(dataset: Vec<i32>) -> anyhow::Result<()> {
         stage2
             .add(move |_| async move {
                 // Filter even numbers and sum
-                let filtered_sum: i32 = chunk.iter()
-                    .filter(|&&x| x % 2 == 0)
-                    .sum();
+                let filtered_sum: i32 = chunk.iter().filter(|&&x| x % 2 == 0).sum();
 
                 tokio::time::sleep(Duration::from_millis(20)).await;
 
@@ -161,7 +159,11 @@ async fn example_batch_processing(dataset: Vec<i32>) -> anyhow::Result<()> {
         .map(|chunk| chunk.to_vec())
         .collect();
 
-    println!("  Processing {} batches of {} items each", batches.len(), batch_size);
+    println!(
+        "  Processing {} batches of {} items each",
+        batches.len(),
+        batch_size
+    );
 
     let mut workers = Workers::new();
     workers.with_limit(3); // Process 3 batches concurrently
@@ -206,7 +208,11 @@ async fn example_batch_processing(dataset: Vec<i32>) -> anyhow::Result<()> {
     let results = results.lock().await;
     let total_processed: usize = results.iter().sum();
 
-    println!("  Processed {} items in {}ms\n", total_processed, start.elapsed().as_millis());
+    println!(
+        "  Processed {} items in {}ms\n",
+        total_processed,
+        start.elapsed().as_millis()
+    );
 
     Ok(())
 }

crates/noworkers/examples/wait_lock.rs

@@ -0,0 +1,62 @@
+#![feature(random)]
+use std::{random, sync::atomic::AtomicUsize};
+
+use tokio_util::sync::CancellationToken;
+
+static SLOW_IN_PROGRESS: AtomicUsize = AtomicUsize::new(0);
+static FAST_IN_PROGRESS: AtomicUsize = AtomicUsize::new(0);
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    tokio::spawn(async move {
+        loop {
+            println!(
+                "slow: {}, fast: {}",
+                SLOW_IN_PROGRESS.load(std::sync::atomic::Ordering::Relaxed),
+                FAST_IN_PROGRESS.load(std::sync::atomic::Ordering::Relaxed)
+            );
+            tokio::time::sleep(std::time::Duration::from_millis(500)).await;
+        }
+    });
+    let mut workers = noworkers::Workers::new();
+
+    workers.with_limit(30);
+
+    for _ in 0..1000 {
+        let range: u16 = random::random(..);
+        if range < (u16::MAX / 4) {
+            workers.add(slow).await?;
+            continue;
+        }
+
+        workers.add(fast).await?;
+    }
+
+    workers.wait().await?;
+
+    Ok(())
+}
+
+async fn fast(_cancel: CancellationToken) -> anyhow::Result<()> {
+    FAST_IN_PROGRESS.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+    // println!("{}: running fast", now());
+    tokio::time::sleep(std::time::Duration::from_millis(200)).await;
+
+    FAST_IN_PROGRESS.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
+    Ok(())
+}
+async fn slow(_cancel: CancellationToken) -> anyhow::Result<()> {
+    SLOW_IN_PROGRESS.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+    // println!("{}: running slow", now());
+    tokio::time::sleep(std::time::Duration::from_secs(3)).await;
+    // println!("{}: completed slow", now());
+    SLOW_IN_PROGRESS.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
+    Ok(())
+}
+
+fn now() -> u128 {
+    std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .unwrap()
+        .as_millis()
+}

crates/noworkers/src/lib.rs

@@ -191,7 +191,10 @@
 
 use std::{future::Future, sync::Arc};
 
-use tokio::{sync::Mutex, task::JoinHandle};
+use tokio::{
+    sync::{Mutex, OwnedSemaphorePermit, Semaphore},
+    task::JoinHandle,
+};
 use tokio_util::sync::CancellationToken;
 
 /// Extension traits for common patterns.
@@ -322,8 +325,7 @@ enum WorkerLimit {
     #[default]
     NoLimit,
     Amount {
-        queue: tokio::sync::mpsc::Sender<()>,
-        done: Arc<Mutex<tokio::sync::mpsc::Receiver<()>>>,
+        done: Arc<tokio::sync::Semaphore>,
     },
 }
 
@@ -331,18 +333,21 @@ impl WorkerLimit {
     pub async fn queue_worker(&self) -> WorkerGuard {
         match self {
             WorkerLimit::NoLimit => {}
-            WorkerLimit::Amount { queue, .. } => {
+            WorkerLimit::Amount { done } => {
                 // Queue work, if the channel is limited, we will block until there is enough room
-                queue
-                    .send(())
+                let permit = done
+                    .clone()
+                    .acquire_owned()
                     .await
-                    .expect("tried to queue work on a closed worker channel");
+                    .expect("to be able to acquire permit");
+
+                return WorkerGuard {
+                    _permit: Some(permit),
+                };
             }
         }
 
-        WorkerGuard {
-            limit: self.clone(),
-        }
+        WorkerGuard { _permit: None }
     }
 }
 
@@ -353,24 +358,7 @@ impl WorkerLimit {
 ///
 /// This type is not directly constructible by users.
 pub struct WorkerGuard {
-    limit: WorkerLimit,
-}
-
-impl Drop for WorkerGuard {
-    fn drop(&mut self) {
-        match &self.limit {
-            WorkerLimit::NoLimit => { /* no limit on dequeue */ }
-            WorkerLimit::Amount { done, .. } => {
-                let done = done.clone();
-                tokio::spawn(async move {
-                    let mut done = done.lock().await;
-
-                    // dequeue an item, leave room for the next
-                    done.recv().await
-                });
-            }
-        }
-    }
+    _permit: Option<OwnedSemaphorePermit>,
 }
 
 impl Workers {
@@ -542,11 +530,8 @@ impl Workers {
     /// # }
     /// ```
     pub fn with_limit(&mut self, limit: usize) -> &mut Self {
-        let (tx, rx) = tokio::sync::mpsc::channel(limit);
-
         self.limit = WorkerLimit::Amount {
-            queue: tx,
-            done: Arc::new(Mutex::new(rx)),
+            done: Arc::new(Semaphore::new(limit)),
         };
         self
     }