subscriber: change registry exit to decrement local span ref only (#3331)

## Motivation

Using with_subscriber will trigger a panic if the async block creates
a span & enters it on a different thread that will target the default
subscriber. I believe both the with_subscriber AND the enter off-thread
scenario are idiomatic. The reason for the crash is that the span.enter()
try's recording the span exit into the default subscriber which isn't
aware of it since it's registered by the scoped subscriber.

## Solution

From my understanding, it's sufficient for `Registry::exit` just to call
`self.try_close()` and do local bookkeeping:
- The attempt to use the thread dispatcher leads directly to the issues
  observed in #3223, and
- I don't think it's correct to call `try_close` on the whole `Layer` stack
  at this point anyway, a span being exited is not yet ready to close. All
  that is needed is to decrement the reference count within the current
  registry.

I've added a test which spawns a thread and enters (and exits, and drops)
a span created on a dispatcher not registered to that thread.

Before this patch, the test fails with the span never being closed
(because there is no thread dispatcher when the span is exited, and so a
reference is leaked in `Registry::exit`).

With this patch, the bookkeeping demonstrated in the test seems correct
to me.

Fixes: #3223

Author: davidhewitt

tracing-subscriber/src/registry/sharded.rs

@@ -301,7 +301,7 @@ impl Subscriber for Registry {
     fn exit(&self, id: &span::Id) {
         if let Some(spans) = self.current_spans.get() {
             if spans.borrow_mut().pop(id) {
-                dispatcher::get_default(|dispatch| dispatch.try_close(id.clone()));
+                self.try_close(id.clone());
             }
         }
     }

tracing-subscriber/tests/registry_span_instrument_drop.rs

@@ -0,0 +1,190 @@
+#![cfg(feature = "registry")]
+
+use std::sync::{Arc, Mutex};
+
+use tracing::{
+    span::{self, Id},
+    Dispatch, Event, Metadata, Subscriber,
+};
+use tracing_core::{Interest, LevelFilter};
+use tracing_subscriber::{
+    layer::{Context, SubscriberExt},
+    Layer, Registry,
+};
+
+#[test]
+fn span_entered_on_different_thread_from_subscriber() {
+    /// Counters for various lifecycle events we want to track.
+    #[derive(Default)]
+    struct LifecycleCounts {
+        layer_new_count: usize,
+        layer_enter_count: usize,
+        layer_exit_count: usize,
+        layer_close_count: usize,
+
+        sub_new_count: usize,
+        sub_clone_count: usize,
+        sub_enter_count: usize,
+        sub_exit_count: usize,
+        sub_close_count: usize,
+    }
+
+    /// Wraps `tracing_subscriber::Registry` and adds some accounting
+    /// to verify that the subscriber is receiving the expected number of calls.
+    struct CountingSubscriber {
+        inner: Registry,
+        counts: Arc<Mutex<LifecycleCounts>>,
+    }
+
+    // Forward all subscriber methods to the inner registry, adding counts where appropriate.
+    impl Subscriber for CountingSubscriber {
+        fn on_register_dispatch(&self, subscriber: &Dispatch) {
+            self.inner.on_register_dispatch(subscriber);
+        }
+
+        fn register_callsite(&self, metadata: &'static Metadata<'static>) -> Interest {
+            self.inner.register_callsite(metadata)
+        }
+
+        fn max_level_hint(&self) -> Option<LevelFilter> {
+            self.inner.max_level_hint()
+        }
+
+        fn event_enabled(&self, event: &Event<'_>) -> bool {
+            self.inner.event_enabled(event)
+        }
+
+        fn clone_span(&self, id: &span::Id) -> span::Id {
+            self.counts.lock().unwrap().sub_clone_count += 1;
+            self.inner.clone_span(id)
+        }
+
+        fn drop_span(&self, id: span::Id) {
+            self.counts.lock().unwrap().sub_close_count += 1;
+            #[allow(deprecated)]
+            self.inner.drop_span(id);
+        }
+
+        fn try_close(&self, id: span::Id) -> bool {
+            self.counts.lock().unwrap().sub_close_count += 1;
+            self.inner.try_close(id)
+        }
+
+        fn current_span(&self) -> tracing_core::span::Current {
+            self.inner.current_span()
+        }
+
+        unsafe fn downcast_raw(&self, id: std::any::TypeId) -> Option<*const ()> {
+            self.inner.downcast_raw(id)
+        }
+
+        fn enabled(&self, metadata: &Metadata<'_>) -> bool {
+            self.inner.enabled(metadata)
+        }
+
+        fn new_span(&self, span: &span::Attributes<'_>) -> span::Id {
+            self.counts.lock().unwrap().sub_new_count += 1;
+            self.inner.new_span(span)
+        }
+
+        fn record(&self, span: &span::Id, values: &span::Record<'_>) {
+            self.inner.record(span, values);
+        }
+
+        fn record_follows_from(&self, span: &span::Id, follows: &span::Id) {
+            self.inner.record_follows_from(span, follows);
+        }
+
+        fn event(&self, event: &Event<'_>) {
+            self.inner.event(event);
+        }
+
+        fn enter(&self, span: &span::Id) {
+            self.inner.enter(span);
+            self.counts.lock().unwrap().sub_enter_count += 1;
+        }
+
+        fn exit(&self, span: &span::Id) {
+            self.inner.exit(span);
+            self.counts.lock().unwrap().sub_exit_count += 1;
+        }
+    }
+
+    /// Similar to the above, but for a `Layer` which sits atop the subscriber.
+    struct CountingLayer {
+        counts: Arc<Mutex<LifecycleCounts>>,
+    }
+
+    // Just does bookkeeping where relevant.
+    impl Layer<CountingSubscriber> for CountingLayer {
+        fn on_new_span(
+            &self,
+            _attrs: &span::Attributes<'_>,
+            _id: &span::Id,
+            _ctx: Context<'_, CountingSubscriber>,
+        ) {
+            self.counts.lock().unwrap().layer_new_count += 1;
+        }
+
+        fn on_enter(&self, _id: &span::Id, _ctx: Context<'_, CountingSubscriber>) {
+            self.counts.lock().unwrap().layer_enter_count += 1;
+        }
+
+        fn on_exit(&self, _id: &span::Id, _ctx: Context<'_, CountingSubscriber>) {
+            self.counts.lock().unwrap().layer_exit_count += 1;
+        }
+
+        fn on_close(&self, _id: Id, _ctx: Context<'_, CountingSubscriber>) {
+            self.counts.lock().unwrap().layer_close_count += 1;
+        }
+    }
+
+    // Setup subscriber and layer.
+
+    let counts = Arc::new(Mutex::new(LifecycleCounts::default()));
+
+    let layer = CountingLayer {
+        counts: counts.clone(),
+    };
+
+    let subscriber = CountingSubscriber {
+        inner: tracing_subscriber::registry(),
+        counts: counts.clone(),
+    };
+    let subscriber = Arc::new(subscriber.with(layer));
+
+    // Create a span using the subscriber
+    let span = tracing::subscriber::with_default(subscriber.clone(), move || {
+        tracing::span!(tracing::Level::INFO, "span")
+    });
+
+    // Enter the span in a thread which doesn't have a direct relationship to the subscriber.
+    std::thread::spawn(move || {
+        let _ = span.entered();
+    })
+    .join()
+    .unwrap();
+
+    // layer should have seen exactly one new span & close
+    // should be one enter / exit cycle
+
+    let counts = counts.lock().unwrap();
+
+    assert_eq!(counts.layer_new_count, 1);
+    assert_eq!(counts.layer_enter_count, 1);
+    assert_eq!(counts.layer_exit_count, 1);
+    assert_eq!(counts.layer_close_count, 1);
+
+    // subscriber should have seen one new span
+    // new + any clones should equal number of closes
+    // enter and exit should match layer counts
+
+    assert_eq!(counts.sub_new_count, 1);
+    assert_eq!(
+        counts.sub_new_count + counts.sub_clone_count,
+        counts.sub_close_count
+    );
+
+    assert_eq!(counts.sub_enter_count, counts.layer_enter_count);
+    assert_eq!(counts.sub_exit_count, counts.layer_exit_count);
+}