tessera_shard/
task_handles.rs

1use std::future::Future;
2
3use futures_util::future::{AbortHandle, Abortable};
4use parking_lot::Mutex;
5
6#[cfg(not(target_family = "wasm"))]
7use tokio::task::JoinHandle;
8
9#[cfg(target_family = "wasm")]
10use wasm_bindgen_futures::spawn_local;
11
12#[cfg(not(target_family = "wasm"))]
13pub type TaskRuntimeHandle = JoinHandle<()>;
14#[cfg(target_family = "wasm")]
15pub type TaskRuntimeHandle = ();
16
17pub struct TaskHandle {
18    pub handle: TaskRuntimeHandle,
19    cancel: AbortHandle,
20}
21
22pub struct TaskHandles {
23    tasks: Mutex<Vec<TaskHandle>>,
24}
25
26impl Default for TaskHandles {
27    fn default() -> Self {
28        Self::new()
29    }
30}
31
32impl Drop for TaskHandles {
33    fn drop(&mut self) {
34        self.cancel_all();
35    }
36}
37
38impl TaskHandles {
39    pub fn new() -> Self {
40        Self {
41            tasks: Mutex::new(Vec::new()),
42        }
43    }
44
45    pub fn spawn<F>(&self, fut: F)
46    where
47        F: Future<Output = ()> + Send + 'static,
48    {
49        let (handle, cancel) = spawn_task(fut);
50        self.tasks.lock().push(TaskHandle { handle, cancel });
51    }
52
53    pub fn cancel_all(&self) {
54        let mut tasks = self.tasks.lock();
55        for task in tasks.drain(..) {
56            task.cancel.abort();
57        }
58    }
59}
60
61#[cfg(not(target_family = "wasm"))]
62fn spawn_task<F>(fut: F) -> (TaskRuntimeHandle, AbortHandle)
63where
64    F: Future<Output = ()> + Send + 'static,
65{
66    let (cancel, registration) = AbortHandle::new_pair();
67    let wrapped = Abortable::new(fut, registration);
68    let handle = crate::tokio_runtime::get().spawn(async move {
69        let _ = wrapped.await;
70    });
71    (handle, cancel)
72}
73
74#[cfg(target_family = "wasm")]
75fn spawn_task<F>(fut: F) -> (TaskRuntimeHandle, AbortHandle)
76where
77    F: Future<Output = ()> + Send + 'static,
78{
79    let (cancel, registration) = AbortHandle::new_pair();
80    let wrapped = Abortable::new(fut, registration);
81    spawn_local(async move {
82        let _ = wrapped.await;
83    });
84    ((), cancel)
85}