greentic_telemetry/
layer.rs

1use crate::context::TelemetryCtx;
2use crate::tasklocal::with_current_telemetry_ctx;
3use std::sync::Arc;
4use tracing::Subscriber;
5use tracing_subscriber::{
6    layer::{Context, Layer},
7    registry::LookupSpan,
8};
9
10#[derive(Clone)]
11struct ContextLayer {
12    provider: Arc<dyn Fn() -> Option<TelemetryCtx> + Send + Sync>,
13}
14
15impl ContextLayer {
16    fn new(provider: Arc<dyn Fn() -> Option<TelemetryCtx> + Send + Sync>) -> Self {
17        Self { provider }
18    }
19}
20
21impl<S> Layer<S> for ContextLayer
22where
23    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
24{
25    fn on_new_span(
26        &self,
27        _attrs: &tracing::span::Attributes<'_>,
28        id: &tracing::span::Id,
29        ctx: Context<'_, S>,
30    ) {
31        if let Some(span) = ctx.span(id)
32            && let Some(tctx) = (self.provider)()
33        {
34            let _ = span.extensions_mut().replace(tctx);
35        }
36    }
37
38    fn on_enter(&self, id: &tracing::span::Id, ctx: Context<'_, S>) {
39        if let Some(span) = ctx.span(id) {
40            let telemetry = span
41                .extensions()
42                .get::<TelemetryCtx>()
43                .cloned()
44                .or_else(|| (self.provider)());
45
46            if let Some(tctx) = telemetry {
47                // Store the context for capture/inspection layers only. Do NOT
48                // record gt.* fields here: `Span::record` is a silent no-op for
49                // fields the callsite didn't declare, and `Span::current()` is
50                // the parent span, not the one being entered — so the old record
51                // loop attributed nothing. OTLP attribute export goes through
52                // `annotate_span`.
53                let _ = span.extensions_mut().replace(tctx);
54            }
55        }
56    }
57}
58
59/// Attach a [`TelemetryCtx`]'s attributes to `span` as real OpenTelemetry span
60/// attributes (the `gt.*` keys from [`TelemetryCtx::kv`]), so they reach OTLP
61/// export regardless of whether the span's callsite declared those fields.
62///
63/// This is the **export primitive** for context attribution. `Span::record`
64/// only affects fields declared at the callsite, and `tracing_opentelemetry`'s
65/// `OtelData` builder is private to that crate — so the supported way to attach
66/// callsite-unknown attributes is `OpenTelemetrySpanExt::set_attribute` on an
67/// owned [`tracing::Span`]. A producer that has a span handle (e.g. the runner
68/// creating an invocation span) calls this right after creating the span.
69#[cfg(any(feature = "otlp", feature = "azure", feature = "gcp"))]
70pub fn annotate_span(span: &tracing::Span, tctx: &TelemetryCtx) {
71    use tracing_opentelemetry::OpenTelemetrySpanExt;
72    for (key, value) in tctx.kv() {
73        if let Some(v) = value {
74            span.set_attribute(key, v.to_string());
75        }
76    }
77}
78
79/// [`annotate_span`] using the task-local [`TelemetryCtx`] (set via
80/// `set_current_telemetry_ctx` / the tenant-ctx bridge). No-op when none is set.
81#[cfg(any(feature = "otlp", feature = "azure", feature = "gcp"))]
82pub fn annotate_current_span(span: &tracing::Span) {
83    match with_current_telemetry_ctx(|ctx| ctx.cloned()) {
84        Some(tctx) => annotate_span(span, &tctx),
85        // No task-local TelemetryCtx in scope (called outside `with_task_local`
86        // or before `set_current_telemetry_ctx`). The span then exports with no
87        // `gt.*` attribution — surface it at debug so unattributed spans are
88        // diagnosable rather than silently lost.
89        None => tracing::debug!(
90            "annotate_current_span: no task-local TelemetryCtx in scope; \
91             span exported without gt.* attributes"
92        ),
93    }
94}
95
96/// Layer that attaches the task-local [`TelemetryCtx`] to spans (stored in span
97/// extensions, and recorded onto any declared `gt.*` fields). For OTLP
98/// attribute export of callsite-unknown context, use [`annotate_span`] —
99/// `tracing_opentelemetry` 0.32 keeps its span builder private, so a layer
100/// cannot inject attributes after the fact.
101pub fn layer_from_task_local() -> impl Layer<tracing_subscriber::Registry> + Clone {
102    let provider = Arc::new(|| with_current_telemetry_ctx(|ctx| ctx.cloned()));
103    ContextLayer::new(provider)
104}
105
106/// Like [`layer_from_task_local`] but with a caller-supplied context provider.
107pub fn layer_with_provider(
108    provider: impl Fn() -> Option<TelemetryCtx> + Send + Sync + 'static,
109) -> impl Layer<tracing_subscriber::Registry> + Clone {
110    ContextLayer::new(Arc::new(provider))
111}
greentic_telemetry/layer.rs

greentic_telemetry/
layer.rs
