feat: allow OpenTelemetry context access from SpanRef

Author: bantonsson

examples/span_ref_ext.rs

@@ -0,0 +1,278 @@
+//! This example demonstrates how to use `OpenTelemetrySpanRefExt` from a separate layer
+//! to access OpenTelemetry context data from tracing spans.
+
+use opentelemetry::trace::{TraceContextExt, TracerProvider as _};
+use opentelemetry_sdk::trace::SdkTracerProvider;
+use opentelemetry_stdout as stdout;
+use std::sync::{Arc, Mutex, RwLock};
+use tracing::{debug, info, span, warn, Subscriber};
+use tracing::{dispatcher::WeakDispatch, level_filters::LevelFilter, Dispatch};
+use tracing_opentelemetry::{layer, OpenTelemetrySpanRefExt};
+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 OpenTelemetrySpanRefExt
+/// to extract OpenTelemetry contexts from span references.
+#[derive(Clone)]
+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>>>,
+}
+
+#[derive(Debug, Clone)]
+struct SpanAnalysis {
+    span_name: String,
+    trace_id: String,
+    span_id: String,
+    is_root_span: bool,
+}
+
+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()
+    }
+
+    fn analyze_span_context(&self, span_name: &str, otel_context: &opentelemetry::Context) {
+        let span = otel_context.span();
+        let span_context = span.span_context();
+
+        if span_context.is_valid() {
+            let analysis = SpanAnalysis {
+                span_name: span_name.to_string(),
+                trace_id: format!("{:032x}", span_context.trace_id()),
+                span_id: format!("{:016x}", span_context.span_id()),
+                is_root_span: span_context.trace_flags().is_sampled(),
+            };
+
+            println!(
+                "🔍 Analyzing span '{}': trace_id={}, span_id={}, sampled={}",
+                analysis.span_name,
+                analysis.trace_id,
+                analysis.span_id,
+                span_context.trace_flags().is_sampled()
+            );
+
+            if let Ok(mut results) = self.analysis_results.lock() {
+                results.push(analysis);
+            }
+        }
+    }
+}
+
+impl<S> Layer<S> for SpanAnalysisLayer
+where
+    S: Subscriber + for<'span> LookupSpan<'span>,
+{
+    fn on_new_span(
+        &self,
+        attrs: &tracing::span::Attributes<'_>,
+        id: &tracing::span::Id,
+        ctx: Context<'_, S>,
+    ) {
+        // Get the weak dispatch reference
+        let weak_dispatch = {
+            let read_guard = self.dispatch.read().unwrap();
+            match read_guard.as_ref() {
+                Some(weak_dispatch) => 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 => {
+                    drop(read_guard);
+                    let mut dispatch = self.dispatch.write().unwrap();
+                    let weak_dispatch = Dispatch::default().downgrade();
+                    *dispatch = Some(weak_dispatch.clone());
+                    weak_dispatch
+                }
+            }
+        };
+
+        // Get the span reference and extract OpenTelemetry context
+        if let Some(span_ref) = ctx.span(id) {
+            // This is the key functionality: using OpenTelemetrySpanRefExt
+            // to extract the OpenTelemetry context from a span reference
+            if let Some(otel_context) = span_ref.context(&weak_dispatch.upgrade()) {
+                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.dispatch.read().unwrap().as_ref().map(|wd| wd.clone()) {
+            if let Some(span_ref) = ctx.span(id) {
+                if let Some(otel_context) = span_ref.context(&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()
+                        );
+                    }
+                }
+            }
+        }
+    }
+
+    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) {
+    // Create OpenTelemetry tracer that outputs to stdout
+    let provider = SdkTracerProvider::builder()
+        .with_simple_exporter(stdout::SpanExporter::default())
+        .build();
+    let tracer = provider.tracer("span_ref_ext_example");
+
+    // Create our custom analysis layer
+    let analysis_layer = SpanAnalysisLayer::new();
+
+    // Build the subscriber with multiple layers:
+    // 1. OpenTelemetry layer for trace export
+    // 2. Our custom analysis layer that uses SpanRefExt
+    // 3. Formatting layer for console output
+    let subscriber = tracing_subscriber::registry()
+        .with(layer().with_tracer(tracer).with_filter(LevelFilter::DEBUG))
+        .with(analysis_layer.clone())
+        .with(
+            tracing_subscriber::fmt::layer()
+                .with_target(false)
+                .with_filter(LevelFilter::INFO),
+        );
+
+    (subscriber, provider, analysis_layer)
+}
+
+fn simulate_application_work() {
+    // Create a root span for the main application work
+    let root_span = span!(tracing::Level::INFO, "application_main", version = "1.0.0");
+    let _root_guard = root_span.enter();
+
+    info!("Starting application");
+
+    // Simulate some business logic with nested spans
+    {
+        let auth_span = span!(tracing::Level::DEBUG, "authenticate_user", user_id = 12345);
+        let _auth_guard = auth_span.enter();
+
+        debug!("Validating user credentials");
+
+        // Simulate authentication work
+        std::thread::sleep(std::time::Duration::from_millis(10));
+
+        info!("User authenticated successfully");
+    }
+
+    // Simulate database operations
+    {
+        let db_span = span!(
+            tracing::Level::DEBUG,
+            "database_query",
+            query = "SELECT * FROM users",
+            table = "users"
+        );
+        let _db_guard = db_span.enter();
+
+        debug!("Executing database query");
+
+        // Nested span for connection management
+        {
+            let conn_span = span!(tracing::Level::DEBUG, "acquire_connection", pool_size = 10);
+            let _conn_guard = conn_span.enter();
+
+            debug!("Acquiring database connection from pool");
+            std::thread::sleep(std::time::Duration::from_millis(5));
+        }
+
+        std::thread::sleep(std::time::Duration::from_millis(20));
+        info!("Database query completed");
+    }
+
+    // Simulate some processing work
+    {
+        let process_span = span!(
+            tracing::Level::DEBUG,
+            "process_data",
+            records_count = 150,
+            batch_size = 50
+        );
+        let _process_guard = process_span.enter();
+
+        debug!("Processing user data");
+
+        for batch in 1..=3 {
+            let batch_span = span!(
+                tracing::Level::DEBUG,
+                "process_batch",
+                batch_number = batch,
+                batch_size = 50
+            );
+            let _batch_guard = batch_span.enter();
+
+            debug!("Processing batch {}", batch);
+            std::thread::sleep(std::time::Duration::from_millis(8));
+        }
+
+        info!("Data processing completed");
+    }
+
+    warn!("Application work completed");
+}
+
+fn main() {
+    println!("🚀 SpanRefExt Example: Extracting OpenTelemetry Contexts from Separate Layer");
+    println!("{}", "=".repeat(80));
+
+    // Setup tracing with our custom layer
+    let (subscriber, provider, analysis_layer) = setup_tracing();
+
+    tracing::subscriber::with_default(subscriber, || {
+        // Simulate application work that generates spans
+        simulate_application_work();
+    });
+
+    // Ensure all spans are flushed
+    drop(provider);
+
+    // Display the analysis results
+    println!("\n📊 Span Analysis Results:");
+    println!("{}", "-".repeat(80));
+
+    let results = analysis_layer.get_analysis_results();
+    for (i, analysis) in results.iter().enumerate() {
+        println!(
+            "{}. Span: '{}'\n   Trace ID: {}\n   Span ID: {}\n   Root Span: {}\n",
+            i + 1,
+            analysis.span_name,
+            analysis.trace_id,
+            analysis.span_id,
+            analysis.is_root_span
+        );
+    }
+
+    println!(
+        "✅ Example completed! Total spans analyzed: {}",
+        results.len()
+    );
+}

