grafos_observe/

trace.rs

//! Distributed trace context types for W3C traceparent propagation.
//!
//! All types are `no_std` compatible. The [`TraceContext`] struct carries a
//! 16-byte trace ID, 8-byte span ID, and 1-byte flags field conforming to
//! the W3C Trace Context specification.
//!
//! # Binary encoding
//!
//! ```text
//! [version: 1B] [trace_id: 16B] [span_id: 8B] [flags: 1B] [reserved: 6B]
//! ```
//!
//! Total: 32 bytes (fits in a single cache line, suitable for shared-memory
//! RPC header extension).
//!
//! # W3C string format
//!
//! ```text
//! 00-{trace_id:032x}-{span_id:016x}-{flags:02x}
//! ```

use core::fmt;

/// 16-byte trace identifier, unique per trace.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub struct TraceId(pub u128);

impl TraceId {
    /// The zero (invalid) trace ID.
    pub const INVALID: TraceId = TraceId(0);

    /// Whether this trace ID is the zero/invalid value.
    pub fn is_valid(&self) -> bool {
        self.0 != 0
    }
}

impl fmt::Debug for TraceId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "TraceId({:032x})", self.0)
    }
}

impl fmt::Display for TraceId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:032x}", self.0)
    }
}

impl fmt::LowerHex for TraceId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::LowerHex::fmt(&self.0, f)
    }
}

/// 8-byte span identifier, unique per span within a trace.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub struct SpanId(pub u64);

impl SpanId {
    /// The zero (invalid) span ID.
    pub const INVALID: SpanId = SpanId(0);

    /// Whether this span ID is the zero/invalid value.
    pub fn is_valid(&self) -> bool {
        self.0 != 0
    }
}

impl fmt::Debug for SpanId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "SpanId({:016x})", self.0)
    }
}

impl fmt::Display for SpanId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:016x}", self.0)
    }
}

impl fmt::LowerHex for SpanId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::LowerHex::fmt(&self.0, f)
    }
}

/// W3C trace context in binary form.
///
/// Carries the identifiers needed to correlate spans across service boundaries.
/// The `flags` field uses bit 0 as the "sampled" flag (W3C traceparent spec).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct TraceContext {
    /// Trace-wide identifier (shared by all spans in the trace).
    pub trace_id: TraceId,
    /// Span identifier for the current span.
    pub span_id: SpanId,
    /// Parent span identifier (zero if this is the root span).
    pub parent_span_id: SpanId,
    /// W3C trace flags. Bit 0 = sampled.
    pub flags: u8,
}

/// W3C traceparent version byte.
const W3C_VERSION: u8 = 0x00;

/// Flag bit: trace is sampled.
pub const FLAG_SAMPLED: u8 = 0x01;

/// Binary encoding size.
pub const TRACE_CONTEXT_BYTES: usize = 32;

impl TraceContext {
    /// Create a new root trace context with the given random source.
    ///
    /// Generates fresh trace_id and span_id from the provided random bytes.
    /// `random_bytes` must be at least 24 bytes (16 for trace_id + 8 for span_id).
    pub fn new_root(random_bytes: &[u8; 24]) -> Self {
        let trace_id = u128::from_be_bytes(random_bytes[0..16].try_into().unwrap());
        let span_id = u64::from_be_bytes(random_bytes[16..24].try_into().unwrap());
        TraceContext {
            trace_id: TraceId(trace_id),
            span_id: SpanId(span_id),
            parent_span_id: SpanId::INVALID,
            flags: FLAG_SAMPLED,
        }
    }

    /// Create a child context within the same trace.
    ///
    /// The child inherits `trace_id` and `flags`, gets a new `span_id` from
    /// the provided random bytes, and records the current `span_id` as its parent.
    pub fn child(&self, random_span_bytes: &[u8; 8]) -> Self {
        let new_span_id = u64::from_be_bytes(*random_span_bytes);
        TraceContext {
            trace_id: self.trace_id,
            span_id: SpanId(new_span_id),
            parent_span_id: self.span_id,
            flags: self.flags,
        }
    }

