greentic_telemetry/
tasklocal.rs

1use crate::context::TelemetryCtx;
2use std::{cell::RefCell, future::Future};
3
4tokio::task_local! {
5    static GT_TELEMETRY_CTX: RefCell<Option<TelemetryCtx>>;
6}
7
8/// Set the task-local telemetry context. No-op if outside a Tokio task.
9pub fn set_current_telemetry_ctx(ctx: TelemetryCtx) {
10    let _ = GT_TELEMETRY_CTX.try_with(|slot| {
11        *slot.borrow_mut() = Some(ctx);
12    });
13}
14
15/// Execute `f` with the telemetry context currently stored on this task, if any.
16pub fn with_current_telemetry_ctx<R>(f: impl FnOnce(Option<&TelemetryCtx>) -> R) -> R {
17    let mut f = Some(f);
18    let result = GT_TELEMETRY_CTX.try_with(|slot| {
19        let guard = slot.borrow();
20        let func = f
21            .take()
22            .expect("telemetry context closure already consumed");
23        func(guard.as_ref())
24    });
25
26    match result {
27        Ok(value) => value,
28        Err(_) => {
29            let func = f
30                .take()
31                .expect("telemetry context closure already consumed");
32            func(None)
33        }
34    }
35}
36
37/// Run `fut` with a task-local telemetry context slot initialized.
38pub async fn with_task_local<Fut, R>(fut: Fut) -> R
39where
40    Fut: Future<Output = R>,
41{
42    GT_TELEMETRY_CTX.scope(RefCell::new(None), fut).await
43}
greentic_telemetry/tasklocal.rs

greentic_telemetry/
tasklocal.rs
