fix(subscriber): make Layered propagate on_register_dispatch

Author: bantonsson

tracing-subscriber/src/layer/layered.rs

@@ -91,6 +91,11 @@ where
     L: Layer<S>,
     S: Subscriber,
 {
+    fn on_register_dispatch(&self, subscriber: &Dispatch) {
+        self.layer.on_register_dispatch(subscriber);
+        self.inner.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)
@@ -552,3 +557,76 @@ where
             .finish()
     }
 }
+
+#[cfg(all(test, feature = "std"))]
+mod tests {
+    use super::Layered;
+    use crate::Layer;
+    use std::sync::{Arc, Mutex};
+    use tracing::{Dispatch, Metadata, Subscriber};
+    use tracing_core::span::{Attributes, Id, Record};
+
+    #[test]
+    fn on_register_dispatch_is_called() {
+        // Test that the on_register_dispatch method is called when a Layered
+        // becomes a Dispatch.
+        let dispatch_called_1 = Arc::new(Mutex::new(false));
+        let dispatch_called_2 = Arc::new(Mutex::new(false));
+
+        struct TestLayer {
+            dispatch_called: Arc<Mutex<bool>>,
+        }
+
+        impl<S: Subscriber> Layer<S> for TestLayer {
+            fn on_register_dispatch(&self, _dispatch: &Dispatch) {
+                *self.dispatch_called.lock().unwrap() = true;
+            }
+        }
+
+        impl Subscriber for TestLayer {
+            fn on_register_dispatch(&self, _subscriber: &Dispatch) {
+                *self.dispatch_called.lock().unwrap() = true;
+            }
+
+            fn enabled(&self, _metadata: &Metadata<'_>) -> bool {
+                true
+            }
+
+            fn new_span(&self, _span: &Attributes<'_>) -> Id {
+                Id::from_u64(1)
+            }
+
+            fn record(&self, _span: &Id, _values: &Record<'_>) {}
+
+            fn record_follows_from(&self, _span: &Id, _follows: &Id) {}
+
+            fn event(&self, _event: &tracing::Event<'_>) {}
+
+            fn enter(&self, _span: &Id) {}
+
+            fn exit(&self, _span: &Id) {}
+        }
+
+        let subscriber = Layered::new(
+            TestLayer {
+                dispatch_called: dispatch_called_1.clone(),
+            },
+            TestLayer {
+                dispatch_called: dispatch_called_2.clone(),
+            },
+            false,
+        );
+
+        tracing::subscriber::with_default(subscriber, || {});
+
+        // Verify that on_register_dispatch was called
+        assert!(
+            *dispatch_called_1.lock().unwrap(),
+            "on_register_dispatch should have been called for Layered.layer"
+        );
+        assert!(
+            *dispatch_called_2.lock().unwrap(),
+            "on_register_dispatch should have been called for Layered.inner"
+        );
+    }
+}