spin_telemetry/
propagation.rs

1use opentelemetry::{
2    global,
3    propagation::{Extractor, Injector},
4};
5use tracing_opentelemetry::OpenTelemetrySpanExt;
6
7/// Injects the current W3C TraceContext into the provided request.
8pub fn inject_trace_context<'a>(req: impl Into<HeaderInjector<'a>>) {
9    let mut injector = req.into();
10    global::get_text_map_propagator(|propagator| {
11        let context = tracing::Span::current().context();
12        propagator.inject_context(&context, &mut injector);
13    });
14}
15
16/// Extracts the W3C TraceContext from the provided request and sets it as the parent of the
17/// current span.
18pub fn extract_trace_context<'a>(req: impl Into<HeaderExtractor<'a>>) {
19    let extractor = req.into();
20    let parent_context =
21        global::get_text_map_propagator(|propagator| propagator.extract(&extractor));
22    tracing::Span::current().set_parent(parent_context);
23}
24
25pub enum HeaderInjector<'a> {
26    Http0(&'a mut http0::HeaderMap),
27    Http1(&'a mut http1::HeaderMap),
28}
29
30impl Injector for HeaderInjector<'_> {
31    fn set(&mut self, key: &str, value: String) {
32        match self {
33            HeaderInjector::Http0(headers) => {
34                if let Ok(name) = http0::header::HeaderName::from_bytes(key.as_bytes()) {
35                    if let Ok(val) = http0::header::HeaderValue::from_str(&value) {
36                        headers.insert(name, val);
37                    }
38                }
39            }
40            HeaderInjector::Http1(headers) => {
41                if let Ok(name) = http1::header::HeaderName::from_bytes(key.as_bytes()) {
42                    if let Ok(val) = http1::header::HeaderValue::from_str(&value) {
43                        headers.insert(name, val);
44                    }
45                }
46            }
47        }
48    }
49}
50
51impl<'a, T> From<&'a mut http0::Request<T>> for HeaderInjector<'a> {
52    fn from(req: &'a mut http0::Request<T>) -> Self {
53        Self::Http0(req.headers_mut())
54    }
55}
56
57impl<'a, T> From<&'a mut http1::Request<T>> for HeaderInjector<'a> {
58    fn from(req: &'a mut http1::Request<T>) -> Self {
59        Self::Http1(req.headers_mut())
60    }
61}
62
63impl<'a> From<&'a mut http0::HeaderMap> for HeaderInjector<'a> {
64    fn from(headers: &'a mut http0::HeaderMap) -> Self {
65        Self::Http0(headers)
66    }
67}
68
69impl<'a> From<&'a mut http1::HeaderMap> for HeaderInjector<'a> {
70    fn from(headers: &'a mut http1::HeaderMap) -> Self {
71        Self::Http1(headers)
72    }
73}
74
75pub enum HeaderExtractor<'a> {
76    Http0(&'a http0::HeaderMap),
77    Http1(&'a http1::HeaderMap),
78}
79
80impl Extractor for HeaderExtractor<'_> {
81    fn get(&self, key: &str) -> Option<&str> {
82        match self {
83            HeaderExtractor::Http0(headers) => {
84                headers.get(key).map(|v| v.to_str().unwrap_or_default())
85            }
86            HeaderExtractor::Http1(headers) => {
87                headers.get(key).map(|v| v.to_str().unwrap_or_default())
88            }
89        }
90    }
91
92    fn keys(&self) -> Vec<&str> {
93        match self {
94            HeaderExtractor::Http0(headers) => headers.keys().map(|k| k.as_str()).collect(),
95            HeaderExtractor::Http1(headers) => headers.keys().map(|k| k.as_str()).collect(),
96        }
97    }
98}
99
100impl<'a, T> From<&'a http0::Request<T>> for HeaderExtractor<'a> {
101    fn from(req: &'a http0::Request<T>) -> Self {
102        Self::Http0(req.headers())
103    }
104}
105
106impl<'a, T> From<&'a http1::Request<T>> for HeaderExtractor<'a> {
107    fn from(req: &'a http1::Request<T>) -> Self {
108        Self::Http1(req.headers())
109    }
110}
spin_telemetry/propagation.rs

spin_telemetry/
propagation.rs
