diff --git a/datadog-opentelemetry/src/text_map_propagator.rs b/datadog-opentelemetry/src/text_map_propagator.rs index 0241dd7..ff11e46 100644 --- a/datadog-opentelemetry/src/text_map_propagator.rs +++ b/datadog-opentelemetry/src/text_map_propagator.rs @@ -1,7 +1,7 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 -use std::{collections::HashMap, str::FromStr, sync::Arc}; +use std::{collections::HashMap, sync::Arc}; use dd_trace::{catch_panic, sampling::priority, Config}; use opentelemetry::{ @@ -10,7 +10,7 @@ use opentelemetry::{ }; use dd_trace_propagation::{ - context::{InjectSpanContext, Sampling, SpanContext, SpanLink, Tracestate}, + context::{InjectSpanContext, InjectTraceState, Sampling, SpanContext, SpanLink}, DatadogCompositePropagator, }; @@ -113,7 +113,7 @@ impl DatadogPropagator { let tracestate = if *otel_tracestate == opentelemetry::trace::TraceState::NONE { None } else { - Tracestate::from_str(&otel_tracestate.header()).ok() + Some(InjectTraceState::from_header(otel_tracestate.header())) }; let tags = if let Some(propagation_tags) = &mut propagation_data.tags { @@ -129,7 +129,7 @@ impl DatadogPropagator { sampling, origin: propagation_data.origin.as_deref(), tags, - tracestate: tracestate.as_ref(), + tracestate, }; self.inner.inject(dd_span_context, &mut injector) diff --git a/dd-trace-propagation/benches/inject_benchmark.rs b/dd-trace-propagation/benches/inject_benchmark.rs index 3d5aaf7..28bf17d 100644 --- a/dd-trace-propagation/benches/inject_benchmark.rs +++ b/dd-trace-propagation/benches/inject_benchmark.rs @@ -26,7 +26,7 @@ fn span_context_to_inject(c: &mut SpanContext) -> InjectSpanContext<'_> { origin: c.origin.as_deref(), tags: &mut c.tags, is_remote: c.is_remote, - tracestate: c.tracestate.as_ref(), + tracestate: None, } } diff --git a/dd-trace-propagation/src/context.rs b/dd-trace-propagation/src/context.rs index 15b73ab..4fec8e0 100644 --- a/dd-trace-propagation/src/context.rs +++ b/dd-trace-propagation/src/context.rs @@ -75,7 +75,7 @@ pub struct InjectSpanContext<'a> { // tags needs to be mutable because we insert the error meta field pub tags: &'a mut HashMap, pub is_remote: bool, - pub tracestate: Option<&'a Tracestate>, + pub tracestate: Option, } #[cfg(test)] @@ -89,7 +89,17 @@ pub(crate) fn span_context_to_inject(c: &mut SpanContext) -> InjectSpanContext<' origin: c.origin.as_deref(), tags: &mut c.tags, is_remote: c.is_remote, - tracestate: c.tracestate.as_ref(), + tracestate: c.tracestate.as_ref().map(|ts| { + InjectTraceState::from_header(ts.additional_values.as_ref().map_or( + String::new(), + |v| { + v.iter() + .map(|(k, v)| format!("{k}={v}")) + .collect::>() + .join(",") + }, + )) + }), } } @@ -105,6 +115,29 @@ pub struct SpanContext { pub tracestate: Option, } +/// A tracestate we grab from the parent span +/// +/// Only non-dd keys in the tracestate are injected +pub struct InjectTraceState { + header: String, +} + +impl InjectTraceState { + pub fn from_header(header: String) -> Self { + Self { header } + } + + pub fn additional_values(&self) -> impl Iterator { + self.header.split(',').filter(|part| { + let (key, value) = part.split_once('=').unwrap_or((part, "")); + key != "dd" + && !value.is_empty() + && Tracestate::valid_key(key) + && Tracestate::valid_value(value) + }) + } +} + #[derive(Clone, Debug, PartialEq)] pub struct Traceparent { pub sampling_priority: SamplingPriority, @@ -171,9 +204,7 @@ impl FromStr for Tracestate { let mut additional_values = vec![]; for v in ts_v { - let mut parts = v.splitn(2, '='); - let key = parts.next().unwrap_or_default(); - let value = parts.next().unwrap_or_default(); + let (key, value) = v.split_once('=').unwrap_or(("", "")); if !Tracestate::valid_key(key) || value.is_empty() || !Tracestate::valid_value(value) { dd_debug!("Tracestate: invalid key or header value: {v}"); diff --git a/dd-trace-propagation/src/tracecontext.rs b/dd-trace-propagation/src/tracecontext.rs index 5b125a9..823058e 100644 --- a/dd-trace-propagation/src/tracecontext.rs +++ b/dd-trace-propagation/src/tracecontext.rs @@ -249,14 +249,10 @@ fn inject_tracestate(context: &InjectSpanContext, carrier: &mut dyn Injector) { } // Add additional tracestate values if present - if let Some(ts) = context.tracestate { - if let Some(ref additional) = ts.additional_values { - for (key, value) in additional.iter().take(31) { - tracestate.push_str(TRACESTATE_VALUES_SEPARATOR); - tracestate.push_str(key); - tracestate.push('='); - tracestate.push_str(value); - } + if let Some(ts) = &context.tracestate { + for part in ts.additional_values().take(31) { + tracestate.push_str(TRACESTATE_VALUES_SEPARATOR); + tracestate.push_str(part) } } @@ -509,7 +505,10 @@ pub fn keys() -> &'static [String] { mod test { use dd_trace::{configuration::TracePropagationStyle, sampling::priority, Config}; - use crate::{context::span_context_to_inject, Propagator}; + use crate::{ + context::{span_context_to_inject, InjectTraceState}, + Propagator, + }; use super::*; @@ -716,7 +715,6 @@ mod test { #[test] fn test_inject_traceparent() { - let ts = Tracestate::from_str("other=bleh,atel=test,dd=s:2;o:foo_bar_;t.dm:-4").unwrap(); let mut context = InjectSpanContext { trace_id: u128::from_str_radix("1111aaaa2222bbbb3333cccc4444dddd", 16).unwrap(), span_id: u64::from_str_radix("5555eeee6666ffff", 16).unwrap(), @@ -730,7 +728,9 @@ mod test { "abc~!@#$%^&*()_+`-=".to_string(), )]), is_remote: false, - tracestate: Some(&ts), + tracestate: Some(InjectTraceState::from_header( + "other=bleh,atel=test,dd=s:2;o:foo_bar_;t.dm:-4".to_owned(), + )), }; let mut carrier: HashMap = HashMap::new(); @@ -791,7 +791,7 @@ mod test { for index in 0..35 { tracestate.push(format!("state{index}=value-{index}")); } - let tracestate = Tracestate::from_str(&tracestate.join(",")).unwrap(); + let tracestate = tracestate.join(","); let mut context = InjectSpanContext { trace_id: u128::from_str_radix("1111aaaa2222bbbb3333cccc4444dddd", 16).unwrap(), @@ -803,7 +803,7 @@ mod test { origin: Some("rum"), tags: &mut HashMap::from([("_dd.p.foo".to_string(), "abc".to_string())]), is_remote: false, - tracestate: Some(&tracestate), + tracestate: Some(InjectTraceState::from_header(tracestate)), }; let mut carrier: HashMap = HashMap::new();