libasync_dispatcher/
task.rs

1// Copyright 2025 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5//! Safe bindings for the C libasync async dispatcher library
6
7use core::task::Context;
8use fuchsia_sync::Mutex;
9use std::pin::Pin;
10use std::sync::atomic::{AtomicBool, Ordering};
11use std::sync::{Arc, mpsc};
12use std::task::{Poll, Wake, Waker};
13
14use zx_status::Status;
15
16use futures::future::{BoxFuture, FutureExt};
17use futures::task::AtomicWaker;
18
19use crate::{AsAsyncDispatcherRef, AsyncDispatcher};
20
21/// The future returned by [`crate::OnDispatcher::compute`]. If this is dropped, the task will be
22/// cancelled.
23#[must_use]
24#[derive(Debug)]
25pub struct Task<T> {
26    state: Arc<TaskFutureState>,
27    result_receiver: mpsc::Receiver<Result<T, Status>>,
28    detached: bool,
29}
30
31impl<T: Send + 'static> Task<T> {
32    fn new(
33        future: impl Future<Output = T> + Send + 'static,
34        dispatcher: AsyncDispatcher,
35    ) -> (Self, Arc<TaskWakerState<T>>) {
36        let future_state = Arc::new(TaskFutureState {
37            waker: AtomicWaker::new(),
38            aborted: AtomicBool::new(false),
39        });
40        let (result_sender, result_receiver) = mpsc::sync_channel(1);
41        let state = Arc::new(TaskWakerState {
42            result_sender,
43            future_state: future_state.clone(),
44            future: Mutex::new(Some(future.boxed())),
45            dispatcher,
46        });
47        let future = Task { state: future_state, result_receiver, detached: false };
48        (future, state)
49    }
50
51    /// Constructs a task that never runs, but immediately fails with the given error status.
52    pub fn new_failed(status: Status) -> Self {
53        let state = Arc::new(TaskFutureState {
54            waker: AtomicWaker::new(),
55            aborted: AtomicBool::new(false),
56        });
57        let (result_sender, result_receiver) = mpsc::sync_channel(1);
58        // send the error to the result receiver. This should never fail, since
59        // we just created both ends.
60        result_sender.try_send(Err(status)).unwrap();
61        Task { state, result_receiver, detached: false }
62    }
63
64    pub(crate) fn start(
65        future: impl Future<Output = T> + Send + 'static,
66        dispatcher: AsyncDispatcher,
67    ) -> Self {
68        let (future, state) = Self::new(future, dispatcher);
69
70        // try to queue the task and if it fails short circuit the delivery of failure to the
71        // caller.
72        if let Err(err) = state.queue() {
73            // drop the future we were given
74            drop(state.future.lock().take());
75            // send the error to the result receiver. This should never fail, since
76            // we just created both ends and the task queuing failed.
77            state.result_sender.try_send(Err(err)).unwrap();
78        }
79
80        future
81    }
82}
83
84impl<T> Task<T> {
85    /// Detaches this future from the task so that it will continue executing without waiting
86    /// on the future. If this is not called, and the future is dropped, the task will be aborted
87    /// the next time it is awoken.
88    pub fn detach(self) {
89        drop(self.detach_on_drop());
90    }
91
92    /// Detaches this future from the task so that it will continue executing without waiting
93    /// on the future. If this is not called, and the future is dropped, the task will be aborted
94    /// the next time it is awoken.
95    ///
96    /// Returns a future that can be awaited on or dropped without affecting the task.
97    pub fn detach_on_drop(mut self) -> JoinHandle<T> {
98        self.detached = true;
99        JoinHandle(self)
100    }
101
102    /// Aborts the task and returns a future that can be used to wait for the task to either
103    /// complete or cancel. If the task was canceled the result of the future will be
104    /// [`Status::CANCELED`].
105    pub fn abort(&self) {
106        self.state.aborted.store(true, Ordering::Relaxed);
107    }
108}
109
110impl<T> Drop for Task<T> {
111    fn drop(&mut self) {
112        if !self.detached {
113            self.state.aborted.store(true, Ordering::Relaxed);
114        }
115    }
116}
117
118#[derive(Debug)]
119struct TaskFutureState {
120    waker: AtomicWaker,
121    aborted: AtomicBool,
122}
123
124impl<T> Future for Task<T> {
125    type Output = Result<T, Status>;
126
127    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
128        use std::sync::mpsc::TryRecvError;
129        self.state.waker.register(cx.waker());
130        match self.result_receiver.try_recv() {
131            Ok(res) => Poll::Ready(res),
132            Err(TryRecvError::Disconnected) => Poll::Ready(Err(Status::CANCELED)),
133            Err(TryRecvError::Empty) => Poll::Pending,
134        }
135    }
136}
137
138/// A handle for a task that will detach on drop. Returned by [`OnDispatcher::spawn`].
139#[derive(Debug)]
140pub struct JoinHandle<T>(Task<T>);
141
142impl<T> JoinHandle<T> {
143    /// Aborts the task and returns a future that can be used to wait for the task to either
144    /// complete or cancel. If the task was canceled the result of the future will be
145    /// [`Status::CANCELED`].
146    pub fn abort(&self) {
147        self.0.abort()
148    }
149}
150
151impl<T> Future for JoinHandle<T> {
152    type Output = Result<T, Status>;
153
154    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
155        self.0.poll_unpin(cx)
156    }
157}
158
159struct TaskWakerState<T> {
160    result_sender: mpsc::SyncSender<Result<T, Status>>,
161    future_state: Arc<TaskFutureState>,
162    future: Mutex<Option<BoxFuture<'static, T>>>,
163    dispatcher: AsyncDispatcher,
164}
165
166impl<T: Send + 'static> Wake for TaskWakerState<T> {
167    fn wake(self: Arc<Self>) {
168        self.wake_by_ref();
169    }
170    fn wake_by_ref(self: &Arc<Self>) {
171        match self.queue() {
172            Err(e) if e == Status::BAD_STATE => {
173                // the dispatcher is shutting down so drop the future, if there
174                // is one, to cancel it.
175                let future_slot = self.future.lock().take();
176                drop(future_slot);
177                self.send_result(Err(e));
178            }
179            res => res.expect("Unexpected error waking dispatcher task"),
180        }
181    }
182}
183
184impl<T: Send + 'static> TaskWakerState<T> {
185    /// Sends the result to the future end of this task, if it still exists.
186    fn send_result(&self, res: Result<T, Status>) {
187        // send the result and wake the waker if any has been registered.
188        // We ignore the result here because if the other end has dropped it's
189        // fine for the result to go nowhere.
190        self.result_sender.try_send(res).ok();
191        self.future_state.waker.wake();
192    }
193
194    /// Posts a task to progress the currently stored future. The task will
195    /// consume the future if the future is ready after the next poll.
196    /// Otherwise, the future is kept to be polled again after being woken.
197    pub(crate) fn queue(self: &Arc<Self>) -> Result<(), Status> {
198        let arc_self = self.clone();
199        self.dispatcher
200            .post_task_sync(move |status| {
201                let mut future_slot = arc_self.future.lock();
202                // if the executor is shutting down, drop the future we're waiting on and pass
203                // on the error.
204                if let Err(status) = status {
205                    drop(future_slot.take());
206                    arc_self.send_result(Err(status));
207                    return;
208                }
209
210                // if the future has been dropped without being detached, drop the future and
211                // send an Err(Status::CANCELED) if the caller is still listening.
212                if arc_self.future_state.aborted.load(Ordering::Relaxed) {
213                    drop(future_slot.take());
214                    arc_self.send_result(Err(Status::CANCELED));
215                    return;
216                }
217
218                let Some(mut future) = future_slot.take() else {
219                    return;
220                };
221                let waker = Waker::from(arc_self.clone());
222                let context = &mut Context::from_waker(&waker);
223                match future.as_mut().poll(context) {
224                    Poll::Pending => *future_slot = Some(future),
225                    Poll::Ready(res) => {
226                        arc_self.send_result(Ok(res));
227                    }
228                }
229            })
230            .map(|_| ())
231    }
232}
libasync_dispatcher/task.rs

libasync_dispatcher/
task.rs
