rootle_trace/
context.rs

1//! Thread identity and operation correlation stamped into every envelope.
2//!
3//! [`operation_id`] allocates a correlation identity only while a trace is
4//! enabled, so job-spawn sites pay nothing when capture is off.
5//! [`in_operation`] installs that identity in a thread-local for the
6//! duration of the work — including through provider/HTTP calls that record
7//! on their own — and restores the previous identity even when the work
8//! panics.
9
10use serde::{Deserialize, Serialize};
11use std::cell::{Cell, OnceCell};
12use std::sync::Arc;
13use std::sync::atomic::{AtomicU64, Ordering};
14
15use crate::enabled;
16
17/// Correlation identity linking a job's lifecycle records across threads.
18/// Serializes as its bare number inside envelopes and producer payloads.
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
20#[serde(transparent)]
21pub struct OperationId(pub u64);
22
23static NEXT_OPERATION: AtomicU64 = AtomicU64::new(1);
24static THREAD_SERIAL: AtomicU64 = AtomicU64::new(1);
25
26thread_local! {
27    static OPERATION: Cell<Option<OperationId>> = const { Cell::new(None) };
28    static THREAD_LABEL: OnceCell<Arc<str>> = const { OnceCell::new() };
29}
30
31/// Allocate a fresh correlation identity, but only while a trace is
32/// enabled. Returns `None` (and allocates nothing) for disabled capture;
33/// pass the `Option` straight through to [`in_operation`].
34pub fn operation_id() -> Option<OperationId> {
35    enabled().then(|| OperationId(NEXT_OPERATION.fetch_add(1, Ordering::Relaxed)))
36}
37
38/// The identity installed on this thread, stamped into records admitted
39/// here. Absent outside [`in_operation`] — including on the spawning
40/// thread that only allocated the identity.
41pub(crate) fn current_operation() -> Option<OperationId> {
42    OPERATION.get()
43}
44
45/// Run `work` with `id` as this thread's correlation identity, restoring
46/// whatever was installed before — including on panic unwind, so a
47/// panicking worker cannot leak its identity into unrelated records.
48pub fn in_operation<T>(id: Option<OperationId>, work: impl FnOnce() -> T) -> T {
49    let _restore = OperationGuard::enter(id);
50    work()
51}
52
53struct OperationGuard {
54    previous: Option<OperationId>,
55}
56
57impl OperationGuard {
58    fn enter(id: Option<OperationId>) -> Self {
59        Self {
60            previous: OPERATION.replace(id),
61        }
62    }
63}
64
65impl Drop for OperationGuard {
66    fn drop(&mut self) {
67        OPERATION.set(self.previous);
68    }
69}
70
71/// A stable per-thread label for the envelope's `thread` field: the
72/// platform thread name when one exists, disambiguated by a process-wide
73/// serial so unnamed workers still stay distinguishable.
74pub(crate) fn thread_label() -> Arc<str> {
75    THREAD_LABEL.with(|cell| {
76        Arc::clone(cell.get_or_init(|| {
77            let current = std::thread::current();
78            let serial = THREAD_SERIAL.fetch_add(1, Ordering::Relaxed);
79            format!("{}-{serial}", current.name().unwrap_or("unnamed")).into()
80        }))
81    })
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87
88    #[test]
89    fn guard_restores_through_nested_work() {
90        let outer = OperationId(7);
91        let _restore = OperationGuard::enter(Some(outer));
92        assert_eq!(current_operation(), Some(outer));
93        in_operation(None, || assert_eq!(current_operation(), None));
94        assert_eq!(current_operation(), Some(outer));
95        let inner = OperationId(9);
96        let seen = in_operation(Some(inner), current_operation);
97        assert_eq!(seen, Some(inner));
98        assert_eq!(current_operation(), Some(outer));
99    }
100
101    #[test]
102    fn thread_labels_are_stable_and_distinct() {
103        let first = thread_label();
104        assert!(
105            Arc::ptr_eq(&first, &thread_label()),
106            "label computed once per thread"
107        );
108        let other = std::thread::Builder::new()
109            .name("probe".into())
110            .spawn(thread_label)
111            .unwrap()
112            .join()
113            .unwrap();
114        assert!(first != other, "distinct threads get distinct labels");
115        assert!(
116            other.starts_with("probe-"),
117            "label carries the name: {other}"
118        );
119    }
120}
rootle_trace/context.rs

rootle_trace/
context.rs
