use http::HeaderMap;
pub const TRACEPARENT: &str = "traceparent";
#[must_use]
pub fn get_traceparent(headers: &HeaderMap) -> Option<&str> {
headers.get(TRACEPARENT)?.to_str().ok()
}
#[must_use]
pub fn parse_trace_id(traceparent: &str) -> Option<String> {
let parts: Vec<&str> = traceparent.split('-').collect();
if parts.len() >= 4 && parts[0] == "00" {
Some(parts[1].to_owned())
} else {
None
}
}
#[cfg(feature = "otel")]
mod imp {
use super::{get_traceparent, parse_trace_id};
use http::{HeaderMap, HeaderName, HeaderValue};
use opentelemetry::{
Context, global,
propagation::{Extractor, Injector},
};
use tracing::Span;
use tracing_opentelemetry::OpenTelemetrySpanExt;
struct HeadersExtractor<'a>(&'a HeaderMap);
impl Extractor for HeadersExtractor<'_> {
fn get(&self, key: &str) -> Option<&str> {
self.0.get(key).and_then(|v| v.to_str().ok())
}
fn keys(&self) -> Vec<&str> {
self.0.keys().map(http::HeaderName::as_str).collect()
}
}
struct HeadersInjector<'a>(&'a mut HeaderMap);
impl Injector for HeadersInjector<'_> {
fn set(&mut self, key: &str, value: String) {
if let Ok(name) = HeaderName::from_bytes(key.as_bytes())
&& let Ok(val) = HeaderValue::from_str(&value)
{
self.0.insert(name, val);
}
}
}
pub fn inject_current_span(headers: &mut HeaderMap) {
let cx = Context::current();
global::get_text_map_propagator(|propagator| {
propagator.inject_context(&cx, &mut HeadersInjector(headers));
});
}
pub fn set_parent_from_headers(span: &Span, headers: &HeaderMap) {
let parent_cx = global::get_text_map_propagator(|propagator| {
propagator.extract(&HeadersExtractor(headers))
});
_ = span.set_parent(parent_cx);
if let Some(traceparent) = get_traceparent(headers)
&& let Some(trace_id) = parse_trace_id(traceparent)
{
span.record("trace_id", &trace_id);
span.record("parent.trace_id", &trace_id);
}
}
}
#[cfg(not(feature = "otel"))]
mod imp {
use super::{get_traceparent, parse_trace_id};
use http::HeaderMap;
use tracing::Span;
pub fn inject_current_span(_headers: &mut HeaderMap) {
}
pub fn set_parent_from_headers(span: &Span, headers: &HeaderMap) {
if let Some(traceparent) = get_traceparent(headers)
&& let Some(trace_id) = parse_trace_id(traceparent)
{
span.record("trace_id", &trace_id);
span.record("parent.trace_id", &trace_id);
}
}
}
pub use imp::{inject_current_span, set_parent_from_headers};
#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
mod tests {
use super::*;
use tracing::info_span;
#[test]
fn test_get_traceparent_none() {
let headers = HeaderMap::new();
assert!(get_traceparent(&headers).is_none());
}
#[test]
fn test_get_traceparent_ok() {
let mut headers = HeaderMap::new();
headers.insert(
TRACEPARENT,
"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
.parse()
.expect("valid header"),
);
let tp = get_traceparent(&headers);
assert!(tp.is_some());
assert_eq!(
tp.expect("should be some"),
"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
);
}
#[test]
fn test_parse_trace_id_ok() {
let traceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
let trace_id = parse_trace_id(traceparent);
assert_eq!(
trace_id,
Some("4bf92f3577b34da6a3ce929d0e0e4736".to_owned())
);
}
#[test]
fn test_parse_trace_id_invalid() {
assert!(parse_trace_id("invalid").is_none());
assert!(parse_trace_id("").is_none());
}
#[test]
#[cfg(not(feature = "otel"))]
fn test_inject_current_span_noop() {
let mut headers = HeaderMap::new();
inject_current_span(&mut headers);
assert!(headers.is_empty());
}
#[test]
#[cfg(feature = "otel")]
fn test_inject_current_span_no_panic() {
use opentelemetry::global;
use opentelemetry_sdk::propagation::TraceContextPropagator;
global::set_text_map_propagator(TraceContextPropagator::new());
let mut headers = http::HeaderMap::new();
let _span = tracing::info_span!("test").entered();
inject_current_span(&mut headers);
}
#[test]
fn test_set_parent_from_headers_no_panic() {
let mut headers = HeaderMap::new();
headers.insert(
TRACEPARENT,
"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
.parse()
.expect("valid header"),
);
let span = info_span!(
"test",
trace_id = tracing::field::Empty,
parent.trace_id = tracing::field::Empty
);
set_parent_from_headers(&span, &headers);
}
}