use std::sync::{Arc, Condvar, Mutex};
#[derive(Clone)]
pub(crate) struct WaitGroup {
inner: Arc<Inner>,
}
struct Inner {
var: Condvar,
count: Mutex<usize>,
}
impl Default for WaitGroup {
fn default() -> Self {
Self {
inner: Arc::new(Inner {
var: Condvar::new(),
count: Mutex::new(0),
}),
}
}
}
impl WaitGroup {
pub(crate) fn add(&self, n: usize) {
*self.inner.count.lock().unwrap() += n;
}
pub(crate) fn done(&self) {
let mut count = self.inner.count.lock().unwrap();
*count -= 1;
if *count == 0 {
self.inner.var.notify_all();
}
}
pub(crate) fn wait(self) {
let mut count = self.inner.count.lock().unwrap();
while *count > 0 {
count = self.inner.var.wait(count).unwrap();
}
}
}