1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

use opentelemetry::{
    propagation::{
        text_map_propagator::FieldIter, Extractor, Injector, TextMapPropagator,
    },
    Context,
};
use opentelemetry_sdk::propagation::{
    TextMapCompositePropagator, TraceContextPropagator,
};
use opentelemetry_zipkin::{B3Encoding, Propagator as B3Propagator};
use std::collections::BTreeSet;

#[derive(Debug)]
pub struct TextMapSplitPropagator {
    extract_propagator: Box<dyn TextMapPropagator + Send + Sync>,
    inject_propagator: Box<dyn TextMapPropagator + Send + Sync>,
    fields: Vec<String>,
}

impl TextMapSplitPropagator {
    pub fn new(
        extract_propagator: Box<dyn TextMapPropagator + Send + Sync>,
        inject_propagator: Box<dyn TextMapPropagator + Send + Sync>,
    ) -> Self {
        let mut fields = BTreeSet::from_iter(extract_propagator.fields());
        fields.extend(inject_propagator.fields());
        let fields = fields.into_iter().map(String::from).collect();

        Self {
            extract_propagator,
            inject_propagator,
            fields,
        }
    }
}

impl TextMapPropagator for TextMapSplitPropagator {
    fn inject_context(&self, cx: &Context, injector: &mut dyn Injector) {
        self.inject_propagator.inject_context(cx, injector)
    }

    fn extract_with_context(&self, cx: &Context, extractor: &dyn Extractor) -> Context {
        self.extract_propagator.extract_with_context(cx, extractor)
    }

    fn fields(&self) -> FieldIter<'_> {
        FieldIter::new(self.fields.as_slice())
    }
}

impl Default for TextMapSplitPropagator {
    fn default() -> Self {
        let trace_context_propagator = TraceContextPropagator::new();
        let b3_propagator = B3Propagator::with_encoding(B3Encoding::SingleAndMultiHeader);
        let composite_propagator = TextMapCompositePropagator::new(vec![
            Box::new(trace_context_propagator.clone()),
            Box::new(b3_propagator),
        ]);

        Self::new(
            Box::new(composite_propagator),
            Box::new(trace_context_propagator),
        )
    }
}