use opentelemetry::propagation::{Extractor, Injector, TextMapPropagator};
use opentelemetry_sdk::propagation::TraceContextPropagator;
use std::collections::HashMap;
use tracing_opentelemetry::OpenTelemetrySpanExt;
struct Carrier(HashMap<String, String>);
impl Injector for Carrier {
fn set(&mut self, key: &str, value: String) {
self.0.insert(key.to_string(), value);
}
}
impl Extractor for Carrier {
fn get(&self, key: &str) -> Option<&str> {
self.0.get(key).map(String::as_str)
}
fn keys(&self) -> Vec<&str> {
self.0.keys().map(String::as_str).collect()
}
}
pub(crate) fn current_traceparent() -> Option<String> {
let context = tracing::Span::current().context();
let mut carrier = Carrier(HashMap::new());
TraceContextPropagator::new().inject_context(&context, &mut carrier);
carrier.0.remove("traceparent")
}
pub(crate) fn set_parent(span: &tracing::Span, traceparent: Option<&str>) {
let Some(traceparent) = traceparent else {
return;
};
let carrier = Carrier(HashMap::from([(
"traceparent".to_string(),
traceparent.to_string(),
)]));
let context = TraceContextPropagator::new().extract(&carrier);
span.set_parent(context);
}
#[cfg(test)]
mod tests {
use super::*;
use opentelemetry::trace::{TraceContextExt, TracerProvider as _};
use tracing_subscriber::layer::SubscriberExt;
#[test]
fn child_span_joins_the_publisher_trace() {
let provider = opentelemetry_sdk::trace::TracerProvider::builder().build();
let tracer = provider.tracer("test");
let subscriber = tracing_subscriber::registry()
.with(tracing_opentelemetry::layer().with_tracer(tracer));
tracing::subscriber::with_default(subscriber, || {
let publisher = tracing::info_span!("publisher");
let (traceparent, publisher_trace) = publisher.in_scope(|| {
(
current_traceparent(),
tracing::Span::current().context().span().span_context().trace_id(),
)
});
assert!(traceparent.is_some());
let worker = tracing::info_span!("worker");
set_parent(&worker, traceparent.as_deref());
let worker_trace = worker.context().span().span_context().trace_id();
assert_eq!(publisher_trace, worker_trace);
});
}
}