1use std::cell::Cell;
2use std::ptr;
3
4use crate::task::{LocalsMap, Task, TaskId};
5use crate::utils::abort_on_panic;
6
7thread_local! {
8 /// A pointer to the currently running task.
9 static CURRENT: Cell<*const TaskLocalsWrapper> = Cell::new(ptr::null_mut());
10}
11
12/// A wrapper to store task local data.
13pub(crate) struct TaskLocalsWrapper {
14 /// The actual task details.
15 task: Task,
16
17 /// The map holding task-local values.
18 locals: LocalsMap,
19}
20
21impl TaskLocalsWrapper {
22 /// Creates a new task handle.
23 ///
24 /// If the task is unnamed, the inner representation of the task will be lazily allocated on
25 /// demand.
26 #[inline]
27 pub(crate) fn new(task: Task) -> Self {
28 Self {
29 task,
30 locals: LocalsMap::new(),
31 }
32 }
33
34 /// Gets the task's unique identifier.
35 #[inline]
36 pub fn id(&self) -> TaskId {
37 self.task.id()
38 }
39
40 /// Returns a reference to the inner `Task`.
41 pub(crate) fn task(&self) -> &Task {
42 &self.task
43 }
44
45 /// Returns the map holding task-local values.
46 pub(crate) fn locals(&self) -> &LocalsMap {
47 &self.locals
48 }
49
50 /// Set a reference to the current task.
51 pub(crate) unsafe fn set_current<F, R>(task: *const TaskLocalsWrapper, f: F) -> R
52 where
53 F: FnOnce() -> R,
54 {
55 CURRENT.with(|current| {
56 let old_task = current.replace(task);
57 defer! {
58 current.set(old_task);
59 }
60 f()
61 })
62 }
63
64 /// Gets a reference to the current task.
65 pub(crate) fn get_current<F, R>(f: F) -> Option<R>
66 where
67 F: FnOnce(&TaskLocalsWrapper) -> R,
68 {
69 let res = CURRENT.try_with(|current| unsafe { current.get().as_ref().map(f) });
70 match res {
71 Ok(Some(val)) => Some(val),
72 Ok(None) | Err(_) => None,
73 }
74 }
75}
76
77impl Drop for TaskLocalsWrapper {
78 fn drop(&mut self) {
79 // Abort the process if dropping task-locals panics.
80 abort_on_panic(|| {
81 unsafe { self.locals.clear() };
82 });
83 }
84}
85