memscope_rs/capture/backends/async_tracker/
task_context.rs

1//! Thread-local task scopes with guard lifetimes that can end on another thread.
2
3use std::cell::RefCell;
4use std::sync::atomic::{AtomicBool, Ordering};
5use std::sync::Arc;
6
7thread_local! {
8    static SCOPES: RefCell<Vec<TaskFrame>> = const { RefCell::new(Vec::new()) };
9}
10
11struct TaskFrame {
12    task_id: u64,
13    // None denotes an explicitly installed context, rather than a guard scope.
14    active: Option<Arc<AtomicBool>>,
15}
16
17impl TaskFrame {
18    fn is_active(&self) -> bool {
19        self.active
20            .as_ref()
21            .is_none_or(|active| active.load(Ordering::Acquire))
22    }
23}
24
25/// Restores the enclosing live scope when dropped or explicitly released.
26///
27/// Moving the guard does not install its context on the destination thread. Dropping
28/// it there marks only its own frame as ended; the origin prunes it on its next read.
29pub struct TaskGuard {
30    task_id: u64,
31    active: Arc<AtomicBool>,
32}
33
34impl TaskGuard {
35    pub fn task_id(&self) -> u64 {
36        self.task_id
37    }
38
39    pub fn release(self) {
40        drop(self);
41    }
42}
43
44impl Drop for TaskGuard {
45    fn drop(&mut self) {
46        self.active.store(false, Ordering::Release);
47    }
48}
49
50pub(super) fn current() -> Option<u64> {
51    SCOPES.with(|scopes| {
52        let mut scopes = scopes.borrow_mut();
53        scopes.retain(TaskFrame::is_active);
54        scopes.last().map(|frame| frame.task_id)
55    })
56}
57
58pub(super) fn enter(task_id: u64) -> TaskGuard {
59    let active = Arc::new(AtomicBool::new(true));
60    SCOPES.with(|scopes| {
61        let mut scopes = scopes.borrow_mut();
62        scopes.retain(TaskFrame::is_active);
63        scopes.push(TaskFrame {
64            task_id,
65            active: Some(Arc::clone(&active)),
66        });
67    });
68    TaskGuard { task_id, active }
69}
70
71pub(super) fn set(task_id: u64) {
72    SCOPES.with(|scopes| {
73        let mut scopes = scopes.borrow_mut();
74        // An explicit set replaces the previous explicit context, but does not end guards.
75        scopes.retain(|frame| frame.active.is_some() && frame.is_active());
76        scopes.push(TaskFrame {
77            task_id,
78            active: None,
79        });
80    });
81}
82
83pub(super) fn clear() {
84    SCOPES.with(|scopes| scopes.borrow_mut().clear());
85}
86
87/// Drop the *explicit* context of a task that ended.
88///
89/// Guard-owned frames are left alone: a guard still holds their lifetime, and removing
90/// them here would silently delete a live scope when the same task id is nested. They
91/// disappear on their own when the guard drops (possibly on another thread).
92pub(super) fn end(task_id: u64) {
93    SCOPES.with(|scopes| {
94        scopes.borrow_mut().retain(|frame| {
95            frame.is_active() && !(frame.task_id == task_id && frame.active.is_none())
96        });
97    });
98}
memscope_rs/capture/backends/async_tracker/task_context.rs

memscope_rs/capture/backends/async_tracker/
task_context.rs
