core: add Dispatch::downgrade() and WeakDispatch (#2293)

Allows collectors and subscribers to stash their own `Dispatch` without
causing a memory leak.

## Motivation

Resolves a shortcoming of #2269: that it's impossible for collectors or
subscribers 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 collector. `WeakDispatch` can be created via
`Dispatch::downgrade`, and can be transformed into a `Dispatch` via
`WeakDispatch::upgrade`.

Co-authored-by: Eliza Weisman <eliza@buoyant.io>

Author: jswrenn

tracing-core/src/collect.rs

@@ -79,6 +79,24 @@ use core::ptr::NonNull;
 /// [`event_enabled`]: Collect::event_enabled
 pub trait Collect: 'static {
     /// Invoked when this collector becomes a [`Dispatch`].
+    ///
+    /// ## Avoiding Memory Leaks
+    ///
+    /// Collectors should not store their own [`Dispatch`]. Because the
+    /// `Dispatch` owns the collector, storing the `Dispatch` within the
+    /// collector 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 collector, 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 collector without causing a
+    /// memory leak, and can be [upgraded] into a `Dispatch` temporarily when
+    /// the `Dispatch` must be accessed by the collector.
+    ///
+    /// [`WeakDispatch`]: crate::dispatch::WeakDispatch
+    /// [upgraded]: crate::dispatch::WeakDispatch::upgrade
     fn on_register_dispatch(&self, collector: &Dispatch) {
         let _ = collector;
     }

tracing-core/src/dispatch.rs

@@ -150,11 +150,10 @@ use core::{
 use std::{
     cell::{Cell, RefCell, RefMut},
     error,
-    sync::Weak,
 };
 
 #[cfg(feature = "alloc")]
-use alloc::sync::Arc;
+use alloc::sync::{Arc, Weak};
 
 #[cfg(feature = "alloc")]
 use core::ops::Deref;
@@ -169,6 +168,34 @@ pub struct Dispatch {
     collector: &'static (dyn Collect + Send + Sync),
 }
 
+/// `WeakDispatch` is a version of [`Dispatch`] that holds a non-owning reference
+/// to a [collector].
+///
+/// The collector may be accessed by calling [`WeakDispatch::upgrade`],
+/// which returns an `Option<Dispatch>`. If all [`Dispatch`] clones that point
+/// at the collector 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 collector 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`].
+///
+/// [collector]: Collect
+/// [`Arc`]: std::sync::Arc
+/// [here]: Collect#avoiding-memory-leaks
+#[derive(Clone)]
+pub struct WeakDispatch {
+    #[cfg(feature = "alloc")]
+    collector: Kind<Weak<dyn Collect + Send + Sync>>,
+
+    #[cfg(not(feature = "alloc"))]
+    collector: &'static (dyn Collect + Send + Sync),
+}
+
 #[cfg(feature = "alloc")]
 #[derive(Clone)]
 enum Kind<T> {
@@ -577,6 +604,33 @@ impl Dispatch {
         me
     }
 
+    /// Creates a [`WeakDispatch`] from this `Dispatch`.
+    ///
+    /// A [`WeakDispatch`] is similar to a [`Dispatch`], but it does not prevent
+    /// the underlying [collector] from being dropped. Instead, it only permits
+    /// access while other references to the collector 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 collector
+    /// to hold a cyclical reference to itself without creating a memory leak.
+    /// See [here] for details.
+    ///
+    /// [collector]: Collect
+    /// [`Arc::downgrade`]: std::sync::Arc::downgrade
+    /// [here]: Collect#avoiding-memory-leaks
+    pub fn downgrade(&self) -> WeakDispatch {
+        #[cfg(feature = "alloc")]
+        let collector = match &self.collector {
+            Kind::Global(dispatch) => Kind::Global(*dispatch),
+            Kind::Scoped(dispatch) => Kind::Scoped(Arc::downgrade(dispatch)),
+        };
+        #[cfg(not(feature = "alloc"))]
+        let collector = self.collector;
+
+        WeakDispatch { collector }
+    }
+
     #[cfg(feature = "std")]
     pub(crate) fn registrar(&self) -> Registrar {
         Registrar(match self.collector {
@@ -859,6 +913,63 @@ where
     }
 }
 
+impl WeakDispatch {
+    /// Attempts to upgrade this `WeakDispatch` to a [`Dispatch`].
+    ///
+    /// Returns `None` if the referenced `Dispatch` has already been dropped.
+    ///
+    /// ## Examples
+    ///
+    /// ```
+    /// # use tracing_core::collect::NoCollector;
+    /// # use tracing_core::dispatch::Dispatch;
+    /// static COLLECTOR: NoCollector = NoCollector::new();
+    /// let strong = Dispatch::new(COLLECTOR);
+    /// 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> {
+        #[cfg(feature = "alloc")]
+        let collector = match &self.collector {
+            Kind::Global(dispatch) => Some(Kind::Global(*dispatch)),
+            Kind::Scoped(dispatch) => dispatch.upgrade().map(Kind::Scoped),
+        };
+        #[cfg(not(feature = "alloc"))]
+        let collector = Some(self.collector);
+
+        collector.map(|collector| Dispatch { collector })
+    }
+}
+
+impl fmt::Debug for WeakDispatch {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match &self.collector {
+            #[cfg(feature = "alloc")]
+            Kind::Global(collector) => f
+                .debug_tuple("WeakDispatch::Global")
+                .field(&format_args!("{:p}", collector))
+                .finish(),
+
+            #[cfg(feature = "alloc")]
+            Kind::Scoped(collector) => f
+                .debug_tuple("WeakDispatch::Scoped")
+                .field(&format_args!("{:p}", collector))
+                .finish(),
+
+            #[cfg(not(feature = "alloc"))]
+            collector => f
+                .debug_tuple("WeakDispatch::Global")
+                .field(&format_args!("{:p}", collector))
+                .finish(),
+        }
+    }
+}
+
 #[cfg(feature = "std")]
 impl Registrar {
     pub(crate) fn upgrade(&self) -> Option<Dispatch> {

tracing-subscriber/src/subscribe/mod.rs

@@ -719,6 +719,23 @@ where
     /// Performs late initialization when installing this subscriber as a
     /// [collector].
     ///
+    /// ## Avoiding Memory Leaks
+    ///
+    /// Subscribers should not store the [`Dispatch`] pointing to the collector
+    /// that they are a part of. Because the `Dispatch` owns the collector,
+    /// storing the `Dispatch` within the collector 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`]: tracing_core::dispatch::WeakDispatch
+    /// [upgraded]: tracing_core::dispatch::WeakDispatch::upgrade
     /// [collector]: tracing_core::Collect
     fn on_register_dispatch(&self, collector: &Dispatch) {
         let _ = collector;

tracing/src/dispatch.rs

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