ntex_rt/
task.rs

1use std::{cell::Cell, ptr};
2
3thread_local! {
4    static CB: Cell<*const Callbacks> = const { Cell::new(ptr::null()) };
5}
6
7struct Callbacks {
8    before: Box<dyn Fn() -> Option<*const ()>>,
9    enter: Box<dyn Fn(*const ()) -> *const ()>,
10    exit: Box<dyn Fn(*const ())>,
11    after: Box<dyn Fn(*const ())>,
12}
13
14pub(crate) struct Data {
15    cb: &'static Callbacks,
16    ptr: *const (),
17}
18
19impl Data {
20    pub(crate) fn load() -> Option<Data> {
21        let cb = CB.with(Cell::get);
22
23        if let Some(cb) = unsafe { cb.as_ref() }
24            && let Some(ptr) = (*cb.before)()
25        {
26            return Some(Data { cb, ptr });
27        }
28        None
29    }
30
31    pub(crate) fn run<F, R>(&mut self, f: F) -> R
32    where
33        F: FnOnce() -> R,
34    {
35        let ptr = (*self.cb.enter)(self.ptr);
36        let result = f();
37        (*self.cb.exit)(ptr);
38        result
39    }
40}
41
42impl Drop for Data {
43    fn drop(&mut self) {
44        (*self.cb.after)(self.ptr);
45    }
46}
47
48/// # Safety
49///
50/// The user must ensure that the pointer returned by `before` has a `'static` lifetime.
51/// This pointer will be owned by the spawned task for the duration of that task, and
52/// ownership will be returned to the user at the end of the task via `after`.
53/// The pointer remains opaque to the runtime.
54///
55/// # Panics
56///
57/// Panics if task callbacks have already been set.
58pub unsafe fn task_callbacks<FBefore, FEnter, FExit, FAfter>(
59    before: FBefore,
60    enter: FEnter,
61    exit: FExit,
62    after: FAfter,
63) where
64    FBefore: Fn() -> Option<*const ()> + 'static,
65    FEnter: Fn(*const ()) -> *const () + 'static,
66    FExit: Fn(*const ()) + 'static,
67    FAfter: Fn(*const ()) + 'static,
68{
69    CB.with(|cb| {
70        assert!(cb.get().is_null(), "Spawn callbacks already set");
71
72        let new: *mut Callbacks = Box::leak(Box::new(Callbacks {
73            before: Box::new(before),
74            enter: Box::new(enter),
75            exit: Box::new(exit),
76            after: Box::new(after),
77        }));
78        cb.replace(new);
79    });
80}
81
82/// # Safety
83///
84/// The user must ensure that the pointer returned by `before` has a `'static` lifetime.
85/// This pointer will be owned by the spawned task for the duration of that task, and
86/// ownership will be returned to the user at the end of the task via `after`.
87/// The pointer remains opaque to the runtime.
88///
89/// Returns false if task callbacks have already been set.
90pub unsafe fn task_opt_callbacks<FBefore, FEnter, FExit, FAfter>(
91    before: FBefore,
92    enter: FEnter,
93    exit: FExit,
94    after: FAfter,
95) -> bool
96where
97    FBefore: Fn() -> Option<*const ()> + 'static,
98    FEnter: Fn(*const ()) -> *const () + 'static,
99    FExit: Fn(*const ()) + 'static,
100    FAfter: Fn(*const ()) + 'static,
101{
102    CB.with(|cb| {
103        if cb.get().is_null() {
104            unsafe {
105                task_callbacks(before, enter, exit, after);
106            }
107            true
108        } else {
109            false
110        }
111    })
112}
ntex_rt/task.rs

ntex_rt/
task.rs
