compact_waitgroup/
wait_group.rs

1use core::{
2    pin::Pin,
3    task::{Context, Poll},
4};
5
6use crate::{
7    core_impl::{WaitGroupUtil, WaitGroupWrapper},
8    state::SharedWgInner,
9    twin_ref::{ClonableTwinRef, TwinRef},
10};
11
12#[cfg(feature = "compact-mono")]
13type MonoInner = crate::state::MonoWgInner;
14#[cfg(not(feature = "compact-mono"))]
15type MonoInner = crate::state::SharedWgInner;
16
17/// WaitGroup with clonable worker handles.
18///
19/// # Cancellation safety
20///
21/// This future is cancellation safe.
22///
23/// It is also safe to poll again after completion.
24///
25/// ```rust
26/// # use compact_waitgroup::WaitGroup;
27/// # futures_executor::block_on(async {
28/// let (wg, handle) = WaitGroup::new();
29/// let mut wg = core::pin::pin!(wg);
30///
31/// assert!(!wg.is_done());
32///
33/// handle.done();
34///
35/// wg.as_mut().await;
36/// assert!(wg.is_done());
37///
38/// // It is safe to await again (re-poll)
39/// wg.as_mut().await;
40/// assert!(wg.is_done());
41/// # });
42/// ```
43#[must_use]
44#[derive(Debug)]
45pub struct WaitGroup(WaitGroupWrapper<TwinRef<SharedWgInner>>);
46
47/// WaitGroup with a single non-clonable worker handle.
48///
49/// # Cancellation safety
50///
51/// This future is cancellation safe.
52///
53/// It is also safe to poll again after completion.
54///
55/// ```rust
56/// # use compact_waitgroup::MonoWaitGroup;
57/// # futures_executor::block_on(async {
58/// let (wg, handle) = MonoWaitGroup::new();
59/// let mut wg = core::pin::pin!(wg);
60///
61/// assert!(!wg.is_done());
62///
63/// handle.done();
64///
65/// wg.as_mut().await;
66/// assert!(wg.is_done());
67///
68/// // It is safe to await again (re-poll)
69/// wg.as_mut().await;
70/// assert!(wg.is_done());
71/// # });
72/// ```
73#[must_use]
74#[derive(Debug)]
75pub struct MonoWaitGroup(WaitGroupWrapper<TwinRef<MonoInner>>);
76
77/// Clonable worker handle.
78#[must_use]
79#[derive(Clone, Debug)]
80pub struct WorkerHandle {
81    _handle: ClonableTwinRef<SharedWgInner>,
82}
83
84/// Non-clonable worker handle.
85#[must_use]
86#[derive(Debug)]
87pub struct MonoWorkerHandle(TwinRef<MonoInner>);
88
89impl WaitGroup {
90    /// Creates a new `WaitGroup` and a clonable `WorkerHandle`.
91    ///
92    /// The `WaitGroup` is used to await the completion of tasks. The
93    /// `WorkerHandle` is used to signal task completion.
94    ///
95    /// # Examples
96    ///
97    /// ```
98    /// use compact_waitgroup::WaitGroup;
99    ///
100    /// let (wg, handle) = WaitGroup::new();
101    /// // ... distribute handle ...
102    /// ```
103    pub fn new() -> (Self, WorkerHandle) {
104        let inner = SharedWgInner::new();
105        let (wg, handle) = TwinRef::new_clonable(inner);
106        (
107            Self(WaitGroupWrapper::new(wg)),
108            WorkerHandle { _handle: handle },
109        )
110    }
111
112    /// Checks if the `WaitGroup` has completed.
113    ///
114    /// This returns `true` if all `WorkerHandle`s have been dropped.
115    ///
116    /// # Examples
117    ///
118    /// ```
119    /// use compact_waitgroup::WaitGroup;
120    ///
121    /// let (wg, handle) = WaitGroup::new();
122    /// assert!(!wg.is_done());
123    ///
124    /// drop(handle);
125    /// assert!(wg.is_done());
126    /// ```
127    #[inline]
128    pub fn is_done(&self) -> bool {
129        self.0.is_done()
130    }
131}
132
133impl MonoWaitGroup {
134    /// Creates a new `MonoWaitGroup` and a single `MonoWorkerHandle`.
135    ///
136    /// This variant is optimized for scenarios where there is exactly one
137    /// worker task. The handle cannot be cloned.
138    ///
139    /// # Examples
140    ///
141    /// ```
142    /// use compact_waitgroup::MonoWaitGroup;
143    ///
144    /// let (wg, handle) = MonoWaitGroup::new();
145    /// ```
146    pub fn new() -> (Self, MonoWorkerHandle) {
147        let inner = MonoInner::new();
148        let (wg, handle) = TwinRef::new_mono(inner);
149        (Self(WaitGroupWrapper::new(wg)), MonoWorkerHandle(handle))
150    }
151
152    /// Checks if the `MonoWaitGroup` has completed.
153    ///
154    /// This returns `true` if the `MonoWorkerHandle` has been dropped.
155    ///
156    /// # Examples
157    ///
158    /// ```
159    /// use compact_waitgroup::MonoWaitGroup;
160    ///
161    /// let (wg, handle) = MonoWaitGroup::new();
162    /// assert!(!wg.is_done());
163    ///
164    /// drop(handle);
165    /// assert!(wg.is_done());
166    /// ```
167    #[inline]
168    pub fn is_done(&self) -> bool {
169        self.0.is_done()
170    }
171}
172
173impl Future for WaitGroup {
174    type Output = ();
175
176    #[inline]
177    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
178        Pin::new(&mut self.0).poll(cx)
179    }
180}
181
182impl Future for MonoWaitGroup {
183    type Output = ();
184
185    #[inline]
186    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
187        Pin::new(&mut self.0).poll(cx)
188    }
189}
190
191#[cfg(feature = "futures-core")]
192impl futures_core::FusedFuture for WaitGroup {
193    #[inline]
194    fn is_terminated(&self) -> bool {
195        self.0.is_terminated()
196    }
197}
198
199#[cfg(feature = "futures-core")]
200impl futures_core::FusedFuture for MonoWaitGroup {
201    #[inline]
202    fn is_terminated(&self) -> bool {
203        self.0.is_terminated()
204    }
205}
206
207impl WorkerHandle {
208    /// Consumes the handle.
209    ///
210    /// This is equivalent to dropping the handle.
211    #[inline]
212    pub fn done(self) {
213        drop(self);
214    }
215}
216
217impl MonoWorkerHandle {
218    /// Consumes the handle.
219    ///
220    /// This is equivalent to dropping the handle.
221    #[inline]
222    pub fn done(self) {
223        drop(self);
224    }
225}
226
227impl Drop for MonoWorkerHandle {
228    #[inline]
229    fn drop(&mut self) {
230        unsafe {
231            self.0.send_done();
232        }
233    }
234}
compact_waitgroup/wait_group.rs

compact_waitgroup/
wait_group.rs
