main
use std::collections::HashMap;
use opentelemetry::propagation::{
Extractor, Injector, TextMapCompositePropagator, TextMapPropagator,
};
use opentelemetry::trace::{SpanContext, TraceContextExt};
use opentelemetry::{Context, global};
use opentelemetry_sdk::propagation::{BaggagePropagator, TraceContextPropagator};
use tracing_opentelemetry::OpenTelemetrySpanExt;
pub const TRACEPARENT: &str = "traceparent";
pub const TRACESTATE: &str = "tracestate";
pub const BAGGAGE: &str = "baggage";
#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct TextMap(pub HashMap<String, String>);
impl TextMap {
pub fn new() -> Self {
Self(HashMap::new())
}
pub fn get(&self, key: &str) -> Option<&str> {
self.0.get(key).map(String::as_str)
}
pub fn insert(&mut self, key: impl Into<String>, value: impl Into<String>) {
self.0.insert(key.into(), value.into());
}
}
impl Injector for TextMap {
fn set(&mut self, key: &str, value: String) {
self.0.insert(key.to_owned(), value);
}
}
impl Extractor for TextMap {
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 fn install() {
global::set_text_map_propagator(TextMapCompositePropagator::new(vec![
Box::new(TraceContextPropagator::new()),
Box::new(BaggagePropagator::new()),
]));
}
pub fn current_context() -> Context {
let context = tracing::Span::current().context();
if context.span().span_context().is_valid() {
return context;
}
Context::current()
}
pub fn inject(into: &mut dyn Injector) {
let context = current_context();
global::get_text_map_propagator(|propagator| {
propagator.inject_context(&context, &mut *into);
});
}
pub fn inject_map(into: &mut TextMap) {
inject(into);
}
pub fn extract(from: &dyn Extractor) -> Context {
global::get_text_map_propagator(|propagator| propagator.extract(from))
}
pub fn traceparent() -> Option<String> {
let mut map = TextMap::new();
inject_map(&mut map);
map.0.get(TRACEPARENT).cloned()
}
pub fn context_from_traceparent(traceparent: &str) -> Context {
let mut map = TextMap::new();
map.insert(TRACEPARENT, traceparent);
TraceContextPropagator::new().extract(&map)
}
pub fn current() -> Option<SpanContext> {
let span_context = current_context().span().span_context().clone();
span_context.is_valid().then_some(span_context)
}
pub fn set_parent(span: &tracing::Span, parent: &Context) -> bool {
span.set_parent(parent.clone()).is_ok()
}
pub fn set_parent_from_traceparent(span: &tracing::Span, traceparent: Option<&str>) -> bool {
match traceparent.filter(|value| !value.is_empty()) {
Some(value) => set_parent(span, &context_from_traceparent(value)),
None => false,
}
}