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

libasync_dispatcher/
task.rs
