From b1123e431584a6b75f795d93d4eec87f06314d02 Mon Sep 17 00:00:00 2001
From: Eliza Weisman <eliza@buoyant.io>
Date: Fri, 18 Mar 2022 12:33:27 -0700
Subject: [PATCH] core: don't use NoSubscriber as local placeholder  (#2001)

## Motivation

Currently, it is not actually possible to use `set_default(NoSubscriber)`
or similar to temporarily disable the global default subscriber (see
#1999).

This is because `NoSubscriber` is currently used as a placeholder value
when the thread-local cell that stores the current scoped default
subscriber is initialized. Therefore, we currently check if the current
scoped subscriber is `NoSubscriber`, and if it is, we fall back to
returning the global default instead.

This was fine, _when `NoSubscriber` was a private internal type only_.
However, PR #1549 makes `NoSubscriber` into a public API type. When users
can publicly construct `NoSubscriber` instances, it makes sense to want
to be able to use `NoSubscriber` to disable the current subscriber. This
is not possible when there is a global default set, because the local
default being `NoSubscriber` will cause the global default to be
returned.

## Solution

This branch changes the thread-local cell to store an `Option<Dispatch>`
instead, and use the `None` case to indicate no local default is set.
This way, when the local default is explicitly set to `NoSubscriber`, we
will return `NoSubscriber` rather than falling back.

This may also be a slight performance improvement, because we now check
if there's no global default by checking if the `Option` is `None`,
rather than downcasting it to a `NoSubscriber`.

I've also added a test reproducing #1999.

Fixes #1999

Signed-off-by: Eliza Weisman <eliza@buoyant.io>
---
 tracing-core/src/dispatcher.rs | 32 +++++++++++++++++---------------
 tracing/tests/no_subscriber.rs | 17 +++++++++++++++++
 2 files changed, 34 insertions(+), 15 deletions(-)
 create mode 100644 tracing/tests/no_subscriber.rs

diff --git a/tracing-core/src/dispatcher.rs b/tracing-core/src/dispatcher.rs
index 278f79df26..efda4ebbb9 100644
--- a/tracing-core/src/dispatcher.rs
+++ b/tracing-core/src/dispatcher.rs
@@ -160,7 +160,7 @@ pub struct Dispatch {
 #[cfg(feature = "std")]
 thread_local! {
     static CURRENT_STATE: State = State {
-        default: RefCell::new(Dispatch::none()),
+        default: RefCell::new(None),
         can_enter: Cell::new(true),
     };
 }
@@ -178,7 +178,7 @@ static mut GLOBAL_DISPATCH: Option<Dispatch> = None;
 #[cfg(feature = "std")]
 struct State {
     /// This thread's current default dispatcher.
-    default: RefCell<Dispatch>,
+    default: RefCell<Option<Dispatch>>,
     /// Whether or not we can currently begin dispatching a trace event.
     ///
     /// This is set to `false` when functions such as `enter`, `exit`, `event`,
@@ -641,7 +641,9 @@ impl Default for Dispatch {
 
 impl fmt::Debug for Dispatch {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.pad("Dispatch(...)")
+        f.debug_tuple("Dispatch")
+            .field(&format_args!("{:p}", self.subscriber))
+            .finish()
     }
 }
 
@@ -682,7 +684,13 @@ impl State {
         let prior = CURRENT_STATE
             .try_with(|state| {
                 state.can_enter.set(true);
-                state.default.replace(new_dispatch)
+                state
+                    .default
+                    .replace(Some(new_dispatch))
+                    // if the scoped default was not set on this thread, set the
+                    // `prior` default to the global default to populate the
+                    // scoped default when unsetting *this* default
+                    .unwrap_or_else(|| get_global().cloned().unwrap_or_else(Dispatch::none))
             })
             .ok();
         EXISTS.store(true, Ordering::Release);
@@ -705,16 +713,10 @@ impl State {
 impl<'a> Entered<'a> {
     #[inline]
     fn current(&self) -> RefMut<'a, Dispatch> {
-        let mut default = self.0.default.borrow_mut();
-
-        if default.is::<NoSubscriber>() {
-            if let Some(global) = get_global() {
-                // don't redo this call on the next check
-                *default = global.clone();
-            }
-        }
-
-        default
+        let default = self.0.default.borrow_mut();
+        RefMut::map(default, |default| {
+            default.get_or_insert_with(|| get_global().cloned().unwrap_or_else(Dispatch::none))
+        })
     }
 }
 
@@ -738,7 +740,7 @@ impl Drop for DefaultGuard {
             // lead to the drop of a subscriber which, in the process,
             // could then also attempt to access the same thread local
             // state -- causing a clash.
-            let prev = CURRENT_STATE.try_with(|state| state.default.replace(dispatch));
+            let prev = CURRENT_STATE.try_with(|state| state.default.replace(Some(dispatch)));
             drop(prev)
         }
     }
diff --git a/tracing/tests/no_subscriber.rs b/tracing/tests/no_subscriber.rs
new file mode 100644
index 0000000000..f7ff2485b2
--- /dev/null
+++ b/tracing/tests/no_subscriber.rs
@@ -0,0 +1,17 @@
+#![cfg(feature = "std")]
+
+use tracing::subscriber::{self, NoSubscriber};
+
+mod support;
+
+#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
+#[test]
+fn no_subscriber_disables_global() {
+    // Reproduces https://github.com/tokio-rs/tracing/issues/1999
+    let (subscriber, handle) = support::subscriber::mock().done().run_with_handle();
+    subscriber::set_global_default(subscriber).expect("setting global default must succeed");
+    subscriber::with_default(NoSubscriber::default(), || {
+        tracing::info!("this should not be recorded");
+    });
+    handle.assert_finished();
+}