diff --git a/Cargo.lock b/Cargo.lock index ecc15530..c832f5a5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -487,6 +487,7 @@ dependencies = [ "opentelemetry-otlp", "opentelemetry-semantic-conventions", "opentelemetry_sdk", + "percent-encoding", "pretty_assertions", "proptest", "rand 0.8.5", diff --git a/datadog-opentelemetry/Cargo.toml b/datadog-opentelemetry/Cargo.toml index b18c14b3..807efc3b 100644 --- a/datadog-opentelemetry/Cargo.toml +++ b/datadog-opentelemetry/Cargo.toml @@ -64,6 +64,7 @@ libdd-trace-utils = { workspace = true } libdd-telemetry = { workspace = true } libdd-common = { workspace = true } libdd-tinybytes = { workspace = true } +percent-encoding = "2.3.1" [dev-dependencies] assert_unordered = "0.3" diff --git a/datadog-opentelemetry/examples/propagator/src/server.rs b/datadog-opentelemetry/examples/propagator/src/server.rs index 81656a22..2ed563d0 100644 --- a/datadog-opentelemetry/examples/propagator/src/server.rs +++ b/datadog-opentelemetry/examples/propagator/src/server.rs @@ -119,7 +119,10 @@ async fn send_request( ) -> std::result::Result<(), Box> { let client = Client::builder(TokioExecutor::new()).build_http(); - let cx = Context::current(); + let cx = Context::current().with_baggage(vec![ + KeyValue::new("request-id", "xyz-123"), + KeyValue::new("caller", "rust-propagator-example"), + ]); let mut req = hyper::Request::builder().uri(url); global::get_text_map_propagator(|propagator| { diff --git a/datadog-opentelemetry/src/core/configuration/configuration.rs b/datadog-opentelemetry/src/core/configuration/configuration.rs index 248faac1..808a97e2 100644 --- a/datadog-opentelemetry/src/core/configuration/configuration.rs +++ b/datadog-opentelemetry/src/core/configuration/configuration.rs @@ -773,6 +773,8 @@ pub enum TracePropagationStyle { Datadog, /// W3C Trace Context propagation format using `traceparent` and `tracestate` headers. TraceContext, + /// W3C Baggage propagation format using the `baggage` header. + Baggage, /// No propagation - trace context is not propagated. None, } @@ -804,6 +806,7 @@ impl FromStr for TracePropagationStyle { match s.trim().to_lowercase().as_str() { "datadog" => Ok(TracePropagationStyle::Datadog), "tracecontext" => Ok(TracePropagationStyle::TraceContext), + "baggage" => Ok(TracePropagationStyle::Baggage), "none" => Ok(TracePropagationStyle::None), _ => Err(format!("Unknown trace propagation style: '{s}'")), } @@ -815,6 +818,7 @@ impl Display for TracePropagationStyle { let style = match self { TracePropagationStyle::Datadog => "datadog", TracePropagationStyle::TraceContext => "tracecontext", + TracePropagationStyle::Baggage => "baggage", TracePropagationStyle::None => "none", }; write!(f, "{style}") @@ -1822,6 +1826,7 @@ fn default_config() -> Config { Some(vec![ TracePropagationStyle::Datadog, TracePropagationStyle::TraceContext, + TracePropagationStyle::Baggage, ]), ), trace_propagation_style_extract: ConfigItem::new( @@ -2726,6 +2731,7 @@ mod tests { Some(vec![ TracePropagationStyle::Datadog, TracePropagationStyle::TraceContext, + TracePropagationStyle::Baggage, ]) .as_deref() ); @@ -2737,6 +2743,47 @@ mod tests { assert!(config.trace_propagation_extract_first()); } + #[test] + fn test_propagation_style_baggage_parsed_from_env() { + // "baggage" is recognised as a valid style value (case-insensitive) in all three env vars. + let mut sources = CompositeSource::new(); + sources.add_source(HashMapSource::from_iter( + [ + ("DD_TRACE_PROPAGATION_STYLE", "datadog,tracecontext,baggage"), + ("DD_TRACE_PROPAGATION_STYLE_EXTRACT", "Baggage,datadog"), + ("DD_TRACE_PROPAGATION_STYLE_INJECT", "BAGGAGE,tracecontext"), + ], + ConfigSourceOrigin::EnvVar, + )); + let config = Config::builder_with_sources(&sources).build(); + + assert_eq!( + config.trace_propagation_style(), + Some(vec![ + TracePropagationStyle::Datadog, + TracePropagationStyle::TraceContext, + TracePropagationStyle::Baggage, + ]) + .as_deref() + ); + assert_eq!( + config.trace_propagation_style_extract(), + Some(vec![ + TracePropagationStyle::Baggage, + TracePropagationStyle::Datadog, + ]) + .as_deref() + ); + assert_eq!( + config.trace_propagation_style_inject(), + Some(vec![ + TracePropagationStyle::Baggage, + TracePropagationStyle::TraceContext, + ]) + .as_deref() + ); + } + #[test] fn test_stats_computation_enabled_config() { let mut sources = CompositeSource::new(); diff --git a/datadog-opentelemetry/src/propagation/baggage.rs b/datadog-opentelemetry/src/propagation/baggage.rs new file mode 100644 index 00000000..98582dd6 --- /dev/null +++ b/datadog-opentelemetry/src/propagation/baggage.rs @@ -0,0 +1,270 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! W3C Baggage propagation (`baggage` header). +//! +//! This module exposes the header key so the composite propagator can include it in its `fields()` +//! list. It contains only extraction code. +//! +//! Injection is performed by [`opentelemetry_sdk::propagation::BaggagePropagator`] +//! at the [`DatadogPropagator`](crate::text_map_propagator::DatadogPropagator) layer, which has +//! access to the OTel [`Context`](opentelemetry::Context) that carries baggage. +use std::sync::LazyLock; + +use opentelemetry::baggage::{Baggage, KeyValueMetadata}; +use percent_encoding::percent_decode_str; + +use crate::{dd_warn, propagation::carrier::Extractor}; + +/// The W3C `baggage` header name. +pub const BAGGAGE_KEY: &str = "baggage"; + +/// Cap on the number of baggage members extracted. Entries beyond this are dropped. +const MAX_BAGGAGE_MEMBERS: usize = 32; + +/// Cap on the cumulative byte length of baggage members parsed. Once exceeded, remaining entries +/// are dropped. +const MAX_BAGGAGE_LENGTH: usize = 1024; + +static BAGGAGE_HEADER_KEYS: LazyLock<[String; 1]> = LazyLock::new(|| [BAGGAGE_KEY.to_owned()]); + +/// Returns the header keys used by the W3C baggage propagator. +pub fn keys() -> &'static [String] { + BAGGAGE_HEADER_KEYS.as_slice() +} + +fn parse_baggage_member(baggage_member: &str) -> Option { + let mut member = baggage_member.split(';'); + let Some(name_and_value) = member.next() else { + dd_warn!("Propagator (baggage): invalid format"); + return None; + }; + let mut iter = name_and_value.split('='); + let (Some(name), Some(value)) = (iter.next(), iter.next()) else { + dd_warn!("Propagator (baggage): invalid key-value format"); + return None; + }; + let decode_name = percent_decode_str(name).decode_utf8(); + let decode_value = percent_decode_str(value).decode_utf8(); + + let (Ok(name), Ok(value)) = (decode_name, decode_value) else { + dd_warn!("Propagator (baggage): invalid percent encoded UTF8 string in key values"); + return None; + }; + + // decode and trim metadata entries associated with the key-value + let decoded_props = member + .flat_map(|prop| percent_decode_str(prop).decode_utf8()) + .enumerate() + .fold(String::new(), |mut acc, (i, prop)| { + if i != 0 { + acc.push(';'); + } + acc.push_str(prop.trim()); + acc + }); + + Some(KeyValueMetadata::new( + name.trim().to_owned(), + value.trim().to_string(), + decoded_props, + )) +} + +pub(crate) fn extract_baggage(extractor: &dyn Extractor) -> Option { + let header_value = extractor.get(BAGGAGE_KEY)?; + let mut members = 0; + let mut allocated_size = 0; + let baggage = header_value.split(',') + .take_while(|member| { + allocated_size += member.len(); + let drop_entry = allocated_size > MAX_BAGGAGE_LENGTH; + if drop_entry { + dd_warn!("Propagator (baggage): ignored baggage key-values, only first {} bytes propagated", MAX_BAGGAGE_LENGTH) + } + !drop_entry + }) + .filter_map(parse_baggage_member) + .take_while(|_| { + members +=1; + let drop_entry = members > MAX_BAGGAGE_MEMBERS; + if drop_entry { + dd_warn!("Propagator (baggage): ignored baggage key-values, only first {} propagated", MAX_BAGGAGE_MEMBERS) + } + !drop_entry + }); + Some(Baggage::from(baggage)) +} + +#[cfg(test)] +mod tests { + use super::*; + use opentelemetry::{baggage::BaggageMetadata, Key, StringValue}; + use std::collections::HashMap; + + fn valid_extract_data() -> Vec<(&'static str, HashMap)> { + vec![ + // "valid w3cHeader" + ( + "key1=val1,key2=val2", + vec![ + (Key::new("key1"), StringValue::from("val1")), + (Key::new("key2"), StringValue::from("val2")), + ] + .into_iter() + .collect(), + ), + // "valid w3cHeader with spaces" + ( + "key1 = val1, key2 =val2 ", + vec![ + (Key::new("key1"), StringValue::from("val1")), + (Key::new("key2"), StringValue::from("val2")), + ] + .into_iter() + .collect(), + ), + // "valid header with url-escaped comma" + ( + "key1=val1,key2=val2%2Cval3", + vec![ + (Key::new("key1"), StringValue::from("val1")), + (Key::new("key2"), StringValue::from("val2,val3")), + ] + .into_iter() + .collect(), + ), + // "valid header with an invalid header" + ( + "key1=val1,key2=val2,a,val3", + vec![ + (Key::new("key1"), StringValue::from("val1")), + (Key::new("key2"), StringValue::from("val2")), + ] + .into_iter() + .collect(), + ), + // "valid header with no value" + ( + "key1=,key2=val2", + vec![ + (Key::new("key1"), StringValue::from("")), + (Key::new("key2"), StringValue::from("val2")), + ] + .into_iter() + .collect(), + ), + ] + } + + #[allow(clippy::type_complexity)] + fn valid_extract_data_with_metadata( + ) -> Vec<(&'static str, HashMap)> { + vec![ + // "valid w3cHeader with properties" + ("key1=val1,key2=val2;prop=1", vec![(Key::new("key1"), (StringValue::from("val1"), BaggageMetadata::default())), (Key::new("key2"), (StringValue::from("val2"), BaggageMetadata::from("prop=1")))].into_iter().collect()), + // prop can don't need to be key value pair + ("key1=val1,key2=val2;prop1", vec![(Key::new("key1"), (StringValue::from("val1"), BaggageMetadata::default())), (Key::new("key2"), (StringValue::from("val2"), BaggageMetadata::from("prop1")))].into_iter().collect()), + ("key1=value1;property1;property2, key2 = value2, key3=value3; propertyKey=propertyValue", + vec![ + (Key::new("key1"), (StringValue::from("value1"), BaggageMetadata::from("property1;property2"))), + (Key::new("key2"), (StringValue::from("value2"), BaggageMetadata::default())), + (Key::new("key3"), (StringValue::from("value3"), BaggageMetadata::from("propertyKey=propertyValue"))), + ].into_iter().collect()), + ] + } + + #[test] + fn test_extract_baggage() { + for (header_value, kvs) in valid_extract_data() { + let mut extractor: HashMap = HashMap::new(); + extractor.insert(BAGGAGE_KEY.to_string(), header_value.to_string()); + let baggage = extract_baggage(&extractor).expect("baggage extracted"); + + assert_eq!(kvs.len(), baggage.len()); + for (key, (value, _metadata)) in &baggage { + assert_eq!(Some(value), kvs.get(key)) + } + } + } + + #[test] + fn test_extract_baggage_with_metadata() { + for (header_value, kvm) in valid_extract_data_with_metadata() { + let mut extractor: HashMap = HashMap::new(); + extractor.insert(BAGGAGE_KEY.to_string(), header_value.to_string()); + let baggage = extract_baggage(&extractor).expect("baggage extracted"); + + assert_eq!(kvm.len(), baggage.len()); + for (key, value_and_prop) in &baggage { + assert_eq!(Some(value_and_prop), kvm.get(key)) + } + } + } + + #[test] + fn test_extract_baggage_respects_max_members() { + let total = MAX_BAGGAGE_MEMBERS + 8; + let header_value = (0..total) + .map(|i| format!("k{i}=v{i}")) + .collect::>() + .join(","); + + let mut extractor: HashMap = HashMap::new(); + extractor.insert(BAGGAGE_KEY.to_string(), header_value); + let baggage = extract_baggage(&extractor).expect("baggage extracted"); + + assert_eq!(baggage.len(), MAX_BAGGAGE_MEMBERS); + for i in 0..MAX_BAGGAGE_MEMBERS { + let key = Key::new(format!("k{i}")); + assert_eq!( + baggage.get(&key).map(|v| v.as_str()), + Some(format!("v{i}").as_str()) + ); + } + for i in MAX_BAGGAGE_MEMBERS..total { + let key = Key::new(format!("k{i}")); + assert!(baggage.get(&key).is_none()); + } + } + + #[test] + fn test_extract_baggage_respects_max_length() { + // Each member is exactly 50 bytes: "k{NN}={padding...}". + // 21 members * 50 = 1050 bytes (separators excluded by the + // implementation), so the 21st entry pushes allocated_size past + // MAX_BAGGAGE_LENGTH (1024) and must be dropped. + let member_size = 50; + let total = 25; + let members: Vec = (0..total) + .map(|i| { + let prefix = format!("k{i:02}="); + let pad = "x".repeat(member_size - prefix.len()); + format!("{prefix}{pad}") + }) + .collect(); + assert!(members.iter().map(|m| m.len()).sum::() > MAX_BAGGAGE_LENGTH); + let header_value = members.join(","); + + let mut extractor: HashMap = HashMap::new(); + extractor.insert(BAGGAGE_KEY.to_string(), header_value); + let baggage = extract_baggage(&extractor).expect("baggage extracted"); + + let expected_kept = MAX_BAGGAGE_LENGTH / member_size; + assert_eq!(baggage.len(), expected_kept); + for i in 0..expected_kept { + let key = Key::new(format!("k{i:02}")); + assert!( + baggage.get(&key).is_some(), + "expected key k{i:02} to be present" + ); + } + for i in expected_kept..total { + let key = Key::new(format!("k{i:02}")); + assert!( + baggage.get(&key).is_none(), + "expected key k{i:02} to be dropped" + ); + } + } +} diff --git a/datadog-opentelemetry/src/propagation/config.rs b/datadog-opentelemetry/src/propagation/config.rs index ea31576f..334110a7 100644 --- a/datadog-opentelemetry/src/propagation/config.rs +++ b/datadog-opentelemetry/src/propagation/config.rs @@ -5,7 +5,7 @@ use crate::configuration::TracePropagationStyle; use crate::propagation::PropagationConfig; -pub(super) fn get_extractors(config: &impl PropagationConfig) -> &[TracePropagationStyle] { +pub(crate) fn get_extractors(config: &impl PropagationConfig) -> &[TracePropagationStyle] { if let Some(extractors) = config.trace_propagation_style_extract() { extractors } else { @@ -13,7 +13,7 @@ pub(super) fn get_extractors(config: &impl PropagationConfig) -> &[TracePropagat } } -pub(super) fn get_injectors(config: &impl PropagationConfig) -> &[TracePropagationStyle] { +pub(crate) fn get_injectors(config: &impl PropagationConfig) -> &[TracePropagationStyle] { if let Some(injectors) = config.trace_propagation_style_inject() { injectors } else { diff --git a/datadog-opentelemetry/src/propagation/mod.rs b/datadog-opentelemetry/src/propagation/mod.rs index 07a19ca2..66106e3a 100644 --- a/datadog-opentelemetry/src/propagation/mod.rs +++ b/datadog-opentelemetry/src/propagation/mod.rs @@ -14,6 +14,8 @@ use config::{get_extractors, get_injectors}; use datadog::DATADOG_LAST_PARENT_ID_KEY; use tracecontext::TRACESTATE_KEY; +/// W3C Baggage propagation (`baggage` header). +pub mod baggage; pub mod carrier; pub(crate) mod config; pub mod context; @@ -70,7 +72,7 @@ pub struct DatadogCompositePropagator { config: Arc, extractors: Vec, injectors: Vec, - keys: Vec, + fields: Vec, } impl DatadogCompositePropagator { @@ -99,18 +101,18 @@ impl DatadogCompositePropagator { .copied() .collect(); - let keys = extractors.iter().fold(Vec::new(), |mut keys, extractor| { - >::keys(extractor) + let fields = injectors.iter().fold(Vec::new(), |mut fields, injector| { + >::keys(injector) .iter() - .for_each(|key| keys.push(key.clone())); - keys + .for_each(|key| fields.push(key.clone())); + fields }); Self { config, extractors, injectors, - keys, + fields, } } @@ -135,9 +137,9 @@ impl DatadogCompositePropagator { .for_each(|propagator| propagator.inject(context, carrier, self.config.as_ref())); } - /// Returns the header keys used by the configured extractors. - pub fn keys(&self) -> &[String] { - &self.keys + /// Returns the header keys used by the configured injectors. + pub fn fields(&self) -> &[String] { + &self.fields } fn extract_available_contexts( @@ -927,10 +929,11 @@ pub(crate) mod tests { fn get_config( extract: Option>, - _: Option>, + inject: Option>, ) -> Arc { let mut builder = Config::builder(); builder.set_trace_propagation_style_extract(extract.unwrap_or_default()); + builder.set_trace_propagation_style_inject(inject.unwrap_or_default()); Arc::new(builder.build()) } @@ -1331,12 +1334,12 @@ pub(crate) mod tests { } #[test] - fn test_default_keys() { - let extract = Some(vec![ + fn test_default_fields() { + let inject = Some(vec![ TracePropagationStyle::Datadog, TracePropagationStyle::TraceContext, ]); - let config = get_config(extract, None); + let config = get_config(None, inject); let propagator = DatadogCompositePropagator::new(config); @@ -1350,24 +1353,24 @@ pub(crate) mod tests { "traceparent", "tracestate" ], - propagator.keys() + propagator.fields() ) } #[test] - fn test_tracecontext_keys() { - let extract = Some(vec![TracePropagationStyle::TraceContext]); - let config = get_config(extract, None); + fn test_tracecontext_fields() { + let inject = Some(vec![TracePropagationStyle::TraceContext]); + let config = get_config(None, inject); let propagator = DatadogCompositePropagator::new(config); - assert_eq!(vec!["traceparent", "tracestate"], propagator.keys()) + assert_eq!(vec!["traceparent", "tracestate"], propagator.fields()) } #[test] - fn test_datadog_keys() { - let extract = Some(vec![TracePropagationStyle::Datadog]); - let config = get_config(extract, None); + fn test_datadog_fields() { + let inject = Some(vec![TracePropagationStyle::Datadog]); + let config = get_config(None, inject); let propagator = DatadogCompositePropagator::new(config); @@ -1379,7 +1382,7 @@ pub(crate) mod tests { "x-datadog-sampling-priority", "x-datadog-tags", ], - propagator.keys() + propagator.fields() ) } } diff --git a/datadog-opentelemetry/src/propagation/trace_propagation_style.rs b/datadog-opentelemetry/src/propagation/trace_propagation_style.rs index 735424fa..9f7deb8b 100644 --- a/datadog-opentelemetry/src/propagation/trace_propagation_style.rs +++ b/datadog-opentelemetry/src/propagation/trace_propagation_style.rs @@ -5,6 +5,7 @@ use crate::core::configuration::TracePropagationStyle; use serde::{Deserialize, Deserializer}; use crate::propagation::{ + baggage, carrier::{Extractor, Injector}, context::{InjectSpanContext, SpanContext}, datadog, tracecontext, PropagationConfig, Propagator, @@ -17,7 +18,8 @@ impl Propagator for TracePropagationStyle { match self { Self::Datadog => datadog::extract(carrier, config), Self::TraceContext => tracecontext::extract(carrier), - _ => None, + // Baggage extraction operates on OTel Context and is handled by DatadogPropagator. + Self::Baggage | Self::None => None, } } @@ -25,7 +27,8 @@ impl Propagator for TracePropagationStyle { match self { Self::Datadog => datadog::inject(context, carrier, config), Self::TraceContext => tracecontext::inject(context, carrier), - _ => {} + // Baggage injection operates on OTel Context and is handled by DatadogPropagator. + Self::Baggage | Self::None => {} } } @@ -33,7 +36,8 @@ impl Propagator for TracePropagationStyle { match self { Self::Datadog => datadog::keys(), Self::TraceContext => tracecontext::keys(), - _ => &NONE_KEYS, + Self::Baggage => baggage::keys(), + Self::None => &NONE_KEYS, } } } diff --git a/datadog-opentelemetry/src/text_map_propagator.rs b/datadog-opentelemetry/src/text_map_propagator.rs index cfb4d278..1c9f366f 100644 --- a/datadog-opentelemetry/src/text_map_propagator.rs +++ b/datadog-opentelemetry/src/text_map_propagator.rs @@ -7,14 +7,18 @@ use crate::{ catch_panic, core::{configuration::Config, sampling::priority}, propagation::{ + baggage::extract_baggage, + config::{get_extractors, get_injectors}, context::{InjectSpanContext, InjectTraceState, Sampling, SpanContext, SpanLink}, - DatadogCompositePropagator, + DatadogCompositePropagator, TracePropagationStyle, }, }; use opentelemetry::{ + baggage::BaggageExt as _, propagation::{text_map_propagator::FieldIter, TextMapPropagator}, trace::TraceContextExt, }; +use opentelemetry_sdk::propagation::BaggagePropagator; use crate::TraceRegistry; @@ -64,14 +68,25 @@ pub struct DatadogPropagator { inner: DatadogCompositePropagator, registry: TraceRegistry, cfg: Arc, + baggage_propagator: BaggagePropagator, + baggage_extract: bool, + baggage_inject: bool, } impl DatadogPropagator { pub(crate) fn new(config: Arc, registry: TraceRegistry) -> Self { + let baggage_extract = + get_extractors(config.as_ref()).contains(&TracePropagationStyle::Baggage); + let baggage_inject = + get_injectors(config.as_ref()).contains(&TracePropagationStyle::Baggage); + let inner = DatadogCompositePropagator::new(config.clone()); DatadogPropagator { - inner: DatadogCompositePropagator::new(config.clone()), + inner, registry, cfg: config, + baggage_propagator: BaggagePropagator::new(), + baggage_extract, + baggage_inject, } } @@ -83,6 +98,10 @@ impl DatadogPropagator { cx: &opentelemetry::Context, mut injector: &mut dyn opentelemetry::propagation::Injector, ) { + if self.baggage_inject { + self.baggage_propagator.inject_context(cx, injector); + } + let span = cx.span(); let otel_span_context = span.span_context(); @@ -150,7 +169,8 @@ impl DatadogPropagator { cx: &opentelemetry::Context, extractor: &dyn opentelemetry::propagation::Extractor, ) -> opentelemetry::Context { - self.inner + let cx = self + .inner .extract(&extractor) .map(|dd_span_context| { let trace_flags = extract_trace_flags(&dd_span_context); @@ -167,7 +187,17 @@ impl DatadogPropagator { cx.with_remote_span_context(otel_span_context) .with_value(DatadogExtractData::from_span_context(dd_span_context)) }) - .unwrap_or_else(|| cx.clone()) + .unwrap_or_else(|| cx.clone()); + + if self.baggage_extract { + if let Some(baggage) = extract_baggage(&extractor) { + cx.with_baggage(baggage) + } else { + cx + } + } else { + cx + } } } @@ -199,8 +229,8 @@ impl TextMapPropagator for DatadogPropagator { } fn fields(&self) -> opentelemetry::propagation::text_map_propagator::FieldIter<'_> { - let fields = if self.cfg.enabled() { - self.inner.keys() + let fields: &[String] = if self.cfg.enabled() { + self.inner.fields() } else { &[] }; @@ -245,6 +275,7 @@ pub mod tests { }; use assert_unordered::assert_eq_unordered; use opentelemetry::{ + baggage::BaggageExt, propagation::{Extractor, TextMapPropagator}, trace::{Span, SpanContext as OtelSpanContext, Status, TraceContextExt, TraceState}, Context, KeyValue, SpanId, TraceFlags, TraceId, @@ -679,4 +710,326 @@ pub mod tests { } } } + + use crate::propagation::baggage::BAGGAGE_KEY; + + fn get_propagator_with_separate_styles( + extract: Vec, + inject: Vec, + ) -> DatadogPropagator { + let config = Arc::new( + Config::builder() + .set_trace_propagation_style_extract(extract) + .set_trace_propagation_style_inject(inject) + .build(), + ); + DatadogPropagator::new(config.clone(), TraceRegistry::new(config)) + } + + #[test] + fn baggage_extract_only_when_not_in_inject_styles() { + // Extract has baggage, inject does not. + let propagator = get_propagator_with_separate_styles( + vec![ + TracePropagationStyle::Baggage, + TracePropagationStyle::Datadog, + ], + vec![TracePropagationStyle::TraceContext], + ); + + // Extraction: baggage header present → lands in OTel Context. + let mut carrier = HashMap::new(); + carrier.insert(BAGGAGE_KEY.to_string(), "user=alice".to_string()); + let cx = propagator.extract(&carrier); + assert!( + cx.baggage().get("user").is_some(), + "baggage should be extracted when Baggage is in extract styles" + ); + + // Injection: same Context → baggage header must NOT be written. + let mut injector: HashMap = HashMap::new(); + propagator.inject_context(&cx, &mut injector); + assert!( + !injector.contains_key(BAGGAGE_KEY), + "baggage header must not be injected when Baggage is absent from inject styles" + ); + } + + #[test] + fn baggage_inject_only_when_not_in_extract_styles() { + // Inject has baggage, extract does not. + let propagator = get_propagator_with_separate_styles( + vec![TracePropagationStyle::Datadog], + vec![ + TracePropagationStyle::Baggage, + TracePropagationStyle::TraceContext, + ], + ); + + // Extraction: baggage header present → must NOT land in OTel Context. + let mut carrier = HashMap::new(); + carrier.insert(BAGGAGE_KEY.to_string(), "user=alice".to_string()); + let cx = propagator.extract(&carrier); + assert!( + cx.baggage().get("user").is_none(), + "baggage must not be extracted when Baggage is absent from extract styles" + ); + + // Injection: context carrying baggage → header must be written. + let cx_with_baggage = Context::current_with_baggage(vec![KeyValue::new("server", "42")]); + let span_cx = OtelSpanContext::new( + TraceId::from(0x4bf9_2f35_77b3_4da6_a3ce_929d_0e0e_4736_u128), + SpanId::from(0x00f0_67aa_0ba9_02b7_u64), + TraceFlags::SAMPLED, + true, + TraceState::NONE, + ); + let cx_with_both = cx_with_baggage.with_remote_span_context(span_cx); + + let mut injector: HashMap = HashMap::new(); + propagator.inject_context(&cx_with_both, &mut injector); + let header = injector + .get(BAGGAGE_KEY) + .expect("baggage header must be injected"); + assert!( + header.contains("server=42"), + "injected baggage header should contain the baggage entry" + ); + } + + // ── Helper: build a propagator from a single combined style list ────────── + + fn get_propagator_with_combined_style(styles: Vec) -> DatadogPropagator { + let config = Arc::new( + Config::builder() + .set_trace_propagation_style(styles) + .build(), + ); + DatadogPropagator::new(config.clone(), TraceRegistry::new(config)) + } + + // ── Unit: composite propagator inclusion/exclusion ──────────────────────── + + #[test] + fn baggage_included_in_fields_when_configured() { + let propagator = get_propagator_with_combined_style(vec![TracePropagationStyle::Baggage]); + let fields: Vec<&str> = propagator.fields().collect(); + assert!( + fields.contains(&BAGGAGE_KEY), + "fields() must include 'baggage' when Baggage style is configured; got {fields:?}" + ); + } + + #[test] + fn baggage_excluded_from_fields_when_not_configured() { + let propagator = get_propagator_with_combined_style(vec![ + TracePropagationStyle::Datadog, + TracePropagationStyle::TraceContext, + ]); + let fields: Vec<&str> = propagator.fields().collect(); + assert!( + !fields.contains(&BAGGAGE_KEY), + "fields() must not include 'baggage' when Baggage style is absent; got {fields:?}" + ); + } + + #[test] + fn baggage_included_in_fields_when_only_in_inject_styles() { + let propagator = get_propagator_with_separate_styles( + vec![TracePropagationStyle::Datadog], + vec![ + TracePropagationStyle::Baggage, + TracePropagationStyle::TraceContext, + ], + ); + let fields: Vec<&str> = propagator.fields().collect(); + assert!( + fields.contains(&BAGGAGE_KEY), + "fields() must include 'baggage' when Baggage is only in inject styles; got {fields:?}" + ); + } + + #[test] + fn baggage_excluded_from_fields_when_only_in_extract_styles() { + // fields() represents what the propagator *writes* (injection only, per OTel spec). + // When Baggage is only in the extract list it should not appear in fields(). + let propagator = get_propagator_with_separate_styles( + vec![ + TracePropagationStyle::Baggage, + TracePropagationStyle::TraceContext, + ], + vec![TracePropagationStyle::Datadog], + ); + let fields: Vec<&str> = propagator.fields().collect(); + assert!( + !fields.contains(&BAGGAGE_KEY), + "fields() must not include 'baggage' when Baggage is only in extract styles; got {fields:?}" + ); + } + + #[test] + fn baggage_not_injected_when_only_in_extract_styles() { + let propagator = get_propagator_with_separate_styles( + vec![ + TracePropagationStyle::Baggage, + TracePropagationStyle::TraceContext, + ], + vec![TracePropagationStyle::Datadog], + ); + + let cx = Context::current_with_baggage(vec![KeyValue::new("user", "alice")]); + let mut carrier: HashMap = HashMap::new(); + propagator.inject_context(&cx, &mut carrier); + + assert!( + !carrier.contains_key(BAGGAGE_KEY), + "baggage header must not be injected when Baggage is only in extract styles; got {carrier:?}" + ); + } + + // ── Integration: inject / extract / round-trip ──────────────────────────── + + #[test] + fn baggage_inject_encodes_multiple_entries() { + let propagator = get_propagator_with_combined_style(vec![TracePropagationStyle::Baggage]); + + let cx = Context::current_with_baggage(vec![ + KeyValue::new("user", "alice"), + KeyValue::new("tenant", "acme"), + ]); + + let mut carrier: HashMap = HashMap::new(); + propagator.inject_context(&cx, &mut carrier); + + let header = carrier + .get(BAGGAGE_KEY) + .expect("baggage header must be present"); + assert!( + header.contains("user=alice"), + "header must contain user=alice; got {header}" + ); + assert!( + header.contains("tenant=acme"), + "header must contain tenant=acme; got {header}" + ); + } + + #[test] + fn baggage_extract_populates_context() { + let propagator = get_propagator_with_combined_style(vec![TracePropagationStyle::Baggage]); + + let mut carrier: HashMap = HashMap::new(); + carrier.insert( + BAGGAGE_KEY.to_string(), + "user=alice,tenant=acme".to_string(), + ); + + let cx = propagator.extract(&carrier); + + let baggage = cx.baggage(); + assert_eq!( + baggage.get("user").map(|v| v.as_str()), + Some("alice"), + "extracted baggage must contain user=alice" + ); + assert_eq!( + baggage.get("tenant").map(|v| v.as_str()), + Some("acme"), + "extracted baggage must contain tenant=acme" + ); + } + + #[test] + fn baggage_round_trip() { + let propagator = get_propagator_with_combined_style(vec![TracePropagationStyle::Baggage]); + + // Inject + let original_cx = Context::current_with_baggage(vec![ + KeyValue::new("request-id", "xyz-123"), + KeyValue::new("region", "us-east-1"), + ]); + let mut carrier: HashMap = HashMap::new(); + propagator.inject_context(&original_cx, &mut carrier); + + // Extract into a fresh context + let extracted_cx = propagator.extract(&carrier); + let baggage = extracted_cx.baggage(); + + assert_eq!( + baggage.get("request-id").map(|v| v.as_str()), + Some("xyz-123"), + "round-trip must preserve request-id" + ); + assert_eq!( + baggage.get("region").map(|v| v.as_str()), + Some("us-east-1"), + "round-trip must preserve region" + ); + } + + // ── Env-var style scenarios ─────────────────────────────────────────────── + + #[test] + fn baggage_not_injected_when_style_is_datadog_tracecontext() { + // DD_TRACE_PROPAGATION_STYLE=datadog,tracecontext — no Baggage variant. + let propagator = get_propagator_with_combined_style(vec![ + TracePropagationStyle::Datadog, + TracePropagationStyle::TraceContext, + ]); + + let cx = Context::current_with_baggage(vec![KeyValue::new("user", "alice")]); + let mut carrier: HashMap = HashMap::new(); + propagator.inject_context(&cx, &mut carrier); + + assert!( + !carrier.contains_key(BAGGAGE_KEY), + "baggage header must NOT be present when Baggage is not in propagation style" + ); + } + + #[test] + fn baggage_only_style_injects_only_baggage_header() { + // DD_TRACE_PROPAGATION_STYLE=baggage — no Datadog or TraceContext headers. + let propagator = get_propagator_with_combined_style(vec![TracePropagationStyle::Baggage]); + + let cx = Context::current_with_baggage(vec![KeyValue::new("user", "alice")]); + let mut carrier: HashMap = HashMap::new(); + propagator.inject_context(&cx, &mut carrier); + + assert!( + carrier.contains_key(BAGGAGE_KEY), + "baggage header must be present" + ); + assert!( + !carrier.contains_key("traceparent"), + "traceparent must NOT be present when TraceContext is not in propagation style" + ); + assert!( + !carrier.contains_key("x-datadog-trace-id"), + "x-datadog-trace-id must NOT be present when Datadog is not in propagation style" + ); + } + + #[test] + fn baggage_injected_by_default_config() { + // Default config (no explicit style set) includes Baggage. + let config = Arc::new(Config::builder().build()); + let propagator = DatadogPropagator::new(config.clone(), TraceRegistry::new(config)); + + // Verify via fields() + let fields: Vec<&str> = propagator.fields().collect(); + assert!( + fields.contains(&BAGGAGE_KEY), + "fields() must include 'baggage' with default config; got {fields:?}" + ); + + // Verify via actual injection + let cx = Context::current_with_baggage(vec![KeyValue::new("env", "prod")]); + let mut carrier: HashMap = HashMap::new(); + propagator.inject_context(&cx, &mut carrier); + assert!( + carrier.contains_key(BAGGAGE_KEY), + "baggage header must be injected with default config" + ); + } }