A basic Rust implementation of Golang WaitGroup (atomic & spin-loop)

wait_group_atomic.rs

use std::{sync::{ Arc, atomic::{AtomicUsize, Ordering}}, hint::spin_loop};

#[derive(Clone)]
pub struct WaitGroup(Arc<WaitGroupInner>);

impl WaitGroup {
    pub fn new(count: usize) -> Self {
        Self(Arc::new(WaitGroupInner::new(count)))
    }

    pub fn add(&self, count: usize) {
        self.0.add(count)
    }

    pub fn done(&self) -> WaitGroupFinalizer {
        self.0.done()
    }

    pub fn wait(&self) {
        self.0.wait()
    }

}

struct WaitGroupInner {
    counter: AtomicUsize,
}

impl WaitGroupInner {

    pub fn new(count: usize) -> Self {
        Self { 
            counter: AtomicUsize::new(count), 
        }
    }

    pub fn add(&self, count: usize) {
        self.counter.fetch_add(count, Ordering::Release);
    }

    pub fn done(&self) -> WaitGroupFinalizer {
        return WaitGroupFinalizer(self)
    }

    pub fn wait(&self) {
        while self.counter.load(Ordering::Acquire) > 0 {
            spin_loop();
        }
    }

    
    fn sub(&self) {
        self.counter.fetch_sub(1, Ordering::Release);
    }
}


pub struct WaitGroupFinalizer<'a>(&'a WaitGroupInner);

impl<'a> Drop for WaitGroupFinalizer<'_> {
    fn drop(&mut self) {
        self.0.sub()
    }
}


#[cfg(test)]
mod tests {

    use std::{thread, time::Duration};
    use crate::WaitGroup;

    #[test]
    fn test() {
        let wg = WaitGroup::new(10);
        for id in 0..10 {
            let wg =  wg.clone();
            thread::spawn(move || {
                let _finalizer = wg.done();
                println!("thread_id: {:?}", id);
                thread::sleep(Duration::from_secs(1));
            });
        }
        wg.wait();
    }
}