    /// Whether the sampled flag is set.
    pub fn is_sampled(&self) -> bool {
        self.flags & FLAG_SAMPLED != 0
    }

    /// Set or clear the sampled flag.
    pub fn set_sampled(&mut self, sampled: bool) {
        if sampled {
            self.flags |= FLAG_SAMPLED;
        } else {
            self.flags &= !FLAG_SAMPLED;
        }
    }

    /// Encode to a 32-byte binary representation.
    ///
    /// Layout: `[version:1][trace_id:16][span_id:8][flags:1][reserved:6]`
    pub fn encode(&self) -> [u8; TRACE_CONTEXT_BYTES] {
        let mut buf = [0u8; TRACE_CONTEXT_BYTES];
        buf[0] = W3C_VERSION;
        buf[1..17].copy_from_slice(&self.trace_id.0.to_be_bytes());
        buf[17..25].copy_from_slice(&self.span_id.0.to_be_bytes());
        buf[25] = self.flags;
        // bytes 26..32 reserved (zeros)
        buf
    }

    /// Decode from a 32-byte binary representation.
    pub fn decode(buf: &[u8; TRACE_CONTEXT_BYTES]) -> Result<Self, TraceContextError> {
        if buf[0] != W3C_VERSION {
            return Err(TraceContextError::UnsupportedVersion(buf[0]));
        }
        let trace_id = u128::from_be_bytes(buf[1..17].try_into().unwrap());
        let span_id = u64::from_be_bytes(buf[17..25].try_into().unwrap());
        let flags = buf[25];

        Ok(TraceContext {
            trace_id: TraceId(trace_id),
            span_id: SpanId(span_id),
            parent_span_id: SpanId::INVALID, // not stored in wire format
            flags,
        })
    }

    /// Format as a W3C traceparent string: `00-{trace_id}-{span_id}-{flags}`.
    ///
    /// Returns a fixed-length 55-character string.
    pub fn to_w3c_string(&self) -> alloc::string::String {
        alloc::format!(
            "{:02x}-{:032x}-{:016x}-{:02x}",
            W3C_VERSION,
            self.trace_id.0,
            self.span_id.0,
            self.flags,
        )
    }

    /// Parse a W3C traceparent string.
    ///
    /// Expected format: `{version:2}-{trace_id:32}-{span_id:16}-{flags:2}`
    /// (55 characters total).
    pub fn from_w3c_string(s: &str) -> Result<Self, TraceContextError> {
        if s.len() != 55 {
            return Err(TraceContextError::InvalidFormat);
        }
        let parts: alloc::vec::Vec<&str> = s.split('-').collect();
        if parts.len() != 4 {
            return Err(TraceContextError::InvalidFormat);
        }
        if parts[0].len() != 2
            || parts[1].len() != 32
            || parts[2].len() != 16
            || parts[3].len() != 2
        {
            return Err(TraceContextError::InvalidFormat);
        }
        let version =
            u8::from_str_radix(parts[0], 16).map_err(|_| TraceContextError::InvalidFormat)?;
        if version != W3C_VERSION {
            return Err(TraceContextError::UnsupportedVersion(version));
        }
        let trace_id =
            u128::from_str_radix(parts[1], 16).map_err(|_| TraceContextError::InvalidFormat)?;
        let span_id =
            u64::from_str_radix(parts[2], 16).map_err(|_| TraceContextError::InvalidFormat)?;
        let flags =
            u8::from_str_radix(parts[3], 16).map_err(|_| TraceContextError::InvalidFormat)?;

        Ok(TraceContext {
            trace_id: TraceId(trace_id),
            span_id: SpanId(span_id),
            parent_span_id: SpanId::INVALID,
            flags,
        })
    }

    /// Check if this context has all-zero IDs (i.e. no trace context was set).
    pub fn is_empty(&self) -> bool {
        !self.trace_id.is_valid() && !self.span_id.is_valid()
    }
}

