alloy-telemetry 0.3.0

//! Owned trace context types for cross-actor trace propagation.
//!
//! This module provides [`OwnedTraceContext`], [`Traceparent`], and [`Traced`] for
//! embedding W3C-compatible trace context in actor message envelopes.

/// Owned trace context suitable for embedding in actor message envelopes.
///
/// Unlike [`crate::api::TraceContext`], all fields are owned `String` values, allowing
/// `OwnedTraceContext` to be cloned and sent across actor boundaries without
/// lifetime constraints.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct OwnedTraceContext {
    /// Stable 32-hex-character trace identifier (W3C trace-id format).
    pub trace_id: String,
    /// Current span identifier (16 hex characters, W3C parent-id format).
    pub span_id: String,
    /// Parent span identifier, if this context was derived from a parent span.
    pub parent_span_id: Option<String>,
    /// Logical operation name for the span that will be created from this context.
    pub operation: String,
    /// Optional higher-level correlation or scenario identifier.
    pub correlation_id: Option<String>,
}

impl OwnedTraceContext {
    /// Builds a root context with no parent span.
    #[must_use]
    pub fn root(
        trace_id: impl Into<String>,
        span_id: impl Into<String>,
        operation: impl Into<String>,
    ) -> Self {
        Self {
            trace_id: trace_id.into(),
            span_id: span_id.into(),
            parent_span_id: None,
            operation: operation.into(),
            correlation_id: None,
        }
    }

    /// Builds a child context derived from this context's span.
    ///
    /// The new context inherits `trace_id` and `correlation_id`, sets `parent_span_id`
    /// to this context's `span_id`, and takes a fresh `span_id` and `operation`.
    #[must_use]
    pub fn child(&self, span_id: impl Into<String>, operation: impl Into<String>) -> Self {
        Self {
            trace_id: self.trace_id.clone(),
            span_id: span_id.into(),
            parent_span_id: Some(self.span_id.clone()),
            operation: operation.into(),
            correlation_id: self.correlation_id.clone(),
        }
    }

    /// Returns a zero-copy [`ChildSpanSpec`] for framework-level span creation.
    #[must_use]
    pub fn as_child_span_spec(&self) -> ChildSpanSpec<'_> {
        ChildSpanSpec {
            trace_id: &self.trace_id,
            parent_span_id: &self.span_id,
            operation: &self.operation,
            correlation_id: self.correlation_id.as_deref(),
        }
    }
}

/// Zero-copy descriptor used by framework code to create a child span.
///
/// Produced by [`OwnedTraceContext::as_child_span_spec`]. Provides the fields
/// needed to construct a child span without cloning strings.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct ChildSpanSpec<'a> {
    /// Trace identifier inherited from the parent context.
    pub trace_id: &'a str,
    /// Parent span identifier (the `span_id` of the sending actor's context).
    pub parent_span_id: &'a str,
    /// Operation name for the new child span.
    pub operation: &'a str,
    /// Propagated correlation identifier, if present.
    pub correlation_id: Option<&'a str>,
}

/// Reason a traceparent header string failed to parse.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum TraceparentParseError {
    /// Header did not contain exactly four dash-separated segments.
    InvalidSegmentCount,
    /// Version segment was not the two-character hex string `"00"`.
    InvalidVersion,
    /// Trace-id segment was not exactly 32 lowercase hex characters.
    InvalidTraceId,
    /// Parent-id segment was not exactly 16 lowercase hex characters.
    InvalidParentId,
    /// Flags segment was not exactly two lowercase hex characters.
    InvalidFlags,
}

impl TraceparentParseError {
    /// Returns a stable string description of the error variant.
    #[must_use]
    pub const fn as_str(&self) -> &'static str {
        match self {
            Self::InvalidSegmentCount => {
                "traceparent must have exactly four dash-separated segments"
            }
            Self::InvalidVersion => "traceparent version must be '00'",
            Self::InvalidTraceId => "trace-id must be exactly 32 lowercase hex characters",
            Self::InvalidParentId => "parent-id must be exactly 16 lowercase hex characters",
            Self::InvalidFlags => "trace-flags must be exactly two lowercase hex characters",
        }
    }
}

fn is_lowercase_hex(s: &str) -> bool {
    s.chars().all(|c| matches!(c, '0'..='9' | 'a'..='f'))
}

/// Parsed W3C traceparent header value.
///
/// Wire format: `{version}-{trace-id}-{parent-id}-{flags}`
///
/// This implementation accepts version `00` only. Future version support can be
/// added without breaking this API.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Traceparent {
    /// Version byte. Must be `0x00` for the current W3C specification.
    pub version: u8,
    /// 32 lowercase hex character trace identifier.
    pub trace_id: String,
    /// 16 lowercase hex character parent-id (the sending span's identifier).
    pub parent_id: String,
    /// Trace flags byte (bit 0 = sampled).
    pub trace_flags: u8,
}

