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