scoped_join_set/
lib.rs

1use std::{
2    any::Any,
3    collections::HashMap,
4    error::Error,
5    fmt,
6    future::Future,
7    pin::Pin,
8    ptr::drop_in_place,
9    sync::atomic::{AtomicBool, Ordering},
10    task::{Context, Poll},
11};
12
13use tokio::{
14    runtime::{Handle, RuntimeFlavor},
15    task::{block_in_place, AbortHandle, Id, JoinSet},
16};
17
18type TokioJoinError = tokio::task::JoinError;
19
20/// A scoped wrapper around [`tokio::task::JoinSet`] for dynamically spawning
21/// and awaiting multiple asynchronous tasks.
22///
23/// `ScopedJoinSet` allows you to spawn futures that are bound to a specific lifetime `'scope`,
24/// ensuring all tasks outlive the scope in which they are spawned. Unlike a standard
25/// `JoinSet`, this structure internally tracks uniquely-owned pinned futures via `FutureHolder`,
26/// while exposing a weak reference (`WeakFuture`) for safe polling.
27///
28/// # Features
29/// - **Scoped Futures**: Ensures futures cannot outlive the provided `'scope`.
30/// - **Dynamic Spawning**: Tasks can be spawned dynamically at runtime.
31/// - **Cancellation Safe**: `join_next()` is cancellation safe; tasks remain in the set
32///   until they complete and are awaited.
33/// - **Automatic Abort**: Dropping a `FutureHolder` aborts its associated task.
34///
35/// # Safety
36/// - Each spawned task has exactly one `FutureHolder` (owning the `Pin<Box>` of the future)
37///   and one `WeakFuture` (polling handle).  
38/// - `FutureHolder` tracks whether the future is alive using an `AtomicBool`.  
39#[derive(Default)]
40pub struct ScopedJoinSet<'scope, T>
41where
42    T: 'static,
43{
44    join_set: JoinSet<Option<T>>,
45    holders: HashMap<Id, FutureHolder<'scope, T>>,
46}
47
48impl<'scope, T> ScopedJoinSet<'scope, T>
49where
50    T: 'static,
51{
52    pub fn new() -> Self {
53        Self {
54            join_set: JoinSet::new(),
55            holders: HashMap::new(),
56        }
57    }
58
59    pub fn spawn<F>(&mut self, task: F)
60    where
61        F: Future<Output = T> + Send + 'scope,
62        T: Send,
63    {
64        let alive: &'static AtomicBool = Box::leak(Box::new(AtomicBool::new(true)));
65        let strong = Box::pin(task);
66        let weak_future = WeakFuture {
67            future: unsafe {
68                std::mem::transmute::<
69                    *mut (dyn Future<Output = T> + Send + 'scope),
70                    *mut dyn Future<Output = T>,
71                >((&raw const *strong).cast_mut())
72            },
73            alive,
74        };
75        let handle = self.join_set.spawn(weak_future);
76        let holder = FutureHolder {
77            abort_handle: handle.clone(),
78            _future: strong,
79            alive,
80        };
81        self.holders.insert(handle.id(), holder);
82    }
83
84    pub fn is_empty(&self) -> bool {
85        self.join_set.is_empty()
86    }
87
88    pub async fn join_next(&mut self) -> Option<Result<T, JoinError>> {
89        match self.join_set.join_next_with_id().await? {
90            Ok((id, Some(value))) => {
91                self.holders.remove(&id);
92                Some(Ok(value))
93            }
94            Ok((id, None)) => {
95                self.holders.remove(&id);
96                Some(Err(JoinError::Cancelled))
97            }
98            Err(error) => {
99                self.holders.remove(&error.id());
100                Some(Err(error.into()))
101            }
102        }
103    }
104}
105
106unsafe impl<'scope, T> Send for ScopedJoinSet<'scope, T> {}
107
108struct FutureHolder<'scope, T>
109where
110    T: 'static,
111{
112    abort_handle: AbortHandle,
113    _future: Pin<Box<dyn Future<Output = T> + Send + 'scope>>,
114    alive: &'static AtomicBool,
115}
116
117impl<'scope, T> Drop for FutureHolder<'scope, T> {
118    fn drop(&mut self) {
119        self.abort_handle.abort();
120
121        if Handle::current().runtime_flavor() == RuntimeFlavor::CurrentThread {
122            return;
123        }
124
125        // Wait until the future is not being polled
126        if !self.alive.load(Ordering::Acquire) {
127            return;
128        }
129
130        block_in_place(|| {
131            while self.alive.load(Ordering::Acquire) {
132                std::hint::spin_loop();
133            }
134        });
135
136        unsafe {
137            drop_in_place(self.alive.as_ptr());
138        }
139    }
140}
141
142struct WeakFuture<T> {
143    future: *mut dyn Future<Output = T>,
144    alive: &'static AtomicBool,
145}
146
147impl<T> Drop for WeakFuture<T> {
148    fn drop(&mut self) {
149        // Release the polling flag
150        self.alive.store(false, Ordering::Release);
151    }
152}
153
154unsafe impl<T: Send> Send for WeakFuture<T> {}
155
156impl<T> Future for WeakFuture<T> {
157    type Output = Option<T>;
158
159    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
160        let future_ref = unsafe { self.future.as_mut().unwrap_unchecked() };
161        let future_pin = unsafe { Pin::new_unchecked(future_ref) };
162        Future::poll(future_pin, cx).map(Some)
163    }
164}
165
166#[derive(Debug)]
167pub enum JoinError {
168    Cancelled,
169    Panicked(Box<dyn Any + Send + 'static>),
170}
171
172impl fmt::Display for JoinError {
173    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
174        match self {
175            JoinError::Cancelled => write!(f, "task was cancelled"),
176            JoinError::Panicked(_) => write!(f, "task panicked"),
177        }
178    }
179}
180
181impl Error for JoinError {}
182
183impl JoinError {
184    pub fn is_cancelled(&self) -> bool {
185        matches!(self, JoinError::Cancelled)
186    }
187
188    pub fn is_panic(&self) -> bool {
189        matches!(self, JoinError::Panicked(_))
190    }
191}
192
193impl From<TokioJoinError> for JoinError {
194    fn from(err: TokioJoinError) -> Self {
195        if err.is_cancelled() {
196            JoinError::Cancelled
197        } else if err.is_panic() {
198            JoinError::Panicked(err.into_panic())
199        } else {
200            JoinError::Cancelled
201        }
202    }
203}
scoped_join_set/lib.rs

scoped_join_set/
lib.rs
