Skip to content

Commit

Permalink
core: add Dispatch::downgrade() and WeakDispatch (#2293)
Browse files Browse the repository at this point in the history
Allows `Subscriber`s and `Layer`s to stash their own `Dispatch` without
causing a memory leak.

## Motivation

Resolves a shortcoming of #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 <eliza@buoyant.io>
  • Loading branch information
jswrenn and hawkw committed Sep 30, 2022
1 parent b740179 commit 37c2434
Show file tree
Hide file tree
Showing 4 changed files with 133 additions and 14 deletions.
102 changes: 93 additions & 9 deletions tracing-core/src/dispatcher.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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),
}
}

Expand Down Expand Up @@ -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> {
Expand Down
18 changes: 18 additions & 0 deletions tracing-core/src/subscriber.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
25 changes: 21 additions & 4 deletions tracing-subscriber/src/layer/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion tracing/src/dispatcher.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down

0 comments on commit 37c2434

Please sign in to comment.