provide_telemetry/
context.rs

1// SPDX-FileCopyrightText: Copyright (C) 2026 provide.io llc
2// SPDX-License-Identifier: Apache-2.0
3// SPDX-Comment: Part of provide-telemetry.
4//
5
6use std::cell::RefCell;
7use std::collections::{BTreeMap, HashMap};
8use std::sync::atomic::{AtomicU64, Ordering};
9use std::sync::{Mutex, OnceLock};
10
11use serde_json::Value;
12
13#[derive(Clone, Debug, Default, PartialEq)]
14pub struct ContextSnapshot {
15    pub fields: BTreeMap<String, Value>,
16    pub session_id: Option<String>,
17    pub trace_id: Option<String>,
18    pub span_id: Option<String>,
19}
20
21#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
22enum ContextScopeKey {
23    Task(tokio::task::Id),
24    Thread,
25}
26
27pub struct ContextGuard {
28    key: ContextScopeKey,
29    previous: ContextSnapshot,
30    epoch: u64,
31}
32
33thread_local! {
34    static THREAD_CONTEXT: RefCell<ContextSnapshot> = RefCell::new(ContextSnapshot::default());
35}
36
37static TASK_CONTEXTS: OnceLock<Mutex<HashMap<tokio::task::Id, ContextSnapshot>>> = OnceLock::new();
38static CONTEXT_EPOCH: AtomicU64 = AtomicU64::new(0);
39
40fn task_contexts() -> &'static Mutex<HashMap<tokio::task::Id, ContextSnapshot>> {
41    TASK_CONTEXTS.get_or_init(|| Mutex::new(HashMap::new()))
42}
43
44fn current_scope_key() -> ContextScopeKey {
45    tokio::task::try_id()
46        .map(ContextScopeKey::Task)
47        .unwrap_or(ContextScopeKey::Thread)
48}
49
50fn current_snapshot() -> ContextSnapshot {
51    match current_scope_key() {
52        ContextScopeKey::Task(task_id) => task_contexts()
53            .lock()
54            .expect("task context lock poisoned")
55            .get(&task_id)
56            .cloned()
57            .unwrap_or_default(),
58        ContextScopeKey::Thread => THREAD_CONTEXT.with(|ctx| ctx.borrow().clone()),
59    }
60}
61
62fn set_snapshot_for_key(key: ContextScopeKey, snapshot: ContextSnapshot) {
63    match key {
64        ContextScopeKey::Task(task_id) => {
65            let mut map = task_contexts().lock().expect("task context lock poisoned");
66            if snapshot == ContextSnapshot::default() {
67                // Remove the entry when restoring to default to prevent unbounded growth
68                map.remove(&task_id);
69            } else {
70                map.insert(task_id, snapshot);
71            }
72        }
73        ContextScopeKey::Thread => {
74            THREAD_CONTEXT.with(|ctx| {
75                *ctx.borrow_mut() = snapshot;
76            });
77        }
78    }
79}
80
81fn replace_snapshot(next: ContextSnapshot) -> ContextGuard {
82    let key = current_scope_key();
83    let previous = current_snapshot();
84    let epoch = CONTEXT_EPOCH.load(Ordering::SeqCst);
85    set_snapshot_for_key(key, next);
86    ContextGuard {
87        key,
88        previous,
89        epoch,
90    }
91}
92
93pub fn get_context() -> BTreeMap<String, Value> {
94    current_snapshot().fields
95}
96
97pub fn bind_context<I, K>(fields: I) -> ContextGuard
98where
99    I: IntoIterator<Item = (K, Value)>,
100    K: Into<String>,
101{
102    let mut next = current_snapshot();
103    for (key, value) in fields {
104        next.fields.insert(key.into(), value);
105    }
106    replace_snapshot(next)
107}
108
109pub fn unbind_context(keys: &[&str]) -> ContextGuard {
110    let mut next = current_snapshot();
111    for key in keys {
112        next.fields.remove(*key);
113    }
114    replace_snapshot(next)
115}
116
117pub fn clear_context() -> ContextGuard {
118    let mut next = current_snapshot();
119    next.fields.clear();
120    replace_snapshot(next)
121}
122
123pub fn bind_session_context(session_id: impl Into<String>) -> ContextGuard {
124    let mut next = current_snapshot();
125    let session_id = session_id.into();
126    next.session_id = Some(session_id.clone());
127    next.fields
128        .insert("session_id".to_string(), Value::String(session_id));
129    replace_snapshot(next)
130}
131
132pub fn get_session_id() -> Option<String> {
133    current_snapshot().session_id
134}
135
136pub fn clear_session_context() -> ContextGuard {
137    let mut next = current_snapshot();
138    next.session_id = None;
139    next.fields.remove("session_id");
140    replace_snapshot(next)
141}
142
143pub(crate) fn set_trace_context_internal(
144    trace_id: Option<String>,
145    span_id: Option<String>,
146) -> ContextGuard {
147    let mut next = current_snapshot();
148    next.trace_id = trace_id;
149    next.span_id = span_id;
150    replace_snapshot(next)
151}
152
153pub(crate) fn trace_snapshot() -> ContextSnapshot {
154    current_snapshot()
155}
156
157pub(crate) fn reset_context_for_tests() {
158    CONTEXT_EPOCH.fetch_add(1, Ordering::SeqCst);
159    THREAD_CONTEXT.with(|ctx| {
160        *ctx.borrow_mut() = ContextSnapshot::default();
161    });
162    task_contexts()
163        .lock()
164        .expect("task context lock poisoned")
165        .clear();
166}
167
168pub(crate) fn reset_trace_context_for_tests() {
169    CONTEXT_EPOCH.fetch_add(1, Ordering::SeqCst);
170    THREAD_CONTEXT.with(|ctx| {
171        let mut snapshot = ctx.borrow_mut();
172        snapshot.trace_id = None;
173        snapshot.span_id = None;
174    });
175    let mut tasks = task_contexts().lock().expect("task context lock poisoned");
176    for snapshot in tasks.values_mut() {
177        snapshot.trace_id = None;
178        snapshot.span_id = None;
179    }
180}
181
182impl Drop for ContextGuard {
183    fn drop(&mut self) {
184        if CONTEXT_EPOCH.load(Ordering::SeqCst) == self.epoch {
185            set_snapshot_for_key(self.key, self.previous.clone());
186        }
187    }
188}
provide_telemetry/context.rs

provide_telemetry/
context.rs
