liminal/tracing/
context.rs1use 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}