chore: update SpanRef context access based on feedback

bantonsson · bantonsson · commit c18a70e4aa1f · 2025-12-02T11:49:25.000+01:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -29,7 +29,7 @@ opentelemetry_sdk = { version = "0.31.0", default-features = false, features = [
 ] }
 tracing = { version = "0.1.35", default-features = false, features = ["std"] }
 tracing-core = "0.1.28"
-tracing-subscriber = { version = "0.3.0", default-features = false, features = [
+tracing-subscriber = { version = "0.3.22", default-features = false, features = [
     "registry",
     "std",
 ] }
diff --git a/examples/otel_context.rs b/examples/otel_context.rs
@@ -4,23 +4,23 @@
 use opentelemetry::trace::{TraceContextExt, TracerProvider as _};
 use opentelemetry_sdk::trace::SdkTracerProvider;
 use opentelemetry_stdout as stdout;
-use std::sync::{Arc, Mutex, RwLock};
+use std::sync::{Arc, Mutex, OnceLock};
 use tracing::{debug, info, span, warn, Subscriber};
 use tracing::{dispatcher::WeakDispatch, level_filters::LevelFilter, Dispatch};
-use tracing_opentelemetry::{layer, OpenTelemetryContext};
+use tracing_opentelemetry::{get_otel_context, layer};
 use tracing_subscriber::layer::Context;
 use tracing_subscriber::prelude::*;
 use tracing_subscriber::registry::LookupSpan;
 use tracing_subscriber::Layer;
 
 /// A custom layer that demonstrates how to use OpenTelemetryContext
 /// to extract OpenTelemetry contexts from span extensions.
-#[derive(Clone)]
+#[derive(Clone, Default)]
 struct SpanAnalysisLayer {
     /// Store span analysis results for demonstration
     analysis_results: Arc<Mutex<Vec<SpanAnalysis>>>,
     /// Weak reference to the dispatcher for context extraction
-    dispatch: Arc<RwLock<Option<WeakDispatch>>>,
+    dispatch: Arc<OnceLock<WeakDispatch>>,
 }
 
 #[derive(Debug, Clone)]
@@ -32,13 +32,6 @@ struct SpanAnalysis {
 }
 
 impl SpanAnalysisLayer {
-    fn new() -> Self {
-        Self {
-            analysis_results: Arc::new(Mutex::new(Vec::new())),
-            dispatch: Arc::new(RwLock::new(None)),
-        }
-    }
-
     fn get_analysis_results(&self) -> Vec<SpanAnalysis> {
         self.analysis_results.lock().unwrap().clone()
     }
@@ -68,46 +61,23 @@ impl SpanAnalysisLayer {
             }
         }
     }
-
-    fn get_weak_dispatch(&self, get_default: bool) -> Option<WeakDispatch> {
-        let read_guard = self.dispatch.read().unwrap();
-        match read_guard.as_ref() {
-            Some(weak_dispatch) => Some(weak_dispatch.clone()),
-            // Note: This workaround is needed until https://github.com/tokio-rs/tracing/pull/3379
-            // is merged and released. It should really be handled in on_register_dispatch
-            None => {
-                if !get_default {
-                    None
-                } else {
-                    drop(read_guard);
-                    let mut dispatch = self.dispatch.write().unwrap();
-                    let weak_dispatch = Dispatch::default().downgrade();
-                    *dispatch = Some(weak_dispatch.clone());
-                    Some(weak_dispatch)
-                }
-            }
-        }
-    }
 }
 
 impl<S> Layer<S> for SpanAnalysisLayer
 where
     S: Subscriber + for<'span> LookupSpan<'span>,
 {
+    fn on_register_dispatch(&self, subscriber: &Dispatch) {
+        let _ = self.dispatch.set(subscriber.downgrade());
+    }
+
     fn on_new_span(
         &self,
         attrs: &tracing::span::Attributes<'_>,
         id: &tracing::span::Id,
         ctx: Context<'_, S>,
     ) {
-        // Get the weak dispatch reference.
-        //
-        // Note: We can't use the Dispatch::default() workaround described above here since this
-        // method is called from inside a dispatcher::get_default block, and such calls can't be
-        // nested so we would get the global dispatcher instead, which can't downcast to the right
-        // types when extracting the OpenTelemetry context. This also means that we will miss
-        // analyzing the first span that is created 🤷🏼‍♂️
-        let Some(weak_dispatch) = self.get_weak_dispatch(false) else {
+        let Some(weak_dispatch) = self.dispatch.get() else {
             return;
         };
 
@@ -116,46 +86,39 @@ where
             // This is the key functionality: using OpenTelemetryContext
             // to extract the OpenTelemetry context from span extensions
             let mut extensions = span_ref.extensions_mut();
-            if let Some(otel_context) =
-                OpenTelemetryContext::context(&mut extensions, &weak_dispatch.upgrade())
-            {
-                self.analyze_span_context(attrs.metadata().name(), &otel_context);
-            } else {
-                println!(
-                    "⚠️  Could not extract OpenTelemetry context for span '{}'",
-                    attrs.metadata().name()
-                );
+            if let Some(dispatch) = weak_dispatch.upgrade() {
+                if let Some(otel_context) = get_otel_context(&mut extensions, &dispatch) {
+                    self.analyze_span_context(attrs.metadata().name(), &otel_context);
+                } else {
+                    println!(
+                        "⚠️  Could not extract OpenTelemetry context for span '{}'",
+                        attrs.metadata().name()
+                    );
+                }
             }
         }
     }
 
     fn on_enter(&self, id: &tracing::span::Id, ctx: Context<'_, S>) {
-        if let Some(weak_dispatch) = self.get_weak_dispatch(true) {
+        if let Some(weak_dispatch) = self.dispatch.get() {
             if let Some(span_ref) = ctx.span(id) {
                 let mut extensions = span_ref.extensions_mut();
-                if let Some(otel_context) =
-                    OpenTelemetryContext::context(&mut extensions, &weak_dispatch.upgrade())
-                {
-                    let span = otel_context.span();
-                    let span_context = span.span_context();
-                    if span_context.is_valid() {
-                        println!(
-                            "📍 Entering span with trace_id: {:032x}, span_id: {:016x}",
-                            span_context.trace_id(),
-                            span_context.span_id()
-                        );
+                if let Some(dispatch) = weak_dispatch.upgrade() {
+                    if let Some(otel_context) = get_otel_context(&mut extensions, &dispatch) {
+                        let span = otel_context.span();
+                        let span_context = span.span_context();
+                        if span_context.is_valid() {
+                            println!(
+                                "📍 Entering span with trace_id: {:032x}, span_id: {:016x}",
+                                span_context.trace_id(),
+                                span_context.span_id()
+                            );
+                        }
                     }
                 }
             }
         }
     }
-
-    fn on_register_dispatch(&self, subscriber: &tracing::Dispatch) {
-        // Note: This does not work for Layer until https://github.com/tokio-rs/tracing/pull/3379
-        // is merged and released, since `on_register_dispatch` is never called.
-        let mut dispatch = self.dispatch.write().unwrap();
-        *dispatch = Some(subscriber.clone().downgrade());
-    }
 }
 
 fn setup_tracing() -> (impl Subscriber, SdkTracerProvider, SpanAnalysisLayer) {
@@ -166,7 +129,7 @@ fn setup_tracing() -> (impl Subscriber, SdkTracerProvider, SpanAnalysisLayer) {
     let tracer = provider.tracer("span_ref_ext_example");
 
     // Create our custom analysis layer
-    let analysis_layer = SpanAnalysisLayer::new();
+    let analysis_layer = SpanAnalysisLayer::default();
 
     // Build the subscriber with multiple layers:
     // 1. OpenTelemetry layer for trace export
diff --git a/src/layer.rs b/src/layer.rs
@@ -112,11 +112,11 @@ pub(crate) struct WithContext {
 
     ///
     /// Ensures the given SpanId has been activated - that is, created in the OTel side of things,
-    /// and had its SpanBuilder consumed - and then provides access to the OtelData associated with it.
+    /// and had its SpanBuilder consumed - and then provides access to the OTel Context associated with it.
     ///
     #[allow(clippy::type_complexity)]
-    pub(crate) with_activated_context_extensions:
-        fn(&tracing::Dispatch, &mut ExtensionsMut<'_>, f: &mut dyn FnMut(&mut OtelData)),
+    pub(crate) with_activated_otel_context:
+        fn(&tracing::Dispatch, &mut ExtensionsMut<'_>, f: &mut dyn FnMut(&OtelContext)),
 }
 
 impl WithContext {
@@ -152,13 +152,13 @@ impl WithContext {
     /// and had its SpanBuilder consumed - and then provides access to the OtelData associated with it.
     ///
     #[allow(clippy::type_complexity)]
-    pub(crate) fn with_activated_context_extensions(
+    pub(crate) fn with_activated_otel_context(
         &self,
         dispatch: &tracing::Dispatch,
         extensions: &mut ExtensionsMut<'_>,
-        mut f: impl FnMut(&mut OtelData),
+        mut f: impl FnMut(&OtelContext),
     ) {
-        (self.with_activated_context_extensions)(dispatch, extensions, &mut f)
+        (self.with_activated_otel_context)(dispatch, extensions, &mut f)
     }
 }
 
@@ -652,7 +652,7 @@ where
             with_context: WithContext {
                 with_context: Self::get_context,
                 with_activated_context: Self::get_activated_context,
-                with_activated_context_extensions: Self::get_activated_context_extensions,
+                with_activated_otel_context: Self::get_activated_otel_context,
             },
             _registry: marker::PhantomData,
         }
@@ -708,8 +708,8 @@ where
             with_context: WithContext {
                 with_context: OpenTelemetryLayer::<S, Tracer>::get_context,
                 with_activated_context: OpenTelemetryLayer::<S, Tracer>::get_activated_context,
-                with_activated_context_extensions:
-                    OpenTelemetryLayer::<S, Tracer>::get_activated_context_extensions,
+                with_activated_otel_context:
+                    OpenTelemetryLayer::<S, Tracer>::get_activated_otel_context,
             },
             _registry: self._registry,
             // cannot use ``..self` here due to different generics
@@ -967,6 +967,18 @@ where
         }
     }
 
+    /// Retrieves the OpenTelemetry data for a span and activates its context before calling
+    /// the provided function.
+    ///
+    /// This function retrieves the span from the subscriber's registry using the provided
+    /// span ID, activates the OTel `Context` in the span's `OtelData` if present, and then
+    /// applies the callback function `f` to the `OtelData`.
+    ///
+    /// # Parameters
+    ///
+    /// * `dispatch` - The tracing dispatch to downcast and retrieve the span from
+    /// * `id` - The span ID to look up in the registry
+    /// * `f` - The closure to invoke with mutable access to the span's `OtelData`
     fn get_activated_context(
         dispatch: &tracing::Dispatch,
         id: &span::Id,
@@ -984,6 +996,34 @@ where
         Self::get_activated_context_extensions(dispatch, &mut extensions, f)
     }
 
+    /// Retrieves the activated OpenTelemetry context from a span's extensions and passes it
+    /// to the provided function.
+    ///
+    /// This method activates the context and extracts the current OTel `Context` from the
+    /// `OtelData` state if present, and then applies the callback function `f` to a reference
+    /// to the OTel `Context`.
+    ///
+    /// # Parameters
+    ///
+    /// * `dispatch` - The tracing dispatch to downcast to the `OpenTelemetryLayer`
+    /// * `extensions` - Mutable reference to the span's extensions containing `OtelData`
+    /// * `f` - The closure to invoke with a reference to the OTel `Context`
+    fn get_activated_otel_context(
+        dispatch: &tracing::Dispatch,
+        extensions: &mut ExtensionsMut<'_>,
+        f: &mut dyn FnMut(&OtelContext),
+    ) {
+        Self::get_activated_context_extensions(
+            dispatch,
+            extensions,
+            &mut |otel_data: &mut OtelData| {
+                if let OtelDataState::Context { current_cx } = &otel_data.state {
+                    f(current_cx)
+                }
+            },
+        );
+    }
+
     fn get_activated_context_extensions(
         dispatch: &tracing::Dispatch,
         extensions: &mut ExtensionsMut<'_>,
diff --git a/src/lib.rs b/src/lib.rs
@@ -115,7 +115,7 @@ mod metrics;
 
 /// Implementation of the trace::Layer as a source of OpenTelemetry data.
 mod layer;
-/// OpenTelemetryContext which enables OpenTelemetry context extraction from span extensions.
+/// Function which enables OpenTelemetry context extraction from span extensions.
 mod otel_context;
 /// Span extension which enables OpenTelemetry context management.
 mod span_ext;
@@ -129,7 +129,7 @@ pub use layer::{layer, FilteredOpenTelemetryLayer, OpenTelemetryLayer};
 #[cfg(feature = "metrics")]
 pub use metrics::MetricsLayer;
 use opentelemetry::trace::TraceContextExt as _;
-pub use otel_context::OpenTelemetryContext;
+pub use otel_context::get_otel_context;
 pub use span_ext::{OpenTelemetrySpanExt, SetParentError};
 
 /// Per-span OpenTelemetry data tracked by this crate.
diff --git a/src/otel_context.rs b/src/otel_context.rs
diff --git a/tests/otel_context.rs b/tests/otel_context.rs
