add LogfireTracingLayer as a public type (#61)

* move `LogfireTracer` into `LogfireTracingLayer`

* add `LogfireTracingLayer` as a public type

Author: davidhewitt

src/bridges/tracing.rs

@@ -1,4 +1,4 @@
-use std::time::SystemTime;
+use std::{any::TypeId, time::SystemTime};
 
 use opentelemetry::{
     KeyValue,
@@ -9,11 +9,35 @@ use tracing::{Subscriber, field::Visit};
 use tracing_opentelemetry::{OpenTelemetrySpanExt, OtelData, PreSampledTracer};
 use tracing_subscriber::{Layer, registry::LookupSpan};
 
-use crate::{LogfireTracer, try_with_logfire_tracer};
+use crate::LogfireTracer;
 
-pub(crate) struct LogfireTracingLayer(pub(crate) opentelemetry_sdk::trace::Tracer);
+/// A `tracing` layer that bridges `tracing` spans to OpenTelemetry spans using the Logfire tracer.
+///
+/// This layer is a wrapper around `tracing_opentelemetry::OpenTelemetryLayer` that adds additional
+/// Logfire-specific metadata.
+///
+/// See [`ShutdownHandler::tracing_layer`][crate::ShutdownHandler::tracing_layer] for how to use
+/// this layer.
+pub struct LogfireTracingLayer<S> {
+    tracer: LogfireTracer,
+    otel_layer: tracing_opentelemetry::OpenTelemetryLayer<S, opentelemetry_sdk::trace::Tracer>,
+}
 
-impl<S> Layer<S> for LogfireTracingLayer
+impl<S> LogfireTracingLayer<S>
+where
+    S: Subscriber + for<'span> LookupSpan<'span>,
+{
+    /// Create a new `LogfireTracingLayer` with the given tracer.
+    pub(crate) fn new(tracer: LogfireTracer) -> Self {
+        let otel_layer = tracing_opentelemetry::layer()
+            .with_error_records_to_exceptions(true)
+            .with_tracer(tracer.inner.clone());
+
+        LogfireTracingLayer { tracer, otel_layer }
+    }
+}
+
+impl<S> Layer<S> for LogfireTracingLayer<S>
 where
     S: Subscriber + for<'span> LookupSpan<'span>,
 {
@@ -23,6 +47,10 @@ where
         id: &tracing::span::Id,
         ctx: tracing_subscriber::layer::Context<'_, S>,
     ) {
+        // Delegate to OpenTelemetry layer first
+        self.otel_layer.on_new_span(attrs, id, ctx.clone());
+
+        // Add Logfire-specific attributes
         let span = ctx.span(id).expect("span not found");
         let mut extensions = span.extensions_mut();
         if let Some(otel_data) = extensions.get_mut::<OtelData>() {
@@ -46,6 +74,9 @@ where
     ///
     /// e.g. <https://github.com/davidB/tracing-opentelemetry-instrumentation-sdk/blob/5830c9113b0d42b72167567bf8e5f4c6b20933c8/axum-tracing-opentelemetry/src/middleware/trace_extractor.rs#L132>
     fn on_enter(&self, id: &tracing::span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
+        // Delegate to OpenTelemetry layer first
+        self.otel_layer.on_enter(id, ctx.clone());
+
         let span = ctx.span(id).expect("span not found");
         let mut extensions = span.extensions_mut();
 
@@ -58,7 +89,7 @@ where
         // Guaranteed to be on first entering of the span
         if let Some(otel_data) = extensions.get_mut::<OtelData>() {
             // Emit a pending span, if this span will be sampled.
-            let context = self.0.sampled_context(otel_data);
+            let context = self.tracer.inner.sampled_context(otel_data);
             let sampling_result = otel_data
                 .builder
                 .sampling_result
@@ -110,14 +141,15 @@ where
                     ));
                 }
 
-                pending_span_builder.span_id = Some(self.0.new_span_id());
+                pending_span_builder.span_id = Some(self.tracer.inner.new_span_id());
 
                 let start_time = pending_span_builder
                     .start_time
                     .expect("otel SDK sets start time");
 
                 // emit pending span
-                let mut pending_span = pending_span_builder.start_with_context(&self.0, &context);
+                let mut pending_span =
+                    pending_span_builder.start_with_context(&self.tracer.inner, &context);
                 pending_span.end_with_timestamp(start_time);
             }
         }
@@ -131,39 +163,59 @@ where
         event: &tracing::Event<'_>,
         _ctx: tracing_subscriber::layer::Context<'_, S>,
     ) {
-        try_with_logfire_tracer(|tracer| {
-            // All events are emitted as log spans
-            emit_event_as_log_span(tracer, event, &tracing::Span::current());
-        });
+        // Don't delegate events to OpenTelemetry layer, we emit them as log spans instead.
+        // FIXME: can we get current span from `ctx`?
+        emit_event_as_log_span(&self.tracer, event, &tracing::Span::current());
     }
-}
 
-/// Helper to print spans when dropped; if it was never entered then the pending span
-/// is never sent (the console writer uses pending spans).
-///
-/// This needs to be a separate layer so that it can access the `OtelData` before the
-/// `tracing_opentelemetry` layer removes it.
-pub struct LogfireTracingPendingSpanNotSentLayer;
+    fn on_exit(&self, id: &tracing::span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
+        self.otel_layer.on_exit(id, ctx);
+    }
 
