pgtask_otel/
propagation.rs

1use opentelemetry::{
2    Context, global,
3    propagation::{Extractor, Injector, TextMapCompositePropagator},
4};
5use opentelemetry_sdk::propagation::{BaggagePropagator, TraceContextPropagator};
6use serde_json::{Map, Value};
7use tracing::Span;
8use tracing_opentelemetry::{OpenTelemetrySpanExt, SetParentError};
9
10pub fn configure_propagation() {
11    global::set_text_map_propagator(TextMapCompositePropagator::new(vec![
12        Box::new(TraceContextPropagator::new()),
13        Box::new(BaggagePropagator::new()),
14    ]));
15}
16
17pub fn inject_context(headers: &Map<String, Value>, context: &Context) -> Map<String, Value> {
18    let mut headers = headers.clone();
19    global::get_text_map_propagator(|propagator| propagator.inject_context(context, &mut HeaderInjector(&mut headers)));
20    headers
21}
22
23pub fn inject_span_context(headers: &Map<String, Value>, span: &Span) -> Map<String, Value> {
24    inject_context(headers, &span.context())
25}
26
27pub fn set_parent_from_headers(span: &Span, headers: &Map<String, Value>) -> Result<(), SetParentError> {
28    let context = global::get_text_map_propagator(|propagator| propagator.extract(&HeaderExtractor(headers)));
29    span.set_parent(context)
30}
31
32struct HeaderInjector<'a>(&'a mut Map<String, Value>);
33
34impl Injector for HeaderInjector<'_> {
35    fn set(&mut self, key: &str, value: String) {
36        self.0.insert(key.to_owned(), Value::String(value));
37    }
38}
39
40struct HeaderExtractor<'a>(&'a Map<String, Value>);
41
42impl Extractor for HeaderExtractor<'_> {
43    fn get(&self, key: &str) -> Option<&str> {
44        self.0.get(key).and_then(Value::as_str)
45    }
46
47    fn keys(&self) -> Vec<&str> {
48        self.0.keys().map(String::as_str).collect()
49    }
50}
pgtask_otel/propagation.rs

pgtask_otel/
propagation.rs
