veecle_telemetry/
id.rs

1// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.
2// Copyright 2025 Veecle GmbH.
3//
4// This file has been modified from the original TiKV implementation.
5
6//! Unique identifiers for traces and spans.
7//!
8//! This module provides the core identifier types used throughout the telemetry system
9//! to uniquely identify traces and spans in distributed tracing scenarios.
10//!
11//! # Core Types
12//!
13//! - [`TraceId`]: A 128-bit globally unique identifier that groups related spans together
14//! - [`SpanId`]: A 64-bit identifier that uniquely identifies a span within a trace
15//! - [`SpanContext`]: A combination of trace ID and span ID that uniquely identifies a span
16
17use core::fmt;
18use core::str::FromStr;
19
20use serde::{Deserialize, Serialize};
21
22use crate::collector::get_collector;
23#[cfg(feature = "enable")]
24use crate::span::CURRENT_SPAN;
25
26/// An identifier for a trace, which groups a set of related spans together.
27#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
28pub struct TraceId(pub u128);
29
30impl TraceId {
31    /// Uses the state in the collector to generate a `TraceId`.
32    ///
33    /// Returns 0 if the collector has not been initialized via [`crate::collector::set_exporter`].
34    #[inline]
35    pub fn generate() -> Self {
36        get_collector().generate_trace_id()
37    }
38}
39
40impl fmt::Display for TraceId {
41    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
42        write!(f, "{:032x}", self.0)
43    }
44}
45
46impl FromStr for TraceId {
47    type Err = core::num::ParseIntError;
48
49    fn from_str(s: &str) -> Result<Self, Self::Err> {
50        u128::from_str_radix(s, 16).map(TraceId)
51    }
52}
53
54impl serde::Serialize for TraceId {
55    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
56    where
57        S: serde::Serializer,
58    {
59        let mut hex_bytes = [0u8; size_of::<u128>() * 2];
60        hex::encode_to_slice(self.0.to_le_bytes(), &mut hex_bytes).unwrap();
61
62        serializer.serialize_str(str::from_utf8(&hex_bytes).unwrap())
63    }
64}
65
66impl<'de> serde::Deserialize<'de> for TraceId {
67    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
68    where
69        D: serde::Deserializer<'de>,
70    {
71        let bytes: [u8; size_of::<u128>()] = hex::serde::deserialize(deserializer)?;
72
73        Ok(TraceId(u128::from_le_bytes(bytes)))
74    }
75}
76
77/// An identifier for a span within a trace.
78#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
79pub struct SpanId(pub u64);
80
81#[cfg(feature = "enable")]
82impl SpanId {
83    #[inline]
84    #[doc(hidden)]
85    /// Creates a non-zero `SpanId`
86    pub fn next_id() -> SpanId {
87        #[cfg(feature = "std")]
88        {
89            std::thread_local! {
90                static LOCAL_ID_GENERATOR: core::cell::Cell<u64> = core::cell::Cell::new(rand::random::<u64>() & 0xffffffff00000000);
91            }
92
93            LOCAL_ID_GENERATOR
94                .try_with(|g| {
95                    let id = g.get().wrapping_add(1);
96                    g.set(id);
97
98                    SpanId(id)
99                })
100                .unwrap_or({
101                    // This only gets called if the TLS key has been destroyed, it should be safe to fall back to a 0
102                    // value (noop) `SpanId`
103                    SpanId(0)
104                })
105        }
106
107        #[cfg(not(feature = "std"))]
108        {
109            use core::sync::atomic;
110            // For no_std, use a simple counter approach
111            static COUNTER: atomic::AtomicU64 = atomic::AtomicU64::new(1);
112            SpanId(COUNTER.fetch_add(1, atomic::Ordering::Relaxed))
113        }
114    }
115}
116
117impl fmt::Display for SpanId {
118    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
119        write!(f, "{:016x}", self.0)
120    }
121}
122
123impl FromStr for SpanId {
124    type Err = core::num::ParseIntError;
125
126    fn from_str(s: &str) -> Result<Self, Self::Err> {
127        u64::from_str_radix(s, 16).map(SpanId)
128    }
129}
130
131impl serde::Serialize for SpanId {
132    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
133    where
134        S: serde::Serializer,
135    {
136        let mut hex_bytes = [0u8; size_of::<u64>() * 2];
137        hex::encode_to_slice(self.0.to_le_bytes(), &mut hex_bytes).unwrap();
138
139        serializer.serialize_str(str::from_utf8(&hex_bytes).unwrap())
140    }
141}
142
143impl<'de> serde::Deserialize<'de> for SpanId {
144    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
145    where
146        D: serde::Deserializer<'de>,
147    {
148        let bytes: [u8; size_of::<u64>()] = hex::serde::deserialize(deserializer)?;
149
150        Ok(SpanId(u64::from_le_bytes(bytes)))
151    }
152}
153
154/// A struct representing the context of a span, including its [`TraceId`] and [`SpanId`].
155///
156/// [`TraceId`]: crate::id::TraceId
157/// [`SpanId`]: crate::id::SpanId
158#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Serialize, Deserialize)]
159pub struct SpanContext {
160    /// The trace ID this span belongs to
161    pub trace_id: TraceId,
162    /// The unique ID of this span
163    pub span_id: SpanId,
164}
165
166impl SpanContext {
167    /// Creates a new `SpanContext` with the given [`TraceId`] and [`SpanId`].
168    ///
169    /// # Examples
170    ///
171    /// ```
172    /// use veecle_telemetry::id::*;
173    ///
174    /// let span_context = SpanContext::new(TraceId(12), SpanId(13));
175    /// ```
176    ///
177    /// [`TraceId`]: crate::id::TraceId
178    /// [`SpanId`]: crate::id::SpanId
179    pub fn new(trace_id: TraceId, span_id: SpanId) -> Self {
180        Self { trace_id, span_id }
181    }
182
183    /// Creates a new `SpanContext` with a `TraceId` generated with the state in the collector.
184    ///
185    /// Returns 0 if the collector has not been initialized via [`crate::collector::set_exporter`].
186    ///
187    /// # Examples
188    ///
189    /// ```
190    /// use veecle_telemetry::*;
191    ///
192    /// let random = SpanContext::generate();
193    /// ```
194    pub fn generate() -> Self {
195        Self {
196            trace_id: TraceId::generate(),
197            span_id: SpanId(0),
198        }
199    }
200
201    /// Creates a `SpanContext` from the current local parent span. If there is no
202    /// local parent span, this function will return `None`.
203    ///
204    /// # Examples
205    ///
206    /// ```
207    /// use veecle_telemetry::*;
208    ///
209    /// let span = Span::new("root", &[]);
210    /// let _guard = span.entered();
211    ///
212    /// let span_context = SpanContext::current();
213    /// assert!(span_context.is_some());
214    /// ```
215    pub fn current() -> Option<Self> {
216        #[cfg(not(feature = "enable"))]
217        {
218            None
219        }
220
221        #[cfg(feature = "enable")]
222        {
223            CURRENT_SPAN.get()
224        }
225    }
226}
227
228#[cfg(all(test, feature = "std"))]
229mod tests {
230    use std::collections::HashSet;
231    use std::format;
232    use std::string::String;
233    use std::vec::Vec;
234
235    use super::*;
236
237    #[test]
238    #[cfg(not(miri))] // VERY slow with Miri.
239    #[allow(clippy::needless_collect)]
240    fn unique_id() {
241        let handles = std::iter::repeat_with(|| {
242            std::thread::spawn(|| {
243                std::iter::repeat_with(SpanId::next_id)
244                    .take(1000)
245                    .collect::<Vec<_>>()
246            })
247        })
248        .take(32)
249        .collect::<Vec<_>>();
250
251        let k = handles
252            .into_iter()
253            .flat_map(|h| h.join().unwrap())
254            .collect::<HashSet<_>>();
255
256        assert_eq!(k.len(), 32 * 1000);
257    }
258
259    #[test]
260    fn trace_id_formatting() {
261        assert_eq!(
262            format!("{}", TraceId(0)),
263            "00000000000000000000000000000000"
264        );
265        assert_eq!(
266            format!("{}", TraceId(u128::MAX)),
267            "ffffffffffffffffffffffffffffffff"
268        );
269        assert_eq!(
270            format!("{}", TraceId(0x123456789ABCDEF0FEDCBA9876543210)),
271            "123456789abcdef0fedcba9876543210"
272        );
273        assert_eq!(
274            format!("{}", TraceId(0x123)),
275            "00000000000000000000000000000123"
276        );
277    }
278
279    #[test]
280    fn span_id_formatting() {
281        assert_eq!(format!("{}", SpanId(0)), "0000000000000000");
282        assert_eq!(format!("{}", SpanId(u64::MAX)), "ffffffffffffffff");
283        assert_eq!(
284            format!("{}", SpanId(0xFEDCBA9876543210)),
285            "fedcba9876543210"
286        );
287        assert_eq!(format!("{}", SpanId(0x123)), "0000000000000123");
288    }
289
290    #[test]
291    fn trace_id_from_str() {
292        assert_eq!(
293            "123456789abcdef0fedcba9876543210"
294                .parse::<TraceId>()
295                .unwrap(),
296            TraceId(0x123456789ABCDEF0FEDCBA9876543210)
297        );
298        assert_eq!(
299            "123456789ABCDEF0FEDCBA9876543210"
300                .parse::<TraceId>()
301                .unwrap(),
302            TraceId(0x123456789ABCDEF0FEDCBA9876543210)
303        );
304        assert_eq!(
305            "00000000000000000000000000000000"
306                .parse::<TraceId>()
307                .unwrap(),
308            TraceId(0)
309        );
310        assert_eq!(
311            "ffffffffffffffffffffffffffffffff"
312                .parse::<TraceId>()
313                .unwrap(),
314            TraceId(u128::MAX)
315        );
316        // Shorter hex string works as u128::from_str_radix handles it
317        assert_eq!("123".parse::<TraceId>().unwrap(), TraceId(0x123));
318
319        assert!("xyz".parse::<TraceId>().is_err());
320        assert!("".parse::<TraceId>().is_err());
321    }
322
323    #[test]
324    fn span_id_from_str() {
325        assert_eq!(
326            "fedcba9876543210".parse::<SpanId>().unwrap(),
327            SpanId(0xFEDCBA9876543210)
328        );
329        assert_eq!(
330            "FEDCBA9876543210".parse::<SpanId>().unwrap(),
331            SpanId(0xFEDCBA9876543210)
332        );
333        assert_eq!("0000000000000000".parse::<SpanId>().unwrap(), SpanId(0));
334        assert_eq!(
335            "ffffffffffffffff".parse::<SpanId>().unwrap(),
336            SpanId(u64::MAX)
337        );
338        assert_eq!("123".parse::<SpanId>().unwrap(), SpanId(0x123));
339
340        assert!("xyz".parse::<SpanId>().is_err());
341        assert!("".parse::<SpanId>().is_err());
342    }
343
344    #[test]
345    fn trace_id_format_from_str_roundtrip() {
346        let test_cases = [
347            0u128,
348            1,
349            0x123,
350            0x123456789ABCDEF0FEDCBA9876543210,
351            u128::MAX,
352            u128::MAX - 1,
353        ];
354
355        for value in test_cases {
356            let trace_id = TraceId(value);
357            let formatted = format!("{trace_id}");
358            let parsed = formatted.parse::<TraceId>().unwrap();
359            assert_eq!(trace_id, parsed, "Failed roundtrip for value {value:#x}");
360        }
361    }
362
363    #[test]
364    fn span_id_format_from_str_roundtrip() {
365        let test_cases = [0u64, 1, 0x123, 0xFEDCBA9876543210, u64::MAX, u64::MAX - 1];
366
367        for value in test_cases {
368            let span_id = SpanId(value);
369            let formatted = format!("{span_id}");
370            let parsed = formatted.parse::<SpanId>().unwrap();
371            assert_eq!(span_id, parsed, "Failed roundtrip for value {value:#x}");
372        }
373    }
374
375    #[test]
376    fn trace_id_serde_roundtrip() {
377        let test_cases = [
378            TraceId(0),
379            TraceId(1),
380            TraceId(0x123),
381            TraceId(0x123456789ABCDEF0FEDCBA9876543210),
382            TraceId(u128::MAX),
383            TraceId(u128::MAX - 1),
384        ];
385
386        for original in test_cases {
387            let json = serde_json::to_string(&original).unwrap();
388            let deserialized: TraceId = serde_json::from_str(&json).unwrap();
389            assert_eq!(
390                original, deserialized,
391                "JSON roundtrip failed for {:#x}",
392                original.0
393            );
394        }
395    }
396
397    #[test]
398    fn span_id_serde_roundtrip() {
399        let test_cases = [
400            SpanId(0),
401            SpanId(1),
402            SpanId(0x123),
403            SpanId(0xFEDCBA9876543210),
404            SpanId(u64::MAX),
405            SpanId(u64::MAX - 1),
406        ];
407
408        for original in test_cases {
409            let json = serde_json::to_string(&original).unwrap();
410            let deserialized: SpanId = serde_json::from_str(&json).unwrap();
411            assert_eq!(
412                original, deserialized,
413                "JSON roundtrip failed for {:#x}",
414                original.0
415            );
416        }
417    }
418
419    #[test]
420    fn span_context_serde_roundtrip() {
421        let test_cases = [
422            SpanContext::new(TraceId(0), SpanId(0)),
423            SpanContext::new(
424                TraceId(0x123456789ABCDEF0FEDCBA9876543210),
425                SpanId(0xFEDCBA9876543210),
426            ),
427            SpanContext::new(TraceId(u128::MAX), SpanId(u64::MAX)),
428            SpanContext::new(TraceId(1), SpanId(1)),
429        ];
430
431        for original in test_cases {
432            let json = serde_json::to_string(&original).unwrap();
433            let deserialized: SpanContext = serde_json::from_str(&json).unwrap();
434            assert_eq!(
435                original.trace_id, deserialized.trace_id,
436                "JSON roundtrip failed for trace_id"
437            );
438            assert_eq!(
439                original.span_id, deserialized.span_id,
440                "JSON roundtrip failed for span_id"
441            );
442        }
443    }
444
445    #[test]
446    fn trace_id_serialization_format() {
447        let trace_id = TraceId(0x123456789ABCDEF0FEDCBA9876543210);
448        let json = serde_json::to_string(&trace_id).unwrap();
449
450        // Serialization uses little-endian bytes
451        let expected_le_bytes = 0x123456789ABCDEF0FEDCBA9876543210u128.to_le_bytes();
452        let mut expected_hex = String::new();
453        for byte in &expected_le_bytes {
454            expected_hex.push_str(&format!("{byte:02x}"));
455        }
456        let expected_json = format!("\"{expected_hex}\"");
457
458        assert_eq!(json, expected_json);
459    }
460
461    #[test]
462    fn span_id_serialization_format() {
463        let span_id = SpanId(0xFEDCBA9876543210);
464        let json = serde_json::to_string(&span_id).unwrap();
465
466        let expected_le_bytes = 0xFEDCBA9876543210u64.to_le_bytes();
467        let mut expected_hex = String::new();
468        for byte in &expected_le_bytes {
469            expected_hex.push_str(&format!("{byte:02x}"));
470        }
471        let expected_json = format!("\"{expected_hex}\"");
472
473        assert_eq!(json, expected_json);
474    }
475
476    #[test]
477    fn span_context_new_and_fields() {
478        let trace_id = TraceId(0x123);
479        let span_id = SpanId(0x456);
480        let context = SpanContext::new(trace_id, span_id);
481
482        assert_eq!(context.trace_id, trace_id);
483        assert_eq!(context.span_id, span_id);
484    }
485
486    #[test]
487    fn span_context_generate_produces_non_zero_trace_id() {
488        let context = SpanContext::generate();
489        // span_id should be 0 as per the implementation
490        assert_eq!(context.span_id, SpanId(0));
491        // trace_id value depends on collector initialization
492    }
493
494    #[test]
495    fn span_id_next_id_produces_non_zero_values() {
496        let ids: Vec<SpanId> = (0..100).map(|_| SpanId::next_id()).collect();
497
498        for id in &ids {
499            assert_ne!(id.0, 0, "SpanId::next_id() should not produce zero values");
500        }
501
502        let mut unique_ids = HashSet::new();
503        for id in &ids {
504            assert!(
505                unique_ids.insert(id.0),
506                "SpanId::next_id() should produce unique values"
507            );
508        }
509    }
510}
veecle_telemetry/id.rs

veecle_telemetry/
id.rs
