Skip to main content

polyoxide_rt/
task.rs

1//! `spawn` and the `Task` it returns.
2
3use std::fmt;
4use std::future::Future;
5use std::pin::Pin;
6use std::sync::atomic::{AtomicBool, Ordering};
7use std::sync::Arc;
8use std::task::{Context, Poll};
9
10use futures_channel::oneshot;
11use futures_util::future::{AbortHandle, Abortable};
12
13use crate::Runtime;
14
15/// A spawned future, as a future of its output.
16///
17/// It resolves to `Some(output)`, or to `None` when the task was aborted or
18/// its output never arrived (it panicked, or its runtime dropped it). Dropping
19/// a `Task` detaches the task, as dropping tokio's `JoinHandle` does; only
20/// [`Task::abort`] stops it.
21pub struct Task<T> {
22    abort: AbortHandle,
23    output: oneshot::Receiver<T>,
24    finished: Arc<AtomicBool>,
25}
26
27impl<T> Task<T> {
28    /// Stop the task at its next suspension point. A task that has already
29    /// ended is unaffected.
30    pub fn abort(&self) {
31        self.abort.abort();
32    }
33
34    /// Whether the task has ended: it completed, was aborted and dropped by its
35    /// runtime, or panicked. A task just aborted is not finished until its
36    /// runtime has dropped it.
37    pub fn is_finished(&self) -> bool {
38        self.finished.load(Ordering::Acquire)
39    }
40}
41
42impl<T> Future for Task<T> {
43    type Output = Option<T>;
44
45    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
46        match Pin::new(&mut self.output).poll(cx) {
47            Poll::Ready(output) => {
48                self.finished.store(true, Ordering::Release);
49                Poll::Ready(output.ok())
50            }
51            Poll::Pending => Poll::Pending,
52        }
53    }
54}
55
56impl<T> fmt::Debug for Task<T> {
57    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
58        f.debug_struct("Task")
59            .field("finished", &self.is_finished())
60            .finish_non_exhaustive()
61    }
62}
63
64/// Sets the task's flag when dropped, so a panic, an abort and a runtime that
65/// drops the task unpolled all end it.
66struct Finish(Arc<AtomicBool>);
67
68impl Drop for Finish {
69    fn drop(&mut self) {
70        self.0.store(true, Ordering::Release);
71    }
72}
73
74/// Spawn `future` on `runtime`, which needs only its detached spawn.
75pub fn spawn<T, F>(runtime: &dyn Runtime, future: F) -> Task<T>
76where
77    F: Future<Output = T> + Send + 'static,
78    T: Send + 'static,
79{
80    let (output_tx, output_rx) = oneshot::channel();
81    let (abort, registration) = AbortHandle::new_pair();
82    let finished = Arc::new(AtomicBool::new(false));
83    let finish = Finish(Arc::clone(&finished));
84    let run = Abortable::new(
85        async move {
86            let value = future.await;
87            // Flag the end before the output is visible, so a reader woken by
88            // the output sees the task finished.
89            drop(finish);
90            let _ = output_tx.send(value);
91        },
92        registration,
93    );
94    runtime.spawn(Box::pin(async move {
95        let _ = run.await;
96    }));
97    Task {
98        abort,
99        output: output_rx,
100        finished,
101    }
102}