libasync/
on_signals.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 crate::callback_state::CallbackSharedState;
6use core::fmt;
7use core::future::Future;
8use core::pin::Pin;
9use core::ptr::NonNull;
10use core::sync::atomic::{AtomicI32, AtomicU32, Ordering};
11use core::task::{Context, Poll};
12use futures::task::AtomicWaker;
13use libasync_dispatcher::{DetectDispatcher, GetAsyncDispatcher};
14use libasync_sys::{async_begin_wait, async_cancel_wait, async_dispatcher, async_wait};
15use std::sync::Arc;
16use zx::Status;
17use zx::sys::{ZX_ERR_CANCELED, ZX_OK};
18
19/// Internal state managed for a pending wait callback.
20struct WaitState {
21    async_dispatcher: NonNull<async_dispatcher>,
22    waker: AtomicWaker,
23    status: AtomicI32,
24    observed: AtomicU32,
25}
26
27// SAFETY: It is safe to access an `async_dispatcher_t` from any thread per the libasync C api.
28unsafe impl Send for WaitState {}
29unsafe impl Sync for WaitState {}
30
31type SharedState = CallbackSharedState<async_wait, WaitState>;
32
33impl WaitState {
34    extern "C" fn call(
35        dispatcher: *mut async_dispatcher,
36        wait: *mut async_wait,
37        status: i32,
38        signal: *const zx::sys::zx_packet_signal_t,
39    ) {
40        debug_assert!(
41            status == ZX_OK || status == ZX_ERR_CANCELED,
42            "wait callback called with status other than ok or canceled: {}",
43            status
44        );
45
46        // SAFETY: We called make_raw_ptr when we began the wait.
47        let state = unsafe { SharedState::from_raw_ptr(wait) };
48
49        debug_assert!(
50            dispatcher == state.async_dispatcher.as_ptr(),
51            "dispatcher does not match wait's dispatcher"
52        );
53
54        if status == ZX_OK {
55            debug_assert!(!signal.is_null(), "signal must not be null when status is ZX_OK");
56            // SAFETY: signal is not null and valid when status is ZX_OK, as per the contract of
57            // async_wait_handler_t.
58            let observed_signals = unsafe { (*signal).observed };
59            // Store `observed` first with Relaxed ordering because it is guarded by the
60            // upcoming Release store on `status`.
61            state.observed.store(observed_signals, Ordering::Relaxed);
62        }
63
64        // Store `status` using Release ordering, since we already stored `observed` first, loading
65        // `status` in poll and cancel_wait with Acquire ensures `observed` updates are also seen
66        // by the other thread.
67        state.status.store(status, Ordering::Release);
68
69        state.waker.wake();
70    }
71}
72
73/// Implements methods for setting and waiting on signals on a dispatcher.
74pub trait DispatcherSignalExt {
75    /// Returns a future that will fire when the given signals are asserted on the handle.
76    fn on_signals<H: zx::AsHandleRef>(&self, handle: H, signals: zx::Signals) -> OnSignals<H>;
77
78    /// Returns a future that will fire when the given signals are asserted on the handle. Returns
79    /// `None` if there is no dispatcher found on `self`.
80    fn try_on_signals<H: zx::AsHandleRef>(
81        &self,
82        handle: H,
83        signals: zx::Signals,
84    ) -> Option<OnSignals<H>>;
85}
86
87impl<T> DispatcherSignalExt for T
88where
89    T: GetAsyncDispatcher,
90{
91    fn on_signals<H: zx::AsHandleRef>(&self, handle: H, signals: zx::Signals) -> OnSignals<H> {
92        self.try_on_signals(handle, signals).expect("No current dispatcher")
93    }
94
95    fn try_on_signals<H: zx::AsHandleRef>(
96        &self,
97        handle: H,
98        signals: zx::Signals,
99    ) -> Option<OnSignals<H>> {
100        let dispatcher = self.try_get_async_dispatcher()?;
101        Some(OnSignals::new_on(dispatcher, handle, signals))
102    }
103}
104
105/// A future that completes when specified signals are asserted on a Zircon handle.
106pub struct OnSignals<H: zx::AsHandleRef> {
107    dispatcher: DetectDispatcher,
108    handle: Option<H>,
109    signals: zx::Signals,
110    state: Option<Arc<SharedState>>,
111}
112
113impl<H: zx::AsHandleRef> OnSignals<H> {
114    /// Creates a new `OnSignals` object that will receive notifications when signals occur on
115    /// `handle`.
116    pub fn new(handle: H, signals: zx::Signals) -> Self {
117        Self { dispatcher: DetectDispatcher::default(), handle: Some(handle), signals, state: None }
118    }
119
120    /// Creates a new `OnSignals` object bound to a specific dispatcher.
121    pub fn new_on(dispatcher: impl GetAsyncDispatcher, handle: H, signals: zx::Signals) -> Self {
122        let async_disp = dispatcher.get_async_dispatcher();
123        Self {
124            dispatcher: DetectDispatcher::new(async_disp),
125            handle: Some(handle),
126            signals,
127            state: None,
128        }
129    }
130
131    /// Takes the handle back, cancelling any active wait.
132    pub fn take_handle(&mut self) -> Option<H> {
133        self.cancel_wait();
134        self.handle.take()
135    }
136
137    fn cancel_wait(&mut self) {
138        let Some(state) = self.state.take() else {
139            return;
140        };
141
142        // Load `status` using Acquire ordering to synchronize with `WaitState::call`'s Release
143        // store.
144        if state.status.load(Ordering::Acquire) != Status::SHOULD_WAIT.into_raw() {
145            // Callback has already completed; nothing to cancel.
146            return;
147        }
148
149        let state_ptr = SharedState::as_raw_ptr(&state);
150        // SAFETY: async_cancel_wait is thread-safe per the C API doc.
151        let status = unsafe { async_cancel_wait(state.async_dispatcher.as_ptr(), state_ptr) };
152        if Status::from_raw(status) == Status::OK {
153            // SAFETY: Cancellation succeeded. The callback will not run, so we have to release the
154            // raw pointer here.
155            unsafe { SharedState::release_raw_ptr(state_ptr) };
156        }
157    }
158}
159
160impl<H: zx::AsHandleRef + Unpin> Future for OnSignals<H> {
161    type Output = Result<zx::Signals, Status>;
162
163    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
164        let Some(handle) = self.handle.as_ref() else {
165            return Poll::Ready(Err(Status::BAD_STATE));
166        };
167
168        if let Some(state) = &self.state {
169            // Avoid a lost wake by registering the waker first. This prevents the case that the
170            // callback running in parallel, stores `status` and wakes the previous waker,
171            // after our check but before we register the waker.
172            state.waker.register(cx.waker());
173
174            // Use `Acquire` ordering to ensure that the store in `WaitState::call` is visible
175            // here.
176            let status = state.status.load(Ordering::Acquire);
177            if status != Status::SHOULD_WAIT.into_raw() {
178                if status == ZX_OK {
179                    // We can use `Relaxed` ordering here since we used `Acquire` above, and on the
180                    // store side we store in reverse.
181                    let observed = state.observed.load(Ordering::Relaxed);
182                    return Poll::Ready(Ok(zx::Signals::from_bits_truncate(observed)));
183                } else {
184                    return Poll::Ready(Err(Status::from_raw(status)));
185                }
186            } else {
187                // Optimization: use a non-blocking wait to check if its ready. Alternatively we
188                // could just wait for the callback but this way we don't have to wait for the
189                // two thread hops.
190                match handle
191                    .as_handle_ref()
192                    .wait_one(self.signals, zx::MonotonicInstant::INFINITE_PAST)
193                    .to_result()
194                {
195                    Ok(signals) => {
196                        self.cancel_wait();
197                        return Poll::Ready(Ok(signals));
198                    }
199                    Err(Status::TIMED_OUT) => {
200                        return Poll::Pending;
201                    }
202                    Err(err) => {
203                        self.cancel_wait();
204                        return Poll::Ready(Err(err));
205                    }
206                }
207            }
208        }
209
210        // First poll. Check if its already signaled, if it is we can bypass the async wait setup.
211        match handle
212            .as_handle_ref()
213            .wait_one(self.signals, zx::MonotonicInstant::INFINITE_PAST)
214            .to_result()
215        {
216            Ok(signals) => return Poll::Ready(Ok(signals)),
217            Err(Status::TIMED_OUT) => { /* Signal not yet asserted; proceed to register */ }
218            Err(err) => return Poll::Ready(Err(err)),
219        }
220
221        let dispatcher = self.dispatcher.get_or_detect()?;
222        let async_dispatcher = dispatcher.as_ptr();
223
224        let wait = async_wait {
225            handler: Some(WaitState::call),
226            object: handle.as_handle_ref().raw_handle(),
227            trigger: self.signals.bits(),
228            options: 0,
229            ..Default::default()
230        };
231
232        let state = WaitState {
233            async_dispatcher,
234            waker: AtomicWaker::new(),
235            status: AtomicI32::new(Status::SHOULD_WAIT.into_raw()),
236            observed: AtomicU32::new(0),
237        };
238
239        let shared_state = SharedState::new(wait, state);
240        shared_state.waker.register(cx.waker());
241
242        let state_ptr = SharedState::make_raw_ptr(shared_state.clone());
243
244        // SAFETY: async_begin_wait is thread safe per the C API doc.
245        let res = Status::ok(unsafe { async_begin_wait(async_dispatcher.as_ptr(), state_ptr) });
246        match res {
247            Ok(_) => {
248                self.state = Some(shared_state);
249                Poll::Pending
250            }
251            Err(err) => {
252                // SAFETY: The C callback will not be called on the error case, so we must release
253                // the raw pointer here.
254                unsafe { SharedState::release_raw_ptr(state_ptr) };
255                Poll::Ready(Err(err))
256            }
257        }
258    }
259}
260
261impl<H: zx::AsHandleRef> Drop for OnSignals<H> {
262    fn drop(&mut self) {
263        self.cancel_wait();
264    }
265}
266
267impl<H: zx::AsHandleRef> zx::AsHandleRef for OnSignals<H> {
268    fn as_handle_ref(&self) -> zx::HandleRef<'_> {
269        self.handle.as_ref().expect("OnSignals dereferenced after handle was taken").as_handle_ref()
270    }
271}
272
273impl<H: zx::AsHandleRef> fmt::Debug for OnSignals<H> {
274    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
275        write!(
276            f,
277            "OnSignals {{ handle: {:?}, signals: {:?} }}",
278            self.handle.as_ref().map(|h| h.as_handle_ref()),
279            self.signals
280        )
281    }
282}
283
284/// Alias for the common case where `OnSignals` is used with `zx::HandleRef`.
285pub type OnSignalsRef<'a> = OnSignals<zx::HandleRef<'a>>;
286
287#[cfg(test)]
288mod tests {
289    use super::*;
290    use fdf_env::test::spawn_in_driver;
291    use futures::{FutureExt, poll};
292    use libasync_dispatcher::CurrentDispatcher;
293    use std::sync::{LazyLock, mpsc};
294    use std::task::Waker;
295    use std::thread::sleep;
296    use std::time::Duration;
297
298    #[test]
299    fn wait_on_signals() {
300        spawn_in_driver("testing wait", async move {
301            let event = zx::Event::create();
302            let mut fut = CurrentDispatcher.on_signals(&event, zx::Signals::EVENT_SIGNALED);
303            assert_eq!(poll!(&mut fut), Poll::Pending);
304
305            event.signal(zx::Signals::empty(), zx::Signals::EVENT_SIGNALED).unwrap();
306
307            let res = fut.await;
308            assert!(res.is_ok());
309            assert!(res.unwrap().contains(zx::Signals::EVENT_SIGNALED));
310        });
311    }
312
313    #[test]
314    fn wait_on_already_signaled() {
315        spawn_in_driver("testing wait", async move {
316            let event = zx::Event::create();
317            event.signal(zx::Signals::empty(), zx::Signals::EVENT_SIGNALED).unwrap();
318
319            let fut = CurrentDispatcher.on_signals(&event, zx::Signals::EVENT_SIGNALED);
320
321            let res = fut.await;
322            assert!(res.is_ok());
323            assert!(res.unwrap().contains(zx::Signals::EVENT_SIGNALED));
324        });
325    }
326
327    #[test]
328    fn drop_after_poll() {
329        spawn_in_driver("testing wait", async move {
330            let event = zx::Event::create();
331            let mut fut = CurrentDispatcher.on_signals(&event, zx::Signals::EVENT_SIGNALED);
332            assert_eq!(poll!(&mut fut), Poll::Pending);
333        });
334    }
335
336    #[test]
337    fn dispatcher_shutdown_cancel() {
338        static EVENT: LazyLock<zx::Event> = LazyLock::new(zx::Event::create);
339
340        let (fut_tx, fut_rx) = mpsc::channel();
341        spawn_in_driver("testing wait", async move {
342            let mut fut = CurrentDispatcher.on_signals(&*EVENT, zx::Signals::EVENT_SIGNALED);
343            assert_eq!(poll!(&mut fut), Poll::Pending);
344            fut_tx.send(fut).unwrap();
345        });
346        let mut fut = fut_rx.recv().unwrap();
347        let waker = Waker::noop();
348        let mut context = Context::from_waker(waker);
349        loop {
350            let Poll::Ready(res) = fut.poll_unpin(&mut context) else {
351                sleep(Duration::from_millis(10));
352                continue;
353            };
354            assert_eq!(res, Err(Status::CANCELED));
355            break;
356        }
357    }
358
359    #[test]
360    fn test_take_handle() {
361        spawn_in_driver("testing wait", async move {
362            let event = zx::Event::create();
363            let mut fut = CurrentDispatcher.on_signals(event, zx::Signals::EVENT_SIGNALED);
364            assert_eq!(poll!(&mut fut), Poll::Pending);
365            let event = fut.take_handle().unwrap();
366            event.signal(zx::Signals::empty(), zx::Signals::EVENT_SIGNALED).unwrap();
367
368            assert!(fut.take_handle().is_none());
369            assert_eq!(poll!(&mut fut), Poll::Ready(Err(Status::BAD_STATE)));
370            assert_eq!(
371                event
372                    .wait_one(zx::Signals::EVENT_SIGNALED, zx::MonotonicInstant::INFINITE_PAST)
373                    .unwrap(),
374                zx::Signals::EVENT_SIGNALED
375            );
376        });
377    }
378}
libasync/on_signals.rs

libasync/
on_signals.rs
