From e829430d892b3c494771132ad6673221c926581c Mon Sep 17 00:00:00 2001 From: Zhengchao An Date: Thu, 9 Jul 2026 00:31:06 +0800 Subject: [PATCH] fix(concurrency): rewrite Workers on tokio Semaphore to fix lost wakeups (#4498) fix(concurrency): rewrite Workers on tokio Semaphore to fix lost wakeups (backlog#1023) --- crates/concurrency/src/workers.rs | 173 +++++++++++++++++++----------- 1 file changed, 111 insertions(+), 62 deletions(-) diff --git a/crates/concurrency/src/workers.rs b/crates/concurrency/src/workers.rs index e0eb47913..34bfe466e 100644 --- a/crates/concurrency/src/workers.rs +++ b/crates/concurrency/src/workers.rs @@ -15,14 +15,19 @@ //! Worker slot limiter used by long-running background workflows. use std::sync::Arc; -use tokio::sync::{Mutex, Notify}; +use tokio::sync::{Semaphore, watch}; use tracing::{debug, trace}; /// Cooperative worker-slot controller for async tasks. +/// +/// Slot acquisition is backed by a fair FIFO [`Semaphore`], while drain +/// waiters observe the in-use count through a dedicated [`watch`] channel. +/// The two wakeup sources are deliberately separate so a drain waiter can +/// never consume a wakeup destined for a pending [`take`](Self::take). pub struct Workers { - available: Mutex, // Available working slots - notify: Notify, // Used to notify waiting tasks - limit: usize, // Maximum number of concurrent jobs + semaphore: Semaphore, // Available working slots + in_use: watch::Sender, // Slots currently held; drain waiters watch it + limit: usize, // Maximum number of concurrent jobs } impl Workers { @@ -31,86 +36,70 @@ impl Workers { if n == 0 { return Err("n must be > 0"); } + if n > Semaphore::MAX_PERMITS { + return Err("n exceeds the maximum supported number of slots"); + } Ok(Arc::new(Self { - available: Mutex::new(n), - notify: Notify::new(), + semaphore: Semaphore::new(n), + in_use: watch::Sender::new(0), limit: n, })) } /// Acquire a worker slot, waiting until one becomes available. pub async fn take(&self) { - loop { - let mut available = self.available.lock().await; - if *available == 0 { - trace!( - event = "worker_slot.acquire", - component = "concurrency", - subsystem = "workers", - state = "waiting", - available_slots = *available, - total_slots = self.limit, - "worker slot pending" - ); - drop(available); - self.notify.notified().await; - } else { - *available -= 1; - trace!( - event = "worker_slot.acquire", - component = "concurrency", - subsystem = "workers", - state = "granted", - available_slots = *available, - total_slots = self.limit, - permits_in_use = self.limit.saturating_sub(*available), - "worker slot updated" - ); - break; - } - } + let permit = match self.semaphore.acquire().await { + Ok(permit) => permit, + // The semaphore is owned by `self` and never closed. + Err(_) => unreachable!("worker semaphore is never closed"), + }; + permit.forget(); // returned manually via give() + self.in_use.send_modify(|held| *held += 1); + trace!( + event = "worker_slot.acquire", + component = "concurrency", + subsystem = "workers", + state = "granted", + available_slots = self.semaphore.available_permits(), + total_slots = self.limit, + permits_in_use = *self.in_use.borrow(), + "worker slot updated" + ); } /// Release a worker slot. + /// + /// A `give()` that is not paired with a prior [`take`](Self::take) is + /// clamped: it never inflates capacity beyond `limit`. pub async fn give(&self) { - let mut available = self.available.lock().await; - *available = (*available).saturating_add(1).min(self.limit); // avoid over-release beyond limit + let mut released = false; + self.in_use.send_modify(|held| { + if *held > 0 { + *held -= 1; + released = true; + } + }); + if released { + self.semaphore.add_permits(1); + } trace!( event = "worker_slot.release", component = "concurrency", subsystem = "workers", state = "released", - available_slots = *available, + available_slots = self.semaphore.available_permits(), total_slots = self.limit, - permits_in_use = self.limit.saturating_sub(*available), + permits_in_use = *self.in_use.borrow(), "worker slot updated" ); - self.notify.notify_one(); // Notify a waiting task } /// Wait until all worker slots are released. pub async fn wait(&self) { - loop { - { - let available = self.available.lock().await; - if *available == self.limit { - break; - } - trace!( - event = "worker_slot.wait", - component = "concurrency", - subsystem = "workers", - state = "waiting", - available_slots = *available, - total_slots = self.limit, - permits_in_use = self.limit.saturating_sub(*available), - "worker drain pending" - ); - } - // Wait until all slots are freed - self.notify.notified().await; - } + let mut rx = self.in_use.subscribe(); + // The sender is owned by `self`, so it outlives this borrow. + let _ = rx.wait_for(|&held| held == 0).await; debug!( event = "worker_slot.wait", component = "concurrency", @@ -125,7 +114,7 @@ impl Workers { /// Return the current number of available worker slots. pub async fn available(&self) -> usize { - *self.available.lock().await + self.limit.saturating_sub(*self.in_use.borrow()) } } @@ -133,7 +122,7 @@ impl Workers { mod tests { use super::*; use std::time::Duration; - use tokio::time::sleep; + use tokio::time::{sleep, timeout}; #[tokio::test] async fn test_workers() { @@ -168,4 +157,64 @@ mod tests { assert_eq!(workers.available().await, 2); } + + /// Regression for S08: with limit=1, a pending drain waiter must never + /// steal the wakeup destined for a pending taker. + #[tokio::test] + async fn test_drain_waiter_does_not_steal_take_wakeup() { + let workers = Workers::new(1).unwrap(); + workers.take().await; + + let taker = { + let workers = workers.clone(); + tokio::spawn(async move { workers.take().await }) + }; + let drainer = { + let workers = workers.clone(); + tokio::spawn(async move { workers.wait().await }) + }; + // Let both tasks block before releasing the slot. + sleep(Duration::from_millis(50)).await; + workers.give().await; + + timeout(Duration::from_secs(1), taker) + .await + .expect("taker must be woken by give()") + .unwrap(); + + workers.give().await; + timeout(Duration::from_secs(1), drainer) + .await + .expect("drain waiter must complete once all slots are free") + .unwrap(); + assert_eq!(workers.available().await, 1); + } + + /// Regression for S17: two rapid give() calls must wake two pending + /// takers; permits must not merge into a single wakeup. + #[tokio::test] + async fn test_rapid_gives_wake_all_pending_takers() { + let workers = Workers::new(2).unwrap(); + workers.take().await; + workers.take().await; + + let takers: Vec<_> = (0..2) + .map(|_| { + let workers = workers.clone(); + tokio::spawn(async move { workers.take().await }) + }) + .collect(); + // Let both takers queue up before any slot is released. + sleep(Duration::from_millis(50)).await; + workers.give().await; + workers.give().await; + + for taker in takers { + timeout(Duration::from_secs(1), taker) + .await + .expect("every pending taker must be woken") + .unwrap(); + } + assert_eq!(workers.available().await, 0); + } }