rs_zero/observability/
trace.rs

1use http::{HeaderMap, HeaderValue, header::InvalidHeaderValue};
2
3/// Standard correlation header used by rs-zero REST and RPC adapters.
4pub const REQUEST_ID_HEADER: &str = "x-request-id";
5/// W3C Trace Context header.
6pub const TRACEPARENT_HEADER: &str = "traceparent";
7
8/// Extracts a request id from HTTP headers.
9pub fn request_id_from_headers(headers: &HeaderMap) -> Option<String> {
10    headers
11        .get(REQUEST_ID_HEADER)
12        .and_then(|value| value.to_str().ok())
13        .map(str::trim)
14        .filter(|value| !value.is_empty())
15        .map(ToOwned::to_owned)
16}
17
18/// Extracts a valid W3C traceparent header from HTTP headers.
19pub fn traceparent_from_headers(headers: &HeaderMap) -> Option<String> {
20    headers
21        .get(TRACEPARENT_HEADER)
22        .and_then(|value| value.to_str().ok())
23        .map(str::trim)
24        .filter(|value| is_valid_traceparent(value))
25        .map(ToOwned::to_owned)
26}
27
28/// Inserts a W3C traceparent header into HTTP headers.
29pub fn insert_traceparent_header(
30    headers: &mut HeaderMap,
31    traceparent: &str,
32) -> Result<(), InvalidHeaderValue> {
33    headers.insert(TRACEPARENT_HEADER, HeaderValue::from_str(traceparent)?);
34    Ok(())
35}
36
37/// Extracts a request id from tonic metadata.
38#[cfg(feature = "rpc")]
39pub fn request_id_from_metadata(metadata: &tonic::metadata::MetadataMap) -> Option<String> {
40    metadata
41        .get(REQUEST_ID_HEADER)
42        .and_then(|value| value.to_str().ok())
43        .map(str::trim)
44        .filter(|value| !value.is_empty())
45        .map(ToOwned::to_owned)
46}
47
48/// Extracts a valid W3C traceparent value from tonic metadata.
49#[cfg(feature = "rpc")]
50pub fn traceparent_from_metadata(metadata: &tonic::metadata::MetadataMap) -> Option<String> {
51    metadata
52        .get(TRACEPARENT_HEADER)
53        .and_then(|value| value.to_str().ok())
54        .map(str::trim)
55        .filter(|value| is_valid_traceparent(value))
56        .map(ToOwned::to_owned)
57}
58
59/// Inserts a traceparent value into tonic metadata.
60#[cfg(feature = "rpc")]
61pub fn insert_traceparent_metadata(
62    metadata: &mut tonic::metadata::MetadataMap,
63    traceparent: &str,
64) -> Result<(), tonic::metadata::errors::InvalidMetadataValue> {
65    metadata.insert(TRACEPARENT_HEADER, traceparent.parse()?);
66    Ok(())
67}
68
69/// Returns the current OpenTelemetry trace id when an OTLP context is active.
70pub fn current_trace_id() -> Option<String> {
71    #[cfg(feature = "otlp")]
72    {
73        use opentelemetry::trace::TraceContextExt;
74        use tracing_opentelemetry::OpenTelemetrySpanExt;
75
76        let context = tracing::Span::current().context();
77        let span = context.span();
78        let span_context = span.span_context();
79        if span_context.is_valid() {
80            return Some(span_context.trace_id().to_string());
81        }
82    }
83
84    None
85}
86
87/// Returns the current OpenTelemetry span id when an OTLP context is active.
88pub fn current_span_id() -> Option<String> {
89    #[cfg(feature = "otlp")]
90    {
91        use opentelemetry::trace::TraceContextExt;
92        use tracing_opentelemetry::OpenTelemetrySpanExt;
93
94        let context = tracing::Span::current().context();
95        let span = context.span();
96        let span_context = span.span_context();
97        if span_context.is_valid() {
98            return Some(span_context.span_id().to_string());
99        }
100    }
101
102    None
103}
104
105/// Returns a W3C traceparent value for the current span when OTLP context is active.
106pub fn current_traceparent() -> Option<String> {
107    let trace_id = current_trace_id()?;
108    let span_id = current_span_id()?;
109    Some(format!("00-{trace_id}-{span_id}-01"))
110}
111
112/// Extracts an OpenTelemetry parent context from HTTP headers.
113#[cfg(feature = "otlp")]
114pub fn opentelemetry_context_from_headers(headers: &HeaderMap) -> Option<opentelemetry::Context> {
115    use opentelemetry::{global, trace::TraceContextExt};
116
117    traceparent_from_headers(headers)?;
118    let extractor = HeaderMapExtractor { headers };
119    let context = global::get_text_map_propagator(|propagator| propagator.extract(&extractor));
120    context.span().span_context().is_valid().then_some(context)
121}
122
123/// Extracts an OpenTelemetry parent context from one W3C traceparent value.
124#[cfg(feature = "otlp")]
125pub fn opentelemetry_context_from_traceparent(traceparent: &str) -> Option<opentelemetry::Context> {
126    use opentelemetry::{global, trace::TraceContextExt};
127
128    if !is_valid_traceparent(traceparent) {
129        return None;
130    }
131
132    let extractor = TraceParentExtractor { traceparent };
133    let context = global::get_text_map_propagator(|propagator| propagator.extract(&extractor));
134    context.span().span_context().is_valid().then_some(context)
135}
136
137/// Sets an OpenTelemetry parent context on a tracing span from HTTP headers.
138#[cfg(feature = "otlp")]
139pub fn set_span_parent_from_headers(span: &tracing::Span, headers: &HeaderMap) -> bool {
140    use tracing_opentelemetry::OpenTelemetrySpanExt;
141
142    let Some(context) = opentelemetry_context_from_headers(headers) else {
143        return false;
144    };
145
146    span.set_parent(context).is_ok()
147}
148
149/// Extracts an OpenTelemetry parent context from tonic metadata.
150#[cfg(all(feature = "otlp", feature = "rpc"))]
151pub fn opentelemetry_context_from_metadata(
152    metadata: &tonic::metadata::MetadataMap,
153) -> Option<opentelemetry::Context> {
154    use opentelemetry::{global, trace::TraceContextExt};
155
156    traceparent_from_metadata(metadata)?;
157    let extractor = MetadataMapExtractor { metadata };
158    let context = global::get_text_map_propagator(|propagator| propagator.extract(&extractor));
159    context.span().span_context().is_valid().then_some(context)
160}
161
162/// Sets an OpenTelemetry parent context on a tracing span from tonic metadata.
163#[cfg(all(feature = "otlp", feature = "rpc"))]
164pub fn set_span_parent_from_metadata(
165    span: &tracing::Span,
166    metadata: &tonic::metadata::MetadataMap,
167) -> bool {
168    use tracing_opentelemetry::OpenTelemetrySpanExt;
169
170    let Some(context) = opentelemetry_context_from_metadata(metadata) else {
171        return false;
172    };
173
174    span.set_parent(context).is_ok()
175}
176
177/// Injects the current OpenTelemetry context into tonic metadata.
178#[cfg(all(feature = "otlp", feature = "rpc"))]
179pub fn inject_current_context_metadata(
180    metadata: &mut tonic::metadata::MetadataMap,
181) -> Result<bool, tonic::metadata::errors::InvalidMetadataValue> {
182    use opentelemetry::{global, trace::TraceContextExt};
183    use tracing_opentelemetry::OpenTelemetrySpanExt;
184
185    let context = tracing::Span::current().context();
186    if !context.span().span_context().is_valid() {
187        return Ok(false);
188    }
189
190    let mut injector = MetadataMapInjector {
191        metadata,
192        invalid_value: None,
193    };
194    global::get_text_map_propagator(|propagator| {
195        propagator.inject_context(&context, &mut injector);
196    });
197    if let Some(error) = injector.invalid_value {
198        return Err(error);
199    }
200    Ok(injector.metadata.contains_key(TRACEPARENT_HEADER))
201}
202
203/// Extracts the trace id part from a valid traceparent value.
204pub fn trace_id_from_traceparent(traceparent: &str) -> Option<&str> {
205    if !is_valid_traceparent(traceparent) {
206        return None;
207    }
208    traceparent.split('-').nth(1)
209}
210
211/// Extracts the span id part from a valid traceparent value.
212pub fn span_id_from_traceparent(traceparent: &str) -> Option<&str> {
213    if !is_valid_traceparent(traceparent) {
214        return None;
215    }
216    traceparent.split('-').nth(2)
217}
218
219fn is_valid_traceparent(value: &str) -> bool {
220    let mut parts = value.split('-');
221    let Some(version) = parts.next() else {
222        return false;
223    };
224    let Some(trace_id) = parts.next() else {
225        return false;
226    };
227    let Some(span_id) = parts.next() else {
228        return false;
229    };
230    let Some(flags) = parts.next() else {
231        return false;
232    };
233    parts.next().is_none()
234        && version.len() == 2
235        && trace_id.len() == 32
236        && span_id.len() == 16
237        && flags.len() == 2
238        && trace_id != "00000000000000000000000000000000"
239        && span_id != "0000000000000000"
240        && version.chars().all(|value| value.is_ascii_hexdigit())
241        && trace_id.chars().all(|value| value.is_ascii_hexdigit())
242        && span_id.chars().all(|value| value.is_ascii_hexdigit())
243        && flags.chars().all(|value| value.is_ascii_hexdigit())
244}
245
246#[cfg(feature = "otlp")]
247struct HeaderMapExtractor<'a> {
248    headers: &'a HeaderMap,
249}
250
251#[cfg(feature = "otlp")]
252impl opentelemetry::propagation::Extractor for HeaderMapExtractor<'_> {
253    fn get(&self, key: &str) -> Option<&str> {
254        self.headers.get(key).and_then(|value| value.to_str().ok())
255    }
256
257    fn keys(&self) -> Vec<&str> {
258        self.headers.keys().map(|key| key.as_str()).collect()
259    }
260}
261
262#[cfg(feature = "otlp")]
263struct TraceParentExtractor<'a> {
264    traceparent: &'a str,
265}
266
267#[cfg(feature = "otlp")]
268impl opentelemetry::propagation::Extractor for TraceParentExtractor<'_> {
269    fn get(&self, key: &str) -> Option<&str> {
270        key.eq_ignore_ascii_case(TRACEPARENT_HEADER)
271            .then_some(self.traceparent)
272    }
273
274    fn keys(&self) -> Vec<&str> {
275        vec![TRACEPARENT_HEADER]
276    }
277}
278
279#[cfg(all(feature = "otlp", feature = "rpc"))]
280struct MetadataMapExtractor<'a> {
281    metadata: &'a tonic::metadata::MetadataMap,
282}
283
284#[cfg(all(feature = "otlp", feature = "rpc"))]
285impl opentelemetry::propagation::Extractor for MetadataMapExtractor<'_> {
286    fn get(&self, key: &str) -> Option<&str> {
287        self.metadata.get(key).and_then(|value| value.to_str().ok())
288    }
289
290    fn keys(&self) -> Vec<&str> {
291        self.metadata
292            .keys()
293            .filter_map(|key| match key {
294                tonic::metadata::KeyRef::Ascii(key) => Some(key.as_str()),
295                tonic::metadata::KeyRef::Binary(_) => None,
296            })
297            .collect()
298    }
299}
300
301#[cfg(all(feature = "otlp", feature = "rpc"))]
302struct MetadataMapInjector<'a> {
303    metadata: &'a mut tonic::metadata::MetadataMap,
304    invalid_value: Option<tonic::metadata::errors::InvalidMetadataValue>,
305}
306
307#[cfg(all(feature = "otlp", feature = "rpc"))]
308impl opentelemetry::propagation::Injector for MetadataMapInjector<'_> {
309    fn set(&mut self, key: &str, value: String) {
310        let Ok(key) = key.parse::<tonic::metadata::MetadataKey<tonic::metadata::Ascii>>() else {
311            return;
312        };
313        match tonic::metadata::MetadataValue::try_from(value.as_str()) {
314            Ok(value) => {
315                self.metadata.insert(key, value);
316            }
317            Err(error) => {
318                self.invalid_value = Some(error);
319            }
320        }
321    }
322}
323
324#[cfg(test)]
325mod tests {
326    use http::HeaderMap;
327
328    use super::{
329        REQUEST_ID_HEADER, TRACEPARENT_HEADER, current_trace_id, request_id_from_headers,
330        trace_id_from_traceparent, traceparent_from_headers,
331    };
332
333    #[test]
334    fn extracts_request_id_from_headers() {
335        let mut headers = HeaderMap::new();
336        headers.insert(REQUEST_ID_HEADER, "req-1".parse().expect("header"));
337
338        assert_eq!(request_id_from_headers(&headers).as_deref(), Some("req-1"));
339    }
340
341    #[test]
342    fn trace_id_is_not_forged_without_active_context() {
343        assert!(current_trace_id().is_none());
344    }
345
346    #[test]
347    fn extracts_valid_traceparent_from_headers() {
348        let mut headers = HeaderMap::new();
349        headers.insert(
350            TRACEPARENT_HEADER,
351            "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
352                .parse()
353                .expect("traceparent"),
354        );
355
356        let value = traceparent_from_headers(&headers).expect("traceparent");
357        assert_eq!(
358            trace_id_from_traceparent(&value),
359            Some("4bf92f3577b34da6a3ce929d0e0e4736")
360        );
361    }
362}
rs_zero/observability/trace.rs

rs_zero/observability/
trace.rs