impl Traceparent {
    /// Parses a `traceparent` header string.
    ///
    /// Returns `Err(TraceparentParseError)` if any segment is missing, malformed,
    /// or has the wrong length.
    ///
    /// # Errors
    ///
    /// Returns [`TraceparentParseError`] if the header is not a valid W3C traceparent.
    pub fn parse(header: &str) -> Result<Self, TraceparentParseError> {
        let segments: Vec<&str> = header.split('-').collect();
        if segments.len() != 4 {
            return Err(TraceparentParseError::InvalidSegmentCount);
        }
        let (version_str, trace_id_str, parent_id_str, flags_str) =
            (segments[0], segments[1], segments[2], segments[3]);

        if version_str != "00" {
            return Err(TraceparentParseError::InvalidVersion);
        }
        if trace_id_str.len() != 32 || !is_lowercase_hex(trace_id_str) {
            return Err(TraceparentParseError::InvalidTraceId);
        }
        if parent_id_str.len() != 16 || !is_lowercase_hex(parent_id_str) {
            return Err(TraceparentParseError::InvalidParentId);
        }
        if flags_str.len() != 2 || !is_lowercase_hex(flags_str) {
            return Err(TraceparentParseError::InvalidFlags);
        }
        let trace_flags =
            u8::from_str_radix(flags_str, 16).map_err(|_| TraceparentParseError::InvalidFlags)?;

        Ok(Self {
            version: 0,
            trace_id: trace_id_str.to_owned(),
            parent_id: parent_id_str.to_owned(),
            trace_flags,
        })
    }

    /// Serializes this value to the canonical `traceparent` wire spelling.
    ///
    /// Output format: `00-{trace_id}-{parent_id}-{flags:02x}`
    #[must_use]
    pub fn to_header(&self) -> String {
        format!(
            "00-{}-{}-{:02x}",
            self.trace_id, self.parent_id, self.trace_flags
        )
    }

    /// Converts into an [`OwnedTraceContext`] for embedding in a message envelope.
    ///
    /// The `parent_id` field becomes the `span_id`; `parent_span_id` is `None` because
    /// the traceparent header itself carries the parent reference for the next hop.
    #[must_use]
    pub fn into_owned_trace_context(self, operation: impl Into<String>) -> OwnedTraceContext {
        OwnedTraceContext {
            trace_id: self.trace_id,
            span_id: self.parent_id,
            parent_span_id: None,
            operation: operation.into(),
            correlation_id: None,
        }
    }
}

/// A message of type `M` paired with an [`OwnedTraceContext`].
///
/// Actor dispatch frameworks wrap outbound messages in `Traced<M>` to carry trace
/// context across message boundaries. The receiving handler extracts the context via
/// [`Traced::into_parts`] and creates a child span before processing the inner message.
#[derive(Clone, Debug)]
pub struct Traced<M> {
    /// Trace context from the sending actor.
    pub context: OwnedTraceContext,
    /// The inner message payload.
    pub message: M,
}

impl<M> Traced<M> {
    /// Wraps a message with a trace context.
    #[must_use]
    pub fn new(context: OwnedTraceContext, message: M) -> Self {
        Self { context, message }
    }

    /// Deconstructs into `(OwnedTraceContext, M)` for handler dispatch.
    #[must_use]
    pub fn into_parts(self) -> (OwnedTraceContext, M) {
        (self.context, self.message)
    }

    /// Returns a reference to the embedded trace context.
    #[must_use]
    pub fn as_context(&self) -> &OwnedTraceContext {
        &self.context
    }

