Skip to content

Instantly share code, notes, and snippets.

@evanxg852000

evanxg852000/wait_group.rs

Last active January 9, 2024 12:49
Download ZIP

Revisions

  1. evanxg852000 revised this gist Jan 9, 2024. No changes.
  2. evanxg852000 created this gist Jan 9, 2024.
    93 changes: 93 additions & 0 deletions wait_group.rs
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,93 @@
    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();
    }
    }