use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use event_listener::Event;
struct Inner {
count: AtomicUsize,
event: Event,
}
#[derive(Clone)]
pub struct WaitGroup {
inner: Arc<Inner>,
}
impl WaitGroup {
pub fn new() -> Self {
Self {
inner: Arc::new(Inner {
count: AtomicUsize::new(0),
event: Event::new(),
}),
}
}
pub fn add(&self, n: usize) {
self.inner.count.fetch_add(n, Ordering::AcqRel);
}
pub fn done(&self) {
let prev = self.inner.count.fetch_sub(1, Ordering::AcqRel);
if prev == 1 {
self.inner.event.notify(usize::MAX);
}
}
pub async fn wait(&self) {
loop {
if self.inner.count.load(Ordering::Acquire) == 0 {
return;
}
let listener = self.inner.event.listen();
if self.inner.count.load(Ordering::Acquire) == 0 {
return;
}
listener.await;
}
}
pub fn guard(&self) -> WaitGroupGuard {
WaitGroupGuard { wg: self.clone() }
}
}
impl Default for WaitGroup {
fn default() -> Self {
Self::new()
}
}
pub struct WaitGroupGuard {
wg: WaitGroup,
}
impl Drop for WaitGroupGuard {
fn drop(&mut self) {
self.wg.done();
}
}