grafos_observe/
trace.rs

1//! Distributed trace context types for W3C traceparent propagation.
2//!
3//! All types are `no_std` compatible. The [`TraceContext`] struct carries a
4//! 16-byte trace ID, 8-byte span ID, and 1-byte flags field conforming to
5//! the W3C Trace Context specification.
6//!
7//! # Binary encoding
8//!
9//! ```text
10//! [version: 1B] [trace_id: 16B] [span_id: 8B] [flags: 1B] [reserved: 6B]
11//! ```
12//!
13//! Total: 32 bytes (fits in a single cache line, suitable for shared-memory
14//! RPC header extension).
15//!
16//! # W3C string format
17//!
18//! ```text
19//! 00-{trace_id:032x}-{span_id:016x}-{flags:02x}
20//! ```
21
22use core::fmt;
23
24/// 16-byte trace identifier, unique per trace.
25#[derive(Clone, Copy, PartialEq, Eq, Hash)]
26pub struct TraceId(pub u128);
27
28impl TraceId {
29    /// The zero (invalid) trace ID.
30    pub const INVALID: TraceId = TraceId(0);
31
32    /// Whether this trace ID is the zero/invalid value.
33    pub fn is_valid(&self) -> bool {
34        self.0 != 0
35    }
36}
37
38impl fmt::Debug for TraceId {
39    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
40        write!(f, "TraceId({:032x})", self.0)
41    }
42}
43
44impl fmt::Display for TraceId {
45    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
46        write!(f, "{:032x}", self.0)
47    }
48}
49
50impl fmt::LowerHex for TraceId {
51    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
52        fmt::LowerHex::fmt(&self.0, f)
53    }
54}
55
56/// 8-byte span identifier, unique per span within a trace.
57#[derive(Clone, Copy, PartialEq, Eq, Hash)]
58pub struct SpanId(pub u64);
59
60impl SpanId {
61    /// The zero (invalid) span ID.
62    pub const INVALID: SpanId = SpanId(0);
63
64    /// Whether this span ID is the zero/invalid value.
65    pub fn is_valid(&self) -> bool {
66        self.0 != 0
67    }
68}
69
70impl fmt::Debug for SpanId {
71    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
72        write!(f, "SpanId({:016x})", self.0)
73    }
74}
75
76impl fmt::Display for SpanId {
77    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
78        write!(f, "{:016x}", self.0)
79    }
80}
81
82impl fmt::LowerHex for SpanId {
83    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
84        fmt::LowerHex::fmt(&self.0, f)
85    }
86}
87
88/// W3C trace context in binary form.
89///
90/// Carries the identifiers needed to correlate spans across service boundaries.
91/// The `flags` field uses bit 0 as the "sampled" flag (W3C traceparent spec).
92#[derive(Clone, Copy, Debug, PartialEq, Eq)]
93pub struct TraceContext {
94    /// Trace-wide identifier (shared by all spans in the trace).
95    pub trace_id: TraceId,
96    /// Span identifier for the current span.
97    pub span_id: SpanId,
98    /// Parent span identifier (zero if this is the root span).
99    pub parent_span_id: SpanId,
100    /// W3C trace flags. Bit 0 = sampled.
101    pub flags: u8,
102}
103
104/// W3C traceparent version byte.
105const W3C_VERSION: u8 = 0x00;
106
107/// Flag bit: trace is sampled.
108pub const FLAG_SAMPLED: u8 = 0x01;
109
110/// Binary encoding size.
111pub const TRACE_CONTEXT_BYTES: usize = 32;
112
113impl TraceContext {
114    /// Create a new root trace context with the given random source.
115    ///
116    /// Generates fresh trace_id and span_id from the provided random bytes.
117    /// `random_bytes` must be at least 24 bytes (16 for trace_id + 8 for span_id).
118    pub fn new_root(random_bytes: &[u8; 24]) -> Self {
119        let trace_id = u128::from_be_bytes(random_bytes[0..16].try_into().unwrap());
120        let span_id = u64::from_be_bytes(random_bytes[16..24].try_into().unwrap());
121        TraceContext {
122            trace_id: TraceId(trace_id),
123            span_id: SpanId(span_id),
124            parent_span_id: SpanId::INVALID,
125            flags: FLAG_SAMPLED,
126        }
127    }
128
129    /// Create a child context within the same trace.
130    ///
131    /// The child inherits `trace_id` and `flags`, gets a new `span_id` from
132    /// the provided random bytes, and records the current `span_id` as its parent.
133    pub fn child(&self, random_span_bytes: &[u8; 8]) -> Self {
134        let new_span_id = u64::from_be_bytes(*random_span_bytes);
135        TraceContext {
136            trace_id: self.trace_id,
137            span_id: SpanId(new_span_id),
138            parent_span_id: self.span_id,
139            flags: self.flags,
140        }
141    }
142
143    /// Whether the sampled flag is set.
144    pub fn is_sampled(&self) -> bool {
145        self.flags & FLAG_SAMPLED != 0
146    }
147
148    /// Set or clear the sampled flag.
149    pub fn set_sampled(&mut self, sampled: bool) {
150        if sampled {
151            self.flags |= FLAG_SAMPLED;
152        } else {
153            self.flags &= !FLAG_SAMPLED;
154        }
155    }
156
157    /// Encode to a 32-byte binary representation.
158    ///
159    /// Layout: `[version:1][trace_id:16][span_id:8][flags:1][reserved:6]`
160    pub fn encode(&self) -> [u8; TRACE_CONTEXT_BYTES] {
161        let mut buf = [0u8; TRACE_CONTEXT_BYTES];
162        buf[0] = W3C_VERSION;
163        buf[1..17].copy_from_slice(&self.trace_id.0.to_be_bytes());
164        buf[17..25].copy_from_slice(&self.span_id.0.to_be_bytes());
165        buf[25] = self.flags;
166        // bytes 26..32 reserved (zeros)
167        buf
168    }
169
170    /// Decode from a 32-byte binary representation.
171    pub fn decode(buf: &[u8; TRACE_CONTEXT_BYTES]) -> Result<Self, TraceContextError> {
172        if buf[0] != W3C_VERSION {
173            return Err(TraceContextError::UnsupportedVersion(buf[0]));
174        }
175        let trace_id = u128::from_be_bytes(buf[1..17].try_into().unwrap());
176        let span_id = u64::from_be_bytes(buf[17..25].try_into().unwrap());
177        let flags = buf[25];
178
179        Ok(TraceContext {
180            trace_id: TraceId(trace_id),
181            span_id: SpanId(span_id),
182            parent_span_id: SpanId::INVALID, // not stored in wire format
183            flags,
184        })
185    }
186
187    /// Format as a W3C traceparent string: `00-{trace_id}-{span_id}-{flags}`.
188    ///
189    /// Returns a fixed-length 55-character string.
190    pub fn to_w3c_string(&self) -> alloc::string::String {
191        alloc::format!(
192            "{:02x}-{:032x}-{:016x}-{:02x}",
193            W3C_VERSION,
194            self.trace_id.0,
195            self.span_id.0,
196            self.flags,
197        )
198    }
199
200    /// Parse a W3C traceparent string.
201    ///
202    /// Expected format: `{version:2}-{trace_id:32}-{span_id:16}-{flags:2}`
203    /// (55 characters total).
204    pub fn from_w3c_string(s: &str) -> Result<Self, TraceContextError> {
205        if s.len() != 55 {
206            return Err(TraceContextError::InvalidFormat);
207        }
208        let parts: alloc::vec::Vec<&str> = s.split('-').collect();
209        if parts.len() != 4 {
210            return Err(TraceContextError::InvalidFormat);
211        }
212        if parts[0].len() != 2
213            || parts[1].len() != 32
214            || parts[2].len() != 16
215            || parts[3].len() != 2
216        {
217            return Err(TraceContextError::InvalidFormat);
218        }
219        let version =
220            u8::from_str_radix(parts[0], 16).map_err(|_| TraceContextError::InvalidFormat)?;
221        if version != W3C_VERSION {
222            return Err(TraceContextError::UnsupportedVersion(version));
223        }
224        let trace_id =
225            u128::from_str_radix(parts[1], 16).map_err(|_| TraceContextError::InvalidFormat)?;
226        let span_id =
227            u64::from_str_radix(parts[2], 16).map_err(|_| TraceContextError::InvalidFormat)?;
228        let flags =
229            u8::from_str_radix(parts[3], 16).map_err(|_| TraceContextError::InvalidFormat)?;
230
231        Ok(TraceContext {
232            trace_id: TraceId(trace_id),
233            span_id: SpanId(span_id),
234            parent_span_id: SpanId::INVALID,
235            flags,
236        })
237    }
238
239    /// Check if this context has all-zero IDs (i.e. no trace context was set).
240    pub fn is_empty(&self) -> bool {
241        !self.trace_id.is_valid() && !self.span_id.is_valid()
242    }
243}
244
245impl Default for TraceContext {
246    fn default() -> Self {
247        TraceContext {
248            trace_id: TraceId::INVALID,
249            span_id: SpanId::INVALID,
250            parent_span_id: SpanId::INVALID,
251            flags: 0,
252        }
253    }
254}
255
256impl fmt::Display for TraceContext {
257    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
258        write!(f, "{}", self.to_w3c_string())
259    }
260}
261
262/// Errors from trace context encoding/decoding.
263#[derive(Debug, Clone, PartialEq, Eq)]
264pub enum TraceContextError {
265    /// The version byte is not supported.
266    UnsupportedVersion(u8),
267    /// The string format is invalid.
268    InvalidFormat,
269}
270
271impl fmt::Display for TraceContextError {
272    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
273        match self {
274            TraceContextError::UnsupportedVersion(v) => {
275                write!(f, "unsupported traceparent version: {v:#04x}")
276            }
277            TraceContextError::InvalidFormat => write!(f, "invalid traceparent format"),
278        }
279    }
280}
281
282// ---------------------------------------------------------------------------
283// Thread-local current trace context (requires std)
284// ---------------------------------------------------------------------------
285
286#[cfg(feature = "std")]
287mod thread_local_ctx {
288    use super::TraceContext;
289    use std::cell::RefCell;
290
291    thread_local! {
292        static CURRENT: RefCell<TraceContext> = RefCell::new(TraceContext::default());
293    }
294
295    impl TraceContext {
296        /// Get the current thread-local trace context.
297        pub fn current() -> TraceContext {
298            CURRENT.with(|c| *c.borrow())
299        }
300
301        /// Set the current thread-local trace context.
302        pub fn set_current(ctx: TraceContext) {
303            CURRENT.with(|c| *c.borrow_mut() = ctx);
304        }
305    }
306}
307
308#[cfg(test)]
309mod tests {
310    use super::*;
311
312    fn test_random_bytes() -> [u8; 24] {
313        // Deterministic test bytes
314        let mut bytes = [0u8; 24];
315        for (i, b) in bytes.iter_mut().enumerate() {
316            *b = (i as u8).wrapping_add(0x10);
317        }
318        bytes
319    }
320
321    #[test]
322    fn new_root_creates_valid_context() {
323        let ctx = TraceContext::new_root(&test_random_bytes());
324        assert!(ctx.trace_id.is_valid());
325        assert!(ctx.span_id.is_valid());
326        assert_eq!(ctx.parent_span_id, SpanId::INVALID);
327        assert!(ctx.is_sampled());
328    }
329
330    #[test]
331    fn child_inherits_trace_id() {
332        let root = TraceContext::new_root(&test_random_bytes());
333        let child_bytes: [u8; 8] = [0xAA; 8];
334        let child = root.child(&child_bytes);
335
336        assert_eq!(child.trace_id, root.trace_id);
337        assert_ne!(child.span_id, root.span_id);
338        assert_eq!(child.parent_span_id, root.span_id);
339        assert_eq!(child.flags, root.flags);
340    }
341
342    #[test]
343    fn encode_decode_roundtrip() {
344        let ctx = TraceContext::new_root(&test_random_bytes());
345        let encoded = ctx.encode();
346        let decoded = TraceContext::decode(&encoded).unwrap();
347
348        assert_eq!(decoded.trace_id, ctx.trace_id);
349        assert_eq!(decoded.span_id, ctx.span_id);
350        assert_eq!(decoded.flags, ctx.flags);
351    }
352
353    #[test]
354    fn w3c_string_roundtrip() {
355        let ctx = TraceContext::new_root(&test_random_bytes());
356        let s = ctx.to_w3c_string();
357
358        assert_eq!(s.len(), 55);
359        assert!(s.starts_with("00-"));
360
361        let parsed = TraceContext::from_w3c_string(&s).unwrap();
362        assert_eq!(parsed.trace_id, ctx.trace_id);
363        assert_eq!(parsed.span_id, ctx.span_id);
364        assert_eq!(parsed.flags, ctx.flags);
365    }
366
367    #[test]
368    fn w3c_string_format() {
369        let ctx = TraceContext {
370            trace_id: TraceId(0x0102030405060708090a0b0c0d0e0f10),
371            span_id: SpanId(0x1112131415161718),
372            parent_span_id: SpanId::INVALID,
373            flags: 0x01,
374        };
375        let s = ctx.to_w3c_string();
376        assert_eq!(s, "00-0102030405060708090a0b0c0d0e0f10-1112131415161718-01");
377    }
378
379    #[test]
380    fn from_w3c_string_rejects_bad_version() {
381        let s = "01-0102030405060708090a0b0c0d0e0f10-1112131415161718-01";
382        let err = TraceContext::from_w3c_string(s).unwrap_err();
383        assert_eq!(err, TraceContextError::UnsupportedVersion(0x01));
384    }
385
386    #[test]
387    fn from_w3c_string_rejects_bad_length() {
388        assert_eq!(
389            TraceContext::from_w3c_string("too-short"),
390            Err(TraceContextError::InvalidFormat)
391        );
392    }
393
394    #[test]
395    fn sampled_flag_manipulation() {
396        let mut ctx = TraceContext::default();
397        assert!(!ctx.is_sampled());
398        ctx.set_sampled(true);
399        assert!(ctx.is_sampled());
400        ctx.set_sampled(false);
401        assert!(!ctx.is_sampled());
402    }
403
404    #[test]
405    fn default_is_empty() {
406        let ctx = TraceContext::default();
407        assert!(ctx.is_empty());
408    }
409
410    #[test]
411    fn decode_rejects_bad_version() {
412        let mut buf = [0u8; 32];
413        buf[0] = 0xFF;
414        assert_eq!(
415            TraceContext::decode(&buf),
416            Err(TraceContextError::UnsupportedVersion(0xFF))
417        );
418    }
419
420    #[cfg(feature = "std")]
421    #[test]
422    fn thread_local_current_context() {
423        // Save and restore to avoid polluting other tests
424        let saved = TraceContext::current();
425
426        let ctx = TraceContext::new_root(&test_random_bytes());
427        TraceContext::set_current(ctx);
428        let got = TraceContext::current();
429        assert_eq!(got.trace_id, ctx.trace_id);
430        assert_eq!(got.span_id, ctx.span_id);
431
432        TraceContext::set_current(saved);
433    }
434
435    #[test]
436    fn display_uses_w3c_format() {
437        let ctx = TraceContext::new_root(&test_random_bytes());
438        let display = alloc::format!("{ctx}");
439        let w3c = ctx.to_w3c_string();
440        assert_eq!(display, w3c);
441    }
442
443    #[test]
444    fn trace_id_display() {
445        let tid = TraceId(0x0102030405060708090a0b0c0d0e0f10);
446        assert_eq!(alloc::format!("{tid}"), "0102030405060708090a0b0c0d0e0f10");
447    }
448
449    #[test]
450    fn span_id_display() {
451        let sid = SpanId(0x1112131415161718);
452        assert_eq!(alloc::format!("{sid}"), "1112131415161718");
453    }
454}
grafos_observe/trace.rs

grafos_observe/
trace.rs
