libasync/
after_deadline.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
5use std::pin::Pin;
6use std::ptr::NonNull;
7use std::sync::Arc;
8use std::sync::atomic::{AtomicI32, Ordering};
9use std::task::{Context, Poll};
10
11use libasync_dispatcher::{AsyncDispatcher, DetectDispatcher, GetAsyncDispatcher};
12use libasync_sys::{async_cancel_task, async_dispatcher, async_post_task, async_task};
13
14use futures::task::AtomicWaker;
15use zx::Status;
16use zx::sys::{ZX_ERR_CANCELED, ZX_OK};
17
18use crate::callback_state::CallbackSharedState;
19
20type SharedState = CallbackSharedState<async_task, AfterDeadlineState>;
21
22/// Implements methods used for setting and waiting on timers on a dispatcher.
23pub trait DispatcherTimerExt {
24    /// Returns a future that will fire when after the given deadline time.
25    ///
26    /// This can be used instead of the fuchsia-async timer primitives in situations where
27    /// there isn't a currently active fuchsia-async executor running on that dispatcher for some
28    /// reason (ie. the rust code does not own the dispatcher) or for cases where the small overhead
29    /// of fuchsia-async compatibility is too much.
30    ///
31    /// # Panics
32    ///
33    /// If the dispatcher pointed to by `self` is not available right now (like [`CurrentDispatcher`])
34    /// on a thread with no current dispatcher set, this function will panic. You can use
35    /// [`Self::try_after_deadline`] to handle the condition where there is no current dispatcher,
36    /// or if you're trying to run it on the current dispatcher you may want to use [`AfterDeadline::new`]
37    /// instead.
38    fn after_deadline(&self, deadline: zx::MonotonicInstant) -> AfterDeadline;
39
40    /// Returns a future that will fire when after the given deadline time.
41    ///
42    /// This can be used instead of the fuchsia-async timer primitives in situations where
43    /// there isn't a currently active fuchsia-async executor running on that dispatcher for some
44    /// reason (ie. the rust code does not own the dispatcher) or for cases where the small overhead
45    /// of fuchsia-async compatibility is too much.
46    fn try_after_deadline(&self, deadline: zx::MonotonicInstant) -> Option<AfterDeadline>;
47}
48
49impl<T> DispatcherTimerExt for T
50where
51    T: GetAsyncDispatcher,
52{
53    fn after_deadline(&self, deadline: zx::MonotonicInstant) -> AfterDeadline {
54        self.try_after_deadline(deadline).expect("No current dispatcher")
55    }
56
57    fn try_after_deadline(&self, deadline: zx::MonotonicInstant) -> Option<AfterDeadline> {
58        let dispatcher = self.try_get_async_dispatcher()?;
59        Some(AfterDeadline::new_on(dispatcher, deadline))
60    }
61}
62
63struct AfterDeadlineState {
64    async_dispatcher: NonNull<async_dispatcher>,
65    waker: AtomicWaker,
66    /// The status will initially be [`Status::SHOULD_WAIT`]. Once fired it will be the status
67    /// returned by the callback.
68    status: AtomicI32,
69}
70
71// SAFETY: All fields in AfterDeadlineState are either atomic or immutable.
72unsafe impl Send for AfterDeadlineState {}
73unsafe impl Sync for AfterDeadlineState {}
74
75impl AfterDeadlineState {
76    extern "C" fn call(_dispatcher: *mut async_dispatcher, task: *mut async_task, status: i32) {
77        debug_assert!(
78            status == ZX_OK || status == ZX_ERR_CANCELED,
79            "task callback called with status other than ok or canceled"
80        );
81        // SAFETY: This callback's copy of the `async_task` object was refcounted for when we
82        // started the wait.
83        let state = unsafe { SharedState::from_raw_ptr(task) };
84        state.status.store(status, Ordering::Relaxed);
85        state.waker.wake();
86    }
87}
88
89/// A future that represents a deferral to a future time.
90///
91/// See [`OnDispatcher::after_deadline`] for more information.
92pub struct AfterDeadline {
93    dispatcher: DetectDispatcher,
94    state: Option<Arc<SharedState>>,
95    deadline: zx::MonotonicInstant,
96}
97
98impl AfterDeadline {
99    /// Creates a new timer object that will fire when the deadline has passed.
100    ///
101    /// This will get the current dispatcher on first poll. If you want to run it against a
102    /// specific dispatcher, use [`DispatcherTimerExt::after_deadline`].
103    pub fn new(deadline: zx::MonotonicInstant) -> Self {
104        let state = None;
105        let dispatcher = DetectDispatcher::default();
106        Self { dispatcher, state, deadline }
107    }
108
109    fn new_on(dispatcher: AsyncDispatcher, deadline: zx::MonotonicInstant) -> Self {
110        let state = None;
111        let dispatcher = DetectDispatcher::new(dispatcher);
112        Self { dispatcher, state, deadline }
113    }
114}
115
116impl Future for AfterDeadline {
117    type Output = Result<(), Status>;
118
119    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
120        // if we didn't have a dispatcher when the future was created, return BAD_STATE.
121        let dispatcher = self.dispatcher.get_or_detect()?;
122
123        // if we've already spawned a task then return based on the task's state.
124        if let Some(state) = &self.state {
125            let status = state.status.load(Ordering::Relaxed);
126            if status != Status::SHOULD_WAIT.into_raw() {
127                return Poll::Ready(Status::ok(status));
128            } else {
129                state.waker.register(cx.waker());
130                return Poll::Pending;
131            }
132        }
133
134        let deadline = self.deadline;
135        let now = dispatcher.now();
136        if deadline < zx::MonotonicInstant::from_nanos(now) {
137            return Poll::Ready(Ok(()));
138        }
139
140        // otherwise we want to wait for a callback
141        let async_dispatcher = dispatcher.as_ptr();
142
143        let task = async_task {
144            handler: Some(AfterDeadlineState::call),
145            deadline: deadline.into_nanos(),
146            ..Default::default()
147        };
148        let state = AfterDeadlineState {
149            async_dispatcher,
150            waker: AtomicWaker::new(),
151            status: AtomicI32::new(Status::SHOULD_WAIT.into_raw()),
152        };
153        let state = SharedState::new(task, state);
154        state.waker.register(cx.waker());
155
156        let state_ptr = SharedState::make_raw_ptr(state.clone());
157
158        // SAFETY: We know the `async_dispatcher` is valid because we're running inside
159        // `on_dispatcher` and we are giving ownership of the shared state object to the
160        // callback.
161        let res = Status::ok(unsafe { async_post_task(async_dispatcher.as_ptr(), state_ptr) });
162        match res {
163            Ok(_) => {
164                self.state = Some(state);
165                Poll::Pending
166            }
167            Err(err) => {
168                // SAFETY: Posting the task failed, so we now have an outstanding reference to
169                // the state object that will never have a callback called on it.
170                unsafe { SharedState::release_raw_ptr(state_ptr) };
171                Poll::Ready(Err(err))
172            }
173        }
174    }
175}
176
177impl Drop for AfterDeadline {
178    fn drop(&mut self) {
179        let Some(state) = self.state.take() else {
180            // if we never spawned a task we can just return.
181            return;
182        };
183        let Some(dispatcher) = self.dispatcher.get() else {
184            // if we never got a dispatcher or failed to get a dispatcher then we never
185            // registered a wait and we can just return.
186            return;
187        };
188        if state.status.load(Ordering::Relaxed) != Status::SHOULD_WAIT.into_raw() {
189            // the callback has been called so we don't even need to try to cancel it.
190            return;
191        }
192        let async_dispatcher = dispatcher.as_ptr();
193        if async_dispatcher != state.async_dispatcher {
194            panic!(
195                "Dropping a pending `AfterDeadline` future from a different dispatcher than the one it was awaited on."
196            );
197        }
198        let state_ptr = SharedState::as_raw_ptr(&state);
199        // SAFETY: We know that the current async dispatcher is valid because we are running
200        // inside `on_dispatcher`, and we know the `state_ptr` is valid because the `Arc`
201        // holding it is still held.
202        let status = unsafe { async_cancel_task(async_dispatcher.as_ptr(), state_ptr) };
203        if Status::from_raw(status) == Status::OK {
204            // SAFETY: If the cancellation was successful, we know the callback won't be called
205            // so we need to deallocate the copy of the arc that was given to it.
206            unsafe { SharedState::release_raw_ptr(state_ptr) };
207        }
208    }
209}
210
211// TODO(528052543): Migrate these to a specifically test-oriented dispatcher when there is one
212// so they don't require the driver runtime to be involved.
213#[cfg(test)]
214mod tests {
215    use std::sync::mpsc;
216    use std::thread::sleep;
217    use std::time::Duration;
218
219    use super::*;
220
221    use futures::{FutureExt, poll};
222    use std::task::Waker;
223
224    use fdf_env::test::spawn_in_driver;
225    use libasync_dispatcher::CurrentDispatcher;
226
227    fn now() -> zx::MonotonicInstant {
228        zx::MonotonicInstant::from_nanos(CurrentDispatcher.get_async_dispatcher().now())
229    }
230
231    #[test]
232    fn after_the_past() {
233        spawn_in_driver("testing task", async {
234            let fut = CurrentDispatcher.after_deadline(zx::MonotonicInstant::INFINITE_PAST);
235            assert_eq!(poll!(fut), Poll::Ready(Ok(())));
236        });
237    }
238
239    #[test]
240    fn after_now() {
241        spawn_in_driver("testing task", async {
242            let fut = CurrentDispatcher.after_deadline(now());
243            assert_eq!(poll!(fut), Poll::Ready(Ok(())));
244        });
245    }
246
247    #[test]
248    fn after_future() {
249        spawn_in_driver("testing task", async {
250            let deadline = now() + zx::MonotonicDuration::from_seconds(3);
251            let mut fut = CurrentDispatcher.after_deadline(deadline);
252            assert_eq!(poll!(&mut fut), Poll::Pending);
253            assert!(fut.await.is_ok());
254            assert!(now() >= deadline);
255        });
256    }
257
258    #[test]
259    fn drop_after_poll() {
260        spawn_in_driver("testing task", async {
261            let deadline = now() + zx::MonotonicDuration::from_minutes(3);
262            let mut fut = CurrentDispatcher.after_deadline(deadline);
263            assert_eq!(poll!(&mut fut), Poll::Pending);
264        });
265    }
266
267    #[test]
268    fn dispatcher_shutdown_cancel() {
269        let (fut_tx, fut_rx) = mpsc::channel();
270        spawn_in_driver("testing task", async move {
271            let deadline = now() + zx::MonotonicDuration::from_minutes(3);
272            let mut fut = CurrentDispatcher.after_deadline(deadline);
273            assert_eq!(poll!(&mut fut), Poll::Pending);
274            fut_tx.send(fut).unwrap();
275        });
276        let mut fut = fut_rx.recv().unwrap();
277        loop {
278            let Poll::Ready(res) = fut.poll_unpin(&mut Context::from_waker(Waker::noop())) else {
279                sleep(Duration::from_millis(10));
280                continue;
281            };
282            assert_eq!(res, Err(Status::CANCELED));
283            break;
284        }
285    }
286}
libasync/after_deadline.rs

libasync/
after_deadline.rs
