Skip to content

Commit

Permalink
core: don't use NoCollector as local placeholder (#2001)
Browse files Browse the repository at this point in the history
## 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]>
  • Loading branch information
hawkw authored Mar 18, 2022
1 parent 63015eb commit 4adc0a3
Show file tree
Hide file tree
Showing 2 changed files with 57 additions and 17 deletions.
55 changes: 38 additions & 17 deletions tracing-core/src/dispatch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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),
};
}
Expand Down Expand Up @@ -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`,
Expand Down Expand Up @@ -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);
}

Expand Down Expand Up @@ -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(),
}
}
}

Expand Down Expand Up @@ -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);
Expand All @@ -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())
})
}
}

Expand All @@ -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)
}
}
Expand Down
19 changes: 19 additions & 0 deletions tracing/tests/no_collector.rs
Original file line number Diff line number Diff line change
@@ -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();
}

0 comments on commit 4adc0a3

Please sign in to comment.