fix(subscriber): make Layered propagate on_register_dispatch (#3379)

## Motivation

The `Layered` implementation of `Subscriber` does not implement and propagate the `on_register_dispatch` callback. This means that combined layers will never have their `on_register_dispatch` methods called.

## Solution

Implement the missing `on_register_dispatch` method.

---------
Co-authored-by: Hayden Stainsby <[email protected]>

Author: bantonsson

tracing-mock/src/expect.rs

@@ -42,6 +42,7 @@ pub(crate) enum Expect {
     DropSpan(ExpectedSpan),
     Visit(ExpectedSpan, ExpectedFields),
     NewSpan(NewSpan),
+    OnRegisterDispatch,
     Nothing,
 }
 
@@ -324,6 +325,10 @@ impl Expect {
                 "\n[{}] expected {}\n[{}] but instead {}",
                 name, e, name, what
             ),
+            Expect::OnRegisterDispatch => panic!(
+                "\n[{}] expected on_register_dispatch to be called\n[{}] but instead {}",
+                name, name, what
+            ),
             Expect::Nothing => panic!(
                 "\n[{}] expected nothing else to happen\n[{}] but {} instead",
                 name, name, what,

tracing-mock/src/layer.rs

@@ -642,6 +642,41 @@ impl MockLayerBuilder {
         self
     }
 
+    /// Adds an expectation that [`Layer::on_register_dispatch`] will
+    /// be called next.
+    ///
+    /// **Note**: This expectation is usually fulfilled automatically when
+    /// a layer (wrapped in a subscriber) is set as the default via
+    /// [`tracing::subscriber::with_default`] or
+    /// [`tracing::subscriber::set_global_default`], so explicitly expecting
+    /// this is not usually necessary. However, it may be useful when testing
+    /// custom layer implementations that manually call `on_register_dispatch`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tracing_mock::{expect, layer};
+    /// use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, Layer};
+    ///
+    /// let (layer, handle) = layer::mock()
+    ///     .on_register_dispatch()
+    ///     .run_with_handle();
+    ///
+    /// let _subscriber = tracing_subscriber::registry()
+    ///     .with(layer.with_filter(tracing_subscriber::filter::filter_fn(move |_meta| true)))
+    ///     .set_default();
+    ///
+    /// // The layer's on_register_dispatch was called when the subscriber was set as default
+    ///
+    /// handle.assert_finished();
+    /// ```
+    ///
+    /// [`Layer::on_register_dispatch`]: tracing_subscriber::layer::Layer::on_register_dispatch
+    pub fn on_register_dispatch(mut self) -> Self {
+        self.expected.push_back(Expect::OnRegisterDispatch);
+        self
+    }
+
     /// Expects that no further traces are received.
     ///
     /// The call to `only` should appear immediately before the final
@@ -835,6 +870,14 @@ impl<C> Layer<C> for MockLayer
 where
     C: Subscriber + for<'a> LookupSpan<'a>,
 {
+    fn on_register_dispatch(&self, _subscriber: &tracing::Dispatch) {
+        println!("[{}] on_register_dispatch", self.name);
+        let mut expected = self.expected.lock().unwrap();
+        if let Some(Expect::OnRegisterDispatch) = expected.front() {
+            expected.pop_front();
+        }
+    }
+
     fn register_callsite(
         &self,
         metadata: &'static tracing::Metadata<'static>,

tracing-mock/src/subscriber.rs

@@ -802,6 +802,37 @@ where
         self
     }
 
+    /// Adds an expectation that [`Subscriber::on_register_dispatch`] will
+    /// be called next.
+    ///
+    /// **Note**: This expectation is usually fulfilled automatically when
+    /// a subscriber is set as the default via [`tracing::subscriber::with_default`]
+    /// or [`tracing::subscriber::set_global_default`], so explicitly expecting
+    /// this is not usually necessary. However, it may be useful when testing
+    /// custom subscriber implementations that manually call `on_register_dispatch`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tracing_mock::{expect, subscriber};
+    ///
+    /// let (subscriber, handle) = subscriber::mock()
+    ///     .on_register_dispatch()
+    ///     .run_with_handle();
+    ///
+    /// tracing::subscriber::with_default(subscriber, || {
+    ///     // The subscriber's on_register_dispatch was called when it was set as default
+    /// });
+    ///
+    /// handle.assert_finished();
+    /// ```
+    ///
+    /// [`Subscriber::on_register_dispatch`]: tracing::Subscriber::on_register_dispatch
+    pub fn on_register_dispatch(mut self) -> Self {
+        self.expected.push_back(Expect::OnRegisterDispatch);
+        self
+    }
+
     /// Filter the traces evaluated by the `MockSubscriber`.
     ///
     /// The filter will be applied to all traces received before
@@ -1001,6 +1032,14 @@ impl<F> Subscriber for Running<F>
 where
     F: Fn(&Metadata<'_>) -> bool + 'static,
 {
+    fn on_register_dispatch(&self, _subscriber: &tracing::Dispatch) {
+        println!("[{}] on_register_dispatch", self.name);
+        let mut expected = self.expected.lock().unwrap();
+        if let Some(Expect::OnRegisterDispatch) = expected.front() {
+            expected.pop_front();
+        }
+    }
+
     fn enabled(&self, meta: &Metadata<'_>) -> bool {
         println!("[{}] enabled: {:#?}", self.name, meta);
         let enabled = (self.filter)(meta);

tracing-mock/tests/on_register_dispatch.rs

@@ -0,0 +1,63 @@
+//! Tests for `on_register_dispatch` expectations in MockSubscriber and MockLayer.
+
+use tracing_mock::subscriber;
+
+#[test]
+fn subscriber_on_register_dispatch() {
+    let (subscriber, handle) = subscriber::mock().on_register_dispatch().run_with_handle();
+
+    tracing::subscriber::with_default(subscriber, || {
+        // The subscriber's on_register_dispatch is called when set as default
+    });
+
+    handle.assert_finished();
+}
+
+#[cfg(feature = "tracing-subscriber")]
+#[test]
+fn layer_on_register_dispatch() {
+    use tracing_mock::layer;
+    use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
+
+    let (layer, handle) = layer::mock().on_register_dispatch().run_with_handle();
+
+    let _subscriber = tracing_subscriber::registry().with(layer).set_default();
+
+    // The layer's on_register_dispatch is called when the subscriber is set as default
+    drop(_subscriber);
+
+    handle.assert_finished();
+}
+
+#[test]
+fn subscriber_multiple_expectations() {
+    let (subscriber, handle) = subscriber::mock()
+        .on_register_dispatch()
+        .event(tracing_mock::expect::event())
+        .run_with_handle();
+
+    tracing::subscriber::with_default(subscriber, || {
+        tracing::info!("test event");
+    });
+
+    handle.assert_finished();
+}
+
+#[cfg(feature = "tracing-subscriber")]
+#[test]
+fn layer_multiple_expectations() {
+    use tracing_mock::layer;
+    use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
+
+    let (layer, handle) = layer::mock()
+        .on_register_dispatch()
+        .event(tracing_mock::expect::event())
+        .run_with_handle();
+
+    let _subscriber = tracing_subscriber::registry().with(layer).set_default();
+
+    tracing::info!("test event");
+
+    drop(_subscriber);
+    handle.assert_finished();
+}

tracing-subscriber/src/layer/layered.rs

@@ -91,6 +91,11 @@ where
     L: Layer<S>,
     S: Subscriber,
 {
+    fn on_register_dispatch(&self, subscriber: &Dispatch) {
+        self.inner.on_register_dispatch(subscriber);
+        self.layer.on_register_dispatch(subscriber);
+    }
+
     fn register_callsite(&self, metadata: &'static Metadata<'static>) -> Interest {
         self.pick_interest(self.layer.register_callsite(metadata), || {
             self.inner.register_callsite(metadata)

tracing-subscriber/tests/on_register_dispatch_is_called.rs

@@ -0,0 +1,27 @@
+#![cfg(all(feature = "registry", feature = "std"))]
+//! Test that `on_register_dispatch` is called on both layers when a layered
+//! subscriber is set as the default.
+
+use tracing_mock::layer;
+use tracing_subscriber::layer::SubscriberExt;
+
+#[test]
+fn on_register_dispatch_is_called() {
+    let (inner_layer, inner_handle) = layer::named("inner")
+        .on_register_dispatch()
+        .run_with_handle();
+
+    let (outer_layer, outer_handle) = layer::named("outer")
+        .on_register_dispatch()
+        .run_with_handle();
+
+    let subscriber = tracing_subscriber::registry()
+        .with(inner_layer)
+        .with(outer_layer);
+
+    tracing::subscriber::with_default(subscriber, || {});
+
+    // Verify that on_register_dispatch was called on both layers
+    inner_handle.assert_finished();
+    outer_handle.assert_finished();
+}