1use http::HeaderMap;
8
9pub const TRACEPARENT: &str = "traceparent";
11
12#[must_use]
14pub fn get_traceparent(headers: &HeaderMap) -> Option<&str> {
15 headers.get(TRACEPARENT)?.to_str().ok()
16}
17
18#[must_use]
20pub fn parse_trace_id(traceparent: &str) -> Option<String> {
21 let parts: Vec<&str> = traceparent.split('-').collect();
22 if parts.len() >= 4 && parts[0] == "00" {
23 Some(parts[1].to_owned())
24 } else {
25 None
26 }
27}
28
29#[cfg(feature = "otel")]
30mod imp {
31 use super::{get_traceparent, parse_trace_id};
32 use http::{HeaderMap, HeaderName, HeaderValue};
33 use opentelemetry::{
34 Context, global,
35 propagation::{Extractor, Injector},
36 };
37 use tracing::Span;
38 use tracing_opentelemetry::OpenTelemetrySpanExt;
39
40 struct HeadersExtractor<'a>(&'a HeaderMap);
42
43 impl Extractor for HeadersExtractor<'_> {
44 fn get(&self, key: &str) -> Option<&str> {
45 self.0.get(key).and_then(|v| v.to_str().ok())
46 }
47
48 fn keys(&self) -> Vec<&str> {
49 self.0.keys().map(http::HeaderName::as_str).collect()
50 }
51 }
52
53 struct HeadersInjector<'a>(&'a mut HeaderMap);
55
56 impl Injector for HeadersInjector<'_> {
57 fn set(&mut self, key: &str, value: String) {
58 if let Ok(name) = HeaderName::from_bytes(key.as_bytes())
59 && let Ok(val) = HeaderValue::from_str(&value)
60 {
61 self.0.insert(name, val);
62 }
63 }
64 }
65
66 pub fn inject_current_span(headers: &mut HeaderMap) {
69 let cx = Context::current();
70 global::get_text_map_propagator(|propagator| {
71 propagator.inject_context(&cx, &mut HeadersInjector(headers));
72 });
73 }
74
75 pub fn set_parent_from_headers(span: &Span, headers: &HeaderMap) {
78 let parent_cx = global::get_text_map_propagator(|propagator| {
80 propagator.extract(&HeadersExtractor(headers))
81 });
82
83 let _ = span.set_parent(parent_cx);
85
86 if let Some(traceparent) = get_traceparent(headers)
88 && let Some(trace_id) = parse_trace_id(traceparent)
89 {
90 span.record("trace_id", &trace_id);
91 span.record("parent.trace_id", &trace_id);
92 }
93 }
94}
95
96#[cfg(not(feature = "otel"))]
97mod imp {
98 use super::{get_traceparent, parse_trace_id};
99 use http::HeaderMap;
100 use tracing::Span;
101
102 pub fn inject_current_span(_headers: &mut HeaderMap) {
104 }
106
107 pub fn set_parent_from_headers(span: &Span, headers: &HeaderMap) {
110 if let Some(traceparent) = get_traceparent(headers)
112 && let Some(trace_id) = parse_trace_id(traceparent)
113 {
114 span.record("trace_id", &trace_id);
115 span.record("parent.trace_id", &trace_id);
116 }
117 }
118}
119
120pub use imp::{inject_current_span, set_parent_from_headers};
121
122#[cfg(test)]
123#[cfg_attr(coverage_nightly, coverage(off))]
124mod tests {
125 use super::*;
126 use tracing::info_span;
127
128 #[test]
129 fn test_get_traceparent_none() {
130 let headers = HeaderMap::new();
131 assert!(get_traceparent(&headers).is_none());
132 }
133
134 #[test]
135 fn test_get_traceparent_ok() {
136 let mut headers = HeaderMap::new();
137 headers.insert(
138 TRACEPARENT,
139 "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
140 .parse()
141 .expect("valid header"),
142 );
143
144 let tp = get_traceparent(&headers);
145 assert!(tp.is_some());
146 assert_eq!(
147 tp.expect("should be some"),
148 "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
149 );
150 }
151
152 #[test]
153 fn test_parse_trace_id_ok() {
154 let traceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
155 let trace_id = parse_trace_id(traceparent);
156 assert_eq!(
157 trace_id,
158 Some("4bf92f3577b34da6a3ce929d0e0e4736".to_owned())
159 );
160 }
161
162 #[test]
163 fn test_parse_trace_id_invalid() {
164 assert!(parse_trace_id("invalid").is_none());
165 assert!(parse_trace_id("").is_none());
166 }
167
168 #[test]
169 #[cfg(not(feature = "otel"))]
170 fn test_inject_current_span_noop() {
171 let mut headers = HeaderMap::new();
172 inject_current_span(&mut headers);
173 assert!(headers.is_empty());
175 }
176
177 #[test]
178 #[cfg(feature = "otel")]
179 fn test_inject_current_span_no_panic() {
180 use opentelemetry::global;
181 use opentelemetry_sdk::propagation::TraceContextPropagator;
182
183 global::set_text_map_propagator(TraceContextPropagator::new());
184
185 let mut headers = http::HeaderMap::new();
186 let _span = tracing::info_span!("test").entered();
187 inject_current_span(&mut headers);
189 }
190
191 #[test]
192 fn test_set_parent_from_headers_no_panic() {
193 let mut headers = HeaderMap::new();
194 headers.insert(
195 TRACEPARENT,
196 "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
197 .parse()
198 .expect("valid header"),
199 );
200
201 let span = info_span!(
202 "test",
203 trace_id = tracing::field::Empty,
204 parent.trace_id = tracing::field::Empty
205 );
206
207 set_parent_from_headers(&span, &headers);
209 }
210}