polyc_runtime/

propagation.rs

//! W3C `traceparent` propagation across the connectrpc boundary.
//!
//! Pairs with [`crate::observability::init`] (which installs the global
//! [`TraceContextPropagator`](opentelemetry_sdk::propagation::TraceContextPropagator)).
//! Together they close the `OTel` cross-process story: a turn started on the
//! control plane carries its trace context into the harness's `connect` span,
//! so a single user request stitches into one trace tree across the
//! control-plane → harness hop.
//!
//! # Direction
//!
//! - **Outgoing**: [`inject_current_span_into`] writes the current
//!   `tracing::Span`'s `OTel` context into an [`http::HeaderMap`] using the
//!   globally-installed propagator. Callers attach the headers to a
//!   `connectrpc::client::CallOptions` (per-call) or
//!   `connectrpc::client::ClientConfig` (per-client default).
//! - **Incoming**: [`extract_parent_into_current_span`] reads `traceparent`
//!   (and any other propagator-registered fields) off a borrowed
//!   [`http::HeaderMap`] and re-parents the current `tracing::Span` so the
//!   server-side span shares the dialer's trace id.
//!
//! # Design
//!
//! The [`Injector`]/[`Extractor`] adapters are intentionally trivial — they
//! exist because `opentelemetry`'s trait surface is carrier-agnostic and
//! `http::HeaderMap` is not. Keeping them in `polyc-runtime` means the
//! control plane, the harness, and any future server share one implementation;
//! per-call sites only pick which side (inject vs extract) they need.

use opentelemetry::Context;
use opentelemetry::global;
use opentelemetry::propagation::{Extractor, Injector};
use tracing_opentelemetry::OpenTelemetrySpanExt as _;

/// Adapter that injects propagator fields into an [`http::HeaderMap`].
///
/// Lets the global [`TextMapPropagator`](opentelemetry::propagation::TextMapPropagator)
/// write `traceparent` (and any other fields the propagator registers) into
/// an HTTP request's headers.
///
/// Invalid header names/values (the propagator should never produce them for
/// W3C `traceparent`, but the contract is `String`) are silently dropped:
/// shipping a partially-propagated request is strictly better than failing
/// the request, and an absent `traceparent` just degrades to "new trace tree"
/// on the receiving side.
pub struct HttpHeaderInjector<'a>(pub &'a mut http::HeaderMap);

impl Injector for HttpHeaderInjector<'_> {
    fn set(&mut self, key: &str, value: String) {
        if let (Ok(name), Ok(val)) = (
            http::header::HeaderName::try_from(key),
            http::header::HeaderValue::try_from(value),
        ) {
            self.0.insert(name, val);
        }
    }
}

/// Adapter that lets the global `OTel` propagator read `traceparent` (and any
/// other fields) off an [`http::HeaderMap`].
///
/// Header values that aren't valid UTF-8 are skipped — W3C `traceparent` is
/// ASCII by spec, so a non-string value means a misconfigured upstream, and
/// "start a new trace tree" is the right fallback.
pub struct HttpHeaderExtractor<'a>(pub &'a http::HeaderMap);

impl Extractor for HttpHeaderExtractor<'_> {
    fn get(&self, key: &str) -> Option<&str> {
        self.0.get(key).and_then(|v| v.to_str().ok())
    }

    fn keys(&self) -> Vec<&str> {
        self.0
            .keys()
            .map(http::header::HeaderName::as_str)
            .collect()
    }
}

/// Inject the current `tracing::Span`'s `OTel` context into `headers`.
///
/// Uses the propagator installed by [`crate::observability::init`]. If no
/// propagator was installed (e.g. in a test that didn't run `init`), this is
/// a no-op — the global getter returns a no-op propagator by default.
pub fn inject_current_span_into(headers: &mut http::HeaderMap) {
    let cx = tracing::Span::current().context();
    inject_context_into(&cx, headers);
}

/// Inject a specific [`Context`] into `headers`. Lower-level than
/// [`inject_current_span_into`]; prefer that for production call sites.
pub fn inject_context_into(cx: &Context, headers: &mut http::HeaderMap) {
    global::get_text_map_propagator(|propagator| {
        propagator.inject_context(cx, &mut HttpHeaderInjector(headers));
    });
}

