Skip to main content

aether_telemetry/
trace_context.rs

1use aether_core::events::{TRACEPARENT_KEY, TRACESTATE_KEY, TraceContext};
2use opentelemetry::Context;
3use opentelemetry::propagation::TextMapPropagator;
4use opentelemetry::trace::TraceContextExt;
5use opentelemetry_sdk::propagation::TraceContextPropagator;
6use schemars::JsonSchema;
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9
10#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)]
11#[serde(untagged, rename_all = "camelCase", deny_unknown_fields)]
12pub enum AgentTraceContext {
13    /// Continue the trace beneath a remote W3C parent span.
14    Parent(TraceContext),
15    /// Start root spans in the supplied trace without attaching a parent span.
16    Root {
17        /// W3C trace ID: 32 lowercase hexadecimal characters and not all zeros.
18        #[serde(rename = "traceId")]
19        #[schemars(rename = "traceId")]
20        trace_id: String,
21    },
22}
23
24/// The trace context of `context`'s span, for propagation to another process.
25/// `None` when the span is invalid and there is nothing to propagate.
26pub(crate) fn inject_trace_context(context: &Context) -> Option<TraceContext> {
27    if !context.span().span_context().is_valid() {
28        return None;
29    }
30
31    let mut carrier = HashMap::new();
32    TraceContextPropagator::new().inject_context(context, &mut carrier);
33    Some(TraceContext { traceparent: carrier.remove(TRACEPARENT_KEY)?, tracestate: carrier.remove(TRACESTATE_KEY) })
34}
35
36/// The remote span `trace_context` names, for parenting local spans beneath it.
37/// `None` when the traceparent does not name a valid span.
38pub(crate) fn extract_trace_context(trace_context: &TraceContext) -> Option<Context> {
39    let mut carrier = HashMap::from([(TRACEPARENT_KEY.to_string(), trace_context.traceparent.clone())]);
40    carrier.extend(trace_context.tracestate.clone().map(|tracestate| (TRACESTATE_KEY.to_string(), tracestate)));
41    let context = TraceContextPropagator::new().extract(&carrier);
42    context.span().span_context().is_valid().then_some(context)
43}