async_std/task/
task_locals_wrapper.rs

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}
async_std/task/task_locals_wrapper.rs

async_std/task/
task_locals_wrapper.rs
