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

zlayer_observability/
propagation.rs
