ipx/crates/directory/src/otel.rs

//! Traces, for ipx and the directory both (#213): the exporter, and W3C `traceparent` in from a
//! caller and out to a callee, so one trace follows a request from the browser through ipx-web
//! and ipx into the directory.

use axum::http::HeaderMap;

/// Traces over OTLP, to Tempo for one, when OTEL_EXPORTER_OTLP_ENDPOINT names a collector
/// (`http://host:4318`: the exporter speaks OTLP over HTTP and adds `/v1/traces`). The exporter
/// reads that and the other `OTEL_` variables itself.
pub fn provider(service: &'static str) -> anyhow::Result<Option<opentelemetry_sdk::trace::SdkTracerProvider>> {
    if std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").unwrap_or_default().is_empty() {
        return Ok(None);
    }
    let exporter = opentelemetry_otlp::SpanExporter::builder().with_http().build()?;
    let resource = opentelemetry_sdk::Resource::builder().with_service_name(service).build();
    Ok(Some(
        opentelemetry_sdk::trace::SdkTracerProvider::builder()
            .with_batch_exporter(exporter)
            .with_resource(resource)
            .with_id_generator(Ids::default())
            .build(),
    ))
}

/// The trace a request from the page belongs to, from its `traceparent` (#190), so the requests
/// one action sets off are one trace. None for a missing or malformed one. In the query where no
/// header can be set: a position saved by sendBeacon, and the playlist Safari's player fetches
/// itself (#192).
#[must_use]
pub fn remote_parent(headers: &HeaderMap, query: Option<&str>) -> Option<opentelemetry::Context> {
    use opentelemetry::trace::{SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId, TraceState};
    let v = match headers.get("traceparent") {
        Some(h) => h.to_str().ok()?,
        None => query?.split('&').find_map(|kv| kv.strip_prefix("traceparent="))?,
    };
    let mut parts = v.split('-');
    let (version, trace, span) = (parts.next()?, parts.next()?, parts.next()?);
    if version != "00" || trace.len() != 32 || span.len() != 16 {
        return None;
    }
    let sc = SpanContext::new(
        TraceId::from_hex(trace).ok()?,
        SpanId::from_hex(span).ok()?,
        TraceFlags::SAMPLED,
        true,
        TraceState::default(),
    );
    sc.is_valid().then(|| opentelemetry::Context::new().with_remote_span_context(sc))
}

/// The current span as a `traceparent`, for a request to another service, so its spans join
/// this trace. None outside a traced span.
#[must_use]
pub fn traceparent() -> Option<String> {
    use opentelemetry::trace::TraceContextExt;
    use tracing_opentelemetry::OpenTelemetrySpanExt;
    let cx = tracing::Span::current().context();
    let sc = cx.span().span_context().clone();
    sc.is_valid().then(|| format!("00-{}-{}-{:02x}", sc.trace_id(), sc.span_id(), sc.trace_flags().to_u8()))
}

/// Random ids, as the SDK's own, but for a span built inside `with_ids`, which takes the ids it is
/// given: the page names its tap's span before ipx makes it (#376), and the SDK gives a span
/// builder no other way to be told them.
#[derive(Debug, Default)]
struct Ids(opentelemetry_sdk::trace::RandomIdGenerator);

thread_local! {
    static GIVEN: std::cell::Cell<Option<(opentelemetry::trace::TraceId, opentelemetry::trace::SpanId)>> =
        const { std::cell::Cell::new(None) };
}

impl opentelemetry_sdk::trace::IdGenerator for Ids {
    fn new_trace_id(&self) -> opentelemetry::trace::TraceId {
        GIVEN.get().map_or_else(|| self.0.new_trace_id(), |(t, _)| t)
    }
    fn new_span_id(&self) -> opentelemetry::trace::SpanId {
        GIVEN.get().map_or_else(|| self.0.new_span_id(), |(_, s)| s)
    }
}

/// Runs `f`, which builds one root span, with these ids for it. Synchronous, so nothing else on
/// this thread makes a span in between.
pub fn with_ids<R>(
    trace: opentelemetry::trace::TraceId,
    span: opentelemetry::trace::SpanId,
    f: impl FnOnce() -> R,
) -> R {
    GIVEN.set(Some((trace, span)));
    let r = f();
    GIVEN.set(None);
    r
}

/// The JSON log line with the trace and span it belongs to (#91), so a line in Loki leads to its
/// trace in Tempo. The JSON formatter cannot take a field of its own, so the ids go on the end of
/// the object it writes; the directory's lines carry them too (#375). A line outside any traced span, or with no trace exporter, is unchanged.
/// One JSON object a line for Loki, each event's fields its own and the trace and span ids beside
/// them, for ipx and the directory both: the dashboard and the alert rules read these fields.
/// `current_span` adds the span's own fields under `span`, which ipx keeps and the directory
/// dropped (#363).
pub fn json_layer<S, W>(writer: W, current_span: bool) -> impl tracing_subscriber::Layer<S>
where
    S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
    W: for<'w> tracing_subscriber::fmt::MakeWriter<'w> + Send + Sync + 'static,
{
    use tracing_subscriber::Layer as _;
    let registered = Registered::default();
    tracing_subscriber::fmt::layer()
        .fmt_fields(tracing_subscriber::fmt::format::JsonFields::new())
        .event_format(WithTrace(
            tracing_subscriber::fmt::format()
                .json()
                .flatten_event(true)
                .with_current_span(current_span)
                .with_span_list(false),
            registered.clone(),
        ))
        .with_writer(writer)
        .and_then(registered)
}

