core: add {Collect,Subscriber}::on_register_dispatch (#2269)

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

tracing-core/src/callsite.rs

@@ -232,6 +232,7 @@ mod inner {
         let mut dispatchers = REGISTRY.dispatchers.write().unwrap();
         let callsites = &REGISTRY.callsites;
 
+        dispatch.collector().on_register_dispatch(dispatch);
         dispatchers.push(dispatch.registrar());
 
         rebuild_interest(callsites, &mut dispatchers);

tracing-core/src/collect.rs

@@ -1,5 +1,5 @@
 //! Collectors collect and record trace data.
-use crate::{span, Event, LevelFilter, Metadata};
+use crate::{span, Dispatch, Event, LevelFilter, Metadata};
 
 use core::any::{Any, TypeId};
 use core::ptr::NonNull;
@@ -78,6 +78,11 @@ use core::ptr::NonNull;
 /// [`event`]: Collect::event
 /// [`event_enabled`]: Collect::event_enabled
 pub trait Collect: 'static {
+    /// Invoked when this collector becomes a [`Dispatch`].
+    fn on_register_dispatch(&self, collector: &Dispatch) {
+        let _ = collector;
+    }
+
     // === Span registry methods ==============================================
 
     /// Registers a new [callsite] with this collector, returning whether or not

tracing-core/src/dispatch.rs

@@ -587,7 +587,7 @@ impl Dispatch {
 
     #[inline(always)]
     #[cfg(feature = "alloc")]
-    fn collector(&self) -> &(dyn Collect + Send + Sync) {
+    pub(crate) fn collector(&self) -> &(dyn Collect + Send + Sync) {
         match self.collector {
             Kind::Scoped(ref s) => Arc::deref(s),
             Kind::Global(s) => s,
@@ -596,7 +596,7 @@ impl Dispatch {
 
     #[inline(always)]
     #[cfg(not(feature = "alloc"))]
-    fn collector(&self) -> &(dyn Collect + Send + Sync) {
+    pub(crate) fn collector(&self) -> &(dyn Collect + Send + Sync) {
         self.collector
     }
 

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

@@ -44,7 +44,7 @@ use std::{
 };
 use tracing_core::{
     collect::{Collect, Interest},
-    span, Event, Metadata,
+    span, Dispatch, Event, Metadata,
 };
 pub mod combinator;
 
@@ -603,6 +603,10 @@ where
     F: subscribe::Filter<C> + 'static,
     S: Subscribe<C>,
 {
+    fn on_register_dispatch(&self, collector: &Dispatch) {
+        self.subscriber.on_register_dispatch(collector);
+    }
+
     fn on_subscribe(&mut self, collector: &mut C) {
         self.id = MagicPsfDowncastMarker(collector.register_filter());
         self.subscriber.on_subscribe(collector);

tracing-subscriber/src/reload.rs

@@ -30,7 +30,7 @@ use std::{
 use tracing_core::{
     callsite,
     collect::{Collect, Interest},
-    span, Event, LevelFilter, Metadata,
+    span, Dispatch, Event, LevelFilter, Metadata,
 };
 
 /// Wraps a `Filter` or `Subscribe`, allowing it to be reloaded dynamically at runtime.
@@ -71,6 +71,10 @@ where
     S: crate::Subscribe<C> + 'static,
     C: Collect,
 {
+    fn on_register_dispatch(&self, collector: &Dispatch) {
+        try_lock!(self.inner.read()).on_register_dispatch(collector);
+    }
+
     #[inline]
     fn register_callsite(&self, metadata: &'static Metadata<'static>) -> Interest {
         try_lock!(self.inner.read(), else return Interest::sometimes()).register_callsite(metadata)

tracing-subscriber/src/subscribe/layered.rs

@@ -1,7 +1,7 @@
 use tracing_core::{
     collect::{Collect, Interest},
     metadata::Metadata,
-    span, Event, LevelFilter,
+    span, Dispatch, Event, LevelFilter,
 };
 
 use crate::{
@@ -244,6 +244,11 @@ where
     B: Subscribe<C>,
     C: Collect,
 {
+    fn on_register_dispatch(&self, collector: &Dispatch) {
+        self.subscriber.on_register_dispatch(collector);
+        self.inner.on_register_dispatch(collector);
+    }
+
     fn on_subscribe(&mut self, collect: &mut C) {
         self.subscriber.on_subscribe(collect);
         self.inner.on_subscribe(collect);

tracing-subscriber/src/subscribe/mod.rs

@@ -682,7 +682,7 @@ use crate::filter;
 use tracing_core::{
     collect::{Collect, Interest},
     metadata::Metadata,
-    span, Event, LevelFilter,
+    span, Dispatch, Event, LevelFilter,
 };
 
 use core::{any::TypeId, ptr::NonNull};
@@ -716,6 +716,14 @@ where
     C: Collect,
     Self: 'static,
 {
+    /// Performs late initialization when installing this subscriber as a
+    /// [collector].
+    ///
+    /// [collector]: tracing_core::Collect
+    fn on_register_dispatch(&self, collector: &Dispatch) {
+        let _ = collector;
+    }
+
     /// Performs late initialization when attaching a subscriber to a
     /// [collector].
     ///
@@ -1472,6 +1480,12 @@ where
     S: Subscribe<C>,
     C: Collect,
 {
+    fn on_register_dispatch(&self, collector: &Dispatch) {
+        if let Some(ref subscriber) = self {
+            subscriber.on_register_dispatch(collector)
+        }
+    }
+
     fn on_subscribe(&mut self, collector: &mut C) {
         if let Some(ref mut subscriber) = self {
             subscriber.on_subscribe(collector)
@@ -1584,6 +1598,10 @@ where
 #[cfg(any(feature = "std", feature = "alloc"))]
 macro_rules! subscriber_impl_body {
     () => {
+        fn on_register_dispatch(&self, collector: &Dispatch) {
+            self.deref().on_register_dispatch(collector);
+        }
+
         #[inline]
         fn on_subscribe(&mut self, collect: &mut C) {
             self.deref_mut().on_subscribe(collect);
@@ -1681,6 +1699,11 @@ feature! {
         S: Subscribe<C>,
         C: Collect,
     {
+        fn on_register_dispatch(&self, collector: &Dispatch) {
+            for s in self {
+                s.on_register_dispatch(collector);
+            }
+        }
 
         fn on_subscribe(&mut self, collector: &mut C) {
             for s in self {