chore: update SpanRef context access based on feedback

bantonsson · bantonsson · commit bda9f4dab859 · 2025-12-01T11:29:03.000+01:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -29,7 +29,8 @@ 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 = [
+# The 0.3.22 version contains the fix for on_register_dispatch in Layer
+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,8 +86,7 @@ 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())
+            if let Some(otel_context) = get_otel_context(&mut extensions, &weak_dispatch.upgrade())
             {
                 self.analyze_span_context(attrs.metadata().name(), &otel_context);
             } else {
@@ -130,11 +99,11 @@ where
     }
 
     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())
+                    get_otel_context(&mut extensions, &weak_dispatch.upgrade())
                 {
                     let span = otel_context.span();
                     let span_context = span.span_context();
@@ -149,13 +118,6 @@ where
             }
         }
     }
-
-    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 +128,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
@@ -115,8 +115,8 @@ pub(crate) struct WithContext {
     /// and had its SpanBuilder consumed - and then provides access to the OtelData 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
@@ -984,6 +984,22 @@ where
         Self::get_activated_context_extensions(dispatch, &mut extensions, f)
     }
 
+    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
@@ -1,4 +1,4 @@
-use crate::{layer::WithContext, OtelData, OtelDataState};
+use crate::layer::WithContext;
 use tracing::Dispatch;
 use tracing_subscriber::registry::ExtensionsMut;
 
@@ -8,65 +8,61 @@ use tracing_subscriber::registry::ExtensionsMut;
 /// [`ExtensionsMut`]: tracing_subscriber::registry::ExtensionsMut
 /// [OpenTelemetry]: https://opentelemetry.io
 /// [`Context`]: opentelemetry::Context
-pub struct OpenTelemetryContext {}
-impl OpenTelemetryContext {
-    /// Extracts the OpenTelemetry [`Context`] associated with this span extensions.
-    ///
-    /// This method retrieves the OpenTelemetry context data that has been stored
-    /// for the span by the OpenTelemetry layer. The context includes the span's
-    /// OpenTelemetry span context, which contains trace ID, span ID, and other
-    /// trace-related metadata.
-    ///
-    /// [`Context`]: opentelemetry::Context
-    ///
-    /// # Examples
-    ///
-    /// ```rust
-    /// use tracing_opentelemetry::OpenTelemetryContext;
-    /// use tracing::dispatcher::WeakDispatch;
-    /// use tracing_subscriber::registry::LookupSpan;
-    /// use opentelemetry::trace::TraceContextExt;
-    ///
-    /// fn do_things_with_otel_context<'a, D>(
-    ///     span_ref: &tracing_subscriber::registry::SpanRef<'a, D>,
-    ///     weak_dispatch: &WeakDispatch
-    /// ) where
-    ///     D: LookupSpan<'a>,
-    /// {
-    ///     if let Some(otel_context) = OpenTelemetryContext::context(&mut span_ref.extensions_mut(), &weak_dispatch.upgrade()) {
-    ///         // Process the extracted context
-    ///         let span = otel_context.span();
-    ///         let span_context = span.span_context();
-    ///         if span_context.is_valid() {
-    ///             // Handle the valid context...
-    ///         }
-    ///     }
-    /// }
-    /// ```
-    ///
-    /// # Use Cases
-    ///
-    /// - When working with multiple subscriber configurations
-    /// - When implementing advanced tracing middleware that manages multiple dispatches
-    pub fn context(
-        extensions: &mut ExtensionsMut<'_>,
-        dispatch: &Option<Dispatch>,
-    ) -> Option<opentelemetry::Context> {
-        dispatch.as_ref().and_then(|dispatch| {
-            let mut cx = None;
-            if let Some(get_context) = dispatch.downcast_ref::<WithContext>() {
-                // If our span hasn't been built, we should build it and get the context in one call
-                get_context.with_activated_context_extensions(
-                    dispatch,
-                    extensions,
-                    |data: &mut OtelData| {
-                        if let OtelDataState::Context { current_cx } = &data.state {
-                            cx = Some(current_cx.clone());
-                        }
-                    },
-                );
-            }
-            cx
-        })
-    }
+///
+/// Extracts the OpenTelemetry [`Context`] associated with this span extensions.
+///
+/// This method retrieves the OpenTelemetry context data that has been stored
+/// for the span by the OpenTelemetry layer. The context includes the span's
+/// OpenTelemetry span context, which contains trace ID, span ID, and other
+/// trace-related metadata.
+///
+/// [`Context`]: opentelemetry::Context
+///
+/// # Examples
+///
+/// ```rust
+/// use tracing_opentelemetry::get_otel_context;
+/// use tracing::dispatcher::WeakDispatch;
+/// use tracing_subscriber::registry::LookupSpan;
+/// use opentelemetry::trace::TraceContextExt;
+///
+/// fn do_things_with_otel_context<'a, D>(
+///     span_ref: &tracing_subscriber::registry::SpanRef<'a, D>,
+///     weak_dispatch: &WeakDispatch
+/// ) where
+///     D: LookupSpan<'a>,
+/// {
+///     if let Some(otel_context) = get_otel_context(&mut span_ref.extensions_mut(), &weak_dispatch.upgrade()) {
+///         // Process the extracted context
+///         let span = otel_context.span();
+///         let span_context = span.span_context();
+///         if span_context.is_valid() {
+///             // Handle the valid context...
+///         }
+///     }
+/// }
+/// ```
+///
+/// # Use Cases
+///
+/// - When working with multiple subscriber configurations
+/// - When implementing advanced tracing middleware that manages multiple dispatches
+pub fn get_otel_context(
+    extensions: &mut ExtensionsMut<'_>,
+    dispatch: &Option<Dispatch>,
+) -> Option<opentelemetry::Context> {
+    dispatch.as_ref().and_then(|dispatch| {
+        let mut cx = None;
+        if let Some(get_context) = dispatch.downcast_ref::<WithContext>() {
+            // If our span hasn't been built, we should build it and get the context in one call
+            get_context.with_activated_otel_context(
+                dispatch,
+                extensions,
+                |current_cx: &opentelemetry::Context| {
+                    cx = Some(current_cx.clone());
+                },
+            );
+        }
+        cx
+    })
 }
diff --git a/tests/otel_context.rs b/tests/otel_context.rs
