fix(subscriber): panic when reloading Filtered

This commit adds support for reloading filtered layers without breaking
the existing API.

The `Reload` handle now provides a `reload_filtered` function, which
invokes a newly introduced `on_reload_layer` function on the `Layer`.
This function is responsible for registering the filter with the
subscriber via `register_filter_immut`.

Unlike `register_filter`, `register_filter_immut` takes a shared
reference to the subscriber, allowing the filter registration to occur
after the subscriber has been fully initialized.

Author: samugi

tracing-subscriber/src/filter/layer_filters/mod.rs

@@ -731,6 +731,11 @@ where
         self.layer.on_layer(subscriber);
     }
 
+    fn on_reload_layer(&mut self, subscriber: &S) {
+        self.id = MagicPlfDowncastMarker(subscriber.register_filter_immut());
+        self.layer.on_reload_layer(subscriber);
+    }
+
     // TODO(eliza): can we figure out a nice way to make the `Filtered` layer
     // not call `is_enabled_for` in hooks that the inner layer doesn't actually
     // have real implementations of? probably not...

tracing-subscriber/src/layer/layered.rs

@@ -254,6 +254,11 @@ where
         self.inner.on_layer(subscriber);
     }
 
+    fn on_reload_layer(&mut self, subscriber: &S) {
+        self.layer.on_reload_layer(subscriber);
+        self.inner.on_reload_layer(subscriber);   
+    }
+
     fn register_callsite(&self, metadata: &'static Metadata<'static>) -> Interest {
         self.pick_interest(self.layer.register_callsite(metadata), || {
             self.inner.register_callsite(metadata)
@@ -396,6 +401,11 @@ where
     fn register_filter(&mut self) -> FilterId {
         self.inner.register_filter()
     }
+
+    #[cfg(all(feature = "registry", feature = "std"))]
+    fn register_filter_immut(&self) -> FilterId {
+        self.inner.register_filter_immut()
+    }
 }
 
 impl<L, S> Layered<L, S>

tracing-subscriber/src/layer/mod.rs

@@ -786,6 +786,30 @@ where
         let _ = subscriber;
     }
 
+    /// Performs late initialization when a `Layer` is reloaded into a [`Subscriber`].
+    ///
+    /// This is similar to [`Layer::on_layer`], but is intended for use when a
+    /// `Layer` is dynamically reloaded via [`reload`](crate::reload), rather than
+    /// attached to a subscriber during initial construction.
+    ///
+    /// Unlike `on_layer`, this method receives a shared (`&`) reference to the
+    /// [`Subscriber`], allowing it to be called in contexts where the subscriber
+    /// is already in use and cannot be mutated. This is useful for registering
+    /// internal state (e.g., with [`register_filter_immut`]) that would otherwise
+    /// require a mutable subscriber.
+    ///
+    /// Like `on_layer`, `Layer` implementations that wrap other layers (such as
+    /// [`Filtered`] or [`Layered`]) MUST ensure that their inner layers'
+    /// `on_reload_layer` methods are called during reload, if applicable.
+    ///
+    /// [`Subscriber`]: tracing::Subscriber
+    /// [`Filtered`]: crate::filter::Filtered
+    /// [`Layered`]: crate::layer::Layered
+    /// [`register_filter_immut`]: crate::registry::LookupSpan::register_filter_immut
+    fn on_reload_layer(&mut self, subscriber: &S) {
+        let _ = subscriber;
+    }
+
     /// Registers a new callsite with this layer, returning whether or not
     /// the layer is interested in being notified about the callsite, similarly
     /// to [`Subscriber::register_callsite`].
@@ -1568,6 +1592,12 @@ where
         }
     }
 
