Skip to content

Commit 040a625

Browse files
paullegranddcpaullegranddc
authored
perf(propagation): inject borrowed data (#96)
# What does this PR do? Injection needs to read fields such as propagation tags and origin. Currently the SpanContext struct needs to have ownership of these which forces copying when injecting the span context. This PR introduces and InjectSpanContext which just borrows the fields
1 parent d05c0f7 commit 040a625

File tree

7 files changed

+129
-59
lines changed

7 files changed

+129
-59
lines changed

‎datadog-opentelemetry/src/text_map_propagator.rs‎

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
22
// SPDX-License-Identifier: Apache-2.0
33

4-
use std::{collections::HashMap, str::FromStr, sync::Arc, vec};
4+
use std::{collections::HashMap, str::FromStr, sync::Arc};
55

66
use dd_trace::{catch_panic, sampling::priority, Config};
77
use opentelemetry::{
@@ -10,7 +10,7 @@ use opentelemetry::{
1010
};
1111

1212
use dd_trace_propagation::{
13-
context::{Sampling, SpanContext, SpanLink, Tracestate},
13+
context::{InjectSpanContext, Sampling, SpanContext, SpanLink, Tracestate},
1414
DatadogCompositePropagator,
1515
};
1616

@@ -87,7 +87,7 @@ impl DatadogPropagator {
8787
}
8888

8989
let trace_id = otel_span_context.trace_id().to_bytes();
90-
let propagation_data = self.registry.get_trace_propagation_data(trace_id);
90+
let mut propagation_data = self.registry.get_trace_propagation_data(trace_id);
9191

9292
// get Trace's sampling decision and if it is not present obtain it from otel's SpanContext
9393
// flags
@@ -116,24 +116,20 @@ impl DatadogPropagator {
116116
Tracestate::from_str(&otel_tracestate.header()).ok()
117117
};
118118

119-
let mut tags = HashMap::new();
120-
if let Some(propagation_tags) = propagation_data.tags {
121-
propagation_tags.iter().for_each(|(key, value)| {
122-
if key.starts_with("_dd.p.") {
123-
tags.insert(key.clone(), value.clone());
124-
}
125-
});
126-
}
119+
let tags = if let Some(propagation_tags) = &mut propagation_data.tags {
120+
propagation_tags
121+
} else {
122+
&mut HashMap::new()
123+
};
127124

128-
let dd_span_context = &mut SpanContext {
125+
let dd_span_context = &mut InjectSpanContext {
129126
trace_id: u128::from_be_bytes(trace_id),
130127
span_id: u64::from_be_bytes(otel_span_context.span_id().to_bytes()),
131128
is_remote: otel_span_context.is_remote(),
132-
links: vec![], // links don't affect injection
133129
sampling,
134-
origin: propagation_data.origin,
130+
origin: propagation_data.origin.as_deref(),
135131
tags,
136-
tracestate,
132+
tracestate: tracestate.as_ref(),
137133
};
138134

139135
self.inner.inject(dd_span_context, &mut injector)

‎dd-trace-propagation/benches/inject_benchmark.rs‎

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
22
// SPDX-License-Identifier: Apache-2.0
33

4-
use criterion::{black_box, criterion_group, criterion_main, BatchSize, Criterion};
4+
use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
55
use dd_trace::{
66
configuration::TracePropagationStyle,
77
sampling::{mechanism, priority},
@@ -10,14 +10,26 @@ use dd_trace::{
1010
};
1111
use dd_trace_propagation::{
1212
carrier::Injector,
13-
context::{Sampling, SpanContext},
13+
context::{InjectSpanContext, Sampling, SpanContext},
1414
DatadogCompositePropagator,
1515
};
1616
use std::{collections::HashMap, sync::Arc};
1717

1818
#[global_allocator]
1919
static GLOBAL: ReportingAllocator<std::alloc::System> = ReportingAllocator::new(std::alloc::System);
2020

21+
fn span_context_to_inject(c: &mut SpanContext) -> InjectSpanContext<'_> {
22+
InjectSpanContext {
23+
trace_id: c.trace_id,
24+
span_id: c.span_id,
25+
sampling: c.sampling,
26+
origin: c.origin.as_deref(),
27+
tags: &mut c.tags,
28+
is_remote: c.is_remote,
29+
tracestate: c.tracestate.as_ref(),
30+
}
31+
}
32+
2133
// Mock injector for benchmarking
2234
struct BenchInjector {
2335
headers: HashMap<String, String>,
@@ -97,7 +109,7 @@ fn bench_datadog_only_inject<M: criterion::measurement::Measurement + Measuremen
97109
b.iter_batched(
98110
|| (create_simple_span_context(), BenchInjector::new()),
99111
|(mut context, mut carrier)| {
100-
propagator.inject(black_box(&mut context), black_box(&mut carrier))
112+
propagator.inject(&mut span_context_to_inject(&mut context), &mut carrier)
101113
},
102114
BatchSize::LargeInput,
103115
)
@@ -107,7 +119,7 @@ fn bench_datadog_only_inject<M: criterion::measurement::Measurement + Measuremen
107119
b.iter_batched(
108120
|| (create_complex_span_context(), BenchInjector::new()),
109121
|(mut context, mut carrier)| {
110-
propagator.inject(black_box(&mut context), black_box(&mut carrier))
122+
propagator.inject(&mut span_context_to_inject(&mut context), &mut carrier)
111123
},
112124
BatchSize::LargeInput,
113125
)
@@ -130,7 +142,7 @@ fn bench_tracecontext_only_inject<
130142
b.iter_batched(
131143
|| (create_simple_span_context(), BenchInjector::new()),
132144
|(mut context, mut carrier)| {
133-
propagator.inject(black_box(&mut context), black_box(&mut carrier))
145+
propagator.inject(&mut span_context_to_inject(&mut context), &mut carrier)
134146
},
135147
BatchSize::LargeInput,
136148
)
@@ -143,7 +155,7 @@ fn bench_tracecontext_only_inject<
143155
b.iter_batched(
144156
|| (create_complex_span_context(), BenchInjector::new()),
145157
|(mut context, mut carrier)| {
146-
propagator.inject(black_box(&mut context), black_box(&mut carrier))
158+
propagator.inject(&mut span_context_to_inject(&mut context), &mut carrier)
147159
},
148160
BatchSize::LargeInput,
149161
)

‎dd-trace-propagation/src/context.rs‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,33 @@ impl SpanLink {
6666
}
6767
}
6868
}
69+
70+
pub struct InjectSpanContext<'a> {
71+
pub trace_id: u128,
72+
pub span_id: u64,
73+
pub sampling: Sampling,
74+
pub origin: Option<&'a str>,
75+
// tags needs to be mutable because we insert the error meta field
76+
pub tags: &'a mut HashMap<String, String>,
77+
pub is_remote: bool,
78+
pub tracestate: Option<&'a Tracestate>,
79+
}
80+
81+
#[cfg(test)]
82+
/// A helper function because creating synthetic borrowed data is a bit harder
83+
/// than owned data
84+
pub(crate) fn span_context_to_inject(c: &mut SpanContext) -> InjectSpanContext<'_> {
85+
InjectSpanContext {
86+
trace_id: c.trace_id,
87+
span_id: c.span_id,
88+
sampling: c.sampling,
89+
origin: c.origin.as_deref(),
90+
tags: &mut c.tags,
91+
is_remote: c.is_remote,
92+
tracestate: c.tracestate.as_ref(),
93+
}
94+
}
95+
6996
#[derive(Clone, Default, Debug, PartialEq)]
7097
pub struct SpanContext {
7198
pub trace_id: u128,

‎dd-trace-propagation/src/datadog.rs‎

Lines changed: 43 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,8 @@ use std::{collections::HashMap, str::FromStr, sync::LazyLock};
66
use crate::{
77
carrier::{Extractor, Injector},
88
context::{
9-
combine_trace_id, split_trace_id, Sampling, SpanContext, DATADOG_PROPAGATION_TAG_PREFIX,
9+
combine_trace_id, split_trace_id, InjectSpanContext, Sampling, SpanContext,
10+
DATADOG_PROPAGATION_TAG_PREFIX,
1011
},
1112
error::Error,
1213
};
@@ -38,7 +39,7 @@ static DATADOG_HEADER_KEYS: LazyLock<[String; 5]> = LazyLock::new(|| {
3839
]
3940
});
4041

