opentelemetry_spanprocessor_any/sdk/propagation/
composite.rs

1use crate::{
2    propagation::{text_map_propagator::FieldIter, Extractor, Injector, TextMapPropagator},
3    Context,
4};
5use std::collections::HashSet;
6
7/// Composite propagator
8///
9/// A propagator that chains multiple [`TextMapPropagator`] propagators together,
10/// injecting or extracting by their respective HTTP header names.
11///
12/// Injection and extraction from this propagator will preserve the order of the
13/// injectors and extractors passed in during initialization.
14///
15/// [`TextMapPropagator`]: crate::propagation::TextMapPropagator
16///
17/// # Examples
18///
19/// ```
20/// use opentelemetry::{
21///     baggage::BaggageExt,
22///     propagation::TextMapPropagator,
23///     trace::{TraceContextExt, Tracer, TracerProvider},
24///     Context, KeyValue,
25/// };
26/// use opentelemetry::sdk::propagation::{
27///     BaggagePropagator, TextMapCompositePropagator, TraceContextPropagator,
28/// };
29/// use opentelemetry::sdk::trace as sdktrace;
30/// use std::collections::HashMap;
31///
32/// // First create 1 or more propagators
33/// let baggage_propagator = BaggagePropagator::new();
34/// let trace_context_propagator = TraceContextPropagator::new();
35///
36/// // Then create a composite propagator
37/// let composite_propagator = TextMapCompositePropagator::new(vec![
38///     Box::new(baggage_propagator),
39///     Box::new(trace_context_propagator),
40/// ]);
41///
42/// // Then for a given implementation of `Injector`
43/// let mut injector = HashMap::new();
44///
45/// // And a given span
46/// let example_span = sdktrace::TracerProvider::default()
47///     .tracer("example-component")
48///     .start("span-name");
49///
50/// // with the current context, call inject to add the headers
51/// composite_propagator.inject_context(
52///     &Context::current_with_span(example_span)
53///         .with_baggage(vec![KeyValue::new("test", "example")]),
54///     &mut injector,
55/// );
56///
57/// // The injector now has both `baggage` and `traceparent` headers
58/// assert!(injector.get("baggage").is_some());
59/// assert!(injector.get("traceparent").is_some());
60/// ```
61#[derive(Debug)]
62pub struct TextMapCompositePropagator {
63    propagators: Vec<Box<dyn TextMapPropagator + Send + Sync>>,
64    fields: Vec<String>,
65}
66
67impl TextMapCompositePropagator {
68    /// Constructs a new propagator out of instances of [`TextMapPropagator`].
69    ///
70    /// [`TextMapPropagator`]: crate::propagation::TextMapPropagator
71    pub fn new(propagators: Vec<Box<dyn TextMapPropagator + Send + Sync>>) -> Self {
72        let mut fields = HashSet::new();
73        for propagator in &propagators {
74            for field in propagator.fields() {
75                fields.insert(field.to_string());
76            }
77        }
78
79        TextMapCompositePropagator {
80            propagators,
81            fields: fields.into_iter().collect(),
82        }
83    }
84}
85
86impl TextMapPropagator for TextMapCompositePropagator {
87    /// Encodes the values of the `Context` and injects them into the `Injector`.
88    fn inject_context(&self, context: &Context, injector: &mut dyn Injector) {
89        for propagator in &self.propagators {
90            propagator.inject_context(context, injector)
91        }
92    }
93
94    /// Retrieves encoded `Context` information using the `Extractor`. If no data was
95    /// retrieved OR if the retrieved data is invalid, then the current `Context` is
96    /// returned.
97    fn extract_with_context(&self, cx: &Context, extractor: &dyn Extractor) -> Context {
98        self.propagators
99            .iter()
100            .fold(cx.clone(), |current_cx, propagator| {
101                propagator.extract_with_context(&current_cx, extractor)
102            })
103    }
104
105    fn fields(&self) -> FieldIter<'_> {
106        FieldIter::new(self.fields.as_slice())
107    }
108}
109
110#[cfg(all(test, feature = "testing", feature = "trace"))]
111mod tests {
112    use crate::sdk::propagation::{TextMapCompositePropagator, TraceContextPropagator};
113    use crate::testing::trace::TestSpan;
114    use crate::{
115        propagation::{text_map_propagator::FieldIter, Extractor, Injector, TextMapPropagator},
116        trace::{SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId, TraceState},
117        Context,
118    };
119    use std::collections::HashMap;
120    use std::str::FromStr;
121
122    /// Dummy propagator for testing
123    ///
124    /// The format we are using is {trace id(in base10 u128)}-{span id(in base10 u64)}-{flag(in u8)}
125    #[derive(Debug)]
126    struct TestPropagator {
127        fields: [String; 1],
128    }
129
130    impl TestPropagator {
131        #[allow(unreachable_pub)]
132        pub fn new() -> Self {
133            TestPropagator {
134                fields: ["testheader".to_string()],
135            }
136        }
137    }
138
139    impl TextMapPropagator for TestPropagator {
140        fn inject_context(&self, cx: &Context, injector: &mut dyn Injector) {
141            let span = cx.span();
142            let span_context = span.span_context();
143            injector.set(
144                "testheader",
145                format!(
146                    "{:x}-{:x}-{:02x}",
147                    span_context.trace_id(),
148                    span_context.span_id(),
149                    span_context.trace_flags()
150                ),
151            )
152        }
153
154        fn extract_with_context(&self, cx: &Context, extractor: &dyn Extractor) -> Context {
155            let span = if let Some(val) = extractor.get("testheader") {
156                let parts = val.split_terminator('-').collect::<Vec<&str>>();
157                if parts.len() != 3 {
158                    SpanContext::empty_context()
159                } else {
160                    SpanContext::new(
161                        TraceId::from_u128(u128::from_str(parts[0]).unwrap_or(0)),
162                        SpanId::from_u64(u64::from_str(parts[1]).unwrap_or(0)),
163                        TraceFlags::new(u8::from_str(parts[2]).unwrap_or(0)),
164                        true,
165                        TraceState::default(),
166                    )
167                }
168            } else {
169                SpanContext::empty_context()
170            };
171
172            cx.with_remote_span_context(span)
173        }
174
175        fn fields(&self) -> FieldIter<'_> {
176            FieldIter::new(&self.fields)
177        }
178    }
179
180    fn test_data() -> Vec<(&'static str, &'static str)> {
181        vec![
182            ("testheader", "1-1-00"),
183            (
184                "traceparent",
185                "00-00000000000000000000000000000001-0000000000000001-00",
186            ),
187        ]
188    }
189
190    #[test]
191    fn zero_propogators_are_noop() {
192        let composite_propagator = TextMapCompositePropagator::new(vec![]);
193
194        let cx = Context::default().with_span(TestSpan(SpanContext::new(
195            TraceId::from_u128(1),
196            SpanId::from_u64(1),
197            TraceFlags::default(),
198            false,
199            TraceState::default(),
200        )));
201        let mut injector = HashMap::new();
202        composite_propagator.inject_context(&cx, &mut injector);
203
204        assert_eq!(injector.len(), 0);
205
206        for (header_name, header_value) in test_data() {
207            let mut extractor = HashMap::new();
208            extractor.insert(header_name.to_string(), header_value.to_string());
209            assert_eq!(
210                composite_propagator
211                    .extract(&extractor)
212                    .span()
213                    .span_context(),
214                &SpanContext::empty_context()
215            );
216        }
217    }
218
219    #[test]
220    fn inject_multiple_propagators() {
221        let test_propagator = TestPropagator::new();
222        let trace_context = TraceContextPropagator::new();
223        let composite_propagator = TextMapCompositePropagator::new(vec![
224            Box::new(test_propagator),
225            Box::new(trace_context),
226        ]);
227
228        let cx = Context::default().with_span(TestSpan(SpanContext::new(
229            TraceId::from_u128(1),
230            SpanId::from_u64(1),
231            TraceFlags::default(),
232            false,
233            TraceState::default(),
234        )));
235        let mut injector = HashMap::new();
236        composite_propagator.inject_context(&cx, &mut injector);
237
238        for (header_name, header_value) in test_data() {
239            assert_eq!(injector.get(header_name), Some(&header_value.to_string()));
240        }
241    }
242
243    #[test]
244    fn extract_multiple_propagators() {
245        let test_propagator = TestPropagator::new();
246        let trace_context = TraceContextPropagator::new();
247        let composite_propagator = TextMapCompositePropagator::new(vec![
248            Box::new(test_propagator),
249            Box::new(trace_context),
250        ]);
251
252        for (header_name, header_value) in test_data() {
253            let mut extractor = HashMap::new();
254            extractor.insert(header_name.to_string(), header_value.to_string());
255            assert_eq!(
256                composite_propagator
257                    .extract(&extractor)
258                    .span()
259                    .span_context(),
260                &SpanContext::new(
261                    TraceId::from_u128(1),
262                    SpanId::from_u64(1),
263                    TraceFlags::default(),
264                    true,
265                    TraceState::default(),
266                )
267            );
268        }
269    }
270
271    #[test]
272    fn test_get_fields() {
273        let test_propagator = TestPropagator::new();
274        let b3_fields = test_propagator
275            .fields()
276            .map(|s| s.to_string())
277            .collect::<Vec<String>>();
278
279        let trace_context = TraceContextPropagator::new();
280        let trace_context_fields = trace_context
281            .fields()
282            .map(|s| s.to_string())
283            .collect::<Vec<String>>();
284
285        let composite_propagator = TextMapCompositePropagator::new(vec![
286            Box::new(test_propagator),
287            Box::new(trace_context),
288        ]);
289
290        let mut fields = composite_propagator
291            .fields()
292            .map(|s| s.to_string())
293            .collect::<Vec<String>>();
294        fields.sort();
295
296        let mut expected = vec![b3_fields, trace_context_fields]
297            .into_iter()
298            .flatten()
299            .collect::<Vec<String>>();
300        expected.sort();
301        expected.dedup();
302
303        assert_eq!(fields, expected);
304    }
305}
opentelemetry_spanprocessor_any/sdk/propagation/composite.rs

opentelemetry_spanprocessor_any/sdk/propagation/
composite.rs