impl Default for TraceContext {
    fn default() -> Self {
        TraceContext {
            trace_id: TraceId::INVALID,
            span_id: SpanId::INVALID,
            parent_span_id: SpanId::INVALID,
            flags: 0,
        }
    }
}

impl fmt::Display for TraceContext {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.to_w3c_string())
    }
}

/// Errors from trace context encoding/decoding.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TraceContextError {
    /// The version byte is not supported.
    UnsupportedVersion(u8),
    /// The string format is invalid.
    InvalidFormat,
}

impl fmt::Display for TraceContextError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            TraceContextError::UnsupportedVersion(v) => {
                write!(f, "unsupported traceparent version: {v:#04x}")
            }
            TraceContextError::InvalidFormat => write!(f, "invalid traceparent format"),
        }
    }
}

// ---------------------------------------------------------------------------
// slice 224 — bridge between runtime TraceContext (this module) and the
// v1.1 §4.3 graph-layer EdgeTraceContext in grafos-core.
//
// The two types coexist by design (see grafos_core::trace module docs):
// - This `TraceContext` has u8 flags matching the W3C traceparent wire
//   format the observability runtime decodes.
// - `EdgeTraceContext` has u32 flags per v1.1 §4.3 to carry the wider
//   schema-level flag space on every edge.
//
// The bridge gives callers two typed directions:
//   - `From<TraceContext> for EdgeTraceContext`: lossless u8 → u32
//     widening. Always succeeds.
//   - `TryFrom<EdgeTraceContext> for TraceContext`: u32 → u8 narrowing
//     that fails closed when high bits are set (per v1.1 §3.3 "no
//     silent best-effort loss" and the project CLAUDE.md global rule
//     against fallbacks circumventing correct system behavior).
//
// Orphan-rule note: both impls live here because grafos-core cannot
// see grafos_observe::TraceContext (would be cyclic), and From/TryFrom
// must be defined in the crate that owns either the source or target.
// `EdgeTraceContext` is from grafos-core; `TraceContext` is from this
// crate — so this crate is the right home.
// ---------------------------------------------------------------------------

impl From<TraceContext> for grafos_core::EdgeTraceContext {
    fn from(ctx: TraceContext) -> Self {
        grafos_core::EdgeTraceContext {
            trace_id: ctx.trace_id.0,
            span_id: ctx.span_id.0,
            parent_span_id: ctx.parent_span_id.0,
            // u8 → u32: lossless widening. High bits stay zero.
            flags: ctx.flags as u32,
        }
    }
}

/// Error returned when narrowing an [`grafos_core::EdgeTraceContext`]
/// to a runtime [`TraceContext`] would discard set high bits in the
/// `flags` field. The full u32 value is preserved on the error so
/// callers can log/audit what was rejected.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct EdgeTraceContextToObserveError {
    /// The original u32 flags value that did not fit in the
    /// runtime's u8 flags slot.
    pub flags: u32,
}

impl fmt::Display for EdgeTraceContextToObserveError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "EdgeTraceContext → TraceContext: flags 0x{:08x} has high bits set; runtime TraceContext is W3C-shaped with u8 flags",
            self.flags
        )
    }
}

#[cfg(feature = "std")]
impl std::error::Error for EdgeTraceContextToObserveError {}

impl TryFrom<grafos_core::EdgeTraceContext> for TraceContext {
    type Error = EdgeTraceContextToObserveError;

    fn try_from(ctx: grafos_core::EdgeTraceContext) -> Result<Self, Self::Error> {
        // u32 → u8: fail closed if any high bit is set so the bridge
        // never silently drops semantics the schema is allowed to
        // carry but the runtime wire format isn't.
        if ctx.flags > u8::MAX as u32 {
            return Err(EdgeTraceContextToObserveError { flags: ctx.flags });
        }
        Ok(TraceContext {
            trace_id: TraceId(ctx.trace_id),
            span_id: SpanId(ctx.span_id),
            parent_span_id: SpanId(ctx.parent_span_id),
            flags: ctx.flags as u8,
        })
    }
}