src/lib.rs

@@ -117,6 +117,8 @@ mod metrics;
 mod layer;
 /// Span extension which enables OpenTelemetry context management.
 mod span_ext;
+/// SpanRef extension which enables OpenTelemetry context extraction.
+mod span_ref_ext;
 
 mod stack;
 
@@ -127,6 +129,7 @@ pub use layer::{layer, OpenTelemetryLayer};
 #[cfg(feature = "metrics")]
 pub use metrics::MetricsLayer;
 pub use span_ext::OpenTelemetrySpanExt;
+pub use span_ref_ext::OpenTelemetrySpanRefExt;
 
 /// Per-span OpenTelemetry data tracked by this crate.
 #[derive(Debug)]

src/span_ext.rs

@@ -7,6 +7,8 @@ use opentelemetry::{
 };
 use std::{borrow::Cow, time::SystemTime};
 use thiserror::Error;
+use tracing::span::Id;
+use tracing::Dispatch;
 
 /// Utility functions to allow tracing [`Span`]s to accept and return
 /// [OpenTelemetry] [`Context`]s.
@@ -313,20 +315,10 @@ impl OpenTelemetrySpanExt for tracing::Span {
     }
 
     fn context(&self) -> Context {
-        let mut cx = None;
         self.with_subscriber(|(id, subscriber)| {
-            let Some(get_context) = subscriber.downcast_ref::<WithContext>() else {
-                return;
-            };
-            // If our span hasn't been built, we should build it and get the context in one call
-            get_context.with_activated_context(subscriber, id, |data: &mut OtelData| {
-                if let OtelDataState::Context { current_cx } = &data.state {
-                    cx = Some(current_cx.clone());
-                }
-            });
-        });
-
-        cx.unwrap_or_default()
+            get_otel_context(subscriber, id).unwrap_or_default()
+        })
+        .unwrap_or_default()
     }
 
     fn set_attribute(&self, key: impl Into<Key>, value: impl Into<Value>) {
@@ -415,3 +407,16 @@ impl OpenTelemetrySpanExt for tracing::Span {
         });
     }
 }
+
+pub(crate) fn get_otel_context(dispatch: &Dispatch, id: &Id) -> Option<opentelemetry::Context> {
+    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(dispatch, id, |data: &mut OtelData| {
+            if let OtelDataState::Context { current_cx } = &data.state {
+                cx = Some(current_cx.clone());
+            }
+        });
+    }
+    cx
+}