stackable_telemetry/instrumentation/axum/
injector.rs

1use axum::http::{HeaderMap, HeaderName, HeaderValue};
2use opentelemetry::{Context, propagation::Injector};
3
4/// Injects the [`TextMapPropagator`][1] to propagate trace parent information
5/// in HTTP headers.
6///
7/// This propagation is useful when consumers of the HTTP response want to link
8/// up their span data with the data produced in the Tower [`Layer`][2].
9/// A concrete usage example is available in the [`TraceService::call`][3]
10/// implementation for [`TraceService`][4].
11///
12/// This is pretty much a copy-pasted version of the [`HeaderInjector`][5] from
13/// the `opentelemetry_http` crate. However, we cannot use this crate, as it
14/// uses an outdated version of the underlying `http` crate.
15///
16/// [1]: opentelemetry::propagation::TextMapPropagator
17/// [2]: tower::Layer
18/// [3]: tower::Service::call
19/// [4]: crate::instrumentation::axum::TraceService
20/// [5]: https://docs.rs/opentelemetry-http/latest/opentelemetry_http/struct.HeaderInjector.html
21pub struct HeaderInjector<'a>(pub(crate) &'a mut HeaderMap);
22
23impl Injector for HeaderInjector<'_> {
24    fn set(&mut self, key: &str, value: String) {
25        if let Ok(header_name) = HeaderName::from_bytes(key.as_bytes()) {
26            if let Ok(header_value) = HeaderValue::from_str(&value) {
27                self.0.insert(header_name, header_value);
28            }
29        }
30    }
31}
32
33impl<'a> HeaderInjector<'a> {
34    /// Create a new header injecttor from a mutable reference to [`HeaderMap`].
35    pub fn new(headers: &'a mut HeaderMap) -> Self {
36        Self(headers)
37    }
38
39    /// Inject the [`TextMapPropagator`][1] into the HTTP headers.
40    ///
41    /// [1]: opentelemetry::propagation::TextMapPropagator
42    pub fn inject_context(&mut self, cx: &Context) {
43        opentelemetry::global::get_text_map_propagator(|propagator| {
44            propagator.inject_context(cx, self)
45        })
46    }
47}