linera_core/
join_set_ext.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! An extension trait to allow determining at compile time how tasks are spawned on the Tokio
5//! runtime.
6//!
7//! In most cases the [`Future`] task to be spawned should implement [`Send`], but that's
8//! not possible when compiling for the Web. In that case, the task is spawned on the
9//! browser event loop.
10
11use futures::channel::oneshot;
12
13#[cfg(web)]
14mod implementation {
15    pub use futures::future::AbortHandle;
16    use futures::{future, stream, StreamExt as _};
17
18    use super::*;
19
20    /// The set of tasks spawned on the current thread in a Web environment.
21    #[derive(Default)]
22    pub struct JoinSet(Vec<oneshot::Receiver<()>>);
23
24    /// An extension trait for the [`JoinSet`] type.
25    pub trait JoinSetExt: Sized {
26        /// Spawns a `future` task on this [`JoinSet`] using [`JoinSet::spawn_local`].
27        ///
28        /// Returns a [`oneshot::Receiver`] to receive the `future`'s output, and an
29        /// [`AbortHandle`] to cancel execution of the task.
30        fn spawn_task<F: Future + 'static>(&mut self, future: F) -> TaskHandle<F::Output>;
31
32        /// Awaits all tasks spawned in this [`JoinSet`].
33        fn await_all_tasks(&mut self) -> impl Future<Output = ()>;
34
35        /// Reaps tasks that have finished.
36        fn reap_finished_tasks(&mut self);
37    }
38
39    impl JoinSetExt for JoinSet {
40        fn spawn_task<F: Future + 'static>(&mut self, future: F) -> TaskHandle<F::Output> {
41            let (abort_handle, abort_registration) = AbortHandle::new_pair();
42            let (send_done, recv_done) = oneshot::channel();
43            let (send_output, recv_output) = oneshot::channel();
44            let future = async move {
45                // Receiver may have been dropped if the task was aborted.
46                send_output.send(future.await).ok();
47                send_done.send(()).ok();
48            };
49            self.0.push(recv_done);
50            wasm_bindgen_futures::spawn_local(
51                future::Abortable::new(future, abort_registration).map(drop),
52            );
53
54            TaskHandle {
55                output_receiver: recv_output,
56                abort_handle,
57            }
58        }
59
60        async fn await_all_tasks(&mut self) {
61            stream::iter(&mut self.0)
62                .then(|x| x)
63                .map(drop)
64                .collect()
65                .await
66        }
67
68        fn reap_finished_tasks(&mut self) {
69            self.0.retain_mut(|task| task.try_recv() == Ok(None));
70        }
71    }
72}
73
74#[cfg(not(web))]
75mod implementation {
76    pub use tokio::task::AbortHandle;
77
78    use super::*;
79
80    /// The set of tasks spawned on the Tokio runtime.
81    pub type JoinSet = tokio::task::JoinSet<()>;
82
83    /// An extension trait for the [`JoinSet`] type.
84    #[trait_variant::make(Send)]
85    pub trait JoinSetExt: Sized {
86        /// Spawns a `future` task on this [`JoinSet`] using [`JoinSet::spawn`].
87        ///
88        /// Returns a [`oneshot::Receiver`] to receive the `future`'s output, and an
89        /// [`AbortHandle`] to cancel execution of the task.
90        fn spawn_task<F: Future<Output: Send> + Send + 'static>(
91            &mut self,
92            future: F,
93        ) -> TaskHandle<F::Output>;
94
95        /// Awaits all tasks spawned in this [`JoinSet`].
96        async fn await_all_tasks(&mut self);
97
98        /// Reaps tasks that have finished.
99        fn reap_finished_tasks(&mut self);
100    }
101
102    impl JoinSetExt for JoinSet {
103        fn spawn_task<F>(&mut self, future: F) -> TaskHandle<F::Output>
104        where
105            F: Future + Send + 'static,
106            F::Output: Send,
107        {
108            let (output_sender, output_receiver) = oneshot::channel();
109
110            let abort_handle = self.spawn(async move {
111                // Receiver may have been dropped if the task was aborted.
112                output_sender.send(future.await).ok();
113            });
114
115            TaskHandle {
116                output_receiver,
117                abort_handle,
118            }
119        }
120
121        async fn await_all_tasks(&mut self) {
122            while self.join_next().await.is_some() {}
123        }
124
125        fn reap_finished_tasks(&mut self) {
126            while self.try_join_next().is_some() {}
127        }
128    }
129}
130
131use std::{
132    future::Future,
133    pin::Pin,
134    task::{Context, Poll},
135};
136
137use futures::FutureExt as _;
138pub use implementation::*;
139
140/// A handle to a task spawned with [`JoinSetExt`].
141///
142/// Dropping a handle detaches its respective task.
143pub struct TaskHandle<Output> {
144    output_receiver: oneshot::Receiver<Output>,
145    abort_handle: AbortHandle,
146}
147
148impl<Output> Future for TaskHandle<Output> {
149    type Output = Result<Output, oneshot::Canceled>;
150
151    fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
152        self.as_mut().output_receiver.poll_unpin(context)
153    }
154}
155
156impl<Output> TaskHandle<Output> {
157    /// Aborts the task.
158    pub fn abort(&self) {
159        self.abort_handle.abort();
160    }
161
162    /// Returns [`true`] if the task is still running.
163    pub fn is_running(&mut self) -> bool {
164        self.output_receiver.try_recv().is_err()
165    }
166}
linera_core/join_set_ext.rs

linera_core/
join_set_ext.rs
