A basic Rust implementation of Golang WaitGroup (mutex)

wait_group.rs

use std::sync::{Condvar, Mutex, Arc};

#[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: Mutex<usize>,
    condvar: Condvar,
}

impl WaitGroupInner {

    pub fn new(count: usize) -> Self {
        Self { 
            counter: Mutex::new(count),
            condvar: Condvar::new() 
        }
    }

    pub fn add(&self, count: usize) {
        let mut counter = self.counter.lock().unwrap();
        *counter += count;
    }

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

    pub fn wait(&self) {
        let mut counter = self.counter.lock().unwrap();
        while *counter > 0 {
            counter = self.condvar.wait(counter).unwrap();
        }
    }

    fn sub(&self) {
        let mut counter = self.counter.lock().unwrap();
        *counter -= 1;
        if *counter == 0 {
            self.condvar.notify_one();
        }
    }
}


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();
    }
}