tessera_shard/async_support/
task_handles.rs

1//! Background task handles for shard-owned async work.
2//!
3//! ## Usage
4//!
5//! Track spawned jobs in `TaskHandles` so they can be canceled together.
6
7use std::future::Future;
8
9use futures_util::future::{AbortHandle, Abortable};
10use parking_lot::Mutex;
11
12#[cfg(not(target_family = "wasm"))]
13use tokio::task::JoinHandle;
14
15#[cfg(target_family = "wasm")]
16use wasm_bindgen_futures::spawn_local;
17
18#[cfg(not(target_family = "wasm"))]
19/// Platform-specific handle returned when spawning an asynchronous task.
20pub type TaskRuntimeHandle = JoinHandle<()>;
21#[cfg(target_family = "wasm")]
22/// Platform-specific handle returned when spawning an asynchronous task.
23pub type TaskRuntimeHandle = ();
24
25/// Handle pair for a spawned background task.
26pub struct TaskHandle {
27    /// Runtime handle returned by the executor.
28    pub handle: TaskRuntimeHandle,
29    cancel: AbortHandle,
30}
31
32/// Tracks multiple background tasks and aborts them on drop.
33pub struct TaskHandles {
34    tasks: Mutex<Vec<TaskHandle>>,
35    #[cfg(not(target_family = "wasm"))]
36    executor: tokio::runtime::Handle,
37}
38
39impl Default for TaskHandles {
40    fn default() -> Self {
41        Self::new()
42    }
43}
44
45impl Drop for TaskHandles {
46    fn drop(&mut self) {
47        self.cancel_all();
48    }
49}
50
51impl TaskHandles {
52    /// Creates an empty task handle collection.
53    pub fn new() -> Self {
54        Self {
55            tasks: Mutex::new(Vec::new()),
56            #[cfg(not(target_family = "wasm"))]
57            executor:
58                tessera_ui::__private::runtime_resource::<super::tokio_runtime::TaskExecutor>()
59                    .handle()
60                    .clone(),
61        }
62    }
63
64    /// Spawns a task and stores its handle for later cancellation.
65    pub fn spawn<F>(&self, fut: F)
66    where
67        F: Future<Output = ()> + Send + 'static,
68    {
69        #[cfg(not(target_family = "wasm"))]
70        let (handle, cancel) = spawn_task(&self.executor, fut);
71        #[cfg(target_family = "wasm")]
72        let (handle, cancel) = spawn_task(fut);
73        self.tasks.lock().push(TaskHandle { handle, cancel });
74    }
75
76    /// Aborts and removes all tracked tasks.
77    pub fn cancel_all(&self) {
78        let mut tasks = self.tasks.lock();
79        for task in tasks.drain(..) {
80            task.cancel.abort();
81        }
82    }
83}
84
85#[cfg(not(target_family = "wasm"))]
86fn spawn_task<F>(executor: &tokio::runtime::Handle, fut: F) -> (TaskRuntimeHandle, AbortHandle)
87where
88    F: Future<Output = ()> + Send + 'static,
89{
90    let (cancel, registration) = AbortHandle::new_pair();
91    let wrapped = Abortable::new(fut, registration);
92    let handle = executor.spawn(async move {
93        let _ = wrapped.await;
94    });
95    (handle, cancel)
96}
97
98#[cfg(target_family = "wasm")]
99fn spawn_task<F>(fut: F) -> (TaskRuntimeHandle, AbortHandle)
100where
101    F: Future<Output = ()> + Send + 'static,
102{
103    let (cancel, registration) = AbortHandle::new_pair();
104    let wrapped = Abortable::new(fut, registration);
105    spawn_local(async move {
106        let _ = wrapped.await;
107    });
108    ((), cancel)
109}
110
111#[cfg(all(test, not(target_family = "wasm")))]
112mod tests {
113    use super::TaskHandles;
114
115    #[test]
116    fn session_teardown_cancels_jobs_even_when_the_task_collection_survives() {
117        let session = tessera_ui::RuntimeSession::new();
118        let tasks = session.run(TaskHandles::new);
119        tasks.spawn(std::future::pending());
120        let task = tasks.tasks.lock().pop().expect("registered job");
121        drop(session);
122        let executor = tokio::runtime::Builder::new_current_thread()
123            .build()
124            .expect("test executor");
125        let result = executor.block_on(task.handle);
126        assert!(
127            result
128                .expect_err("session must cancel its job")
129                .is_cancelled()
130        );
131    }
132}
tessera_shard/async_support/task_handles.rs

tessera_shard/async_support/
task_handles.rs
