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
//! Golang like [WaitGroup](https://pkg.go.dev/sync#WaitGroup) implementation.
//!
//! ## Usage
//!
//! Add this to your `Cargo.toml`:
//!
//! ```toml
//! [build-dependencies]
//! await-group = "0.1"
//! ```
//!
//! ## Example
//! ```rust
//! use await_group::AwaitGroup;
//!
//! #[tokio::main]
//! async fn main() {
//! let wg = AwaitGroup::new();
//! for _ in 0..10 {
//! let w = wg.clone();
//! tokio::spawn(async move {
//! _ = w;
//! });
//! }
//! wg.await;
//! }
//!
//! ```
extern crate alloc;
use core::{
future::{Future, IntoFuture},
pin::Pin,
task::{Context, Poll},
};
use alloc::{sync::Arc, sync::Weak};
use atomic_waker::AtomicWaker;
#[derive(Clone, Default)]
pub struct AwaitGroup {
inner: Arc<Inner>,
}
impl AwaitGroup {
pub fn new() -> Self {
Self {
inner: Arc::new(Inner {
waker: AtomicWaker::new(),
}),
}
}
}
impl IntoFuture for AwaitGroup {
type Output = ();
type IntoFuture = AwaitGroupFuture;
fn into_future(self) -> Self::IntoFuture {
AwaitGroupFuture {
inner: Arc::downgrade(&self.inner),
}
}
}
pub struct AwaitGroupFuture {
inner: Weak<Inner>,
}
impl Future for AwaitGroupFuture {
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(()),
}
}
}
#[derive(Default)]
struct Inner {
waker: AtomicWaker,
}
impl Drop for Inner {
fn drop(&mut self) {
self.waker.wake();
}
}
#[cfg(test)]
mod test {
use crate::AwaitGroup;
#[tokio::test]
async fn smoke() {
let wg = AwaitGroup::new();
for _ in 0..10 {
let w = wg.clone();
tokio::spawn(async move {
_ = w;
});
}
wg.await;
}
}