tessera_shard/
task_handles.rs1use 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}