Skip to main content

opentelemetry_xray/telemetry/
headers.rs

1/* Copyright © 2025, CosmicMind, Inc. */
2
3use http::{HeaderMap, HeaderName, HeaderValue};
4use opentelemetry::{
5    context::Context,
6    global,
7    propagation::{Extractor, Injector},
8};
9use tracing_opentelemetry::OpenTelemetrySpanExt;
10
11use super::get_x_amzn_trace_id;
12
13/// HeaderInjector is a helper for propagating tracing headers.
14pub struct HeaderInjector<'a>(pub &'a mut HeaderMap);
15
16impl Injector for HeaderInjector<'_> {
17    fn set(&mut self, key: &str, value: String) {
18        if key.eq_ignore_ascii_case("traceparent") {
19            // Instead of inserting "traceparent", use the AWS X‑Ray header key.
20            // Using HeaderName::from_static is simpler and ensures a canonical key.
21            let header_name = HeaderName::from_static(get_x_amzn_trace_id());
22            if let Ok(header_value) = HeaderValue::from_str(&value) {
23                self.0.insert(header_name, header_value);
24            }
25        } else 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
33/// HeaderExtractor is a helper for extracting headers for tracing propagation.
34pub struct HeaderExtractor<'a>(pub &'a HeaderMap);
35
36impl Extractor for HeaderExtractor<'_> {
37    fn get(&self, key: &str) -> Option<&str> {
38        self.0.get(key).and_then(|v| v.to_str().ok())
39    }
40
41    fn keys(&self) -> Vec<&str> {
42        self.0.keys().map(|k| k.as_str()).collect()
43    }
44}
45
46/// inject_headers injects the current OpenTelemetry tracing context into the given headers.
47pub fn inject_headers(headers: &mut HeaderMap) {
48    global::get_text_map_propagator(|propagator| {
49        propagator.inject_context(
50            &tracing::Span::current().context(),
51            &mut HeaderInjector(headers),
52        );
53    });
54}
55
56/// Extracts the current OpenTelemetry tracing context from the given headers.
57/// Returns the extracted Context.
58pub fn extract_headers(headers: &HeaderMap) -> Context {
59    global::get_text_map_propagator(|propagator| propagator.extract(&HeaderExtractor(headers)))
60}