zlayer_observability/
propagation.rs

1//! Context propagation utilities for distributed tracing
2//!
3//! Provides helpers for extracting and injecting trace context
4//! in HTTP headers using W3C Trace Context standard.
5
6use opentelemetry::propagation::{Extractor, Injector};
7use std::collections::HashMap;
8
9/// HTTP header extractor for incoming requests (works with http crate `HeaderMap`)
10pub struct HeaderExtractor<'a, T>(pub &'a T);
11
12impl Extractor for HeaderExtractor<'_, http::HeaderMap> {
13    fn get(&self, key: &str) -> Option<&str> {
14        self.0.get(key).and_then(|v| v.to_str().ok())
15    }
16
17    fn keys(&self) -> Vec<&str> {
18        self.0.keys().map(http::HeaderName::as_str).collect()
19    }
20}
21
22#[allow(clippy::implicit_hasher)]
23impl Extractor for HeaderExtractor<'_, HashMap<String, String>> {
24    fn get(&self, key: &str) -> Option<&str> {
25        self.0.get(key).map(std::string::String::as_str)
26    }
27
28    fn keys(&self) -> Vec<&str> {
29        self.0.keys().map(std::string::String::as_str).collect()
30    }
31}
32
33/// HTTP header injector for outgoing requests
34pub struct HeaderInjector<'a, T>(pub &'a mut T);
35
36impl Injector for HeaderInjector<'_, http::HeaderMap> {
37    fn set(&mut self, key: &str, value: String) {
38        if let Ok(name) = http::header::HeaderName::from_bytes(key.as_bytes()) {
39            if let Ok(val) = http::header::HeaderValue::from_str(&value) {
40                self.0.insert(name, val);
41            }
42        }
43    }
44}
45
46#[allow(clippy::implicit_hasher)]
47impl Injector for HeaderInjector<'_, HashMap<String, String>> {
48    fn set(&mut self, key: &str, value: String) {
49        self.0.insert(key.to_string(), value);
50    }
51}
52
53/// Extract trace context from HTTP headers
54///
55/// Returns a Context that should be used as the parent for spans
56/// handling this request.
57#[cfg(feature = "propagation")]
58#[must_use]
59pub fn extract_context_from_headers(headers: &http::HeaderMap) -> opentelemetry::Context {
60    use opentelemetry::global;
61    global::get_text_map_propagator(|propagator| propagator.extract(&HeaderExtractor(headers)))
62}
63
64/// Inject current trace context into HTTP headers
65///
66/// Used when making outgoing HTTP requests to propagate trace context.
67#[cfg(feature = "propagation")]
68pub fn inject_context_to_headers(headers: &mut http::HeaderMap) {
69    use opentelemetry::global;
70    global::get_text_map_propagator(|propagator| {
71        propagator.inject(&mut HeaderInjector(headers));
72    });
73}
74
75/// Extract trace context from a `HashMap` (useful for gRPC metadata)
76#[cfg(feature = "propagation")]
77#[must_use]
78#[allow(clippy::implicit_hasher)]
79pub fn extract_context_from_map(map: &HashMap<String, String>) -> opentelemetry::Context {
80    use opentelemetry::global;
81    global::get_text_map_propagator(|propagator| propagator.extract(&HeaderExtractor(map)))
82}
83
84/// Inject trace context into a `HashMap`
85#[cfg(feature = "propagation")]
86#[allow(clippy::implicit_hasher)]
87pub fn inject_context_to_map(map: &mut HashMap<String, String>) {
88    use opentelemetry::global;
89    global::get_text_map_propagator(|propagator| {
90        propagator.inject(&mut HeaderInjector(map));
91    });
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97
98    #[test]
99    fn test_hashmap_extractor() {
100        let mut map = HashMap::new();
101        map.insert("traceparent".to_string(), "00-abc123-def456-01".to_string());
102
103        let extractor = HeaderExtractor(&map);
104        assert_eq!(extractor.get("traceparent"), Some("00-abc123-def456-01"));
105        assert_eq!(extractor.get("nonexistent"), None);
106    }
107
108    #[test]
109    fn test_hashmap_injector() {
110        let mut map = HashMap::new();
111        {
112            let mut injector = HeaderInjector(&mut map);
113            injector.set("traceparent", "00-abc123-def456-01".to_string());
114        }
115        assert_eq!(
116            map.get("traceparent"),
117            Some(&"00-abc123-def456-01".to_string())
118        );
119    }
120
121    #[test]
122    fn test_header_map_extractor() {
123        let mut headers = http::HeaderMap::new();
124        headers.insert("traceparent", "00-abc123-def456-01".parse().unwrap());
125
126        let extractor = HeaderExtractor(&headers);
127        assert_eq!(extractor.get("traceparent"), Some("00-abc123-def456-01"));
128    }
129
130    #[test]
131    fn test_header_map_injector() {
132        let mut headers = http::HeaderMap::new();
133        {
134            let mut injector = HeaderInjector(&mut headers);
135            injector.set("traceparent", "00-abc123-def456-01".to_string());
136        }
137        assert_eq!(
138            headers.get("traceparent").and_then(|v| v.to_str().ok()),
139            Some("00-abc123-def456-01")
140        );
141    }
142}
zlayer_observability/propagation.rs

zlayer_observability/
propagation.rs