/// The subscriber the log layer is part of, as it was handed over on registration. The span's
/// OpenTelemetry context is reached through it: `Span::current()` asks the thread's default
/// dispatcher, which inside an event under a scoped subscriber (`with_default`) is hidden by
/// tracing's re-entrancy guard, and every line went out without its ids (#408).
#[derive(Clone, Default)]
struct Registered(std::sync::Arc<std::sync::OnceLock<tracing::dispatcher::WeakDispatch>>);

impl<S: tracing::Subscriber> tracing_subscriber::Layer<S> for Registered {
    fn on_register_dispatch(&self, subscriber: &tracing::Dispatch) {
        let _ = self.0.set(subscriber.downgrade());
    }
}

struct WithTrace<F>(F, Registered);

impl<S, N, F> tracing_subscriber::fmt::FormatEvent<S, N> for WithTrace<F>
where
    F: tracing_subscriber::fmt::FormatEvent<S, N>,
    S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
    N: for<'w> tracing_subscriber::fmt::FormatFields<'w> + 'static,
{
    fn format_event(
        &self,
        ctx: &tracing_subscriber::fmt::FmtContext<'_, S, N>,
        mut w: tracing_subscriber::fmt::format::Writer<'_>,
        ev: &tracing::Event<'_>,
    ) -> std::fmt::Result {
        use opentelemetry::trace::TraceContextExt;
        // Without one (a wrapper that does not pass registration on, as Option does not), the
        // default dispatcher, which a global subscriber's lines still find.
        let registered = self.1.0.get().and_then(tracing::dispatcher::WeakDispatch::upgrade);
        let dispatch = registered.unwrap_or_else(|| tracing::dispatcher::get_default(Clone::clone));
        // The registry's own current span, not the one this layer's filter would show.
        let sc = ctx
            .current_span()
            .id()
            .and_then(|id| tracing_opentelemetry::get_otel_context(id, &dispatch))
            .map(|cx| cx.span().span_context().clone());
        let Some(sc) = sc.filter(|c| c.is_valid()) else {
            return self.0.format_event(ctx, w, ev);
        };
        let mut line = String::new();
        self.0.format_event(ctx, tracing_subscriber::fmt::format::Writer::new(&mut line), ev)?;
        match line.trim_end().strip_suffix('}') {
            Some(body) => writeln!(w, r#"{body},"trace_id":"{}","span_id":"{}"}}"#, sc.trace_id(), sc.span_id()),
            None => w.write_str(&line),
        }
    }
}

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

    #[test]
    fn a_requests_trace_is_the_pages_from_its_traceparent() {
        use opentelemetry::trace::TraceContextExt;
        let with = |v: &str| {
            let mut h = axum::http::HeaderMap::new();
            h.insert("traceparent", v.parse().unwrap());
            remote_parent(&h, None)
        };
        let cx = with("00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01").expect("parsed");
        assert_eq!(cx.span().span_context().trace_id().to_string(), "4bf92f3577b34da6a3ce929d0e0e4736");
        assert!(cx.span().span_context().is_remote());
        // Anything else from the browser starts a trace of its own.
        for bad in [
            "",
            "00-xyz-00f067aa0ba902b7-01",
            "01-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
            "00-00000000000000000000000000000000-00f067aa0ba902b7-01",
            "00-4bf92f3577b34da6a3ce929d0e0e4736",
        ] {
            assert!(with(bad).is_none(), "{bad}");
        }
        assert!(remote_parent(&axum::http::HeaderMap::new(), None).is_none());
        // From the query, where a beacon or Safari's player cannot send a header (#192).
        let q = remote_parent(
            &axum::http::HeaderMap::new(),
            Some("x=1&traceparent=00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"),
        );
        assert_eq!(q.unwrap().span().span_context().trace_id().to_string(), "4bf92f3577b34da6a3ce929d0e0e4736");
        assert!(remote_parent(&axum::http::HeaderMap::new(), Some("traceparent=nonsense")).is_none());
    }

    #[test]
    fn a_span_built_with_given_ids_takes_them_and_the_next_does_not() {
        use opentelemetry::trace::{Span as _, SpanId, TraceId, Tracer as _, TracerProvider as _};
        let tracer = opentelemetry_sdk::trace::SdkTracerProvider::builder()
            .with_id_generator(super::Ids::default())
            .build()
            .tracer("t");
        let (t, s) = (
            TraceId::from_hex("4bf92f3577b34da6a3ce929d0e0e4736").unwrap(),
            SpanId::from_hex("00f067aa0ba902b7").unwrap(),
        );
        let given = super::with_ids(t, s, || tracer.start("tap"));
        assert_eq!((given.span_context().trace_id(), given.span_context().span_id()), (t, s));
        let next = tracer.start("other");
        assert_ne!(next.span_context().trace_id(), t);
    }
}