diff --git a/tracing-core/src/dispatcher.rs b/tracing-core/src/dispatcher.rs index 7f2f4060a..32632abdc 100644 --- a/tracing-core/src/dispatcher.rs +++ b/tracing-core/src/dispatcher.rs @@ -140,7 +140,7 @@ use crate::stdlib::{ #[cfg(feature = "std")] use crate::stdlib::{ - cell::{Cell, RefCell, RefMut}, + cell::{Cell, Ref, RefCell}, error, }; @@ -153,7 +153,7 @@ use core::ops::Deref; /// `Dispatch` trace data to a [`Subscriber`]. #[derive(Clone)] pub struct Dispatch { - subscriber: Arc, + subscriber: Kind>, } /// `WeakDispatch` is a version of [`Dispatch`] that holds a non-owning reference @@ -176,13 +176,12 @@ pub struct Dispatch { /// [here]: Subscriber#avoiding-memory-leaks #[derive(Clone)] pub struct WeakDispatch { - subscriber: Weak, + subscriber: Kind>, } -#[cfg(feature = "alloc")] #[derive(Clone)] enum Kind { - Global(&'static (dyn Collect + Send + Sync)), + Global(&'static (dyn Subscriber + Send + Sync)), Scoped(T), } @@ -197,11 +196,20 @@ thread_local! { static EXISTS: AtomicBool = AtomicBool::new(false); static GLOBAL_INIT: AtomicUsize = AtomicUsize::new(UNINITIALIZED); +#[cfg(feature = "std")] +static SCOPED_COUNT: AtomicUsize = AtomicUsize::new(0); + const UNINITIALIZED: usize = 0; const INITIALIZING: usize = 1; const INITIALIZED: usize = 2; -static mut GLOBAL_DISPATCH: Option = None; +static mut GLOBAL_DISPATCH: Dispatch = Dispatch { + subscriber: Kind::Global(&NO_SUBSCRIBER), +}; +static NONE: Dispatch = Dispatch { + subscriber: Kind::Global(&NO_SUBSCRIBER), +}; +static NO_SUBSCRIBER: NoSubscriber = NoSubscriber::new(); /// The dispatch state of a thread. #[cfg(feature = "std")] @@ -305,8 +313,20 @@ pub fn set_global_default(dispatcher: Dispatch) -> Result<(), SetGlobalDefaultEr ) .is_ok() { + let subscriber = { + let subscriber = match dispatcher.subscriber { + Kind::Global(s) => s, + Kind::Scoped(s) => unsafe { + // safety: this leaks the subscriber onto the heap. the + // reference count will always be at least 1, because the + // global default will never be dropped. + &*Arc::into_raw(s) + }, + }; + Kind::Global(subscriber) + }; unsafe { - GLOBAL_DISPATCH = Some(dispatcher); + GLOBAL_DISPATCH = Dispatch { subscriber }; } GLOBAL_INIT.store(INITIALIZED, Ordering::SeqCst); EXISTS.store(true, Ordering::Release); @@ -365,15 +385,21 @@ pub fn get_default(mut f: F) -> T where F: FnMut(&Dispatch) -> T, { + if SCOPED_COUNT.load(Ordering::Acquire) == 0 { + // fast path if no scoped dispatcher has been set; just use the global + // default. + return f(get_global()); + } + CURRENT_STATE .try_with(|state| { if let Some(entered) = state.enter() { return f(&entered.current()); } - f(&Dispatch::none()) + f(&NONE) }) - .unwrap_or_else(|_| f(&Dispatch::none())) + .unwrap_or_else(|_| f(&NONE)) } /// Executes a closure with a reference to this thread's current [dispatcher]. @@ -387,6 +413,12 @@ where #[doc(hidden)] #[inline(never)] pub fn get_current(f: impl FnOnce(&Dispatch) -> T) -> Option { + if SCOPED_COUNT.load(Ordering::Acquire) == 0 { + // fast path if no scoped dispatcher has been set; just use the global + // default. + return Some(f(get_global())); + } + CURRENT_STATE .try_with(|state| { let entered = state.enter()?; @@ -401,8 +433,7 @@ pub fn get_current(f: impl FnOnce(&Dispatch) -> T) -> Option { #[cfg(not(feature = "std"))] #[doc(hidden)] pub fn get_current(f: impl FnOnce(&Dispatch) -> T) -> Option { - let dispatch = get_global()?; - Some(f(&dispatch)) + Some(f(get_global())) } /// Executes a closure with a reference to the current [dispatcher]. @@ -413,35 +444,30 @@ pub fn get_default(mut f: F) -> T where F: FnMut(&Dispatch) -> T, { - if let Some(d) = get_global() { - f(d) - } else { - f(&Dispatch::none()) - } + f(&get_global()) } -fn get_global() -> Option<&'static Dispatch> { +#[inline] +fn get_global() -> &'static Dispatch { if GLOBAL_INIT.load(Ordering::SeqCst) != INITIALIZED { - return None; + return &NONE; } unsafe { // This is safe given the invariant that setting the global dispatcher // also sets `GLOBAL_INIT` to `INITIALIZED`. - Some(GLOBAL_DISPATCH.as_ref().expect( - "invariant violated: GLOBAL_DISPATCH must be initialized before GLOBAL_INIT is set", - )) + &GLOBAL_DISPATCH } } #[cfg(feature = "std")] -pub(crate) struct Registrar(Weak); +pub(crate) struct Registrar(Kind>); impl Dispatch { /// Returns a new `Dispatch` that discards events and spans. #[inline] pub fn none() -> Self { Dispatch { - subscriber: Arc::new(NoSubscriber::default()), + subscriber: Kind::Global(&NO_SUBSCRIBER), } } @@ -453,7 +479,7 @@ impl Dispatch { S: Subscriber + Send + Sync + 'static, { let me = Dispatch { - subscriber: Arc::new(subscriber), + subscriber: Kind::Scoped(Arc::new(subscriber)), }; callsite::register_dispatch(&me); me @@ -461,7 +487,7 @@ impl Dispatch { #[cfg(feature = "std")] pub(crate) fn registrar(&self) -> Registrar { - Registrar(Arc::downgrade(&self.subscriber)) + Registrar(self.subscriber.downgrade()) } /// Creates a [`WeakDispatch`] from this `Dispatch`. @@ -480,14 +506,16 @@ impl Dispatch { /// [here]: Subscriber#avoiding-memory-leaks pub fn downgrade(&self) -> WeakDispatch { WeakDispatch { - subscriber: Arc::downgrade(&self.subscriber), + subscriber: self.subscriber.downgrade(), } } #[inline(always)] - #[cfg(not(feature = "alloc"))] pub(crate) fn subscriber(&self) -> &(dyn Subscriber + Send + Sync) { - &self.subscriber + match self.subscriber { + Kind::Global(s) => s, + Kind::Scoped(ref s) => s.as_ref(), + } } /// Registers a new callsite with this collector, returning whether or not @@ -500,7 +528,7 @@ impl Dispatch { /// [`register_callsite`]: super::subscriber::Subscriber::register_callsite #[inline] pub fn register_callsite(&self, metadata: &'static Metadata<'static>) -> subscriber::Interest { - self.subscriber.register_callsite(metadata) + self.subscriber().register_callsite(metadata) } /// Returns the highest [verbosity level][level] that this [`Subscriber`] will @@ -516,7 +544,7 @@ impl Dispatch { // TODO(eliza): consider making this a public API? #[inline] pub(crate) fn max_level_hint(&self) -> Option { - self.subscriber.max_level_hint() + self.subscriber().max_level_hint() } /// Record the construction of a new span, returning a new [ID] for the @@ -530,7 +558,7 @@ impl Dispatch { /// [`new_span`]: super::subscriber::Subscriber::new_span #[inline] pub fn new_span(&self, span: &span::Attributes<'_>) -> span::Id { - self.subscriber.new_span(span) + self.subscriber().new_span(span) } /// Record a set of values on a span. @@ -542,7 +570,7 @@ impl Dispatch { /// [`record`]: super::subscriber::Subscriber::record #[inline] pub fn record(&self, span: &span::Id, values: &span::Record<'_>) { - self.subscriber.record(span, values) + self.subscriber().record(span, values) } /// Adds an indication that `span` follows from the span with the id @@ -555,7 +583,7 @@ impl Dispatch { /// [`record_follows_from`]: super::subscriber::Subscriber::record_follows_from #[inline] pub fn record_follows_from(&self, span: &span::Id, follows: &span::Id) { - self.subscriber.record_follows_from(span, follows) + self.subscriber().record_follows_from(span, follows) } /// Returns true if a span with the specified [metadata] would be @@ -569,7 +597,7 @@ impl Dispatch { /// [`enabled`]: super::subscriber::Subscriber::enabled #[inline] pub fn enabled(&self, metadata: &Metadata<'_>) -> bool { - self.subscriber.enabled(metadata) + self.subscriber().enabled(metadata) } /// Records that an [`Event`] has occurred. @@ -582,8 +610,9 @@ impl Dispatch { /// [`event`]: super::subscriber::Subscriber::event #[inline] pub fn event(&self, event: &Event<'_>) { - if self.subscriber.event_enabled(event) { - self.subscriber.event(event); + let subscriber = self.subscriber(); + if subscriber.event_enabled(event) { + subscriber.event(event); } } @@ -595,7 +624,7 @@ impl Dispatch { /// [`Subscriber`]: super::subscriber::Subscriber /// [`enter`]: super::subscriber::Subscriber::enter pub fn enter(&self, span: &span::Id) { - self.subscriber.enter(span); + self.subscriber().enter(span); } /// Records that a span has been exited. @@ -606,7 +635,7 @@ impl Dispatch { /// [`Subscriber`]: super::subscriber::Subscriber /// [`exit`]: super::subscriber::Subscriber::exit pub fn exit(&self, span: &span::Id) { - self.subscriber.exit(span); + self.subscriber().exit(span); } /// Notifies the subscriber that a [span ID] has been cloned. @@ -625,7 +654,7 @@ impl Dispatch { /// [`new_span`]: super::subscriber::Subscriber::new_span #[inline] pub fn clone_span(&self, id: &span::Id) -> span::Id { - self.subscriber.clone_span(id) + self.subscriber().clone_span(id) } /// Notifies the subscriber that a [span ID] has been dropped. @@ -654,7 +683,7 @@ impl Dispatch { #[deprecated(since = "0.1.2", note = "use `Dispatch::try_close` instead")] pub fn drop_span(&self, id: span::Id) { #[allow(deprecated)] - self.subscriber.drop_span(id); + self.subscriber().drop_span(id); } /// Notifies the subscriber that a [span ID] has been dropped, and returns @@ -673,7 +702,7 @@ impl Dispatch { /// [`try_close`]: super::subscriber::Subscriber::try_close /// [`new_span`]: super::subscriber::Subscriber::new_span pub fn try_close(&self, id: span::Id) -> bool { - self.subscriber.try_close(id) + self.subscriber().try_close(id) } /// Returns a type representing this subscriber's view of the current span. @@ -684,21 +713,21 @@ impl Dispatch { /// [`current`]: super::subscriber::Subscriber::current_span #[inline] pub fn current_span(&self) -> span::Current { - self.subscriber.current_span() + self.subscriber().current_span() } /// Returns `true` if this `Dispatch` forwards to a `Subscriber` of type /// `T`. #[inline] pub fn is(&self) -> bool { - ::is::(&self.subscriber) + ::is::(self.subscriber()) } /// Returns some reference to the `Subscriber` this `Dispatch` forwards to /// if it is of type `T`, or `None` if it isn't. #[inline] pub fn downcast_ref(&self) -> Option<&T> { - ::downcast_ref(&self.subscriber) + ::downcast_ref(self.subscriber()) } } @@ -711,9 +740,16 @@ impl Default for Dispatch { impl fmt::Debug for Dispatch { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_tuple("Dispatch") - .field(&format_args!("{:p}", self.subscriber)) - .finish() + match self.subscriber { + Kind::Scoped(ref s) => f + .debug_tuple("Dispatch::Scoped") + .field(&format_args!("{:p}", s)) + .finish(), + Kind::Global(s) => f + .debug_tuple("Dispatch::Global") + .field(&format_args!("{:p}", s)) + .finish(), + } } } @@ -757,12 +793,16 @@ impl WeakDispatch { impl fmt::Debug for WeakDispatch { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let mut tuple = f.debug_tuple("WeakDispatch"); - match self.subscriber.upgrade() { - Some(subscriber) => tuple.field(&format_args!("Some({:p})", subscriber)), - None => tuple.field(&format_args!("None")), - }; - tuple.finish() + match self.subscriber { + Kind::Scoped(ref s) => f + .debug_tuple("WeakDispatch::Scoped") + .field(&format_args!("{:p}", s)) + .finish(), + Kind::Global(s) => f + .debug_tuple("WeakDispatch::Global") + .field(&format_args!("{:p}", s)) + .finish(), + } } } @@ -775,6 +815,26 @@ impl Registrar { // ===== impl State ===== +impl Kind> { + fn downgrade(&self) -> Kind> { + match self { + Kind::Global(s) => Kind::Global(*s), + Kind::Scoped(ref s) => Kind::Scoped(Arc::downgrade(s)), + } + } +} + +impl Kind> { + fn upgrade(&self) -> Option>> { + match self { + Kind::Global(s) => Some(Kind::Global(*s)), + Kind::Scoped(ref s) => Some(Kind::Scoped(s.upgrade()?)), + } + } +} + +// ===== impl State ===== + #[cfg(feature = "std")] impl State { /// Replaces the current default dispatcher on this thread with the provided @@ -792,6 +852,7 @@ impl State { .ok() .flatten(); EXISTS.store(true, Ordering::Release); + SCOPED_COUNT.fetch_add(1, Ordering::Release); DefaultGuard(prior) } @@ -810,10 +871,11 @@ impl State { #[cfg(feature = "std")] impl<'a> Entered<'a> { #[inline] - fn current(&self) -> RefMut<'a, Dispatch> { - let default = self.0.default.borrow_mut(); - RefMut::map(default, |default| { - default.get_or_insert_with(|| get_global().cloned().unwrap_or_else(Dispatch::none)) + fn current(&self) -> Ref<'a, Dispatch> { + let default = self.0.default.borrow(); + Ref::map(default, |default| match default { + Some(default) => default, + None => get_global(), }) } } @@ -838,6 +900,7 @@ impl Drop for DefaultGuard { // could then also attempt to access the same thread local // state -- causing a clash. let prev = CURRENT_STATE.try_with(|state| state.default.replace(self.0.take())); + SCOPED_COUNT.fetch_sub(1, Ordering::Release); drop(prev) } } diff --git a/tracing-core/src/subscriber.rs b/tracing-core/src/subscriber.rs index e8f444119..17b631697 100644 --- a/tracing-core/src/subscriber.rs +++ b/tracing-core/src/subscriber.rs @@ -699,6 +699,14 @@ impl Subscriber for NoSubscriber { fn exit(&self, _span: &span::Id) {} } +impl NoSubscriber { + /// Returns a new `NoSubscriber`. + #[must_use] + pub const fn new() -> Self { + Self(()) + } +} + impl Subscriber for Box where S: Subscriber + ?Sized, diff --git a/tracing-core/tests/local_dispatch_before_init.rs b/tracing-core/tests/local_dispatch_before_init.rs new file mode 100644 index 000000000..71b67f851 --- /dev/null +++ b/tracing-core/tests/local_dispatch_before_init.rs @@ -0,0 +1,43 @@ +mod common; + +use common::*; +use tracing_core::{ + dispatcher::{self, Dispatch}, + subscriber::NoSubscriber, +}; + +/// This test reproduces the following issues: +/// - https://github.com/tokio-rs/tracing/issues/2587 +/// - https://github.com/tokio-rs/tracing/issues/2411 +/// - https://github.com/tokio-rs/tracing/issues/2436 +#[test] +fn local_dispatch_before_init() { + dispatcher::get_default(|current| assert!(dbg!(current).is::())); + + // Temporarily override the default dispatcher with a scoped dispatcher. + // Using a scoped dispatcher makes the thread local state attempt to cache + // the scoped default. + #[cfg(feature = "std")] + { + dispatcher::with_default(&Dispatch::new(TestSubscriberB), || { + dispatcher::get_default(|current| { + assert!( + dbg!(current).is::(), + "overriden subscriber not set", + ); + }) + }) + } + + dispatcher::get_default(|current| assert!(current.is::())); + + dispatcher::set_global_default(Dispatch::new(TestSubscriberA)) + .expect("set global dispatch failed"); + + dispatcher::get_default(|current| { + assert!( + dbg!(current).is::(), + "default subscriber not set" + ); + }); +}