prevent panics when used with per-layer filters (#86)

## Motivation

Currently, if the `tracing-opentelemetry` `Layer` has a per-layer filter
(in these examples, ones that allow only `>=DEBUG`) the following two
cases behave differently:

```
let root = tracing::trace_span!("root");
tracing::debug_span!(parent: &root, "child");
```

```
let root = tracing::trace_span!("root");
oot.in_scope(|| tracing::debug_span!("child"));
```

The former panics, and the latter successfully creates a single-span
trace. Note that this ONLY happens if there is another layer interested
in the `TRACE` level (otherwise, the tracing current-level-filter
fast-path filters the root spans out at their callsites).

This may seem rare, but it becomes more common when the per-layer filter
is a
[reload](https://docs.rs/tracing-subscriber/latest/tracing_subscriber/reload/index.html)
filter, which means the parent and child spans can be differently
filtered if the filter is reloaded between their creation (example:
https://github.com/MaterializeInc/materialize/issues/15223)

## Solution

Handle this case gracefully. I also did the same in `on_follows_from`

Author: guswynn

src/layer.rs

@@ -759,16 +759,30 @@ where
     /// [`span`]: tracing::Span
     /// [`Registry`]: tracing_subscriber::Registry
     fn parent_context(&self, attrs: &Attributes<'_>, ctx: &Context<'_, S>) -> OtelContext {
-        // If a span is specified, it _should_ exist in the underlying `Registry`.
         if let Some(parent) = attrs.parent() {
-            let span = ctx.span(parent).expect("Span not found, this is a bug");
-            let mut extensions = span.extensions_mut();
-            extensions
-                .get_mut::<OtelData>()
-                .map(|builder| self.tracer.sampled_context(builder))
-                .unwrap_or_default()
+            // A span can have an _explicit_ parent that is NOT seen by this `Layer` (for which
+            // `Context::span` returns `None`. This happens if the parent span is filtered away
+            // from the layer by a per-layer filter. In that case, we fall-through to the `else`
+            // case, and consider this span a root span.
+            //
+            // This is likely rare, as most users who use explicit parents will configure their
+            // filters so that children and parents are both seen, but its not guaranteed. Also,
+            // if users configure their filter with a `reload` filter, its possible that a parent
+            // and child have different filters as they are created with a filter change
+            // in-between.
+            //
+            // In these case, we prefer to emit a smaller span tree instead of panicking.
+            if let Some(span) = ctx.span(parent) {
+                let mut extensions = span.extensions_mut();
+                return extensions
+                    .get_mut::<OtelData>()
+                    .map(|builder| self.tracer.sampled_context(builder))
+                    .unwrap_or_default();
+            }
+        }
+
         // Else if the span is inferred from context, look up any available current span.
-        } else if attrs.is_contextual() {
+        if attrs.is_contextual() {
             ctx.lookup_current()
                 .and_then(|span| {
                     let mut extensions = span.extensions_mut();
@@ -948,25 +962,27 @@ where
             .get_mut::<OtelData>()
             .expect("Missing otel data span extensions");
 
-        let follows_span = ctx
-            .span(follows)
-            .expect("Span to follow not found, this is a bug");
-        let mut follows_extensions = follows_span.extensions_mut();
-        let follows_data = follows_extensions
-            .get_mut::<OtelData>()
-            .expect("Missing otel data span extensions");
-
-        let follows_context = self
-            .tracer
-            .sampled_context(follows_data)
-            .span()
-            .span_context()
-            .clone();
-        let follows_link = otel::Link::new(follows_context, Vec::new());
-        if let Some(ref mut links) = data.builder.links {
-            links.push(follows_link);
-        } else {
-            data.builder.links = Some(vec![follows_link]);
+        // The follows span may be filtered away (or closed), from this layer,
+        // in which case we just drop the data, as opposed to panicking. This
+        // uses the same reasoning as `parent_context` above.
+        if let Some(follows_span) = ctx.span(follows) {
+            let mut follows_extensions = follows_span.extensions_mut();
+            let follows_data = follows_extensions
+                .get_mut::<OtelData>()
+                .expect("Missing otel data span extensions");
+
+            let follows_context = self
+                .tracer
+                .sampled_context(follows_data)
+                .span()
+                .span_context()
+                .clone();
+            let follows_link = otel::Link::new(follows_context, Vec::new());
+            if let Some(ref mut links) = data.builder.links {
+                links.push(follows_link);
+            } else {
+                data.builder.links = Some(vec![follows_link]);
+            }
         }
     }
 

tests/filtered.rs

@@ -0,0 +1,67 @@
+use futures_util::future::BoxFuture;
+use opentelemetry::trace::TracerProvider as _;
+use opentelemetry_sdk::{
+    export::trace::{ExportResult, SpanData, SpanExporter},
+    trace::{Tracer, TracerProvider},
+};
+use std::sync::{Arc, Mutex};
+use tracing::level_filters::LevelFilter;
+use tracing::Subscriber;
+use tracing_opentelemetry::layer;
+use tracing_subscriber::prelude::*;
+
+#[derive(Clone, Default, Debug)]
+struct TestExporter(Arc<Mutex<Vec<SpanData>>>);
+
+impl SpanExporter for TestExporter {
+    fn export(&mut self, mut batch: Vec<SpanData>) -> BoxFuture<'static, ExportResult> {
+        let spans = self.0.clone();
+        Box::pin(async move {
+            if let Ok(mut inner) = spans.lock() {
+                inner.append(&mut batch);
+            }
+            Ok(())
+        })
+    }
+}
+
+fn test_tracer() -> (Tracer, TracerProvider, TestExporter, impl Subscriber) {
+    let exporter = TestExporter::default();
+    let provider = TracerProvider::builder()
+        .with_simple_exporter(exporter.clone())
+        .build();
+    let tracer = provider.tracer("test");
+
+    let subscriber = tracing_subscriber::registry()
+        .with(
+            layer()
+                .with_tracer(tracer.clone())
+                // DEBUG-level, so `trace_spans` are skipped.
+                .with_filter(LevelFilter::DEBUG),
+        )
+        // This is REQUIRED so that the tracing fast path doesn't filter
+        // out trace spans at their callsite.
+        .with(tracing_subscriber::fmt::layer().with_filter(LevelFilter::TRACE));
+
+    (tracer, provider, exporter, subscriber)
+}
+
+#[test]
+fn trace_filtered() {
+    let (_tracer, provider, exporter, subscriber) = test_tracer();
+
+    tracing::subscriber::with_default(subscriber, || {
+        // Neither of these should panic
+
+        let root = tracing::trace_span!("root");
+        tracing::debug_span!(parent: &root, "child");
+
+        let root = tracing::trace_span!("root");
+        root.in_scope(|| tracing::debug_span!("child"));
+    });
+
+    drop(provider); // flush all spans
+    let spans = exporter.0.lock().unwrap();
+    // Only the child spans are reported.
+    assert_eq!(spans.len(), 2);
+}

tests/follows_from.rs

@@ -0,0 +1,62 @@
+use futures_util::future::BoxFuture;
+use opentelemetry::trace::TracerProvider as _;
+use opentelemetry_sdk::{
+    export::trace::{ExportResult, SpanData, SpanExporter},
+    trace::{Tracer, TracerProvider},
+};
+use std::sync::{Arc, Mutex};
+use tracing::Subscriber;
+use tracing_opentelemetry::layer;
+use tracing_subscriber::prelude::*;
+
+#[derive(Clone, Default, Debug)]
+struct TestExporter(Arc<Mutex<Vec<SpanData>>>);
+
+impl SpanExporter for TestExporter {
+    fn export(&mut self, mut batch: Vec<SpanData>) -> BoxFuture<'static, ExportResult> {
+        let spans = self.0.clone();
+        Box::pin(async move {
+            if let Ok(mut inner) = spans.lock() {
+                inner.append(&mut batch);
+            }
+            Ok(())
+        })
+    }
+}
+
+fn test_tracer() -> (Tracer, TracerProvider, TestExporter, impl Subscriber) {
+    let exporter = TestExporter::default();
+    let provider = TracerProvider::builder()
+        .with_simple_exporter(exporter.clone())
+        .build();
+    let tracer = provider.tracer("test");
+
+    // Note that if we added a `with_filter` here, the original bug (issue #14) will
+    // not reproduce. This is because the `Filtered` layer will not
+    // call the `tracing-opentelemetry` `Layer`'s `on_follows_from`, as the
+    // closed followed span no longer exists in a way that can checked against
+    // the that `Filtered`'s filter.
+    let subscriber = tracing_subscriber::registry().with(layer().with_tracer(tracer.clone()));
+
+    (tracer, provider, exporter, subscriber)
+}
+
+#[test]
+fn trace_follows_from_closed() {
+    let (_tracer, provider, exporter, subscriber) = test_tracer();
+
+    tracing::subscriber::with_default(subscriber, || {
+        let f = tracing::debug_span!("f");
+        let f_id = f.id().unwrap();
+        drop(f);
+
+        let s = tracing::debug_span!("span");
+        // This should not panic
+        s.follows_from(f_id);
+    });
+
+    drop(provider); // flush all spans
+    let spans = exporter.0.lock().unwrap();
+    // Only the child spans are reported.
+    assert_eq!(spans.len(), 2);
+}