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}