later 0.0.50

Distributed Background jobs manager and runner for Rust
//! Carries the enqueuing span's trace context across the job store, so a
//! worker's span joins the publisher's trace rather than starting its own.

use opentelemetry::propagation::{Extractor, Injector, TextMapPropagator};
use opentelemetry_sdk::propagation::TraceContextPropagator;
use std::collections::HashMap;
use tracing_opentelemetry::OpenTelemetrySpanExt;

struct Carrier(HashMap<String, String>);

impl Injector for Carrier {
    fn set(&mut self, key: &str, value: String) {
        self.0.insert(key.to_string(), value);
    }
}

impl Extractor for Carrier {
    fn get(&self, key: &str) -> Option<&str> {
        self.0.get(key).map(String::as_str)
    }

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

/// `traceparent` of the current span, if it belongs to a valid trace.
pub(crate) fn current_traceparent() -> Option<String> {
    let context = tracing::Span::current().context();
    let mut carrier = Carrier(HashMap::new());
    TraceContextPropagator::new().inject_context(&context, &mut carrier);
    carrier.0.remove("traceparent")
}

/// Makes `span` a child of the span that produced `traceparent`.
pub(crate) fn set_parent(span: &tracing::Span, traceparent: Option<&str>) {
    let Some(traceparent) = traceparent else {
        return;
    };
    let carrier = Carrier(HashMap::from([(
        "traceparent".to_string(),
        traceparent.to_string(),
    )]));
    let context = TraceContextPropagator::new().extract(&carrier);
    span.set_parent(context);
}

#[cfg(test)]
mod tests {
    use super::*;
    use opentelemetry::trace::{TraceContextExt, TracerProvider as _};
    use tracing_subscriber::layer::SubscriberExt;

    #[test]
    fn child_span_joins_the_publisher_trace() {
        let provider = opentelemetry_sdk::trace::TracerProvider::builder().build();
        let tracer = provider.tracer("test");
        let subscriber = tracing_subscriber::registry()
            .with(tracing_opentelemetry::layer().with_tracer(tracer));

        tracing::subscriber::with_default(subscriber, || {
            let publisher = tracing::info_span!("publisher");
            let (traceparent, publisher_trace) = publisher.in_scope(|| {
                (
                    current_traceparent(),
                    tracing::Span::current().context().span().span_context().trace_id(),
                )
            });
            assert!(traceparent.is_some());

            let worker = tracing::info_span!("worker");
            set_parent(&worker, traceparent.as_deref());
            let worker_trace = worker.context().span().span_context().trace_id();
            assert_eq!(publisher_trace, worker_trace);
        });
    }
}