scoped_join_set/
lib.rs

1//! A scoped, lifetime-aware wrapper around [`tokio::task::JoinSet`].
2//!
3//! `ScopedJoinSet` enables dynamically spawning async tasks whose futures are
4//! tied to a specific Rust lifetime `'scope`. This provides a safe alternative
5//! to hand-managed `unsafe` approaches while enabling advanced async task
6//! orchestration where futures do **not** need to be `'static`.
7//!
8//! Each spawned task is split into:
9//!
10//! - A [`FutureHolder`], owning the actual boxed future and ensuring proper drop
11//!   ordering.
12//! - A [`WeakFuture`] (inside the `JoinSet`), which the Tokio runtime polls.
13//!
14//! When the [`ScopedJoinSet`] is dropped, all tasks are aborted and all future
15//! memory is reclaimed safely.
16//!
17//! # Guarantees
18//!
19//! - Futures spawned via [`ScopedJoinSet::spawn`] cannot outlive `'scope`.
20//! - `join_next()` is cancellation-safe.
21//! - Dropping the entire set aborts all tasks.
22//!
23//! # Example
24//!
25//! ```
26//! use scoped_join_set::ScopedJoinSet;
27//!
28//! #[tokio::main]
29//! async fn main() {
30//!     let temporary_value = Box::new(5);
31//!     let value_ref = &temporary_value;
32//!
33//!     let mut set = ScopedJoinSet::<u32>::new();
34//!
35//!     set.spawn(async {
36//!         println!("I have a {value_ref}");
37//!         **value_ref
38//!     });
39//!
40//!     while let Some(result) = set.join_next().await {
41//!         assert_eq!(result.unwrap(), *temporary_value);
42//!     }
43//! }
44//! ```
45
46mod future_wrappers;
47mod lifetime_obfuscation;
48
49use std::{any::Any, collections::HashMap, error::Error, fmt, future::Future, marker::PhantomData};
50
51use tokio::task::{Id, JoinSet};
52
53use crate::{
54    future_wrappers::{split_spawn_task, FutureHolder},
55    lifetime_obfuscation::{from_bytes_unchecked, obfuscate_lifetime},
56};
57
58type TokioJoinError = tokio::task::JoinError;
59
60/// A scoped wrapper around [`tokio::task::JoinSet`] for dynamically spawning
61/// futures with a lifetime `'scope`.
62///
63/// Unlike normal Tokio tasks requiring `'static` futures, this set accepts
64/// futures tied to a temporary stack lifetime, as long as the
65/// `ScopedJoinSet` lives long enough.
66///
67/// Each call to [`spawn`](Self::spawn) places a non-`'static` future into the
68/// runtime via a controlled internal mechanism (`FutureHolder` + `WeakFuture`).
69///
70/// # Type Parameters
71/// - `'scope`: The lifetime during which spawned futures must live.
72/// - `T`: The output of each spawned future.
73#[derive(Default)]
74pub struct ScopedJoinSet<'scope, T>
75where
76    T: 'scope + Send,
77{
78    join_set: JoinSet<Option<Box<[u8]>>>,
79    holders: HashMap<Id, FutureHolder<'scope, Box<[u8]>>>,
80    _marker: PhantomData<T>,
81}
82
83impl<'scope, T> ScopedJoinSet<'scope, T>
84where
85    T: 'scope + Send,
86{
87    /// Creates an empty `ScopedJoinSet`.
88    pub fn new() -> Self {
89        Self {
90            join_set: JoinSet::new(),
91            holders: HashMap::new(),
92            _marker: PhantomData,
93        }
94    }
95
96    /// Spawns a new scoped future into the set.
97    ///
98    /// This function is analogous to [`tokio::spawn`], but accepts futures with
99    /// the given scope lifetime.
100    ///
101    /// # Parameters
102    /// - `task`: The user future to spawn.
103    ///
104    /// # Requirements
105    /// - The returned value `T` must be `Send`.
106    /// - The future must be `Send` and live at least `'scope`.
107    pub fn spawn<F>(&mut self, task: F)
108    where
109        F: Future<Output = T> + Send + 'scope,
110        T: Send,
111    {
112        let task = obfuscate_lifetime(task);
113        let (id, holder) = split_spawn_task(&mut self.join_set, task);
114        self.holders.insert(id, holder);
115    }
116
117    /// Returns whether no tasks remain.
118    ///
119    /// Equivalent to [`JoinSet::is_empty`].
120    pub fn is_empty(&self) -> bool {
121        self.join_set.is_empty()
122    }
123
124    /// Returns the number of tasks currently stored in the set.
125    ///
126    /// Equivalent to [`JoinSet::len`].
127    pub fn len(&self) -> usize {
128        self.holders.len()
129    }
130
131    /// Waits for the next completed task in the set.
132    ///
133    /// Returns:
134    /// - `Some(Ok(T))` on success.
135    /// - `Some(Err(JoinError))` on cancellation or panic.
136    /// - `None` when the set is empty.
137    ///
138    /// This method is **cancellation-safe**: dropping the future returned by
139    /// `.await` does not remove the task from the set.
140    pub async fn join_next(&mut self) -> Option<Result<T, JoinError>> {
141        match self.join_set.join_next_with_id().await? {
142            Ok((id, Some(bytes))) => {
143                self.holders.remove(&id);
144
145                // SAFETY: By having `T` as the type argument, reconstruction needs to be successful.
146                let value = unsafe { from_bytes_unchecked(&bytes) };
147
148                Some(Ok(value))
149            }
150            Ok((id, None)) => {
151                self.holders.remove(&id);
152                Some(Err(JoinError::Cancelled))
153            }
154            Err(error) => {
155                self.holders.remove(&error.id());
156                Some(Err(error.into()))
157            }
158        }
159    }
160
161    /// Attempts to join one completed task without awaiting.
162    ///
163    /// This method is **non-blocking** and returns immediately.
164    ///
165    /// Returns:
166    /// - `Some(Ok(T))` if a task has already completed successfully,
167    /// - `Some(Err(JoinError))` if a completed task was cancelled or panicked,
168    /// - `None` if **no** tasks have completed yet **or** the set is empty.
169    ///
170    /// This differs from [`join_next`](Self::join_next):
171    /// - `join_next()` **awaits** the next completed task.
172    /// - `try_join_next()` returns immediately without awaiting.
173    ///
174    /// # Cancellation Safety
175    ///
176    /// This method is always cancellation-safe, because it does not perform
177    /// any async operations and does not remove tasks unless they have
178    /// already finished.
179    pub fn try_join_next(&mut self) -> Option<Result<T, JoinError>> {
180        match self.join_set.try_join_next_with_id()? {
181            Ok((id, Some(bytes))) => {
182                // Remove the holder before reconstructing the value.
183                self.holders.remove(&id);
184
185                // SAFETY: The bytes originated from a valid instance of T.
186                let value = unsafe { from_bytes_unchecked(&bytes) };
187                Some(Ok(value))
188            }
189            Ok((id, None)) => {
190                self.holders.remove(&id);
191                Some(Err(JoinError::Cancelled))
192            }
193            Err(error) => {
194                self.holders.remove(&error.id());
195                Some(Err(error.into()))
196            }
197        }
198    }
199
200    /// Waits for all tasks in the set to complete, consuming the set and
201    /// returning a list of results.
202    ///
203    /// The returned results are in **completion order**, not spawn order.
204    ///
205    /// # Cancellation Safety
206    /// If this future is dropped, remaining tasks stay inside the set.
207    pub async fn join_all(mut self) -> Vec<Result<T, JoinError>> {
208        let mut results = Vec::with_capacity(self.len());
209        while let Some(res) = self.join_next().await {
210            results.push(res);
211        }
212        results
213    }
214
215    /// Aborts all currently-running tasks.
216    ///
217    /// This does **not** await task completion; it merely issues abort signals.
218    /// To also wait for them, call [`join_remaining`](#method.join_remaining)
219    /// or drop the entire set.
220    ///
221    /// After calling `abort_all`, the set will still contain all tasks until
222    /// they complete their aborts and are polled again.
223    pub fn abort_all(&mut self) {
224        for holder in self.holders.values_mut() {
225            holder.abort();
226        }
227    }
228
229    /// Aborts all tasks *and then waits for them to finish*.
230    ///
231    /// Unlike dropping the set, `shutdown` consumes it gracefully.
232    pub async fn shutdown(mut self) -> Vec<Result<T, JoinError>> {
233        self.abort_all();
234        self.join_all().await
235    }
236}
237
238unsafe impl<'scope, T: Send> Send for ScopedJoinSet<'scope, T> {}
239
240/// Errors returned when joining scoped tasks.
241#[derive(Debug)]
242pub enum JoinError {
243    /// The task was cancelled or aborted.
244    Cancelled,
245    /// The task panicked and its panic payload is returned.
246    Panicked(Box<dyn Any + Send + 'static>),
247}
248
249impl fmt::Display for JoinError {
250    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
251        match self {
252            JoinError::Cancelled => write!(f, "task was cancelled"),
253            JoinError::Panicked(_) => write!(f, "task panicked"),
254        }
255    }
256}
257
258impl Error for JoinError {}
259
260impl JoinError {
261    /// Returns `true` if the task was cancelled.
262    pub fn is_cancelled(&self) -> bool {
263        matches!(self, JoinError::Cancelled)
264    }
265
266    /// Returns `true` if the task panicked.
267    pub fn is_panic(&self) -> bool {
268        matches!(self, JoinError::Panicked(_))
269    }
270}
271
272impl From<TokioJoinError> for JoinError {
273    fn from(err: TokioJoinError) -> Self {
274        if err.is_cancelled() {
275            JoinError::Cancelled
276        } else if err.is_panic() {
277            JoinError::Panicked(err.into_panic())
278        } else {
279            JoinError::Cancelled
280        }
281    }
282}
scoped_join_set/lib.rs

scoped_join_set/
lib.rs
