crossbeam_utils/sync/
wait_group.rs

1use crate::primitive::sync::atomic::{AtomicUsize, Ordering};
2use crate::primitive::sync::{Arc, Condvar, Mutex};
3use core::mem::ManuallyDrop;
4use std::fmt;
5
6/// Enables threads to synchronize the beginning or end of some computation.
7///
8/// # Wait groups vs barriers
9///
10/// `WaitGroup` is very similar to [`Barrier`], but there are a few differences:
11///
12/// * [`Barrier`] needs to know the number of threads at construction, while `WaitGroup` is cloned to
13///   register more threads.
14///
15/// * A [`Barrier`] can be reused even after all threads have synchronized, while a `WaitGroup`
16///   synchronizes threads only once.
17///
18/// * All threads wait for others to reach the [`Barrier`]. With `WaitGroup`, each thread can choose
19///   to either wait for other threads or to continue without blocking.
20///
21/// # Examples
22///
23/// ```
24/// use crossbeam_utils::sync::WaitGroup;
25/// use std::thread;
26///
27/// // Create a new wait group.
28/// let wg = WaitGroup::new();
29///
30/// for _ in 0..4 {
31///     // Create another reference to the wait group.
32///     let wg = wg.clone();
33///
34///     thread::spawn(move || {
35///         // Do some work.
36///
37///         // Drop the reference to the wait group.
38///         drop(wg);
39///     });
40/// }
41///
42/// // Block until all threads have finished their work.
43/// wg.wait();
44/// # std::thread::sleep(std::time::Duration::from_millis(500)); // wait for background threads closed: https://github.com/rust-lang/miri/issues/1371
45/// ```
46///
47/// [`Barrier`]: std::sync::Barrier
48pub struct WaitGroup {
49    inner: Arc<Inner>,
50}
51
52/// Inner state of a `WaitGroup`.
53struct Inner {
54    cvar: Condvar,
55    lock: Mutex<()>,
56    count: AtomicUsize,
57}
58
59impl Default for WaitGroup {
60    fn default() -> Self {
61        Self {
62            inner: Arc::new(Inner {
63                cvar: Condvar::new(),
64                lock: Mutex::new(()),
65                count: AtomicUsize::new(1),
66            }),
67        }
68    }
69}
70
71impl WaitGroup {
72    /// Creates a new wait group and returns the single reference to it.
73    ///
74    /// # Examples
75    ///
76    /// ```
77    /// use crossbeam_utils::sync::WaitGroup;
78    ///
79    /// let wg = WaitGroup::new();
80    /// ```
81    pub fn new() -> Self {
82        Self::default()
83    }
84
85    /// Drops this reference and waits until all other references are dropped.
86    ///
87    /// # Examples
88    ///
89    /// ```
90    /// use crossbeam_utils::sync::WaitGroup;
91    /// use std::thread;
92    ///
93    /// let wg = WaitGroup::new();
94    ///
95    /// thread::spawn({
96    ///     let wg = wg.clone();
97    ///     move || {
98    ///         // Block until both threads have reached `wait()`.
99    ///         wg.wait();
100    ///     }
101    /// });
102    ///
103    /// // Block until both threads have reached `wait()`.
104    /// wg.wait();
105    /// # std::thread::sleep(std::time::Duration::from_millis(500)); // wait for background threads closed: https://github.com/rust-lang/miri/issues/1371
106    /// ```
107    pub fn wait(self) {
108        // SAFETY: this is equivalent to let Self { inner } = self, without calling our Drop.
109        let inner = unsafe {
110            let slf = ManuallyDrop::new(self);
111            core::ptr::read(&slf.inner)
112        };
113
114        if inner.count.fetch_sub(1, Ordering::AcqRel) == 1 {
115            // Acquire lock after updating count, see below.
116            drop(inner.lock.lock().unwrap());
117            inner.cvar.notify_all();
118            return;
119        }
120
121        // We check the counter while holding the lock, and notifiers acquire
122        // the lock between updating the counter and notifying, ensuring we
123        // can not miss the notification.
124        let mut guard = inner.lock.lock().unwrap();
125        while inner.count.load(Ordering::Acquire) != 0 {
126            guard = inner.cvar.wait(guard).unwrap();
127        }
128    }
129}
130
131impl Drop for WaitGroup {
132    fn drop(&mut self) {
133        if self.inner.count.fetch_sub(1, Ordering::Release) == 1 {
134            // Acquire lock after updating count, see wait().
135            drop(self.inner.lock.lock().unwrap());
136            self.inner.cvar.notify_all();
137        }
138    }
139}
140
141impl Clone for WaitGroup {
142    fn clone(&self) -> Self {
143        self.inner.count.fetch_add(1, Ordering::Relaxed);
144        Self {
145            inner: self.inner.clone(),
146        }
147    }
148}
149
150impl fmt::Debug for WaitGroup {
151    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
152        let count = self.inner.count.load(Ordering::Relaxed);
153        f.debug_struct("WaitGroup").field("count", &count).finish()
154    }
155}
crossbeam_utils/sync/wait_group.rs

crossbeam_utils/sync/
wait_group.rs
