opentelemetry/propagation/
composite.rs1use crate::{
9 propagation::{text_map_propagator::FieldIter, Extractor, Injector, TextMapPropagator},
10 Context,
11};
12use std::collections::HashSet;
13
14#[derive(Debug)]
68pub struct TextMapCompositePropagator {
69 propagators: Vec<Box<dyn TextMapPropagator + Send + Sync>>,
70 fields: Vec<String>,
71}
72
73impl TextMapCompositePropagator {
74 pub fn new(propagators: Vec<Box<dyn TextMapPropagator + Send + Sync>>) -> Self {
78 let mut fields = HashSet::new();
79 for propagator in &propagators {
80 fields.extend(propagator.fields().map(ToString::to_string));
81 }
82
83 TextMapCompositePropagator {
84 propagators,
85 fields: fields.into_iter().collect(),
86 }
87 }
88}
89
90impl TextMapPropagator for TextMapCompositePropagator {
91 fn inject_context(&self, context: &Context, injector: &mut dyn Injector) {
93 for propagator in &self.propagators {
94 propagator.inject_context(context, injector)
95 }
96 }
97
98 fn extract_with_context(&self, cx: &Context, extractor: &dyn Extractor) -> Context {
102 self.propagators
103 .iter()
104 .fold(cx.clone(), |current_cx, propagator| {
105 propagator.extract_with_context(¤t_cx, extractor)
106 })
107 }
108
109 fn fields(&self) -> FieldIter<'_> {
110 FieldIter::new(self.fields.as_slice())
111 }
112}
113
114#[cfg(all(test, feature = "trace"))]
115mod tests {
116 use crate::baggage::BaggageExt;
117 use crate::propagation::TextMapCompositePropagator;
118 use crate::testing::trace::TestSpan;
119 use crate::{
120 propagation::{text_map_propagator::FieldIter, Extractor, Injector, TextMapPropagator},
121 trace::{SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId, TraceState},
122 Context, KeyValue,
123 };
124 use std::collections::HashMap;
125
126 #[derive(Debug)]
128 struct TestPropagator {
129 header: &'static str,
130 fields: Vec<String>, }
132
133 impl TestPropagator {
134 #[allow(unreachable_pub)]
135 pub fn new(header: &'static str) -> Self {
136 TestPropagator {
137 header,
138 fields: vec![header.to_string()],
139 }
140 }
141 }
142
143 impl TextMapPropagator for TestPropagator {
144 fn inject_context(&self, cx: &Context, injector: &mut dyn Injector) {
145 let span = cx.span();
146 let span_context = span.span_context();
147 match self.header {
148 "span-id" => injector.set(self.header, format!("{:x}", span_context.span_id())),
149 "baggage" => injector.set(self.header, cx.baggage().to_string()),
150 _ => {}
151 }
152 }
153
154 fn extract_with_context(&self, cx: &Context, extractor: &dyn Extractor) -> Context {
155 match (self.header, extractor.get(self.header)) {
156 ("span-id", Some(val)) => cx.with_remote_span_context(SpanContext::new(
157 TraceId::from_u128(1),
158 SpanId::from_u64(u64::from_str_radix(val, 16).unwrap()),
159 TraceFlags::default(),
160 false,
161 TraceState::default(),
162 )),
163 ("baggage", Some(_)) => cx.with_baggage(vec![KeyValue::new("baggagekey", "value")]),
164 _ => cx.clone(),
165 }
166 }
167
168 fn fields(&self) -> FieldIter<'_> {
169 FieldIter::new(self.fields.as_slice())
170 }
171 }
172
173 fn setup() -> Context {
174 let mut cx = Context::default();
175 cx = cx.with_span(TestSpan(SpanContext::new(
176 TraceId::from_u128(1),
177 SpanId::from_u64(11),
178 TraceFlags::default(),
179 true,
180 TraceState::default(),
181 )));
182 cx.with_baggage(vec![KeyValue::new("baggagekey", "value")])
184 }
185
186 fn test_data() -> Vec<(&'static str, &'static str)> {
187 vec![("span-id", "b"), ("baggage", "baggagekey=value")]
188 }
189
190 #[test]
191 fn zero_propogators_are_noop() {
192 let composite_propagator = TextMapCompositePropagator::new(vec![]);
194 let cx = setup();
195
196 let mut injector = HashMap::new();
197 composite_propagator.inject_context(&cx, &mut injector);
198
199 assert_eq!(injector.len(), 0);
200 for (header_name, header_value) in test_data() {
201 let mut extractor = HashMap::new();
202 extractor.insert(header_name.to_string(), header_value.to_string());
203 assert_eq!(
204 composite_propagator
205 .extract(&extractor)
206 .span()
207 .span_context(),
208 &SpanContext::empty_context()
209 );
210 }
211 }
212
213 #[test]
214 fn inject_multiple_propagators() {
215 let composite_propagator = TextMapCompositePropagator::new(vec![
216 Box::new(TestPropagator::new("span-id")),
217 Box::new(TestPropagator::new("baggage")),
218 ]);
219
220 let cx = setup();
221 let mut injector = HashMap::new();
222 composite_propagator.inject_context(&cx, &mut injector);
223
224 for (header_name, header_value) in test_data() {
225 assert_eq!(injector.get(header_name), Some(&header_value.to_string()));
226 }
227 }
228
229 #[test]
230 fn extract_multiple_propagators() {
231 let composite_propagator = TextMapCompositePropagator::new(vec![
232 Box::new(TestPropagator::new("span-id")),
233 Box::new(TestPropagator::new("baggage")),
234 ]);
235
236 let mut extractor = HashMap::new();
237 for (header_name, header_value) in test_data() {
238 extractor.insert(header_name.to_string(), header_value.to_string());
239 }
240 let cx = composite_propagator.extract(&extractor);
241 assert_eq!(
242 cx.span().span_context(),
243 &SpanContext::new(
244 TraceId::from_u128(1),
245 SpanId::from_u64(11),
246 TraceFlags::default(),
247 false,
248 TraceState::default(),
249 )
250 );
251 assert_eq!(cx.baggage().to_string(), "baggagekey=value",);
252 }
253
254 #[test]
255 fn test_get_fields() {
256 let test_cases = vec![
257 (
260 "multiple propagators with order",
261 vec!["span-id", "baggage"],
262 vec!["baggage", "span-id"],
263 ),
264 ];
265
266 for test_case in test_cases {
267 let test_propagators = test_case
268 .1
269 .into_iter()
270 .map(|name| {
271 Box::new(TestPropagator::new(name)) as Box<dyn TextMapPropagator + Send + Sync>
272 })
273 .collect();
274
275 let composite_propagator = TextMapCompositePropagator::new(test_propagators);
276
277 let mut fields = composite_propagator
278 .fields()
279 .map(|s| s.to_string())
280 .collect::<Vec<String>>();
281 fields.sort();
282
283 assert_eq!(fields, test_case.2);
284 }
285 }
286}