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

scoped_join_set/
lib.rs
