aether_telemetry/
trace_context.rs

1use aether_core::events::{TRACEPARENT_KEY, TRACESTATE_KEY, TraceContext};
2use opentelemetry::Context;
3use opentelemetry::propagation::{Extractor, Injector, TextMapPropagator};
4use opentelemetry::trace::TraceContextExt;
5use opentelemetry_sdk::propagation::TraceContextPropagator;
6use schemars::JsonSchema;
7use serde::{Deserialize, Serialize};
8
9#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
10#[serde(untagged, rename_all = "camelCase", deny_unknown_fields)]
11pub enum AgentTraceContext {
12    /// Continue the trace beneath a remote W3C parent span.
13    Parent(TraceContext),
14    /// Start root spans in the supplied trace without attaching a parent span.
15    Root {
16        /// W3C trace ID: 32 lowercase hexadecimal characters and not all zeros.
17        #[serde(rename = "traceId")]
18        #[schemars(rename = "traceId")]
19        trace_id: String,
20    },
21}
22
23/// The trace context of `context`'s span, for propagation to another process.
24/// `None` when the span is invalid and there is nothing to propagate.
25pub(crate) fn inject_trace_context(context: &Context) -> Option<TraceContext> {
26    if !context.span().span_context().is_valid() {
27        return None;
28    }
29
30    let mut carrier = W3cCarrier::default();
31    TraceContextPropagator::new().inject_context(context, &mut carrier);
32    Some(TraceContext { traceparent: carrier.traceparent?, tracestate: carrier.tracestate })
33}
34
35/// The remote span `trace_context` names, for parenting local spans beneath it.
36/// `None` when the traceparent does not name a valid span.
37pub(crate) fn extract_trace_context(trace_context: &TraceContext) -> Option<Context> {
38    let context = TraceContextPropagator::new().extract(&W3cCarrier::from(trace_context));
39    context.span().span_context().is_valid().then_some(context)
40}
41
42/// Adapts [`TraceContext`] to the propagator's carrier API, which the orphan
43/// rule keeps us from implementing on `TraceContext` itself.
44#[derive(Default)]
45struct W3cCarrier {
46    traceparent: Option<String>,
47    tracestate: Option<String>,
48}
49
50impl From<&TraceContext> for W3cCarrier {
51    fn from(trace_context: &TraceContext) -> Self {
52        Self { traceparent: Some(trace_context.traceparent.clone()), tracestate: trace_context.tracestate.clone() }
53    }
54}
55
56impl Injector for W3cCarrier {
57    fn set(&mut self, key: &str, value: String) {
58        match key {
59            TRACEPARENT_KEY => self.traceparent = Some(value),
60            TRACESTATE_KEY => self.tracestate = Some(value),
61            _ => {}
62        }
63    }
64}
65
66impl Extractor for W3cCarrier {
67    fn get(&self, key: &str) -> Option<&str> {
68        match key {
69            TRACEPARENT_KEY => self.traceparent.as_deref(),
70            TRACESTATE_KEY => self.tracestate.as_deref(),
71            _ => None,
72        }
73    }
74
75    fn keys(&self) -> Vec<&str> {
76        [(TRACEPARENT_KEY, &self.traceparent), (TRACESTATE_KEY, &self.tracestate)]
77            .into_iter()
78            .filter_map(|(key, value)| value.is_some().then_some(key))
79            .collect()
80    }
81}
aether_telemetry/trace_context.rs

aether_telemetry/
trace_context.rs
