apollo-opentelemetry 0.8.0

//! Context propagation configuration.
//!
//! See the [propagator specification](https://opentelemetry.io/docs/specs/otel/context/api-propagators/)
//! for details on context propagation.

use apollo_configuration::configuration;
use opentelemetry::propagation::{TextMapCompositePropagator, TextMapPropagator};
use opentelemetry_aws::trace::XrayPropagator;
use opentelemetry_sdk::propagation::{BaggagePropagator, TraceContextPropagator};

/// Configuration for context propagators.
#[configuration]
pub struct PropagatorConfig {
    /// List of propagators to use for context propagation.
    ///
    /// Common values: `tracecontext`, `baggage`, `b3`, `b3multi`, `jaeger`, `xray`, `ottrace`.
    pub composite: Vec<Propagator>,
}

/// Available propagator types.
///
/// See the [propagator specification](https://opentelemetry.io/docs/specs/otel/context/api-propagators/)
/// for details.
#[configuration]
#[derive(PartialEq, Eq)]
pub enum Propagator {
    /// W3C Trace Context propagator (recommended).
    Tracecontext,
    /// W3C Baggage propagator.
    Baggage,
    /// B3 single-header propagator.
    B3,
    /// B3 multi-header propagator.
    B3multi,
    /// Jaeger propagator.
    Jaeger,
    /// AWS X-Ray propagator.
    Xray,
    /// OpenTracing propagator.
    Ottrace,
}

impl From<&PropagatorConfig> for TextMapCompositePropagator {
    fn from(config: &PropagatorConfig) -> Self {
        TextMapCompositePropagator::new(
            config
                .composite
                .iter()
                .map(|prop| -> Box<dyn TextMapPropagator + Send + Sync> {
                    match prop {
                        Propagator::Tracecontext => Box::new(TraceContextPropagator::new()),
                        Propagator::Baggage => Box::new(BaggagePropagator::new()),
                        Propagator::Xray => Box::new(XrayPropagator::new()),
                        _ => unimplemented!("{prop:?}"),
                    }
                })
                .collect(),
        )
    }
}

impl From<PropagatorConfig> for TextMapCompositePropagator {
    fn from(config: PropagatorConfig) -> Self {
        TextMapCompositePropagator::from(&config)
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use apollo_configuration::parse_yaml;
    use opentelemetry::baggage::BaggageExt;
    use opentelemetry::propagation::{Injector, TextMapCompositePropagator, TextMapPropagator};
    use opentelemetry::trace::{SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId};
    use opentelemetry::{Context, KeyValue};

    use super::{Propagator, PropagatorConfig};
    use crate::config::OpenTelemetryConfig;

    struct TestCarrier(HashMap<String, String>);

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

    #[test]
    fn parse_propagator() {
        let config: OpenTelemetryConfig = parse_yaml(
            indoc::indoc! {"
                propagator:
                  composite:
                    - tracecontext
                    - baggage
            "},
            &Default::default(),
        )
        .unwrap();

        assert_eq!(config.propagator.composite.len(), 2);
    }

    #[test]
    fn tracecontext_propagator_injects_traceparent() {
        let config = PropagatorConfig {
            composite: vec![Propagator::Tracecontext],
        };
        let propagator = TextMapCompositePropagator::from(&config);

        let span_context = SpanContext::new(
            TraceId::from_hex("4bf92f3577b34da6a3ce929d0e0e4736").unwrap(),
            SpanId::from_hex("00f067aa0ba902b7").unwrap(),
            TraceFlags::SAMPLED,
            true,
            Default::default(),
        );
        let cx = Context::current().with_remote_span_context(span_context);

        let mut carrier = TestCarrier(HashMap::new());
        propagator.inject_context(&cx, &mut carrier);

        assert_eq!(
            carrier.0.get("traceparent"),
            Some(&"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01".to_string())
        );
    }

    #[test]
    fn baggage_propagator_injects_baggage() {
        let config = PropagatorConfig {
            composite: vec![Propagator::Baggage],
        };
        let propagator = TextMapCompositePropagator::from(&config);

        let cx = Context::current().with_baggage([KeyValue::new("user_id", "12345")]);

        let mut carrier = TestCarrier(HashMap::new());
        propagator.inject_context(&cx, &mut carrier);

        assert_eq!(carrier.0.get("baggage"), Some(&"user_id=12345".to_string()));
    }

    #[test]
    fn xray_propagator_injects_xray_header() {
        let config = PropagatorConfig {
            composite: vec![Propagator::Xray],
        };
        let propagator = TextMapCompositePropagator::from(&config);

        let span_context = SpanContext::new(
            TraceId::from_hex("4bf92f3577b34da6a3ce929d0e0e4736").unwrap(),
            SpanId::from_hex("00f067aa0ba902b7").unwrap(),
            TraceFlags::SAMPLED,
            true,
            Default::default(),
        );
        let cx = Context::current().with_remote_span_context(span_context);

        let mut carrier = TestCarrier(HashMap::new());
        propagator.inject_context(&cx, &mut carrier);

        let header = carrier
            .0
            .get("x-amzn-trace-id")
            .expect("should inject x-amzn-trace-id");
        assert!(header.contains("Root="), "should contain Root segment");
        assert!(header.contains("Parent="), "should contain Parent segment");
        assert!(header.contains("Sampled=1"), "should contain Sampled=1");
    }

    #[test]
    fn composite_propagator_injects_multiple_formats() {
        let config = PropagatorConfig {
            composite: vec![Propagator::Tracecontext, Propagator::Baggage],
        };
        let propagator = TextMapCompositePropagator::from(&config);

        let span_context = SpanContext::new(
            TraceId::from_hex("4bf92f3577b34da6a3ce929d0e0e4736").unwrap(),
            SpanId::from_hex("00f067aa0ba902b7").unwrap(),
            TraceFlags::SAMPLED,
            true,
            Default::default(),
        );
        let cx = Context::current()
            .with_remote_span_context(span_context)
            .with_baggage([KeyValue::new("key", "value")]);

        let mut carrier = TestCarrier(HashMap::new());
        propagator.inject_context(&cx, &mut carrier);

        assert!(
            carrier.0.contains_key("traceparent"),
            "should inject traceparent"
        );
        assert!(carrier.0.contains_key("baggage"), "should inject baggage");
    }

    #[test]
    fn empty_propagator_injects_nothing() {
        let config = PropagatorConfig { composite: vec![] };
        let propagator = TextMapCompositePropagator::from(&config);

        let span_context = SpanContext::new(
            TraceId::from_hex("4bf92f3577b34da6a3ce929d0e0e4736").unwrap(),
            SpanId::from_hex("00f067aa0ba902b7").unwrap(),
            TraceFlags::SAMPLED,
            true,
            Default::default(),
        );
        let cx = Context::current().with_remote_span_context(span_context);

        let mut carrier = TestCarrier(HashMap::new());
        propagator.inject_context(&cx, &mut carrier);

        assert!(
            carrier.0.is_empty(),
            "empty propagator should not inject anything"
        );
    }
}