foundation_core/synca/
waitgroup.rs

1//! WaitGroup - Thread completion tracking primitive.
2//!
3//! WHY: Provides a counting barrier for tracking N outstanding work items.
4//!
5//! WHAT: `WaitGroup` blocks until count reaches 0. `WaitGroupGuard` is a
6//! RAII guard that calls `done()` on drop, ensuring cleanup even on panic.
7//!
8//! HOW: A `Mutex<usize>` holds the count and a `Condvar` parks waiters. The
9//! count is read and the thread parks ATOMICALLY under the same lock, so a
10//! `done()` that drives the count to 0 between a waiter's check and its park
11//! cannot be lost (no lost-wakeup). An earlier implementation used an atomic
12//! count plus a separate `LockSignal`, which had exactly that race: the
13//! signal could fire in the gap between the count check and the park, and
14//! `lock_and_wait()`'s internal `try_lock()` then clobbered the signalled
15//! state — leaving the waiter parked forever. Surfaced as intermittent
16//! "test running for over 60 seconds" hangs under contention.
17
18use std::sync::{Arc, Condvar, Mutex};
19
20/// Tracks N outstanding work items. `wait()` blocks until count reaches 0.
21///
22/// Each worker thread gets a `WaitGroupGuard` that calls `done()` on drop,
23/// ensuring cleanup even if the thread panics.
24///
25/// # Examples
26///
27/// ```
28/// use foundation_core::synca::WaitGroup;
29/// use std::sync::Arc;
30///
31/// let wg = WaitGroup::new();
32/// wg.add(3); // 3 workers
33///
34/// for i in 0..3 {
35///     let wg = wg.clone();
36///     std::thread::spawn(move || {
37///         let _guard = wg.guard(); // done() called on drop
38///         // do work
39///     });
40/// }
41///
42/// wg.wait(); // block until all workers done
43/// ```
44#[derive(Clone)]
45pub struct WaitGroup {
46    /// (count, condvar). The count lives INSIDE the mutex so that reading it
47    /// and parking on the condvar happen atomically under the same lock.
48    inner: Arc<(Mutex<usize>, Condvar)>,
49}
50
51impl WaitGroup {
52    /// Create a new `WaitGroup` with count 0.
53    #[must_use]
54    pub fn new() -> Self {
55        Self {
56            inner: Arc::new((Mutex::new(0), Condvar::new())),
57        }
58    }
59
60    /// Increment the counter by n.
61    pub fn add(&self, n: usize) {
62        let (lock, _) = &*self.inner;
63        let mut count = lock.lock().unwrap_or_else(|p| p.into_inner());
64        *count += n;
65    }
66
67    /// Decrement the counter. If it reaches 0, wake all waiters.
68    pub fn done(&self) {
69        let (lock, cond) = &*self.inner;
70        let mut count = lock.lock().unwrap_or_else(|p| p.into_inner());
71        debug_assert!(*count > 0, "WaitGroup::done() called more times than add()");
72        *count -= 1;
73        if *count == 0 {
74            // Notify under the lock — a waiter checking the predicate cannot be
75            // between its check and its park, because both happen under this lock.
76            cond.notify_all();
77        }
78    }
79
80    /// Block until count reaches 0.
81    pub fn wait(&self) {
82        let (lock, cond) = &*self.inner;
83        let mut count = lock.lock().unwrap_or_else(|p| p.into_inner());
84        while *count != 0 {
85            count = cond.wait(count).unwrap_or_else(|p| p.into_inner());
86        }
87    }
88
89    /// Block until count reaches 0 or `timeout` elapses.
90    ///
91    /// Returns `true` if the count reached 0 within the deadline, `false` if it
92    /// timed out with work still outstanding. Used for bounded drains (e.g. a
93    /// graceful-shutdown grace window) where the caller force-closes whatever
94    /// remains past the deadline. The deadline is tracked absolutely so
95    /// spurious condvar wakeups shorten the remaining wait rather than resetting
96    /// it.
97    #[must_use]
98    pub fn wait_timeout(&self, timeout: std::time::Duration) -> bool {
99        let (lock, cond) = &*self.inner;
100        let deadline = std::time::Instant::now() + timeout;
101        let mut count = lock.lock().unwrap_or_else(|p| p.into_inner());
102        while *count != 0 {
103            let Some(remaining) = deadline.checked_duration_since(std::time::Instant::now()) else {
104                return false;
105            };
106            let (next, res) = cond
107                .wait_timeout(count, remaining)
108                .unwrap_or_else(|p| p.into_inner());
109            count = next;
110            if res.timed_out() && *count != 0 {
111                return false;
112            }
113        }
114        true
115    }
116
117    /// Create a RAII guard that calls `done()` on drop.
118    #[must_use]
119    pub fn guard(&self) -> WaitGroupGuard {
120        WaitGroupGuard(self.clone())
121    }
122}
123
124impl Default for WaitGroup {
125    fn default() -> Self {
126        Self::new()
127    }
128}
129
130/// Calls `WaitGroup::done()` on drop — ensures threads that panic still decrement.
131pub struct WaitGroupGuard(WaitGroup);
132
133impl Drop for WaitGroupGuard {
134    fn drop(&mut self) {
135        self.0.done();
136    }
137}
foundation_core/synca/waitgroup.rs

foundation_core/synca/
waitgroup.rs
