opentelemetry/propagation/
composite.rs

1//! # Composite Propagator
2//!
3//! A utility over multiple `Propagator`s to group multiple Propagators from different cross-cutting
4//! concerns in order to leverage them as a single entity.
5//!
6//! Each composite Propagator will implement a specific Propagator type, such as TextMapPropagator,
7//! as different Propagator types will likely operate on different data types.
8use crate::{
9    propagation::{text_map_propagator::FieldIter, Extractor, Injector, TextMapPropagator},
10    Context,
11};
12use std::collections::HashSet;
13
14/// Composite propagator for [`TextMapPropagator`]s.
15///
16/// A propagator that chains multiple [`TextMapPropagator`] propagators together,
17/// injecting or extracting by their respective HTTP header names.
18///
19/// Injection and extraction from this propagator will preserve the order of the
20/// injectors and extractors passed in during initialization.
21///
22/// # Examples
23///
24/// ```
25/// use opentelemetry::{
26///     baggage::BaggageExt,
27///     propagation::{TextMapPropagator, TextMapCompositePropagator},
28///
29///     trace::{TraceContextExt, Tracer, TracerProvider},
30///     Context, KeyValue,
31/// };
32/// use opentelemetry_sdk::propagation::{
33///     BaggagePropagator, TraceContextPropagator,
34/// };
35/// use opentelemetry_sdk::trace as sdktrace;
36/// use std::collections::HashMap;
37///
38/// // First create 1 or more propagators
39/// let baggage_propagator = BaggagePropagator::new();
40/// let trace_context_propagator = TraceContextPropagator::new();
41///
42/// // Then create a composite propagator
43/// let composite_propagator = TextMapCompositePropagator::new(vec![
44///     Box::new(baggage_propagator),
45///     Box::new(trace_context_propagator),
46/// ]);
47///
48/// // Then for a given implementation of `Injector`
49/// let mut injector = HashMap::new();
50///
51/// // And a given span
52/// let example_span = sdktrace::TracerProvider::default()
53///     .tracer("example-component")
54///     .start("span-name");
55///
56/// // with the current context, call inject to add the headers
57/// composite_propagator.inject_context(
58///     &Context::current_with_span(example_span)
59///         .with_baggage(vec![KeyValue::new("test", "example")]),
60///     &mut injector,
61/// );
62///
63/// // The injector now has both `baggage` and `traceparent` headers
64/// assert!(injector.get("baggage").is_some());
65/// assert!(injector.get("traceparent").is_some());
66/// ```
67#[derive(Debug)]
68pub struct TextMapCompositePropagator {
69    propagators: Vec<Box<dyn TextMapPropagator + Send + Sync>>,
70    fields: Vec<String>,
71}
72
73impl TextMapCompositePropagator {
74    /// Constructs a new propagator out of instances of [`TextMapPropagator`].
75    ///
76    /// [`TextMapPropagator`]: TextMapPropagator
77    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    /// Encodes the values of the `Context` and injects them into the `Injector`.
94    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    /// Retrieves encoded `Context` information using the `Extractor`. If no data was
101    /// retrieved OR if the retrieved data is invalid, then the current `Context` is
102    /// returned.
103    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(&current_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    /// A test propagator that injects and extracts a single header.
129    #[derive(Debug)]
130    struct TestPropagator {
131        header: &'static str,
132        fields: Vec<String>, // used by fields method
133    }
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        // setup for baggage propagator
185        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        // setup
195        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            // name, header_name, expected_result
260            // ("single propagator", vec!["span-id"], vec!["span-id"]),
261            (
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}