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

libasync/dispatcher/
task.rs
