core: add {Subscriber,Layer}::on_register_dispatch (tokio-rs#2269)

The `on_register_dispatch` method is invoked when a `Subscriber` is
registered as a `Dispatch`. This method should be overridden to
perform actions upon the installation of a subscriber/layer;
for instance, to send a copy of the subscriber's `Dispatch` to a
worker thread.

tracing-core/src/dispatcher.rs

@@ -430,8 +430,23 @@ impl Dispatch {
         Registrar(Arc::downgrade(&self.subscriber))
     }
 
-    /// Registers a new callsite with this subscriber, returning whether or not
-    /// the subscriber is interested in being notified about the callsite.
+    #[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,
+        }
+    }
+
+    #[inline(always)]
+    #[cfg(not(feature = "alloc"))]
+    pub(crate) fn subscriber(&self) -> &(dyn Subscriber + Send + Sync) {
+        &self.subscriber
+    }
+
+    /// Registers a new callsite with this collector, returning whether or not
+    /// the collector is interested in being notified about the callsite.
     ///
     /// This calls the [`register_callsite`] function on the [`Subscriber`]
     /// that this `Dispatch` forwards to.

tracing-core/src/subscriber.rs

@@ -1,5 +1,5 @@
-//! Subscribers collect and record trace data.
-use crate::{span, Event, LevelFilter, Metadata};
+//! Collectors collect and record trace data.
+use crate::{span, Dispatch, Event, LevelFilter, Metadata};
 
 use crate::stdlib::{
     any::{Any, TypeId},
@@ -81,7 +81,10 @@ use crate::stdlib::{
 /// [`event`]: Subscriber::event
 /// [`event_enabled`]: Subscriber::event_enabled
 pub trait Subscriber: 'static {
-    // === Span registry methods ==============================================
+    /// Invoked when this subscriber becomes a [`Dispatch`].
+    fn on_register_dispatch(&self, subscriber: &Dispatch) {
+        let _ = subscriber;
+    }
 
     /// Registers a new [callsite] with this subscriber, returning whether or not
     /// the subscriber is interested in being notified about the callsite.

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

@@ -44,7 +44,7 @@ use std::{
 use tracing_core::{
     span,
     subscriber::{Interest, Subscriber},
-    Event, Metadata,
+    Dispatch, Event, Metadata,
 };
 pub mod combinator;
 
@@ -607,6 +607,10 @@ where
     F: layer::Filter<S> + 'static,
     L: Layer<S>,
 {
+    fn on_register_dispatch(&self, collector: &Dispatch) {
+        self.layer.on_register_dispatch(collector);
+    }
+
     fn on_layer(&mut self, subscriber: &mut S) {
         self.id = MagicPlfDowncastMarker(subscriber.register_filter());
         self.layer.on_layer(subscriber);

tracing-subscriber/src/layer/layered.rs

@@ -1,9 +1,4 @@
-use tracing_core::{
-    metadata::Metadata,
-    span,
-    subscriber::{Interest, Subscriber},
-    Event, LevelFilter,
-};
+use tracing_core::{metadata::Metadata, span, Dispatch, Event, Interest, LevelFilter, Subscriber};
 
 use crate::{
     filter,
@@ -245,6 +240,11 @@ where
     B: Layer<S>,
     S: Subscriber,
 {
+    fn on_register_dispatch(&self, subscriber: &Dispatch) {
+        self.layer.on_register_dispatch(subscriber);
+        self.inner.on_register_dispatch(subscriber);
+    }
+
     fn on_layer(&mut self, subscriber: &mut S) {
         self.layer.on_layer(subscriber);
         self.inner.on_layer(subscriber);

tracing-subscriber/src/layer/mod.rs

@@ -678,7 +678,7 @@ use tracing_core::{
     metadata::Metadata,
     span,
     subscriber::{Interest, Subscriber},
-    Event, LevelFilter,
+    Dispatch, Event, LevelFilter,
 };
 
 use core::any::TypeId;
@@ -711,6 +711,14 @@ where
     S: Subscriber,
     Self: 'static,
 {
+    /// Performs late initialization when installing this layer as a
+    /// [subscriber].
+    ///
+    /// [subscriber]: tracing_core::Subscriber
+    fn on_register_dispatch(&self, subscriber: &Dispatch) {
+        let _ = subscriber;
+    }
+
     /// Performs late initialization when attaching a `Layer` to a
     /// [`Subscriber`].
     ///
@@ -1591,6 +1599,11 @@ feature! {
 
     macro_rules! layer_impl_body {
         () => {
+            #[inline]
+            fn on_register_dispatch(&self, subscriber: &Dispatch) {
+                self.deref().on_register_dispatch(subscriber);
+            }
+
             #[inline]
             fn on_layer(&mut self, subscriber: &mut S) {
                 self.deref_mut().on_layer(subscriber);