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#[derive(Debug)]
44pub struct WaitGroup(WaitGroupWrapper<TwinRef<SharedWgInner>>);
45
46/// WaitGroup with a single non-clonable worker handle.
47///
48/// # Cancellation safety
49///
50/// This future is cancellation safe.
51///
52/// It is also safe to poll again after completion.
53///
54/// ```rust
55/// # use compact_waitgroup::MonoWaitGroup;
56/// # futures_executor::block_on(async {
57/// let (wg, handle) = MonoWaitGroup::new();
58/// let mut wg = core::pin::pin!(wg);
59///
60/// assert!(!wg.is_done());
61///
62/// handle.done();
63///
64/// wg.as_mut().await;
65/// assert!(wg.is_done());
66///
67/// // It is safe to await again (re-poll)
68/// wg.as_mut().await;
69/// assert!(wg.is_done());
70/// # });
71/// ```
72#[derive(Debug)]
73pub struct MonoWaitGroup(WaitGroupWrapper<TwinRef<MonoInner>>);
74
75/// Clonable worker handle.
76#[derive(Clone, Debug)]
77pub struct WorkerHandle {
78    _handle: ClonableTwinRef<SharedWgInner>,
79}
80
81/// Non-clonable worker handle.
82#[derive(Debug)]
83pub struct MonoWorkerHandle(TwinRef<MonoInner>);
84
85impl WaitGroup {
86    pub fn new() -> (Self, WorkerHandle) {
87        let inner = SharedWgInner::new();
88        let (wg, handle) = TwinRef::new_clonable(inner);
89        (
90            Self(WaitGroupWrapper::new(wg)),
91            WorkerHandle { _handle: handle },
92        )
93    }
94
95    #[inline]
96    pub fn is_done(&self) -> bool {
97        self.0.is_done()
98    }
99}
100
101impl MonoWaitGroup {
102    pub fn new() -> (Self, MonoWorkerHandle) {
103        let inner = MonoInner::new();
104        let (wg, handle) = TwinRef::new_mono(inner);
105        (Self(WaitGroupWrapper::new(wg)), MonoWorkerHandle(handle))
106    }
107
108    #[inline]
109    pub fn is_done(&self) -> bool {
110        self.0.is_done()
111    }
112}
113
114impl Future for WaitGroup {
115    type Output = ();
116
117    #[inline]
118    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
119        Pin::new(&mut self.0).poll(cx)
120    }
121}
122
123impl Future for MonoWaitGroup {
124    type Output = ();
125
126    #[inline]
127    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
128        Pin::new(&mut self.0).poll(cx)
129    }
130}
131
132#[cfg(feature = "futures-core")]
133impl futures_core::FusedFuture for WaitGroup {
134    #[inline]
135    fn is_terminated(&self) -> bool {
136        self.0.is_terminated()
137    }
138}
139
140#[cfg(feature = "futures-core")]
141impl futures_core::FusedFuture for MonoWaitGroup {
142    #[inline]
143    fn is_terminated(&self) -> bool {
144        self.0.is_terminated()
145    }
146}
147
148impl WorkerHandle {
149    /// Consumes the handle.
150    ///
151    /// This is equivalent to dropping the handle.
152    #[inline]
153    pub fn done(self) {
154        drop(self);
155    }
156}
157
158impl MonoWorkerHandle {
159    /// Consumes the handle.
160    ///
161    /// This is equivalent to dropping the handle.
162    #[inline]
163    pub fn done(self) {
164        drop(self);
165    }
166}
167
168impl Drop for MonoWorkerHandle {
169    #[inline]
170    fn drop(&mut self) {
171        unsafe {
172            self.0.send_done();
173        }
174    }
175}
compact_waitgroup/wait_group.rs

compact_waitgroup/
wait_group.rs
