geph5_rt/
reaper.rs

1//! A task "reaper" that cancels everything it owns when dropped, yet does not
2//! leak handles to tasks that have already finished. A faithful reimplementation
3//! of `smolscale::reaper` on top of [`crate::Task`].
4//!
5//! Note: a bare `tokio::task::JoinSet` is *not* a drop-in substitute — it needs
6//! `&mut self` to spawn (this `attach` takes `&self`) and never reaps finished
7//! tasks without active `join_next` calls.
8
9use async_channel::{Receiver, Sender};
10use futures_concurrency::future::Race;
11use futures_util::StreamExt;
12use futures_util::stream::FuturesUnordered;
13
14use crate::Task;
15
16/// Owns a set of detached tasks: finished tasks are dropped as they complete,
17/// and all still-running tasks are cancelled when the reaper is dropped.
18pub struct TaskReaper<T> {
19    send_task: Sender<Task<T>>,
20    _reaper: Task<()>,
21}
22
23impl<T: Send + 'static> Default for TaskReaper<T> {
24    fn default() -> Self {
25        Self::new()
26    }
27}
28
29impl<T: Send + 'static> TaskReaper<T> {
30    /// Creates a new reaper with its background driver task.
31    pub fn new() -> Self {
32        let (send_task, recv_task) = async_channel::unbounded();
33        let _reaper = crate::spawn(reaper_loop(recv_task));
34        Self { send_task, _reaper }
35    }
36
37    /// Attaches a task to this reaper, transferring ownership of its handle.
38    pub fn attach(&self, task: Task<T>) {
39        let _ = self.send_task.try_send(task);
40    }
41}
42
43async fn reaper_loop<T: Send + 'static>(recv_task: Receiver<Task<T>>) {
44    let mut inner: FuturesUnordered<Task<T>> = FuturesUnordered::new();
45    loop {
46        // Race receiving a new task against draining finished tasks. The drain
47        // arm never resolves (it parks on `pending` after polling the set); its
48        // only purpose is to drive attached tasks so completed ones are removed.
49        let next = (async { recv_task.recv().await }, async {
50            inner.next().await;
51            std::future::pending::<Result<Task<T>, async_channel::RecvError>>().await
52        })
53            .race()
54            .await;
55        match next {
56            Ok(task) => inner.push(task),
57            // Channel closed: the reaper was dropped. Returning drops `inner`,
58            // which aborts all still-running attached tasks.
59            Err(_) => return,
60        }
61    }
62}
geph5_rt/reaper.rs

geph5_rt/
reaper.rs
