Skip to main content

liminal/tracing/
context.rs

1use rand::RngExt;
2
3#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4pub struct TraceContext {
5    pub trace_id: u128,
6    pub span_id: u64,
7}
8
9impl TraceContext {
10    #[must_use]
11    pub fn new_root() -> Self {
12        Self {
13            trace_id: next_trace_id(),
14            span_id: next_span_id_except(None),
15        }
16    }
17
18    #[must_use]
19    pub fn child(&self) -> Self {
20        Self {
21            trace_id: self.trace_id,
22            span_id: next_span_id_except(Some(self.span_id)),
23        }
24    }
25
26    #[must_use]
27    pub const fn from_ids(trace_id: u128, span_id: u64) -> Self {
28        Self { trace_id, span_id }
29    }
30
31    #[must_use]
32    pub const fn trace_id(&self) -> u128 {
33        self.trace_id
34    }
35
36    #[must_use]
37    pub const fn span_id(&self) -> u64 {
38        self.span_id
39    }
40
41    #[must_use]
42    pub fn to_traceparent(&self) -> String {
43        format!("00-{:032x}-{:016x}-01", self.trace_id, self.span_id)
44    }
45
46    #[must_use]
47    pub fn from_traceparent(header: &str) -> Option<Self> {
48        let mut parts = header.split('-');
49        let version = parts.next()?;
50        let trace_id = parts.next()?;
51        let span_id = parts.next()?;
52        let flags = parts.next()?;
53
54        if parts.next().is_some()
55            || version != "00"
56            || flags != "01"
57            || !is_hex_with_len(trace_id, 32)
58            || !is_hex_with_len(span_id, 16)
59        {
60            return None;
61        }
62
63        let trace_id = u128::from_str_radix(trace_id, 16).ok()?;
64        let span_id = u64::from_str_radix(span_id, 16).ok()?;
65
66        if trace_id == 0 || span_id == 0 {
67            return None;
68        }
69
70        Some(Self::from_ids(trace_id, span_id))
71    }
72}
73
74fn next_trace_id() -> u128 {
75    loop {
76        let candidate = thread_local_random::<u128>();
77        if candidate != 0 {
78            return candidate;
79        }
80    }
81}
82
83fn next_span_id_except(excluded: Option<u64>) -> u64 {
84    loop {
85        let candidate = thread_local_random::<u64>();
86        if candidate != 0 && Some(candidate) != excluded {
87            return candidate;
88        }
89    }
90}
91
92fn thread_local_random<T>() -> T
93where
94    rand::distr::StandardUniform: rand::distr::Distribution<T>,
95{
96    let mut rng = rand::rng();
97    rng.random()
98}
99
100fn is_hex_with_len(value: &str, len: usize) -> bool {
101    value.len() == len && value.bytes().all(|byte| byte.is_ascii_hexdigit())
102}
103
104#[cfg(test)]
105mod tests {
106    use super::TraceContext;
107
108    #[test]
109    fn new_root_generates_non_zero_ids() {
110        let context = TraceContext::new_root();
111
112        assert_ne!(context.trace_id(), 0);
113        assert_ne!(context.span_id(), 0);
114    }
115
116    #[test]
117    fn child_preserves_trace_id_and_changes_span_id() {
118        let parent = TraceContext::new_root();
119        let child = parent.child();
120
121        assert_eq!(child.trace_id(), parent.trace_id());
122        assert_ne!(child.span_id(), parent.span_id());
123        assert_ne!(child.span_id(), 0);
124    }
125
126    #[test]
127    fn roots_have_distinct_trace_ids() {
128        let first = TraceContext::new_root();
129        let second = TraceContext::new_root();
130
131        assert_ne!(first.trace_id(), second.trace_id());
132    }
133
134    #[test]
135    fn from_ids_round_trips() {
136        let context = TraceContext::new_root();
137
138        assert_eq!(
139            TraceContext::from_ids(context.trace_id(), context.span_id()),
140            context
141        );
142    }
143
144    #[test]
145    fn traceparent_round_trips() {
146        let context = TraceContext::new_root();
147        let traceparent = context.to_traceparent();
148
149        assert_eq!(traceparent.len(), 55);
150        assert_eq!(TraceContext::from_traceparent(&traceparent), Some(context));
151    }
152
153    #[test]
154    fn invalid_traceparents_return_none() {
155        assert_eq!(TraceContext::from_traceparent(""), None);
156        assert_eq!(
157            TraceContext::from_traceparent(
158                "01-00000000000000000000000000000001-0000000000000001-01"
159            ),
160            None
161        );
162        assert_eq!(
163            TraceContext::from_traceparent(
164                "00-00000000000000000000000000000001-0000000000000001-00"
165            ),
166            None
167        );
168        assert_eq!(
169            TraceContext::from_traceparent("00-not-0000000000000001-01"),
170            None
171        );
172        assert_eq!(
173            TraceContext::from_traceparent(
174                "00-00000000000000000000000000000000-0000000000000001-01"
175            ),
176            None
177        );
178        assert_eq!(
179            TraceContext::from_traceparent(
180                "00-00000000000000000000000000000001-0000000000000000-01"
181            ),
182            None
183        );
184    }
185}