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 id = meta.as_ptr() as usize;
43        unsafe {
44            meta.as_ref()
45                .scope()
46                .executor()
47                .hooks_map
48                .0
49                .lock()
50                .get(&id)
51                .unwrap()
52                .cast::<Self>()
53                .as_mut()
54        }
55    }
56
57    unsafe fn drop(mut meta: NonNull<Meta>) {
58        let meta_ref = unsafe { meta.as_mut() };
59        // Remove the hooks entry from the map.
60        let id = meta.as_ptr() as usize;
61        let hooks = unsafe {
62            Box::from_raw(
63                meta_ref
64                    .scope()
65                    .executor()
66                    .hooks_map
67                    .0
68                    .lock()
69                    .remove(&id)
70                    .unwrap()
71                    .cast::<Self>()
72                    .as_mut(),
73            )
74        };
75        // Restore the vtable because the drop implementation can call `drop_future` or
76        // `drop_result`, but we've removed the hooks from the map now.
77        meta_ref.vtable = hooks.orig_vtable;
78        unsafe { (hooks.orig_vtable.drop)(meta) };
79    }
80
81    unsafe fn poll(meta: NonNull<Meta>, cx: &mut Context<'_>) -> Poll<()> {
82        let wrapper = unsafe { Self::wrapper(meta) };
83        wrapper.hooks.task_poll_start();
84        let result = unsafe { (wrapper.orig_vtable.poll)(meta, cx) };
85        wrapper.hooks.task_poll_end();
86        if result.is_ready() {
87            wrapper.hooks.task_completed();
88        }
89        result
90    }
91
92    unsafe fn drop_future(meta: NonNull<Meta>) {
93        unsafe { (Self::wrapper(meta).orig_vtable.drop_future)(meta) };
94    }
95
96    unsafe fn get_result(meta: NonNull<Meta>) -> *const () {
97        unsafe { (Self::wrapper(meta).orig_vtable.get_result)(meta) }
98    }
99
100    unsafe fn drop_result(meta: NonNull<Meta>) {
101        unsafe { (Self::wrapper(meta).orig_vtable.drop_result)(meta) };
102    }
103}
104
105impl AtomicFutureHandle<'_> {
106    /// Adds hooks to the future.
107    pub fn add_hooks<H: Hooks>(&mut self, hooks: H) {
108        let id = self.id();
109        // SAFETY: This is safe because we have exclusive access.
110        let meta: &mut Meta = unsafe { self.0.as_mut() };
111        {
112            let mut hooks_map = meta.scope().executor().hooks_map.0.lock();
113            // SAFETY: Safe because `Box::into_raw` is guaranteed to give is a non-null pointer. We
114            // can use `Box::into_non_null` when it's stabilised.
115            assert!(
116                hooks_map
117                    .insert(
118                        id,
119                        NonNull::from_mut(Box::leak(Box::new(HooksWrapper {
120                            orig_vtable: meta.vtable,
121                            hooks,
122                        })))
123                        .cast::<()>()
124                    )
125                    .is_none()
126            );
127        }
128        // Inject our vtable.
129        meta.vtable = &HooksWrapper::<H>::VTABLE;
130    }
131}
132
133#[cfg(test)]
134mod tests {
135    use super::Hooks;
136    use crate::runtime::fuchsia::executor::scope::Spawnable;
137    use crate::{SpawnableFuture, TestExecutor, yield_now};
138    use std::sync::Arc;
139    use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
140
141    #[test]
142    fn test_hooks() {
143        let mut executor = TestExecutor::new();
144        let scope = executor.global_scope();
145        let mut future = SpawnableFuture::new(async {
146            yield_now().await;
147        })
148        .into_task(scope.clone());
149        #[derive(Default)]
150        struct MyHooks {
151            poll_start: AtomicU32,
152            poll_end: AtomicU32,
153            completed: AtomicBool,
154        }
155        impl Hooks for Arc<MyHooks> {
156            fn task_completed(&mut self) {
157                assert!(!self.completed.load(Ordering::Relaxed));
158                self.completed.store(true, Ordering::Relaxed);
159            }
160            fn task_poll_start(&mut self) {
161                self.poll_start.fetch_add(1, Ordering::Relaxed);
162            }
163            fn task_poll_end(&mut self) {
164                self.poll_end.fetch_add(1, Ordering::Relaxed);
165            }
166        }
167        let my_hooks = Arc::new(MyHooks::default());
168        future.add_hooks(my_hooks.clone());
169        scope.insert_task(future, false);
170        assert!(executor.run_until_stalled(&mut std::future::pending::<()>()).is_pending());
171        assert_eq!(my_hooks.poll_start.load(Ordering::Relaxed), 2);
172        assert_eq!(my_hooks.poll_end.load(Ordering::Relaxed), 2);
173        assert!(my_hooks.completed.load(Ordering::Relaxed));
174    }
175}
fuchsia_async/runtime/fuchsia/executor/atomic_future/hooks.rs

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