-impl<S> Layer<S> for LogfireTracingPendingSpanNotSentLayer
-where
-    S: Subscriber + for<'span> LookupSpan<'span>,
-{
     fn on_close(&self, id: tracing::span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
         let span = ctx.span(&id).expect("span not found");
-        let mut extensions = span.extensions_mut();
-
-        if extensions.get_mut::<LogfirePendingSpanSent>().is_some() {
-            return;
-        }
+        let extensions = span.extensions();
 
-        // Guaranteed to be on first entering of the span
-        if let Some(otel_data) = extensions.get_mut::<OtelData>() {
-            try_with_logfire_tracer(|tracer| {
-                if let Some(writer) = &tracer.console_writer {
+        // We write pending spans to the console; if the pending span was never created then
+        // we have to manually write it now.
+        if extensions.get::<LogfirePendingSpanSent>().is_none() {
+            if let Some(otel_data) = extensions.get::<OtelData>() {
+                if let Some(writer) = &self.tracer.console_writer {
                     writer.write_tracing_opentelemetry_data(otel_data);
                 }
-            });
+            }
+        }
+
+        // Delegate to OpenTelemetry layer after handling pending span (it will remove the
+        // `OtelData` so cannot do before).
+        drop(extensions);
+        drop(span);
+        self.otel_layer.on_close(id, ctx);
+    }
+
+    fn on_follows_from(
+        &self,
+        span: &tracing::span::Id,
+        follows: &tracing::span::Id,
+        ctx: tracing_subscriber::layer::Context<'_, S>,
+    ) {
+        self.otel_layer.on_follows_from(span, follows, ctx);
+    }
+
+    fn on_record(
+        &self,
+        span: &tracing::span::Id,
+        values: &tracing::span::Record<'_>,
+        ctx: tracing_subscriber::layer::Context<'_, S>,
+    ) {
+        self.otel_layer.on_record(span, values, ctx);
+    }
+
+    unsafe fn downcast_raw(&self, id: TypeId) -> Option<*const ()> {
+        if id == TypeId::of::<Self>() {
+            Some(std::ptr::from_ref(self).cast())
+        } else {
+            unsafe { self.otel_layer.downcast_raw(id) }
         }
     }
 }

src/lib.rs

@@ -102,7 +102,6 @@ use std::panic::PanicHookInfo;
 use std::sync::{Arc, Once};
 use std::{backtrace::Backtrace, env::VarError, sync::OnceLock, time::Duration};
 
-use bridges::tracing::LogfireTracingPendingSpanNotSentLayer;
 use config::get_base_url_from_token;
 use opentelemetry::trace::TracerProvider;
 use opentelemetry_sdk::metrics::{PeriodicReader, SdkMeterProvider};
@@ -113,7 +112,8 @@ use opentelemetry_sdk::trace::{SdkTracerProvider, Tracer};
 use thiserror::Error;
 use tracing::Subscriber;
 use tracing::level_filters::LevelFilter;
-use tracing_subscriber::layer::{Layer, SubscriberExt};
+use tracing_subscriber::layer::SubscriberExt;
+use tracing_subscriber::registry::LookupSpan;
 
 use crate::bridges::tracing::LogfireTracingLayer;
 use crate::config::{
@@ -571,18 +571,15 @@ impl LogfireConfigBuilder {
             .with_default_directive(default_level_filter.into())
             .from_env()?; // but allow the user to override this with `RUST_LOG`
 
+        let tracer = LogfireTracer {
+            inner: tracer,
+            handle_panics: self.install_panic_handler,
+            console_writer,
+        };
+
         let subscriber = tracing_subscriber::registry()
             .with(filter)
-            .with(LogfireTracingPendingSpanNotSentLayer)
-            .with(
-                tracing_opentelemetry::layer()
-                    .with_error_records_to_exceptions(true)
-                    .with_tracer(tracer.clone())
-                    .with_filter(tracing_subscriber::filter::filter_fn(|metadata| {
-                        !metadata.is_event()
-                    })),
-            )
-            .with(LogfireTracingLayer(tracer.clone()));
+            .with(LogfireTracingLayer::new(tracer.clone()));
 
         let mut meter_provider_builder = SdkMeterProvider::builder();
 
@@ -616,11 +613,7 @@ impl LogfireConfigBuilder {
 
         Ok(LogfireParts {
             local: self.local,
-            tracer: LogfireTracer {
-                inner: tracer,
-                handle_panics: self.install_panic_handler,
-                console_writer,
-            },
+            tracer,
             subscriber: Arc::new(subscriber),
             tracer_provider,
             meter_provider,
@@ -661,6 +654,36 @@ impl ShutdownHandler {
             .map_err(|e| ConfigureError::Other(e.into()))?;
         Ok(())
     }
+
+    /// Get a tracing layer which can be used to embed this `Logfire` instance into a `tracing_subscriber::Registry`.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use tracing_subscriber::{Registry, layer::SubscriberExt};
+    ///
+    /// let shutdown_handler = logfire::configure()
+    ///    .local()  // use local mode to avoid setting global state
+    ///    .finish()
+    ///    .expect("Failed to configure logfire");
+    ///
+    /// let subscriber = tracing_subscriber::registry()
+    ///    .with(shutdown_handler.tracing_layer());
+    ///
+    /// tracing::subscriber::set_global_default(subscriber)
+    ///    .expect("Failed to set global subscriber");
+    ///
+    /// logfire::info!("Hello world");
+    ///
+    /// shutdown_handler.shutdown().expect("Failed to shutdown logfire");
+    /// ```
+    #[must_use]
+    pub fn tracing_layer<S>(&self) -> LogfireTracingLayer<S>
+    where
+        S: Subscriber + for<'span> LookupSpan<'span>,
+    {
+        LogfireTracingLayer::new(self.tracer.clone())
+    }
 }
 
 struct LogfireParts {

tests/test_basic_exports.rs

@@ -1034,7 +1034,7 @@ fn test_basic_span() {
                         "code.lineno",
                     ),
                     value: I64(
-                        693,
+                        716,
                     ),
                 },
                 KeyValue {