+    fn on_reload_layer(&mut self, subscriber: &S) {
+        if let Some(ref mut layer) = self {
+            layer.on_reload_layer(subscriber);
+        }
+    }
+
     #[inline]
     fn on_new_span(&self, attrs: &span::Attributes<'_>, id: &span::Id, ctx: Context<'_, S>) {
         if let Some(ref inner) = self {
@@ -1690,6 +1720,11 @@ feature! {
                 self.deref_mut().on_layer(subscriber);
             }
 
+            #[inline]
+            fn on_reload_layer(&mut self, subscriber: &S) {
+                self.deref_mut().on_reload_layer(subscriber);
+            }
+
             #[inline]
             fn on_new_span(&self, attrs: &span::Attributes<'_>, id: &span::Id, ctx: Context<'_, S>) {
                 self.deref().on_new_span(attrs, id, ctx)
@@ -1787,6 +1822,12 @@ feature! {
             }
         }
 
+        fn on_reload_layer(&mut self, subscriber: &S) {
+            for l in self {
+                l.on_reload_layer(subscriber);
+            }
+        }
+
         fn register_callsite(&self, metadata: &'static Metadata<'static>) -> Interest {
             // Return highest level of interest.
             let mut interest = Interest::never();

tracing-subscriber/src/registry/mod.rs

@@ -151,6 +151,32 @@ pub trait LookupSpan<'a> {
             std::any::type_name::<Self>()
         )
     }
+
+    /// Registers a [`Filter`] for [per-layer filtering] with this
+    /// [`Subscriber`], using a shared reference.
+    ///
+    /// This method is similar to [`register_filter`], but takes `&self`
+    /// instead of `&mut self`. It is intended for use in cases where
+    /// filters must be registered after the subscriber has already been
+    /// constructed, such as when reloading a [`Layer`] dynamically.
+    ///
+    /// # Panics
+    ///
+    /// If this `Subscriber` does not support [per-layer filtering].
+    ///
+    /// [`Filter`]: crate::layer::Filter
+    /// [per-layer filtering]: crate::layer::Layer#per-layer-filtering
+    /// [`Subscriber`]: tracing_core::Subscriber
+    /// [`Layer`]: crate::layer::Layer
+    /// [`register_filter`]: Self::register_filter
+    #[cfg(feature = "registry")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "registry")))]
+    fn register_filter_immut(&self) -> FilterId {
+        panic!(
+            "{} does not currently support filters",
+            std::any::type_name::<Self>()
+        )
+    }
 }
 
 /// A stored representation of data associated with a span.

tracing-subscriber/src/registry/sharded.rs

@@ -12,7 +12,7 @@ use crate::{
 };
 use std::{
     cell::{self, Cell, RefCell},
-    sync::atomic::{fence, AtomicUsize, Ordering},
+    sync::atomic::{fence, AtomicU8, AtomicUsize, Ordering},
 };
 use tracing_core::{
     dispatcher::{self, Dispatch},
@@ -92,7 +92,7 @@ use tracing_core::{
 pub struct Registry {
     spans: Pool<DataInner>,
     current_spans: ThreadLocal<RefCell<SpanStack>>,
-    next_filter_id: u8,
+    next_filter_id: AtomicU8,
 }
 
 /// Span data stored in a [`Registry`].
@@ -137,7 +137,7 @@ impl Default for Registry {
         Self {
             spans: Pool::new(),
             current_spans: ThreadLocal::new(),
-            next_filter_id: 0,
+            next_filter_id: AtomicU8::new(0),
         }
     }
 }
@@ -201,7 +201,7 @@ impl Registry {
     }
 
     pub(crate) fn has_per_layer_filters(&self) -> bool {
-        self.next_filter_id > 0
+        self.next_filter_id.load(Ordering::SeqCst) > 0
     }
 
     pub(crate) fn span_stack(&self) -> cell::Ref<'_, SpanStack> {
@@ -375,9 +375,12 @@ impl<'a> LookupSpan<'a> for Registry {
     }
 
     fn register_filter(&mut self) -> FilterId {
-        let id = FilterId::new(self.next_filter_id);
-        self.next_filter_id += 1;
-        id
+        self.register_filter_immut()
+    }
+
+    fn register_filter_immut(&self) -> FilterId {
+        let filter_id = self.next_filter_id.fetch_add(1, Ordering::SeqCst);
+        FilterId::new(filter_id)
     }
 }
 

tracing-subscriber/src/reload.rs

@@ -32,7 +32,7 @@
 //! info!("This will be logged");
 //! ```
 //!
-//! Reloading a [`Filtered`](crate::filter::Filtered) layer:
+//! Reloading the [`Filter`] of a [`Filtered`](crate::filter::Filtered) layer:
 //!
 //! ```rust
 //! # use tracing::info;
@@ -53,6 +53,53 @@
 //! info!("This will be logged");
 //! ```
 //!
