core: add Dispatch::downgrade() and WeakDispatch (tokio-rs#2293)

Allows `Subscriber`s and `Layer`s to stash their own `Dispatch` without
causing a memory leak.

## Motivation

Resolves a shortcoming of tokio-rs#2269: that it's impossible for `Subscriber`s
or `Layer`s to stash a copy of their own `Dispatch` without creating a
reference cycle that would prevent them from being dropped.

## Solution

Introduces `WeakDispatch` (analogous to `std::sync::Weak`) that holds a
weak pointer to a `Subscriber`. `WeakDispatch` can be created via
`Dispatch::downgrade`, and can be transformed into a `Dispatch` via
`WeakDispatch::upgrade`.

Co-authored-by: Eliza Weisman <[email protected]>

Author: 2 people

tracing-core/src/dispatcher.rs

@@ -134,24 +134,58 @@ use crate::stdlib::{
     fmt,
     sync::{
         atomic::{AtomicBool, AtomicUsize, Ordering},
-        Arc,
+        Arc, Weak,
     },
 };
 
 #[cfg(feature = "std")]
 use crate::stdlib::{
     cell::{Cell, RefCell, RefMut},
     error,
-    sync::Weak,
 };
 
+#[cfg(feature = "alloc")]
+use alloc::sync::{Arc, Weak};
+
+#[cfg(feature = "alloc")]
+use core::ops::Deref;
+
 /// `Dispatch` trace data to a [`Subscriber`].
-///
 #[derive(Clone)]
 pub struct Dispatch {
     subscriber: Arc<dyn Subscriber + Send + Sync>,
 }
 
+/// `WeakDispatch` is a version of [`Dispatch`] that holds a non-owning reference
+/// to a [`Subscriber`].
+///
+/// The Subscriber` may be accessed by calling [`WeakDispatch::upgrade`],
+/// which returns an `Option<Dispatch>`. If all [`Dispatch`] clones that point
+/// at the `Subscriber` have been dropped, [`WeakDispatch::upgrade`] will return
+/// `None`. Otherwise, it will return `Some(Dispatch)`.
+///
+/// A `WeakDispatch` may be created from a [`Dispatch`] by calling the
+/// [`Dispatch::downgrade`] method. The primary use for creating a
+/// [`WeakDispatch`] is to allow a Subscriber` to hold a cyclical reference to
+/// itself without creating a memory leak. See [here] for details.
+///
+/// This type is analogous to the [`std::sync::Weak`] type, but for a
+/// [`Dispatch`] rather than an [`Arc`].
+///
+/// [`Arc`]: std::sync::Arc
+/// [here]: Subscriber#avoiding-memory-leaks
+#[derive(Clone)]
+pub struct WeakDispatch {
+    subscriber: Weak<dyn Subscriber + Send + Sync>,
+}
+
+#[cfg(feature = "alloc")]
+#[derive(Clone)]
+enum Kind<T> {
+    Global(&'static (dyn Collect + Send + Sync)),
+    Scoped(T),
+}
+
 #[cfg(feature = "std")]
 thread_local! {
     static CURRENT_STATE: State = State {
@@ -430,12 +464,23 @@ impl Dispatch {
         Registrar(Arc::downgrade(&self.subscriber))
     }
 
-    #[inline(always)]
-    #[cfg(feature = "alloc")]
-    pub(crate) fn subscriber(&self) -> &(dyn Subscriber + Send + Sync) {
-        match self.subscriber {
-            Kind::Scoped(ref s) => Arc::deref(s),
-            Kind::Global(s) => s,
+    /// Creates a [`WeakDispatch`] from this `Dispatch`.
+    ///
+    /// A [`WeakDispatch`] is similar to a [`Dispatch`], but it does not prevent
+    /// the underlying [`Subscriber`] from being dropped. Instead, it only permits
+    /// access while other references to the `Subscriber` exist. This is equivalent
+    /// to the standard library's [`Arc::downgrade`] method, but for `Dispatch`
+    /// rather than `Arc`.
+    ///
+    /// The primary use for creating a [`WeakDispatch`] is to allow a `Subscriber`
+    /// to hold a cyclical reference to itself without creating a memory leak.
+    /// See [here] for details.
+    ///
+    /// [`Arc::downgrade`]: std::sync::Arc::downgrade
+    /// [here]: Subscriber#avoiding-memory-leaks
+    pub fn downgrade(&self) -> WeakDispatch {
+        WeakDispatch {
+            subscriber: Arc::downgrade(&self.subscriber),
         }
     }
 
@@ -682,6 +727,45 @@ where
     }
 }
 
+// === impl WeakDispatch ===
+
+impl WeakDispatch {
+    /// Attempts to upgrade this `WeakDispatch` to a [`Dispatch`].
+    ///
+    /// Returns `None` if the referenced `Dispatch` has already been dropped.
+    ///
+    /// ## Examples
+    ///
+    /// ```
+    /// # use tracing_core::subscriber::NoSubscriber;
+    /// # use tracing_core::dispatcher::Dispatch;
+    /// let strong = Dispatch::new(NoSubscriber::default());
+    /// let weak = strong.downgrade();
+    ///
+    /// // The strong here keeps it alive, so we can still access the object.
+    /// assert!(weak.upgrade().is_some());
+    ///
+    /// drop(strong); // But not any more.
+    /// assert!(weak.upgrade().is_none());
+    /// ```
+    pub fn upgrade(&self) -> Option<Dispatch> {
+        self.subscriber
+            .upgrade()
+            .map(|subscriber| Dispatch { subscriber })
+    }
+}
+
+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()
+    }
+}
+
 #[cfg(feature = "std")]
 impl Registrar {
     pub(crate) fn upgrade(&self) -> Option<Dispatch> {

tracing-core/src/subscriber.rs

@@ -82,6 +82,24 @@ use crate::stdlib::{
 /// [`event_enabled`]: Subscriber::event_enabled
 pub trait Subscriber: 'static {
     /// Invoked when this subscriber becomes a [`Dispatch`].
+    ///
+    /// ## Avoiding Memory Leaks
+    ///
+    /// `Subscriber`s should not store their own [`Dispatch`]. Because the
+    /// `Dispatch` owns the `Subscriber`, storing the `Dispatch` within the
+    /// `Subscriber` will create a reference count cycle, preventing the `Dispatch`
+    /// from ever being dropped.
+    ///
+    /// Instead, when it is necessary to store a cyclical reference to the
+    /// `Dispatch` within a `Subscriber`, use [`Dispatch::downgrade`] to convert a
+    /// `Dispatch` into a [`WeakDispatch`]. This type is analogous to
+    /// [`std::sync::Weak`], and does not create a reference count cycle. A
+    /// [`WeakDispatch`] can be stored within a `Subscriber` without causing a
+    /// memory leak, and can be [upgraded] into a `Dispatch` temporarily when
+    /// the `Dispatch` must be accessed by the `Subscriber`.
+    ///
+    /// [`WeakDispatch`]: crate::dispatcher::WeakDispatch
+    /// [upgraded]: crate::dispatcher::WeakDispatch::upgrade
     fn on_register_dispatch(&self, subscriber: &Dispatch) {
         let _ = subscriber;
     }

tracing-subscriber/src/layer/mod.rs

@@ -712,11 +712,28 @@ where
     Self: 'static,
 {
     /// Performs late initialization when installing this layer as a
-    /// [subscriber].
+    /// [`Subscriber`].
     ///
-    /// [subscriber]: tracing_core::Subscriber
-    fn on_register_dispatch(&self, subscriber: &Dispatch) {
-        let _ = subscriber;
+    /// ## Avoiding Memory Leaks
+    ///
+    /// `Layer`s should not store the [`Dispatch`] pointing to the [`Subscriber`]
+    /// that they are a part of. Because the `Dispatch` owns the `Subscriber`,
+    /// storing the `Dispatch` within the `Subscriber` will create a reference
+    /// count cycle, preventing the `Dispatch` from ever being dropped.
+    ///
+    /// Instead, when it is necessary to store a cyclical reference to the
+    /// `Dispatch` within a `Layer`, use [`Dispatch::downgrade`] to convert a
+    /// `Dispatch` into a [`WeakDispatch`]. This type is analogous to
+    /// [`std::sync::Weak`], and does not create a reference count cycle. A
+    /// [`WeakDispatch`] can be stored within a subscriber without causing a
+    /// memory leak, and can be [upgraded] into a `Dispatch` temporarily when
+    /// the `Dispatch` must be accessed by the subscriber.
+    ///
+    /// [`WeakDispatch`]: tracing_core::dispatcher::WeakDispatch
+    /// [upgraded]: tracing_core::dispatcher::WeakDispatch::upgrade
+    /// [`Subscriber`]: tracing_core::Subscriber
+    fn on_register_dispatch(&self, collector: &Dispatch) {
+        let _ = collector;
     }
 
     /// Performs late initialization when attaching a `Layer` to a

tracing/src/dispatcher.rs

@@ -133,7 +133,7 @@ pub use tracing_core::dispatcher::with_default;
 #[cfg_attr(docsrs, doc(cfg(feature = "std")))]
 pub use tracing_core::dispatcher::DefaultGuard;
 pub use tracing_core::dispatcher::{
-    get_default, set_global_default, Dispatch, SetGlobalDefaultError,
+    get_default, set_global_default, Dispatch, SetGlobalDefaultError, WeakDispatch,
 };
 
 /// Private API for internal use by tracing's macros.