Skip to main content

libasync_scope_dispatcher/
scope_dispatcher.rs

1// Copyright 2026 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 core::ptr::{NonNull, null};
6use core::sync::atomic::{self, AtomicBool};
7use fuchsia_async::{EHandle, MonotonicInstant, Scope, Timer, WakeupTime};
8use fuchsia_sync::Mutex;
9use futures::task::AtomicWaker;
10use libasync_dispatcher::{AsAsyncDispatcherRef, AsyncDispatcherRef};
11use libasync_sys::async_dispatcher_t;
12use pin_project_lite::pin_project;
13use std::pin::Pin;
14use std::sync::Arc;
15use std::task::{Context, Poll};
16use zx::Status;
17
18use crate::ops;
19use crate::ops::v1::{PendingWaits, Task, TaskQueue};
20
21/// Implements a C++-compatible [`async_dispatcher_t`] around a [`fuchsia_async::Scope`].
22#[derive(Debug)]
23#[repr(C)]
24pub struct ScopeDispatcher {
25    // Safety Note: this must go first in this struct for the callbacks to work correctly.
26    dispatcher: async_dispatcher_t,
27    task_queue: Mutex<TaskQueue>,
28    pub(crate) pending_waits: Mutex<PendingWaits>,
29    shutting_down: AtomicBool,
30    shutdown_complete_waker: AtomicWaker,
31    service_waker: AtomicWaker,
32    shutdown_guard: AtomicBool,
33    executor: EHandle,
34    scope: Scope,
35}
36
37// SAFETY: The API of async_dispatcher_t is expected to be thread safe.
38unsafe impl Send for ScopeDispatcher {}
39// SAFETY: The API of async_dispatcher_t is expected to be thread safe.
40unsafe impl Sync for ScopeDispatcher {}
41
42impl ScopeDispatcher {
43    /// Creates a new [`ScopeDispatcher`] on the currently running fuchsia-async executor.
44    ///
45    /// # Panics
46    ///
47    /// Panics if this is not run on a fuchsia-async executor context.
48    pub fn new() -> Arc<Self> {
49        Self::new_on_executor(EHandle::local())
50    }
51
52    /// Creates a new [`ScopeDispatcher`] with a new [`Scope`] on the given `executor`.
53    pub fn new_on_executor(executor: EHandle) -> Arc<Self> {
54        let scope = executor.global_scope().new_child();
55        let scope_handle = scope.as_handle().clone();
56        let dispatcher = async_dispatcher_t { ops: &ops::ASYNC_OPS };
57        let task_queue = Mutex::new(TaskQueue::default());
58        let pending_waits = Mutex::new(PendingWaits::default());
59        let shutting_down = AtomicBool::new(false);
60        let service_waker = AtomicWaker::new();
61        let shutdown_complete_waker = AtomicWaker::new();
62        let shutdown_guard = AtomicBool::new(false);
63        let this = Arc::new(Self {
64            dispatcher,
65            task_queue,
66            pending_waits,
67            shutting_down,
68            shutdown_complete_waker,
69            service_waker,
70            shutdown_guard,
71            executor,
72            scope,
73        });
74
75        scope_handle.spawn(this.clone().service_loop());
76
77        this
78    }
79
80    /// Get the pointer to the dispatcher callback struct for passing through FFI layers.
81    pub fn as_ptr(&self) -> *const async_dispatcher_t {
82        (self as *const Self).cast()
83    }
84
85    /// Gets the global executor handle for this dispatcher.
86    pub fn global_handle(&self) -> &EHandle {
87        &self.executor
88    }
89
90    /// Gets the [`fuchsia_async::Scope`] of this dispatcher.
91    pub fn as_scope(&self) -> &Scope {
92        &self.scope
93    }
94
95    /// Returns true if the dispatcher is currently shutting down.
96    pub fn is_shutting_down(&self) -> bool {
97        self.shutting_down.load(atomic::Ordering::Acquire)
98    }
99
100    /// Starts the dispatcher shutdown. Resolves when all outstanding tasks have been completed or
101    /// canceled.
102    pub fn shutdown(&self) -> ShutdownCompletionFuture<'_> {
103        // note: we might want to do more to prevent multiple attempts to shut the dispatcher down.
104        self.shutting_down.store(true, atomic::Ordering::Release);
105        self.service_waker.wake();
106        ShutdownCompletionFuture(self)
107    }
108
109    /// Gets the Scope from a dispatcher pointer. Used in the callbacks.
110    ///
111    /// # Safety
112    ///
113    /// The caller must ensure that the dispatcher pointer is a valid pointer originally obtained
114    /// through [`ScopeDispatcher::as_ptr`], and is still alive.
115    pub(crate) unsafe fn from_ptr<'a>(dispatcher_ptr: *mut async_dispatcher_t) -> &'a Self {
116        let this = dispatcher_ptr.cast::<ScopeDispatcher>();
117        // Safety: the caller promises that this is a valid pointer to what was originally a
118        // ScopeDispatcher object.
119        unsafe { this.as_ref() }.expect("null dispatcher pointer")
120    }
121
122    /// Gets the Scope from a dispatcher pointer. Used in the callbacks.
123    ///
124    /// # Safety
125    ///
126    /// The caller must ensure that the dispatcher pointer is a valid pointer originally obtained
127    /// through [`ScopeDispatcher::as_ptr`].
128    pub(crate) unsafe fn arc_from_ptr(dispatcher_ptr: *mut async_dispatcher_t) -> Arc<Self> {
129        let this = dispatcher_ptr.cast::<ScopeDispatcher>();
130        // Safety: the caller promises that this is a valid pointer to what was originally a
131        // ScopeDispatcher object.
132        unsafe {
133            Arc::increment_strong_count(this);
134            Arc::from_raw(this)
135        }
136    }
137
138    /// Posts a task to the dispatcher
139    pub(crate) fn post_task(&self, task: Task) -> Result<(), Status> {
140        // don't queue new tasks if we're shutting down.
141        if self.is_shutting_down() {
142            return Err(Status::BAD_STATE);
143        }
144        self.task_queue.lock().queue_task(task);
145        self.service_waker.wake();
146        Ok(())
147    }
148
149    /// Cancels a task queued on the dispatcher
150    pub(crate) fn cancel_task(&self, task: Task) -> Result<(), Status> {
151        // If we succeed at cancelling, we won't call the callback so this can be fairly simple.
152        if self.task_queue.lock().take_pending_task(&task).is_some() {
153            self.service_waker.wake();
154            Ok(())
155        } else {
156            Err(Status::NOT_FOUND)
157        }
158    }
159
160    async fn service_loop(self: Arc<Self>) {
161        while let Some(next_task) = NextTaskFuture::new(&self).await {
162            next_task.run(self.clone(), Ok(()));
163        }
164        // we're shutting down, so drain the queues of all outstanding tasks with
165        // a status of CANCELED. Note that we don't really care about fanning these out to all
166        // threads, so we just run them directly here.
167        // Note also that we will not allow any new items to be added to the queues after the
168        // shutdown flag has been set, so we don't have to worry about new things being added at
169        // this point.
170        for next_wait in self.pending_waits.lock().get_all_waits() {
171            next_wait.run(self.clone(), null(), Err(Status::CANCELED));
172        }
173        while let Some(next_task) = self.task_queue.lock().next_task(MonotonicInstant::INFINITE) {
174            next_task.run(self.clone(), Err(Status::CANCELED));
175        }
176        self.shutdown_guard.store(true, atomic::Ordering::Release);
177        self.shutdown_complete_waker.wake();
178    }
179}
180
181impl AsAsyncDispatcherRef for ScopeDispatcher {
182    fn as_async_dispatcher_ref(&self) -> AsyncDispatcherRef<'_> {
183        // SAFETY: We know this pointer is valid because it is a member of `&self`, which is a valid
184        // reference.
185        let ptr = unsafe { NonNull::new_unchecked(self.as_ptr().cast_mut()) };
186        // SAFETY: The dispatcher ref's lifetime is tied to `self`, of which the dispatcher
187        // structure and callbacks are members, so will not outlive them.
188        unsafe { AsyncDispatcherRef::from_raw(ptr) }
189    }
190}
191
192impl Drop for ScopeDispatcher {
193    fn drop(&mut self) {
194        assert!(
195            self.shutdown_guard.load(atomic::Ordering::Acquire),
196            "Dispatcher not properly shut down before dropping. Call ScopeDispatcher::shutdown()."
197        );
198    }
199}
200
201/// A future which resolves when the dispatcher has been shut down by [`ScopeDispatcher::shutdown`].
202#[must_use = "a future that is never awaited on will never run"]
203pub struct ShutdownCompletionFuture<'a>(&'a ScopeDispatcher);
204
205impl<'a> Future for ShutdownCompletionFuture<'a> {
206    type Output = ();
207
208    fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
209        self.0.shutdown_complete_waker.register(ctx.waker());
210        if self.0.shutdown_guard.load(atomic::Ordering::Acquire) {
211            Poll::Ready(())
212        } else {
213            Poll::Pending
214        }
215    }
216}
217
218pin_project! {
219    #[must_use = "a future that is never awaited on will never run"]
220    struct NextTaskFuture<'a> {
221        dispatcher: &'a Arc<ScopeDispatcher>,
222        #[pin]
223        next_timeout: Timer,
224    }
225}
226
227impl<'a> NextTaskFuture<'a> {
228    fn new(dispatcher: &'a Arc<ScopeDispatcher>) -> Self {
229        let next_timeout = MonotonicInstant::INFINITE.into_timer();
230        Self { dispatcher, next_timeout }
231    }
232}
233
234impl<'a> Future for NextTaskFuture<'a> {
235    type Output = Option<Task>;
236
237    fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
238        let mut task_queue = self.dispatcher.task_queue.lock();
239        // if we are shutting down, the service handler will do the work of canceling the remaining
240        // tasks, so return None to indicate that it should start doing that.
241        if self.dispatcher.shutting_down.load(atomic::Ordering::Acquire) {
242            return Poll::Ready(None);
243        }
244        let now = self.dispatcher.executor.now();
245        if let Some(task) = task_queue.next_task(now) {
246            Poll::Ready(Some(task))
247        } else {
248            let next_deadline = if let Some(task) = task_queue.peek_next_task() {
249                task.deadline().unwrap_or(MonotonicInstant::INFINITE)
250            } else {
251                MonotonicInstant::INFINITE
252            };
253            self.dispatcher.service_waker.register(ctx.waker());
254            let mut this = self.project();
255            this.next_timeout.as_mut().reset(next_deadline);
256            // Note that we don't really care about resolving the timer, we just want to use its
257            // waker to re-awaken this future when we're ready.
258            let _: Poll<()> = this.next_timeout.poll(ctx);
259            Poll::Pending
260        }
261    }
262}