stackable_telemetry/instrumentation/axum/
extractor.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
use axum::http::{HeaderMap, HeaderName};
use opentelemetry::{Context, propagation::Extractor};

/// Extracts the [`TextMapPropagator`][1] to access trace parent information in
/// HTTP headers.
///
/// This propagation is useful when an HTTP request already has a trace parent
/// which can be picked up by the Tower [`Layer`][2] to link both spans together.
/// A concrete usage example is available in [`SpanExt::from_request`][3].
///
/// This is pretty much a copy-pasted version of the [`HeaderExtractor`][4] from
/// the `opentelemetry_http` crate. However, we cannot use this crate, as it
/// uses an outdated version of the underlying `http` crate.
///
/// [1]: opentelemetry::propagation::TextMapPropagator
/// [2]: tower::Layer
/// [3]: crate::instrumentation::axum::SpanExt::from_request
/// [4]: https://docs.rs/opentelemetry-http/latest/opentelemetry_http/struct.HeaderExtractor.html
pub struct HeaderExtractor<'a>(pub(crate) &'a HeaderMap);

impl Extractor for HeaderExtractor<'_> {
    fn get(&self, key: &str) -> Option<&str> {
        self.0
            .get(key)
            .and_then(|header_value| header_value.to_str().ok())
    }

    fn keys(&self) -> Vec<&str> {
        self.0.keys().map(HeaderName::as_str).collect()
    }
}

impl<'a> HeaderExtractor<'a> {
    /// Create a new header extractor from a reference to a [`HeaderMap`].
    pub fn new(headers: &'a HeaderMap) -> Self {
        Self(headers)
    }

    /// Extracts the [`TextMapPropagator`][1] from the HTTP headers.
    ///
    /// [1]: opentelemetry::propagation::TextMapPropagator
    pub fn extract_context(&self) -> Context {
        opentelemetry::global::get_text_map_propagator(|propagator| propagator.extract(self))
    }
}