rs_zero/observability/

correlation.rs

use http::HeaderMap;

use crate::observability::{
    current_span_id, current_trace_id, request_id_from_headers, span_id_from_traceparent,
    trace_id_from_traceparent, traceparent_from_headers,
};

/// Stable, low-cardinality correlation fields for logs, spans, and metrics.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CorrelationContext {
    service: String,
    transport: &'static str,
    route: String,
    method: String,
    request_id: Option<String>,
    traceparent: Option<String>,
    trace_id: Option<String>,
    span_id: Option<String>,
    status: Option<String>,
}

impl CorrelationContext {
    /// Builds HTTP correlation from headers and a route pattern.
    ///
    /// `route` must be a framework route pattern such as `/users/:id`; raw
    /// paths are intentionally not accepted to avoid high-cardinality logs.
    pub fn from_http_headers(
        service: Option<&str>,
        method: impl Into<String>,
        route: Option<&str>,
        headers: &HeaderMap,
    ) -> Self {
        let traceparent = traceparent_from_headers(headers);
        Self::new(
            service.unwrap_or("unknown"),
            "http",
            route.unwrap_or("unknown"),
            method,
        )
        .with_request_id(request_id_from_headers(headers))
        .with_traceparent(traceparent)
        .with_current_span_context()
    }

    /// Builds gRPC correlation from tonic metadata and an RPC method pattern.
    #[cfg(feature = "rpc")]
    pub fn from_rpc_metadata(
        service: impl Into<String>,
        method: impl Into<String>,
        metadata: &tonic::metadata::MetadataMap,
    ) -> Self {
        let method = method.into();
        let traceparent = crate::observability::traceparent_from_metadata(metadata);
        Self::new(service, "grpc", method.clone(), method)
            .with_request_id(crate::observability::request_id_from_metadata(metadata))
            .with_traceparent(traceparent)
            .with_current_span_context()
    }

    /// Builds correlation from explicit RPC request and trace context.
    pub fn from_rpc_parts(
        service: impl Into<String>,
        method: impl Into<String>,
        request_id: Option<&str>,
        traceparent: Option<&str>,
    ) -> Self {
        let method = method.into();
        Self::new(service, "grpc", method.clone(), method)
            .with_request_id(request_id.map(ToOwned::to_owned))
            .with_traceparent(traceparent.map(ToOwned::to_owned))
            .with_current_span_context()
    }

    /// Creates a low-cardinality context from explicit parts.
    pub fn new(
        service: impl Into<String>,
        transport: &'static str,
        route: impl Into<String>,
        method: impl Into<String>,
    ) -> Self {
        Self {
            service: service.into(),
            transport,
            route: route.into(),
            method: method.into(),
            request_id: None,
            traceparent: None,
            trace_id: None,
            span_id: None,
            status: None,
        }
    }

    /// Sets status or code.
    pub fn with_status(mut self, status: impl Into<String>) -> Self {
        self.status = Some(status.into());
        self
    }

    /// Returns the service name.
    pub fn service(&self) -> &str {
        &self.service
    }

    /// Returns the transport name.
    pub fn transport(&self) -> &'static str {
        self.transport
    }

    /// Returns the route pattern or `unknown`.
    pub fn route(&self) -> &str {
        &self.route
    }

    /// Returns the HTTP or RPC method.
    pub fn method(&self) -> &str {
        &self.method
    }

    /// Returns the traceparent value when one is available.
    pub fn traceparent(&self) -> Option<&str> {
        self.traceparent.as_deref()
    }

    /// Returns the request id when one is available.
    pub fn request_id(&self) -> Option<&str> {
        self.request_id.as_deref()
    }

    /// Returns the trace id when one is available.
    pub fn trace_id(&self) -> Option<&str> {
        self.trace_id.as_deref()
    }

    /// Returns the span id when one is available.
    pub fn span_id(&self) -> Option<&str> {
        self.span_id.as_deref()
    }

    /// Converts the context into stable logging fields.
    #[cfg(feature = "core")]
    pub fn into_log_fields(self) -> crate::core::logging::LogFields {
        use crate::core::logging::LogFields;

        let mut fields = LogFields::new(self.service)
            .with_transport(self.transport)
            .with_route(self.route)
            .with_method(self.method);
        if let Some(request_id) = self.request_id {
            fields = fields.with_request_id(request_id);
        }
        if let Some(trace_id) = self.trace_id {
            fields = fields.with_trace_id(trace_id);
        }
        if let Some(span_id) = self.span_id {
            fields = fields.with_span_id(span_id);
        }
        if let Some(status) = self.status {
            fields = fields.with_status(status);
        }
        fields
    }

    /// Returns all populated fields as key/value pairs.
    pub fn as_pairs(&self) -> Vec<(String, String)> {
        let mut pairs = vec![("service".to_string(), self.service.clone())];
        push_optional(&mut pairs, "transport", Some(self.transport));
        push_optional(&mut pairs, "route", Some(&self.route));
        push_optional(&mut pairs, "method", Some(&self.method));
        push_optional(&mut pairs, "request_id", self.request_id.as_deref());
        push_optional(&mut pairs, "trace_id", self.trace_id.as_deref());
        push_optional(&mut pairs, "span_id", self.span_id.as_deref());
        push_optional(&mut pairs, "status", self.status.as_deref());
        pairs
    }

    fn with_request_id(mut self, request_id: Option<String>) -> Self {
        self.request_id = request_id;
        self
    }

    fn with_traceparent(mut self, traceparent: Option<String>) -> Self {
        if let Some(value) = traceparent {
            self.trace_id = trace_id_from_traceparent(&value).map(ToOwned::to_owned);
            self.span_id = span_id_from_traceparent(&value).map(ToOwned::to_owned);
            self.traceparent = Some(value);
        }
        self
    }

    fn with_current_span_context(mut self) -> Self {
        if let Some(trace_id) = current_trace_id() {
            self.trace_id = Some(trace_id);
        }
        if let Some(span_id) = current_span_id() {
            self.span_id = Some(span_id);
        }
        self
    }
}

fn push_optional(pairs: &mut Vec<(String, String)>, key: &str, value: Option<&str>) {
    if let Some(value) = value
        && !value.is_empty()
    {
        pairs.push((key.to_string(), value.to_string()));
    }
}