ntex_rt/
task.rs

1use std::{sync::Arc, sync::atomic::AtomicUsize, sync::atomic::Ordering};
2
3// The Callbacks static holds a pointer to the global callbacks. It is protected by
4// the STATE static which determines whether `CBS` has been initialized yet.
5static mut CBS: Option<Arc<dyn CallbacksApi>> = None;
6
7static STATE: AtomicUsize = AtomicUsize::new(0);
8
9// There are three different states that we care about: the callback's
10// uninitialized, the callback's initializing (set_cbs's been called but
11// CBS hasn't actually been set yet), or the callbacks's active.
12const UNINITIALIZED: usize = 0;
13const INITIALIZING: usize = 1;
14const INITIALIZED: usize = 2;
15
16trait CallbacksApi {
17    fn before(&self) -> Option<*const ()>;
18    fn enter(&self, _: *const ()) -> *const ();
19    fn exit(&self, _: *const ());
20    fn after(&self, _: *const ());
21}
22
23#[allow(clippy::struct_field_names)]
24struct Callbacks<A, B, C, D> {
25    f_before: A,
26    f_enter: B,
27    f_exit: C,
28    f_after: D,
29}
30
31impl<A, B, C, D> CallbacksApi for Callbacks<A, B, C, D>
32where
33    A: Fn() -> Option<*const ()> + 'static,
34    B: Fn(*const ()) -> *const () + 'static,
35    C: Fn(*const ()) + 'static,
36    D: Fn(*const ()) + 'static,
37{
38    fn before(&self) -> Option<*const ()> {
39        (self.f_before)()
40    }
41    fn enter(&self, d: *const ()) -> *const () {
42        (self.f_enter)(d)
43    }
44    fn exit(&self, d: *const ()) {
45        (self.f_exit)(d);
46    }
47    fn after(&self, d: *const ()) {
48        (self.f_after)(d);
49    }
50}
51
52pub(crate) struct Data {
53    cb: &'static dyn CallbacksApi,
54    ptr: *const (),
55}
56
57impl Data {
58    #[allow(clippy::if_not_else)]
59    pub(crate) fn load() -> Option<Data> {
60        // `Acquire` pairs with the `Release` store in `set_cbs`, so `CBS` is visible
61        let cb = if STATE.load(Ordering::Acquire) != INITIALIZED {
62            None
63        } else {
64            #[allow(static_mut_refs)]
65            unsafe {
66                Some(CBS.as_ref().map(AsRef::as_ref).unwrap())
67            }
68        };
69
70        if let Some(cb) = cb
71            && let Some(ptr) = cb.before()
72        {
73            return Some(Data { cb, ptr });
74        }
75        None
76    }
77
78    pub(crate) fn run<F, R>(&mut self, f: F) -> R
79    where
80        F: FnOnce() -> R,
81    {
82        let ptr = self.cb.enter(self.ptr);
83        let result = f();
84        self.cb.exit(ptr);
85        result
86    }
87}
88
89impl Drop for Data {
90    fn drop(&mut self) {
91        self.cb.after(self.ptr);
92    }
93}
94
95/// # Safety
96///
97/// The user must ensure that the pointer returned by `before` has a `'static` lifetime.
98/// This pointer will be owned by the spawned task for the duration of that task, and
99/// ownership will be returned to the user at the end of the task via `after`.
100/// The pointer remains opaque to the runtime.
101///
102/// Does nothing if task callbacks have already been set, use
103/// [`task_opt_callbacks`] to check whether they were set.
104pub unsafe fn task_callbacks<FBefore, FEnter, FExit, FAfter>(
105    f_before: FBefore,
106    f_enter: FEnter,
107    f_exit: FExit,
108    f_after: FAfter,
109) where
110    FBefore: Fn() -> Option<*const ()> + 'static + Sync,
111    FEnter: Fn(*const ()) -> *const () + 'static + Sync,
112    FExit: Fn(*const ()) + 'static + Sync,
113    FAfter: Fn(*const ()) + 'static + Sync,
114{
115    let new = Arc::new(Callbacks {
116        f_before,
117        f_enter,
118        f_exit,
119        f_after,
120    });
121    let _ = set_cbs(new);
122}
123
124/// # Safety
125///
126/// The user must ensure that the pointer returned by `before` has a `'static` lifetime.
127/// This pointer will be owned by the spawned task for the duration of that task, and
128/// ownership will be returned to the user at the end of the task via `after`.
129/// The pointer remains opaque to the runtime.
130///
131/// Returns false if task callbacks have already been set.
132pub unsafe fn task_opt_callbacks<FBefore, FEnter, FExit, FAfter>(
133    f_before: FBefore,
134    f_enter: FEnter,
135    f_exit: FExit,
136    f_after: FAfter,
137) -> bool
138where
139    FBefore: Fn() -> Option<*const ()> + Sync + 'static,
140    FEnter: Fn(*const ()) -> *const () + Sync + 'static,
141    FExit: Fn(*const ()) + Sync + 'static,
142    FAfter: Fn(*const ()) + Sync + 'static,
143{
144    let new = Arc::new(Callbacks {
145        f_before,
146        f_enter,
147        f_exit,
148        f_after,
149    });
150    set_cbs(new).is_ok()
151}
152
153fn set_cbs(cbs: Arc<dyn CallbacksApi>) -> Result<(), ()> {
154    match STATE.compare_exchange(
155        UNINITIALIZED,
156        INITIALIZING,
157        Ordering::Acquire,
158        Ordering::Relaxed,
159    ) {
160        Ok(UNINITIALIZED) => {
161            unsafe {
162                CBS = Some(cbs);
163            }
164            STATE.store(INITIALIZED, Ordering::Release);
165            Ok(())
166        }
167        Err(INITIALIZING) => {
168            while STATE.load(Ordering::Relaxed) == INITIALIZING {
169                std::hint::spin_loop();
170            }
171            Err(())
172        }
173        _ => Err(()),
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180
181    #[test]
182    fn load_callbacks_set_by_other_thread() {
183        let hnd = std::thread::spawn(|| unsafe {
184            task_opt_callbacks(|| Some(std::ptr::null()), |p| p, |_| {}, |_| {})
185        });
186        // observes callbacks set by another thread without synchronization
187        while Data::load().is_none() {
188            std::hint::spin_loop();
189        }
190        assert!(hnd.join().unwrap());
191    }
192}
ntex_rt/task.rs

ntex_rt/
task.rs
