spin_telemetry/
propagation.rs

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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
use opentelemetry::{
    global,
    propagation::{Extractor, Injector},
};
use tracing_opentelemetry::OpenTelemetrySpanExt;

/// Injects the current W3C TraceContext into the provided request.
pub fn inject_trace_context<'a>(req: impl Into<HeaderInjector<'a>>) {
    let mut injector = req.into();
    global::get_text_map_propagator(|propagator| {
        let context = tracing::Span::current().context();
        propagator.inject_context(&context, &mut injector);
    });
}

/// Extracts the W3C TraceContext from the provided request and sets it as the parent of the
/// current span.
pub fn extract_trace_context<'a>(req: impl Into<HeaderExtractor<'a>>) {
    let extractor = req.into();
    let parent_context =
        global::get_text_map_propagator(|propagator| propagator.extract(&extractor));
    tracing::Span::current().set_parent(parent_context);
}

pub enum HeaderInjector<'a> {
    Http0(&'a mut http0::HeaderMap),
    Http1(&'a mut http1::HeaderMap),
}

impl<'a> Injector for HeaderInjector<'a> {
    fn set(&mut self, key: &str, value: String) {
        match self {
            HeaderInjector::Http0(headers) => {
                if let Ok(name) = http0::header::HeaderName::from_bytes(key.as_bytes()) {
                    if let Ok(val) = http0::header::HeaderValue::from_str(&value) {
                        headers.insert(name, val);
                    }
                }
            }
            HeaderInjector::Http1(headers) => {
                if let Ok(name) = http1::header::HeaderName::from_bytes(key.as_bytes()) {
                    if let Ok(val) = http1::header::HeaderValue::from_str(&value) {
                        headers.insert(name, val);
                    }
                }
            }
        }
    }
}

impl<'a, T> From<&'a mut http0::Request<T>> for HeaderInjector<'a> {
    fn from(req: &'a mut http0::Request<T>) -> Self {
        Self::Http0(req.headers_mut())
    }
}

impl<'a, T> From<&'a mut http1::Request<T>> for HeaderInjector<'a> {
    fn from(req: &'a mut http1::Request<T>) -> Self {
        Self::Http1(req.headers_mut())
    }
}

impl<'a> From<&'a mut http0::HeaderMap> for HeaderInjector<'a> {
    fn from(headers: &'a mut http0::HeaderMap) -> Self {
        Self::Http0(headers)
    }
}

impl<'a> From<&'a mut http1::HeaderMap> for HeaderInjector<'a> {
    fn from(headers: &'a mut http1::HeaderMap) -> Self {
        Self::Http1(headers)
    }
}

pub enum HeaderExtractor<'a> {
    Http0(&'a http0::HeaderMap),
    Http1(&'a http1::HeaderMap),
}

impl<'a> Extractor for HeaderExtractor<'a> {
    fn get(&self, key: &str) -> Option<&str> {
        match self {
            HeaderExtractor::Http0(headers) => {
                headers.get(key).map(|v| v.to_str().unwrap_or_default())
            }
            HeaderExtractor::Http1(headers) => {
                headers.get(key).map(|v| v.to_str().unwrap_or_default())
            }
        }
    }

    fn keys(&self) -> Vec<&str> {
        match self {
            HeaderExtractor::Http0(headers) => headers.keys().map(|k| k.as_str()).collect(),
            HeaderExtractor::Http1(headers) => headers.keys().map(|k| k.as_str()).collect(),
        }
    }
}

impl<'a, T> From<&'a http0::Request<T>> for HeaderExtractor<'a> {
    fn from(req: &'a http0::Request<T>) -> Self {
        Self::Http0(req.headers())
    }
}

impl<'a, T> From<&'a http1::Request<T>> for HeaderExtractor<'a> {
    fn from(req: &'a http1::Request<T>) -> Self {
        Self::Http1(req.headers())
    }
}