/// Extract the parent `OTel` context from `headers` and attach it to the
/// current `tracing::Span`, so server-side spans share the dialer's trace id.
///
/// A missing or malformed `traceparent` produces an empty context, which
/// `set_parent` accepts as a no-op — the local span just stays a root span.
///
/// Calling this from a `#[tracing::instrument]`'d handler must happen
/// **before** the first call to [`tracing::Span::context`] on the same span
/// (or anything that triggers a span entry that consults the parent), per
/// `tracing-opentelemetry`'s ordering contract. The two server-side
/// `connect` handlers wire it first thing in the function body for exactly
/// this reason.
pub fn extract_parent_into_current_span(headers: &http::HeaderMap) {
    let parent_cx = extract_context_from(headers);
    // `set_parent` returns `Err` only if the span has already been entered /
    // its context observed; for our use (called first thing in the handler)
    // that's an internal invariant violation, so log and continue with the
    // local root span rather than fail the request.
    if let Err(err) = tracing::Span::current().set_parent(parent_cx) {
        tracing::debug!(error = %err, "could not set parent trace context");
    }
}

/// Extract a parent [`Context`] from `headers` using the globally-installed
/// propagator. Lower-level than [`extract_parent_into_current_span`].
#[must_use]
pub fn extract_context_from(headers: &http::HeaderMap) -> Context {
    global::get_text_map_propagator(|propagator| propagator.extract(&HttpHeaderExtractor(headers)))
}

#[cfg(test)]
mod tests {
    #![allow(clippy::pedantic, clippy::nursery, missing_docs)]

    use super::*;
    use opentelemetry::trace::{
        SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId, TraceState,
    };
    use opentelemetry_sdk::propagation::TraceContextPropagator;

    /// Build a synthetic `OTel` context with a known trace + span id. We use
    /// `Context::new().with_remote_span_context(...)` so the propagator's
    /// `inject_context` finds a sampleable span context — this mirrors what
    /// `tracing-opentelemetry` does internally.
    fn ctx_with(trace_id: TraceId, span_id: SpanId) -> Context {
        let span_ctx = SpanContext::new(
            trace_id,
            span_id,
            TraceFlags::SAMPLED,
            true, // remote
            TraceState::default(),
        );
        Context::new().with_remote_span_context(span_ctx)
    }

    /// Install the W3C propagator locally for the duration of the test
    /// process. `set_text_map_propagator` is process-global; multiple tests
    /// installing the same propagator is harmless (idempotent replacement).
    fn install_w3c() {
        global::set_text_map_propagator(TraceContextPropagator::new());
    }

    #[test]
    fn injector_extractor_round_trip_preserves_trace_id() {
        install_w3c();
        let trace_id = TraceId::from_hex("0af7651916cd43dd8448eb211c80319c").unwrap();
        let span_id = SpanId::from_hex("b7ad6b7169203331").unwrap();
        let cx = ctx_with(trace_id, span_id);

        let mut headers = http::HeaderMap::new();
        inject_context_into(&cx, &mut headers);

        // The W3C propagator writes a `traceparent` header.
        assert!(
            headers.get("traceparent").is_some(),
            "traceparent should be injected"
        );

        let extracted = extract_context_from(&headers);
        let extracted_span = extracted.span();
        let extracted_ctx = extracted_span.span_context();
        assert_eq!(
            extracted_ctx.trace_id(),
            trace_id,
            "trace id must round-trip"
        );
        assert_eq!(extracted_ctx.span_id(), span_id, "span id must round-trip");
    }

    #[test]
    fn extract_from_empty_headers_yields_invalid_context() {
        install_w3c();
        let headers = http::HeaderMap::new();
        let extracted = extract_context_from(&headers);
        // No `traceparent` → the extracted span context isn't valid.
        assert!(
            !extracted.span().span_context().is_valid(),
            "empty headers should yield an invalid (root) context"
        );
    }

    #[test]
    fn injector_silently_drops_invalid_header_names() {
        // Direct unit test of the adapter: an invalid header name (e.g.
        // containing whitespace) is dropped rather than panicking. This
        // matches the production contract — better to ship a partial
        // propagation than fail the request.
        let mut map = http::HeaderMap::new();
        let mut injector = HttpHeaderInjector(&mut map);
        injector.set("invalid name with spaces", "v".to_owned());
        injector.set("x-good", "ok".to_owned());
        assert!(map.get("invalid name with spaces").is_none());
        assert_eq!(map.get("x-good").unwrap(), "ok");
    }

    #[test]
    fn extractor_keys_returns_header_names() {
        let mut map = http::HeaderMap::new();
        map.insert("traceparent", http::HeaderValue::from_static("x"));
        map.insert("x-custom", http::HeaderValue::from_static("y"));
        let extractor = HttpHeaderExtractor(&map);
        let mut keys: Vec<&str> = extractor.keys();
        keys.sort_unstable();
        assert_eq!(keys, vec!["traceparent", "x-custom"]);
        assert_eq!(extractor.get("traceparent"), Some("x"));
        assert_eq!(extractor.get("missing"), None);
    }
}