// ---------------------------------------------------------------------------
// Thread-local current trace context (requires std)
// ---------------------------------------------------------------------------

#[cfg(feature = "std")]
mod thread_local_ctx {
    use super::TraceContext;
    use std::cell::RefCell;

    thread_local! {
        static CURRENT: RefCell<TraceContext> = RefCell::new(TraceContext::default());
    }

    impl TraceContext {
        /// Get the current thread-local trace context.
        pub fn current() -> TraceContext {
            CURRENT.with(|c| *c.borrow())
        }

        /// Set the current thread-local trace context.
        pub fn set_current(ctx: TraceContext) {
            CURRENT.with(|c| *c.borrow_mut() = ctx);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_random_bytes() -> [u8; 24] {
        // Deterministic test bytes
        let mut bytes = [0u8; 24];
        for (i, b) in bytes.iter_mut().enumerate() {
            *b = (i as u8).wrapping_add(0x10);
        }
        bytes
    }

    #[test]
    fn new_root_creates_valid_context() {
        let ctx = TraceContext::new_root(&test_random_bytes());
        assert!(ctx.trace_id.is_valid());
        assert!(ctx.span_id.is_valid());
        assert_eq!(ctx.parent_span_id, SpanId::INVALID);
        assert!(ctx.is_sampled());
    }

    #[test]
    fn child_inherits_trace_id() {
        let root = TraceContext::new_root(&test_random_bytes());
        let child_bytes: [u8; 8] = [0xAA; 8];
        let child = root.child(&child_bytes);

        assert_eq!(child.trace_id, root.trace_id);
        assert_ne!(child.span_id, root.span_id);
        assert_eq!(child.parent_span_id, root.span_id);
        assert_eq!(child.flags, root.flags);
    }

    #[test]
    fn encode_decode_roundtrip() {
        let ctx = TraceContext::new_root(&test_random_bytes());
        let encoded = ctx.encode();
        let decoded = TraceContext::decode(&encoded).unwrap();

        assert_eq!(decoded.trace_id, ctx.trace_id);
        assert_eq!(decoded.span_id, ctx.span_id);
        assert_eq!(decoded.flags, ctx.flags);
    }

    #[test]
    fn w3c_string_roundtrip() {
        let ctx = TraceContext::new_root(&test_random_bytes());
        let s = ctx.to_w3c_string();

        assert_eq!(s.len(), 55);
        assert!(s.starts_with("00-"));

        let parsed = TraceContext::from_w3c_string(&s).unwrap();
        assert_eq!(parsed.trace_id, ctx.trace_id);
        assert_eq!(parsed.span_id, ctx.span_id);
        assert_eq!(parsed.flags, ctx.flags);
    }

    #[test]
    fn w3c_string_format() {
        let ctx = TraceContext {
            trace_id: TraceId(0x0102030405060708090a0b0c0d0e0f10),
            span_id: SpanId(0x1112131415161718),
            parent_span_id: SpanId::INVALID,
            flags: 0x01,
        };
        let s = ctx.to_w3c_string();
        assert_eq!(s, "00-0102030405060708090a0b0c0d0e0f10-1112131415161718-01");
    }

    #[test]
    fn from_w3c_string_rejects_bad_version() {
        let s = "01-0102030405060708090a0b0c0d0e0f10-1112131415161718-01";
        let err = TraceContext::from_w3c_string(s).unwrap_err();
        assert_eq!(err, TraceContextError::UnsupportedVersion(0x01));
    }

    #[test]
    fn from_w3c_string_rejects_bad_length() {
        assert_eq!(
            TraceContext::from_w3c_string("too-short"),
            Err(TraceContextError::InvalidFormat)
        );
    }

    #[test]
    fn sampled_flag_manipulation() {
        let mut ctx = TraceContext::default();
        assert!(!ctx.is_sampled());
        ctx.set_sampled(true);
        assert!(ctx.is_sampled());
        ctx.set_sampled(false);
        assert!(!ctx.is_sampled());
    }

    #[test]
    fn default_is_empty() {
        let ctx = TraceContext::default();
        assert!(ctx.is_empty());
    }

    #[test]
    fn decode_rejects_bad_version() {
        let mut buf = [0u8; 32];
        buf[0] = 0xFF;
        assert_eq!(
            TraceContext::decode(&buf),
            Err(TraceContextError::UnsupportedVersion(0xFF))
        );
    }

    #[cfg(feature = "std")]
    #[test]
    fn thread_local_current_context() {
        // Save and restore to avoid polluting other tests
        let saved = TraceContext::current();

        let ctx = TraceContext::new_root(&test_random_bytes());
        TraceContext::set_current(ctx);
        let got = TraceContext::current();
        assert_eq!(got.trace_id, ctx.trace_id);
        assert_eq!(got.span_id, ctx.span_id);

        TraceContext::set_current(saved);
    }

    #[test]
    fn display_uses_w3c_format() {
        let ctx = TraceContext::new_root(&test_random_bytes());
        let display = alloc::format!("{ctx}");
        let w3c = ctx.to_w3c_string();
        assert_eq!(display, w3c);
    }

    #[test]
    fn trace_id_display() {
        let tid = TraceId(0x0102030405060708090a0b0c0d0e0f10);
        assert_eq!(alloc::format!("{tid}"), "0102030405060708090a0b0c0d0e0f10");
    }

    #[test]
    fn span_id_display() {
        let sid = SpanId(0x1112131415161718);
        assert_eq!(alloc::format!("{sid}"), "1112131415161718");
    }

    // -----------------------------------------------------------------
    // slice 224 — EdgeTraceContext ↔ TraceContext bridge pin tests.
    //
    // The runtime `TraceContext` has u8 flags (W3C wire) while the
    // graph-layer `grafos_core::EdgeTraceContext` has u32 flags (v1.1
    // §4.3 schema). The bridge fails closed on the narrowing direction
    // when high bits are set, and is lossless on the widening
    // direction.
    // -----------------------------------------------------------------

    #[test]
    fn from_trace_context_widens_flags_losslessly() {
        let ctx = TraceContext {
            trace_id: TraceId(0xdeadbeef_cafe_babe_0123_4567_89ab_cdef),
            span_id: SpanId(0x1234_5678_9abc_def0),
            parent_span_id: SpanId(0xfeed_face_dead_beef),
            flags: FLAG_SAMPLED,
        };
        let edge: grafos_core::EdgeTraceContext = ctx.into();
        assert_eq!(edge.trace_id, ctx.trace_id.0);
        assert_eq!(edge.span_id, ctx.span_id.0);
        assert_eq!(edge.parent_span_id, ctx.parent_span_id.0);
        assert_eq!(edge.flags, FLAG_SAMPLED as u32);
        // High bits are zero — u8 → u32 widening can never set them.
        assert_eq!(edge.flags & 0xffff_ff00, 0);
    }

    #[test]
    fn from_trace_context_preserves_zero_flags() {
        let ctx = TraceContext::default();
        let edge: grafos_core::EdgeTraceContext = ctx.into();
        assert_eq!(edge.flags, 0);
        assert!(edge.is_zero());
    }

    #[test]
    fn from_trace_context_preserves_max_u8_flags() {
        let ctx = TraceContext {
            trace_id: TraceId(1),
            span_id: SpanId(2),
            parent_span_id: SpanId(3),
            flags: u8::MAX,
        };
        let edge: grafos_core::EdgeTraceContext = ctx.into();
        assert_eq!(edge.flags, u8::MAX as u32);
        assert_eq!(edge.flags, 0x0000_00ff);
    }

    #[test]
    fn try_from_edge_trace_context_narrows_u8_range_losslessly() {
        let edge = grafos_core::EdgeTraceContext {
            trace_id: 0xa5a5_a5a5_a5a5_a5a5_a5a5_a5a5_a5a5_a5a5,
            span_id: 0x0123_4567_89ab_cdef,
            parent_span_id: 0xfedc_ba98_7654_3210,
            flags: grafos_core::EDGE_TRACE_FLAG_SAMPLED,
        };
        let ctx = TraceContext::try_from(edge).expect("u8-fitting flags must narrow");
        assert_eq!(ctx.trace_id.0, edge.trace_id);
        assert_eq!(ctx.span_id.0, edge.span_id);
        assert_eq!(ctx.parent_span_id.0, edge.parent_span_id);
        assert_eq!(ctx.flags, FLAG_SAMPLED);
    }

    #[test]
    fn try_from_edge_trace_context_fails_closed_on_high_bits() {
        // High bit (32) set — would silently drop in a u32→u8 cast.
        // The bridge must reject this rather than narrowing to 0.
        let edge = grafos_core::EdgeTraceContext {
            flags: 0x8000_0000,
            ..grafos_core::EdgeTraceContext::default()
        };
        let err = TraceContext::try_from(edge).expect_err("high bits must fail closed");
        assert_eq!(err, EdgeTraceContextToObserveError { flags: 0x8000_0000 });
    }

    #[test]
    fn try_from_edge_trace_context_fails_closed_on_mid_bits() {
        // Bit 8 set — just above u8::MAX. This is the canonical
        // "silently truncates to 0" case the bridge protects against.
        let edge = grafos_core::EdgeTraceContext {
            flags: 0x0000_0100,
            ..grafos_core::EdgeTraceContext::default()
        };
        let err = TraceContext::try_from(edge).expect_err("bit 8 must fail closed");
        assert_eq!(err.flags, 0x0000_0100);
    }

    #[test]
    fn try_from_edge_trace_context_at_u8_max_boundary() {
        // 0xff fits exactly; 0x100 does not. Pin the boundary.
        let at_max = grafos_core::EdgeTraceContext {
            flags: 0xff,
            ..grafos_core::EdgeTraceContext::default()
        };
        let just_over = grafos_core::EdgeTraceContext {
            flags: 0x100,
            ..grafos_core::EdgeTraceContext::default()
        };
        assert!(TraceContext::try_from(at_max).is_ok());
        assert!(TraceContext::try_from(just_over).is_err());
    }

    #[test]
    fn bridge_round_trip_is_identity_in_u8_range() {
        // Any TraceContext can round-trip TC → EdgeTraceContext → TC
        // because the widening direction is lossless and the narrowing
        // direction succeeds for any u32 value that fits in u8 (which
        // every u8 value does by construction).
        let original = TraceContext {
            trace_id: TraceId(0x1111_2222_3333_4444_5555_6666_7777_8888),
            span_id: SpanId(0x9999_aaaa_bbbb_cccc),
            parent_span_id: SpanId(0xdddd_eeee_ffff_0000),
            flags: 0x42,
        };
        let edge: grafos_core::EdgeTraceContext = original.into();
        let round: TraceContext =
            TraceContext::try_from(edge).expect("u8-sourced round-trip must succeed");
        assert_eq!(round.trace_id, original.trace_id);
        assert_eq!(round.span_id, original.span_id);
        assert_eq!(round.parent_span_id, original.parent_span_id);
        assert_eq!(round.flags, original.flags);
    }

    #[test]
    fn bridge_error_display_is_human_readable() {
        let err = EdgeTraceContextToObserveError { flags: 0xdead_beef };
        let s = alloc::format!("{err}");
        assert!(
            s.contains("0xdeadbeef"),
            "error must surface the rejected flags value: {s}"
        );
        assert!(
            s.contains("u8") || s.contains("W3C"),
            "error must explain why narrowing failed: {s}"
        );
    }
}