core: don't use NoCollector as local placeholder (#2001)

## Motivation

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

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

This was fine, _when `NoCollector` was a private internal type only_.
However, PR #1549 makes `NoCollector` into a public API type. When users
can publicly construct `NoCollector` instances, it makes sense to want
to be able to use `NoCollector` to disable the current collector. This
is not possible when there is a global default set, because the local
default being `NoCollector` 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 `NoCollector`, we
will return `NoCollector` 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 `NoCollector`.

I've also added a test reproducing #1999.

Fixes #1999

Signed-off-by: Eliza Weisman <[email protected]>

Author: hawkw

tracing-core/src/dispatch.rs

@@ -179,7 +179,7 @@ enum Kind<T> {
 #[cfg(feature = "std")]
 thread_local! {
     static CURRENT_STATE: State = State {
-        default: RefCell::new(Dispatch::none()),
+        default: RefCell::new(None),
         can_enter: Cell::new(true),
     };
 }
@@ -212,7 +212,7 @@ static NO_COLLECTOR: NoCollector = NoCollector::new();
 #[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`,
@@ -409,11 +409,12 @@ where
                 let _guard = Entered(&state.can_enter);
 
                 let mut default = state.default.borrow_mut();
+                let default = default
+                    // if the local default for this thread has never been set,
+                    // populate it with the global default, so we don't have to
+                    // keep getting the global on every `get_default_slow` call.
+                    .get_or_insert_with(|| get_global().clone());
 
-                if default.is::<NoCollector>() {
-                    // don't redo this call on the next check
-                    *default = get_global().clone();
-                }
                 return f(&*default);
             }
 
@@ -811,7 +812,25 @@ impl Default for Dispatch {
 
 impl fmt::Debug for Dispatch {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.pad("Dispatch(...)")
+        match &self.collector {
+            #[cfg(feature = "alloc")]
+            Kind::Global(collector) => f
+                .debug_tuple("Dispatch::Global")
+                .field(&format_args!("{:p}", collector))
+                .finish(),
+
+            #[cfg(feature = "alloc")]
+            Kind::Scoped(collector) => f
+                .debug_tuple("Dispatch::Scoped")
+                .field(&format_args!("{:p}", collector))
+                .finish(),
+
+            #[cfg(not(feature = "alloc"))]
+            collector => f
+                .debug_tuple("Dispatch::Global")
+                .field(&format_args!("{:p}", collector))
+                .finish(),
+        }
     }
 }
 
@@ -854,7 +873,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().clone())
             })
             .ok();
         EXISTS.store(true, Ordering::Release);
@@ -878,14 +903,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::<NoCollector>() {
-            // don't redo this call on the next check
-            *default = get_global().clone();
-        }
-
-        default
+        let default = self.0.default.borrow_mut();
+        RefMut::map(default, |default| {
+            default.get_or_insert_with(|| get_global().clone())
+        })
     }
 }
 
@@ -910,7 +931,7 @@ impl Drop for DefaultGuard {
             // lead to the drop of a collector 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)
         }
     }

tracing/tests/no_collector.rs

@@ -0,0 +1,19 @@
+#![cfg(feature = "std")]
+
+use tracing::collect;
+
+mod support;
+
+use self::support::*;
+
+#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
+#[test]
+fn no_collector_disables_global() {
+    // Reproduces https://github.com/tokio-rs/tracing/issues/1999
+    let (collector, handle) = collector::mock().done().run_with_handle();
+    collect::set_global_default(collector).expect("setting global default must succeed");
+    collect::with_default(collect::NoCollector::default(), || {
+        tracing::info!("this should not be recorded");
+    });
+    handle.assert_finished();
+}