41-
pub fn inject(context: &mut SpanContext, carrier: &mut dyn Injector, config: &Config) {
42+
pub fn inject(context: &mut InjectSpanContext, carrier: &mut dyn Injector, config: &Config) {
4243
let tags = &mut context.tags;
4344

4445
inject_trace_id(context.trace_id, carrier, tags);
@@ -366,7 +367,10 @@ mod test {
366367
sampling::{mechanism, priority},
367368
};
368369

369-
use crate::{context::split_trace_id, Propagator};
370+
use crate::{
371+
context::{span_context_to_inject, split_trace_id},
372+
Propagator,
373+
};
370374

371375
use super::*;
372376

@@ -557,7 +561,11 @@ mod test {
557561
let propagator = TracePropagationStyle::Datadog;
558562

559563
let mut carrier = HashMap::new();
560-
propagator.inject(&mut context, &mut carrier, &Config::builder().build());
564+
propagator.inject(
565+
&mut span_context_to_inject(&mut context),
566+
&mut carrier,
567+
&Config::builder().build(),
568+
);
561569

562570
assert_eq!(carrier[DATADOG_TRACE_ID_KEY], "1234");
563571
assert_eq!(carrier[DATADOG_PARENT_ID_KEY], "5678");
@@ -596,7 +604,11 @@ mod test {
596604
let propagator = TracePropagationStyle::Datadog;
597605

598606
let mut carrier = HashMap::new();
599-
propagator.inject(&mut context, &mut carrier, &Config::builder().build());
607+
propagator.inject(
608+
&mut span_context_to_inject(&mut context),
609+
&mut carrier,
610+
&Config::builder().build(),
611+
);
600612

601613
assert_eq!(carrier[DATADOG_TRACE_ID_KEY], lower.to_string());
602614
assert_eq!(carrier[DATADOG_ORIGIN_KEY], "synthetics");
@@ -617,7 +629,11 @@ mod test {
617629
let propagator = TracePropagationStyle::Datadog;
618630

619631
let mut carrier = HashMap::new();
620-
propagator.inject(&mut context, &mut carrier, &Config::builder().build());
632+
propagator.inject(
633+
&mut span_context_to_inject(&mut context),
634+
&mut carrier,
635+
&Config::builder().build(),
636+
);
621637

622638
assert_eq!(carrier[DATADOG_TAGS_KEY], "_dd.p.dm=-4");
623639
}
@@ -632,7 +648,11 @@ mod test {
632648
let propagator = TracePropagationStyle::Datadog;
633649

634650
let mut carrier = HashMap::new();
635-
propagator.inject(&mut context, &mut carrier, &Config::builder().build());
651+
propagator.inject(
652+
&mut span_context_to_inject(&mut context),
653+
&mut carrier,
654+
&Config::builder().build(),
655+
);
636656

637657
assert_eq!(carrier.get(DATADOG_TAGS_KEY), None);
638658
}
@@ -648,7 +668,11 @@ mod test {
648668
let propagator = TracePropagationStyle::Datadog;
649669

650670
let mut carrier = HashMap::new();
651-
propagator.inject(&mut context, &mut carrier, &Config::builder().build());
671+
propagator.inject(
672+
&mut span_context_to_inject(&mut context),
673+
&mut carrier,
674+
&Config::builder().build(),
675+
);
652676

653677
assert_eq!(carrier.get(DATADOG_TAGS_KEY), None);
654678
}
@@ -663,7 +687,7 @@ mod test {
663687

664688
let mut carrier = HashMap::new();
665689
propagator.inject(
666-
&mut context,
690+
&mut span_context_to_inject(&mut context),
667691
&mut carrier,
668692
&Config::builder().set_datadog_tags_max_length(0).build(),
669693
);
@@ -680,7 +704,11 @@ mod test {
680704
let propagator = TracePropagationStyle::Datadog;
681705

682706
let mut carrier = HashMap::new();
683-
propagator.inject(&mut context, &mut carrier, &Config::builder().build());
707+
propagator.inject(
708+
&mut span_context_to_inject(&mut context),
709+
&mut carrier,
710+
&Config::builder().build(),
711+
);
684712

685713
assert_eq!(carrier.get(DATADOG_TAGS_KEY), None);
686714
}
@@ -695,7 +723,11 @@ mod test {
695723
let propagator = TracePropagationStyle::Datadog;
696724

697725
let mut carrier = HashMap::new();
698-
propagator.inject(&mut context, &mut carrier, &Config::builder().build());
726+
propagator.inject(
727+
&mut span_context_to_inject(&mut context),
728+
&mut carrier,
729+
&Config::builder().build(),
730+
);
699731

700732
assert_eq!(carrier[DATADOG_TAGS_KEY], "_dd.p.other=test");
701733
}

‎dd-trace-propagation/src/lib.rs‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
use std::sync::Arc;
55

6-
use crate::context::{SpanContext, SpanLink};
6+
use crate::context::{InjectSpanContext, SpanContext, SpanLink};
77
use carrier::{Extractor, Injector};
88
use config::{get_extractors, get_injectors};
99
use datadog::DATADOG_LAST_PARENT_ID_KEY;
@@ -20,7 +20,7 @@ pub mod tracecontext;
2020

2121
pub trait Propagator {
2222
fn extract(&self, carrier: &dyn Extractor, config: &Config) -> Option<SpanContext>;
23-
fn inject(&self, context: &mut SpanContext, carrier: &mut dyn Injector, config: &Config);
23+
fn inject(&self, context: &mut InjectSpanContext, carrier: &mut dyn Injector, config: &Config);
2424
fn keys(&self) -> &[String];
2525
}
2626

@@ -81,7 +81,7 @@ impl DatadogCompositePropagator {
8181
Some(context)
8282
}
8383

84-
pub fn inject(&self, context: &mut SpanContext, carrier: &mut dyn Injector) {
84+
pub fn inject(&self, context: &mut InjectSpanContext, carrier: &mut dyn Injector) {
8585
self.injectors
8686
.iter()
8787
.for_each(|propagator| propagator.inject(context, carrier, &self.config));
@@ -990,7 +990,9 @@ pub mod tests {
990990
$(
991991
#[test]
992992
fn $name() {
993-
let (styles, context, expected) = $value;
993+
use crate::context::span_context_to_inject;
994+
995+
let (styles, mut context, expected) = $value;
994996

995997
let builder = if let Some(styles) = styles {
996998
let mut b = Config::builder();
@@ -1003,8 +1005,9 @@ pub mod tests {
10031005
let config = Arc::new(builder.build());
10041006
let propagator = DatadogCompositePropagator::new(config);
10051007

1008+
let mut inject_context = span_context_to_inject(&mut context);
10061009
let mut carrier = HashMap::new();
1007-
propagator.inject(context, &mut carrier);
1010+
propagator.inject(&mut inject_context, &mut carrier);
10081011

10091012
assert_hashmap_keys(&expected, &carrier);
10101013
assert_hashmap_keys(&carrier, &expected);

‎dd-trace-propagation/src/trace_propagation_style.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ use serde::{Deserialize, Deserializer};
77

88
use crate::{
99
carrier::{Extractor, Injector},
10-
context::SpanContext,
10+
context::{InjectSpanContext, SpanContext},
1111
datadog, tracecontext, Propagator,
1212
};
1313

@@ -22,7 +22,7 @@ impl Propagator for TracePropagationStyle {
2222
}
2323
}
2424

25-
fn inject(&self, context: &mut SpanContext, carrier: &mut dyn Injector, config: &Config) {
25+
fn inject(&self, context: &mut InjectSpanContext, carrier: &mut dyn Injector, config: &Config) {
2626
match self {
2727
Self::Datadog => datadog::inject(context, carrier, config),
2828
Self::TraceContext => tracecontext::inject(context, carrier),

0 commit comments

Comments
 (0)