fuchsia_async/runtime/fuchsia/executor/atomic_future/
hooks.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 crate::instrument::Hooks;
6
7use super::{AtomicFutureHandle, Meta, VTable};
8use fuchsia_sync::Mutex;
9use std::collections::HashMap;
10use std::ptr::NonNull;
11use std::task::{Context, Poll};
12
13/// We don't want to pay a cost if there are no hooks, so we store a mapping from task ID to
14/// HooksWrapper in the executor.
15#[derive(Default)]
16pub struct HooksMap(Mutex<HashMap<usize, NonNull<()>>>);
17
18unsafe impl Send for HooksMap {}
19unsafe impl Sync for HooksMap {}
20
21struct HooksWrapper<H> {
22    orig_vtable: &'static VTable,
23    hooks: H,
24}
25
26impl<H: Hooks> HooksWrapper<H> {
27    // # Safety
28    //
29    // We rely on the fact that all these functions are called whilst we have exclusive
30    // access to the underlying future and the associated Hooks object.
31    const VTABLE: VTable = VTable {
32        drop: Self::drop,
33        drop_future: Self::drop_future,
34        poll: Self::poll,
35        get_result: Self::get_result,
36        drop_result: Self::drop_result,
37    };
38
39    // Returns a mutable reference to the wrapper. This will be safe from the functions below
40    // because they are all called when we have exclusive access.
41    unsafe fn wrapper<'a>(meta: NonNull<Meta>) -> &'a mut Self {
42        let meta = meta.as_ref();
43        meta.scope().executor().hooks_map.0.lock().get(&meta.id).unwrap().cast::<Self>().as_mut()
44    }
45
46    unsafe fn drop(mut meta: NonNull<Meta>) {
47        let meta_ref = meta.as_mut();
48        // Remove the hooks entry from the map.
49        let hooks = Box::from_raw(
50            meta_ref
51                .scope()
52                .executor()
53                .hooks_map
54                .0
55                .lock()
56                .remove(&meta_ref.id)
57                .unwrap()
58                .cast::<Self>()
59                .as_mut(),
60        );
61        // Restore the vtable because the drop implementation can call `drop_future` or
62        // `drop_result`, but we've removed the hooks from the map now.
63        meta_ref.vtable = hooks.orig_vtable;
64        (hooks.orig_vtable.drop)(meta);
65    }
66
67    unsafe fn poll(meta: NonNull<Meta>, cx: &mut Context<'_>) -> Poll<()> {
68        let wrapper = Self::wrapper(meta);
69        wrapper.hooks.task_poll_start();
70        let result = (wrapper.orig_vtable.poll)(meta, cx);
71        wrapper.hooks.task_poll_end();
72        if result.is_ready() {
73            wrapper.hooks.task_completed();
74        }
75        result
76    }
77
78    unsafe fn drop_future(meta: NonNull<Meta>) {
79        (Self::wrapper(meta).orig_vtable.drop_future)(meta);
80    }
81
82    unsafe fn get_result(meta: NonNull<Meta>) -> *const () {
83        (Self::wrapper(meta).orig_vtable.get_result)(meta)
84    }
85
86    unsafe fn drop_result(meta: NonNull<Meta>) {
87        (Self::wrapper(meta).orig_vtable.drop_result)(meta);
88    }
89}
90
91impl AtomicFutureHandle<'_> {
92    /// Adds hooks to the future.
93    pub fn add_hooks<H: Hooks>(&mut self, hooks: H) {
94        // SAFETY: This is safe because we have exclusive access.
95        let meta: &mut Meta = unsafe { self.0.as_mut() };
96        {
97            let mut hooks_map = meta.scope().executor().hooks_map.0.lock();
98            // SAFETY: Safe because `Box::into_raw` is guaranteed to give is a non-null pointer. We
99            // can use `Box::into_non_null` when it's stabilised.
100            assert!(hooks_map
101                .insert(meta.id, unsafe {
102                    NonNull::new_unchecked(Box::into_raw(Box::new(HooksWrapper {
103                        orig_vtable: meta.vtable,
104                        hooks,
105                    })))
106                    .cast::<()>()
107                })
108                .is_none());
109        }
110        // Inject our vtable.
111        meta.vtable = &HooksWrapper::<H>::VTABLE;
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::Hooks;
118    use crate::runtime::fuchsia::executor::scope::Spawnable;
119    use crate::{yield_now, SpawnableFuture, TestExecutor};
120    use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
121    use std::sync::Arc;
122
123    #[test]
124    fn test_hooks() {
125        let mut executor = TestExecutor::new();
126        let scope = executor.global_scope();
127        let mut future = SpawnableFuture::new(async {
128            yield_now().await;
129        })
130        .into_task(scope.clone());
131        #[derive(Default)]
132        struct MyHooks {
133            poll_start: AtomicU32,
134            poll_end: AtomicU32,
135            completed: AtomicBool,
136        }
137        impl Hooks for Arc<MyHooks> {
138            fn task_completed(&mut self) {
139                assert!(!self.completed.load(Ordering::Relaxed));
140                self.completed.store(true, Ordering::Relaxed);
141            }
142            fn task_poll_start(&mut self) {
143                self.poll_start.fetch_add(1, Ordering::Relaxed);
144            }
145            fn task_poll_end(&mut self) {
146                self.poll_end.fetch_add(1, Ordering::Relaxed);
147            }
148        }
149        let my_hooks = Arc::new(MyHooks::default());
150        future.add_hooks(my_hooks.clone());
151        scope.insert_task(future, false);
152        assert!(executor.run_until_stalled(&mut std::future::pending::<()>()).is_pending());
153        assert_eq!(my_hooks.poll_start.load(Ordering::Relaxed), 2);
154        assert_eq!(my_hooks.poll_end.load(Ordering::Relaxed), 2);
155        assert!(my_hooks.completed.load(Ordering::Relaxed));
156    }
157}
fuchsia_async/runtime/fuchsia/executor/atomic_future/hooks.rs

fuchsia_async/runtime/fuchsia/executor/atomic_future/
hooks.rs