    /// Maps the inner message type, preserving the trace context.
    #[must_use]
    pub fn map<N>(self, f: impl FnOnce(M) -> N) -> Traced<N> {
        Traced {
            context: self.context,
            message: f(self.message),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{
        ChildSpanSpec, OwnedTraceContext, Traced, Traceparent, TraceparentParseError,
    };

    const TRACE_ID: &str = "4bf92f3577b34da6a3ce929d0e0e4736";
    const PARENT_ID: &str = "00f067aa0ba902b7";

    // --- OwnedTraceContext ---

    #[test]
    fn owned_context_root_has_no_parent() {
        let ctx = OwnedTraceContext::root(TRACE_ID, PARENT_ID, "session.handle");
        assert_eq!(ctx.trace_id, TRACE_ID);
        assert_eq!(ctx.span_id, PARENT_ID);
        assert_eq!(ctx.operation, "session.handle");
        assert!(ctx.parent_span_id.is_none());
        assert!(ctx.correlation_id.is_none());
    }

    #[test]
    fn owned_context_child_inherits_trace_id() {
        let parent = OwnedTraceContext::root(TRACE_ID, PARENT_ID, "session.handle");
        let child = parent.child("abcdef0123456789", "policy.enforce");
        assert_eq!(child.trace_id, TRACE_ID);
        assert_eq!(child.span_id, "abcdef0123456789");
        assert_eq!(child.parent_span_id.as_deref(), Some(PARENT_ID));
        assert_eq!(child.operation, "policy.enforce");
    }

    #[test]
    fn child_span_spec_fields_match_context() {
        let ctx = OwnedTraceContext::root(TRACE_ID, PARENT_ID, "audit.write");
        let spec: ChildSpanSpec<'_> = ctx.as_child_span_spec();
        assert_eq!(spec.trace_id, TRACE_ID);
        assert_eq!(spec.parent_span_id, PARENT_ID);
        assert_eq!(spec.operation, "audit.write");
        assert!(spec.correlation_id.is_none());
    }

    // --- Traceparent ---

    #[test]
    fn traceparent_parses_valid_header() {
        let header = format!("00-{TRACE_ID}-{PARENT_ID}-01");
        let tp = Traceparent::parse(&header).expect("valid header must parse");
        assert_eq!(tp.version, 0);
        assert_eq!(tp.trace_id, TRACE_ID);
        assert_eq!(tp.parent_id, PARENT_ID);
        assert_eq!(tp.trace_flags, 1);
    }

    #[test]
    fn traceparent_to_header_round_trips() {
        let original = format!("00-{TRACE_ID}-{PARENT_ID}-01");
        let tp = Traceparent::parse(&original).expect("valid header must parse");
        assert_eq!(tp.to_header(), original);
    }

    #[test]
    fn traceparent_into_owned_context_preserves_ids() {
        let header = format!("00-{TRACE_ID}-{PARENT_ID}-01");
        let tp = Traceparent::parse(&header).expect("valid header must parse");
        let ctx = tp.into_owned_trace_context("session.forward");
        assert_eq!(ctx.trace_id, TRACE_ID);
        assert_eq!(ctx.span_id, PARENT_ID);
        assert_eq!(ctx.operation, "session.forward");
        assert!(ctx.parent_span_id.is_none());
    }

    #[test]
    fn traceparent_rejects_too_few_segments() {
        let err = Traceparent::parse("00-abc").unwrap_err();
        assert_eq!(err, TraceparentParseError::InvalidSegmentCount);
    }

    #[test]
    fn traceparent_rejects_non_zero_version() {
        let header = format!("ff-{TRACE_ID}-{PARENT_ID}-00");
        let err = Traceparent::parse(&header).unwrap_err();
        assert_eq!(err, TraceparentParseError::InvalidVersion);
    }

    #[test]
    fn traceparent_rejects_short_trace_id() {
        let short = &TRACE_ID[..31];
        let header = format!("00-{short}-{PARENT_ID}-00");
        let err = Traceparent::parse(&header).unwrap_err();
        assert_eq!(err, TraceparentParseError::InvalidTraceId);
    }

    #[test]
    fn traceparent_rejects_short_parent_id() {
        let short = &PARENT_ID[..15];
        let header = format!("00-{TRACE_ID}-{short}-00");
        let err = Traceparent::parse(&header).unwrap_err();
        assert_eq!(err, TraceparentParseError::InvalidParentId);
    }

    #[test]
    fn traceparent_rejects_invalid_flags() {
        let header = format!("00-{TRACE_ID}-{PARENT_ID}-1");
        let err = Traceparent::parse(&header).unwrap_err();
        assert_eq!(err, TraceparentParseError::InvalidFlags);
    }

    // --- Traced<M> ---

    #[test]
    fn traced_new_and_into_parts() {
        let ctx = OwnedTraceContext::root(TRACE_ID, PARENT_ID, "actor.send");
        let traced = Traced::new(ctx.clone(), 42u32);
        let (extracted_ctx, extracted_msg) = traced.into_parts();
        assert_eq!(extracted_ctx, ctx);
        assert_eq!(extracted_msg, 42u32);
    }

    #[test]
    fn traced_as_context_borrows() {
        let ctx = OwnedTraceContext::root(TRACE_ID, PARENT_ID, "actor.recv");
        let traced = Traced::new(ctx, "payload");
        assert_eq!(traced.as_context().trace_id, TRACE_ID);
        let _ = traced.message;
    }

    #[test]
    fn traced_map_preserves_context() {
        let ctx = OwnedTraceContext::root(TRACE_ID, PARENT_ID, "actor.map");
        let traced: Traced<u32> = Traced::new(ctx.clone(), 10u32);
        let mapped: Traced<String> = traced.map(|n| format!("value:{n}"));
        assert_eq!(mapped.context, ctx);
        assert_eq!(mapped.message, "value:10");
    }

    // --- TraceparentParseError ---

    #[test]
    fn parse_error_as_str_non_empty() {
        let variants = [
            TraceparentParseError::InvalidSegmentCount,
            TraceparentParseError::InvalidVersion,
            TraceparentParseError::InvalidTraceId,
            TraceparentParseError::InvalidParentId,
            TraceparentParseError::InvalidFlags,
        ];
        for variant in &variants {
            assert!(!variant.as_str().is_empty(), "{variant:?} has empty as_str");
        }
    }
}