1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
use atomic_waker::AtomicWaker;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Weak};
use std::task::{Context, Poll};
pub struct WaitGroup {
inner: Arc<Inner>,
}
#[derive(Clone)]
pub struct Worker {
inner: Arc<Inner>,
}
pub struct WaitGroupFuture {
inner: Weak<Inner>,
}
struct Inner {
waker: AtomicWaker,
}
impl Drop for Inner {
fn drop(&mut self) {
self.waker.wake();
}
}
impl WaitGroup {
pub fn new() -> Self {
Self {
inner: Arc::new(Inner {
waker: AtomicWaker::new(),
}),
}
}
pub fn worker(&self) -> Worker {
Worker {
inner: self.inner.clone(),
}
}
pub fn wait(self) -> WaitGroupFuture {
WaitGroupFuture {
inner: Arc::downgrade(&self.inner),
}
}
}
impl Default for WaitGroup {
fn default() -> Self {
Self::new()
}
}
impl Future for WaitGroupFuture {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match self.inner.upgrade() {
Some(inner) => {
inner.waker.register(cx.waker());
Poll::Pending
}
None => Poll::Ready(()),
}
}
}
#[cfg(test)]
mod test {
use super::*;
use async_std::task;
#[async_std::test]
async fn smoke() {
let wg = WaitGroup::new();
for _ in 0..100 {
let w = wg.worker();
task::spawn(async move {
drop(w);
});
}
wg.wait().await;
}
}