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::SdkTracerProvider::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            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    /// Encodes the values of the `Context` and injects them into the `Injector`.
92    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    /// Retrieves encoded `Context` information using the `Extractor`. If no data was
99    /// retrieved OR if the retrieved data is invalid, then the current `Context` is
100    /// returned.
101    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(&current_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    /// A test propagator that injects and extracts a single header.
127    #[derive(Debug)]
128    struct TestPropagator {
129        header: &'static str,
130        fields: Vec<String>, // used by fields method
131    }
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        // setup for baggage propagator
183        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        // setup
193        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            // name, header_name, expected_result
258            // ("single propagator", vec!["span-id"], vec!["span-id"]),
259            (
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}