ntex_rt/
task.rs

1use std::{cell::Cell, panic, 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(|cb| cb.get());
22
23        if let Some(cb) = unsafe { cb.as_ref() } {
24            if let Some(ptr) = (*cb.before)() {
25                return Some(Data { cb, ptr });
26            }
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        if !cb.get().is_null() {
71            panic!("Spawn callbacks already set");
72        }
73
74        let new: *mut Callbacks = Box::leak(Box::new(Callbacks {
75            before: Box::new(before),
76            enter: Box::new(enter),
77            exit: Box::new(exit),
78            after: Box::new(after),
79        }));
80        cb.replace(new);
81    });
82}
83
84/// # Safety
85///
86/// The user must ensure that the pointer returned by `before` has a `'static` lifetime.
87/// This pointer will be owned by the spawned task for the duration of that task, and
88/// ownership will be returned to the user at the end of the task via `after`.
89/// The pointer remains opaque to the runtime.
90///
91/// Returns false if task callbacks have already been set.
92pub unsafe fn task_opt_callbacks<FBefore, FEnter, FExit, FAfter>(
93    before: FBefore,
94    enter: FEnter,
95    exit: FExit,
96    after: FAfter,
97) -> bool
98where
99    FBefore: Fn() -> Option<*const ()> + 'static,
100    FEnter: Fn(*const ()) -> *const () + 'static,
101    FExit: Fn(*const ()) + 'static,
102    FAfter: Fn(*const ()) + 'static,
103{
104    CB.with(|cb| {
105        if !cb.get().is_null() {
106            false
107        } else {
108            unsafe {
109                task_callbacks(before, enter, exit, after);
110            }
111            true
112        }
113    })
114}
ntex_rt/task.rs

ntex_rt/
task.rs