+//! Reloading a [`Filtered`](crate::filter::Filtered) layer using [`Handle::reload_filtered`]:
+//!
+//! ```rust
+//! use tracing::info;
+//! use tracing_subscriber::{filter, fmt, reload, prelude::*};
+//! let filtered_layer = fmt::Layer::default().with_filter(filter::LevelFilter::WARN);
+//! let (filtered_layer, reload_handle) = reload::Layer::new(filtered_layer);
+//! let dispatcher =
+//!     tracing_core::Dispatch::new(tracing_subscriber::registry().with(filtered_layer));
+//!
+//! tracing_core::dispatcher::with_default(&dispatcher, || {
+//!     info!("This will be ignored");
+//!     let new_layer = fmt::Layer::default().with_filter(filter::LevelFilter::INFO);
+//!     let subscriber = dispatcher
+//!         .downcast_ref::<tracing_subscriber::Registry>()
+//!         .unwrap();
+//!     reload_handle
+//!         .reload_filtered(new_layer, subscriber)
+//!         .unwrap();
+//!     info!("This will be logged");
+//! });
+//! ```
+//!
+//! Reloading a [`Filtered`](crate::filter::Filtered) layer using [`Handle::modify`]:
+//!
+//! ```rust
+//! use tracing::info;
+//! use tracing_subscriber::{filter, fmt, reload, prelude::*};
+//! let filtered_layer = fmt::Layer::default().with_filter(filter::LevelFilter::WARN);
+//! let (filtered_layer, reload_handle) = reload::Layer::new(filtered_layer);
+//! let dispatcher =
+//!     tracing_core::Dispatch::new(tracing_subscriber::registry().with(filtered_layer));
+//!
+//! tracing_core::dispatcher::with_default(&dispatcher, || {
+//!    info!("This will be ignored");
+//!    let new_layer = fmt::Layer::default().with_filter(filter::LevelFilter::INFO);
+//!    let subscriber = dispatcher
+//!        .downcast_ref::<tracing_subscriber::Registry>()
+//!        .unwrap();
+//!    reload_handle.modify(|layer| {
+//!        *layer = new_layer;
+//!        layer.on_reload_layer(subscriber);
+//!    }).unwrap();
+//!    info!("This will be logged");
+//! });
+//! ```
+//!
 //! ## Note
 //!
 //! The [`Layer`] implementation is unable to implement downcasting functionality,
@@ -62,6 +109,7 @@
 //! `Filter` on a layer, prefer wrapping that `Filter` in the `reload::Layer`.
 //!
 //! [`Filter` trait]: crate::layer::Filter
+//! [`Filter`]: crate::layer::Filter
 //! [`Layer` type]: Layer
 //! [`Layer` trait]: super::layer::Layer
 use crate::layer;
@@ -124,6 +172,10 @@ where
         try_lock!(self.inner.write(), else return).on_layer(subscriber);
     }
 
+    fn on_reload_layer(&mut self, subscriber: &S) {
+        try_lock!(self.inner.write(), else return).on_reload_layer(subscriber);
+    }
+
     #[inline]
     fn register_callsite(&self, metadata: &'static Metadata<'static>) -> Interest {
         try_lock!(self.inner.read(), else return Interest::sometimes()).register_callsite(metadata)
@@ -287,7 +339,9 @@ impl<L, S> Handle<L, S> {
     /// Replace the current [`Layer`] or [`Filter`] with the provided `new_value`.
     ///
     /// [`Handle::reload`] cannot be used with the [`Filtered`] layer; use
-    /// [`Handle::modify`] instead (see [this issue] for additional details).
+    /// [`Handle::reload_filtered`] instead, or [`Handle::modify`] to replace
+    /// the [`Filter`] of a [`Filtered`] layer (see [this issue] for additional
+    /// details).
     ///
     /// However, if the _only_ the [`Filter`]  needs to be modified, use
     /// `reload::Layer` to wrap the `Filter` directly.
@@ -303,6 +357,36 @@ impl<L, S> Handle<L, S> {
         })
     }
 
+    /// Replaces the current [`Filtered`] layer with a new one and performs
+    /// late initialization using [`Layer::on_reload_layer`].
+    ///
+    /// Unlike [`Handle::reload`], this method is specifically designed to
+    /// support [per-layer filtering] by ensuring that the [`Filter`] within
+    /// a [`Filtered`] layer is correctly registered with the [`Subscriber`].
+    /// It achieves this by invoking the [`Layer::on_reload_layer`] callback
+    /// after replacing the layer, allowing the new filter to be initialized
+    /// using only a shared reference to the `Subscriber`.
+    ///
+    /// This method should be used instead of [`reload`] when reloading a
+    /// [`Filtered`] layer.
+    ///
+    /// [`Filter`]: crate::layer::Filter
+    /// [`Filtered`]: crate::filter::Filtered
+    /// [`reload`]: Self::reload
+    /// [`Subscriber`]: tracing::Subscriber
+    /// [`Layer::on_reload_layer`]: crate::layer::Layer::on_reload_layer
+    /// [per-layer filtering]: crate::layer::Layer#per-layer-filtering
+    pub fn reload_filtered(&self, new_value: impl Into<L>, subscriber: &S) -> Result<(), Error>
+    where
+        L: crate::layer::Layer<S> + 'static,
+        S: Subscriber,
+    {
+        self.modify(|layer| {
+            *layer = new_value.into();
+            layer.on_reload_layer(subscriber);
+        })
+    }
+
     /// Invokes a closure with a mutable reference to the current layer or filter,
     /// allowing it to be modified in place.
     pub fn modify(&self, f: impl FnOnce(&mut L)) -> Result<(), Error> {