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

rs_zero/observability/
trace.rs
