From ad609d5a55d8f92a453f90753cd8b5e58a438ad2 Mon Sep 17 00:00:00 2001
From: CAD97 <cad97@cad97.com>
Date: Wed, 23 Mar 2022 17:31:53 -0500
Subject: [PATCH] Let Subscribe disable specific events

---
 tracing-subscriber/src/reload.rs    |  5 +++
 tracing-subscriber/src/subscribe.rs | 61 ++++++++++++++++++++++++++---
 2 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/tracing-subscriber/src/reload.rs b/tracing-subscriber/src/reload.rs
index d245731d54..e75cff5526 100644
--- a/tracing-subscriber/src/reload.rs
+++ b/tracing-subscriber/src/reload.rs
@@ -93,6 +93,11 @@ where
         try_lock!(self.inner.read()).on_follows_from(span, follows, ctx)
     }
 
+    #[inline]
+    fn event_enabled(&self, event: &Event<'_>, ctx: subscribe::Context<'_, C>) -> bool {
+        try_lock!(self.inner.read(), else return false).event_enabled(event, ctx)
+    }
+
     #[inline]
     fn on_event(&self, event: &Event<'_>, ctx: subscribe::Context<'_, C>) {
         try_lock!(self.inner.read()).on_event(event, ctx)
diff --git a/tracing-subscriber/src/subscribe.rs b/tracing-subscriber/src/subscribe.rs
index 85d96d42c5..c680edbc2c 100644
--- a/tracing-subscriber/src/subscribe.rs
+++ b/tracing-subscriber/src/subscribe.rs
@@ -185,10 +185,10 @@ feature! {
 /// recorded: a filtering subscriber can be applied to other subscribers or
 /// collectors. A `Subscriber` that implements a filtering strategy should override the
 /// [`register_callsite`] and/or [`enabled`] methods. It may also choose to implement
-/// methods such as [`on_enter`], if it wishes to filter trace events based on
-/// the current span context.
+/// methods such as [`on_enter`] or [`event_enabled`], if it wishes to filter events
+/// based on the current span context.
 ///
-/// Note that the [`Subscribe::register_callsite`] and [`Subscribe::enabled`] methods
+/// Note that the [`register_callsite`], [`enabled`], and [`event_enabled`] methods
 /// determine whether a span or event is enabled *globally*. Thus, they should
 /// **not** be used to indicate whether an individual subscriber wishes to record a
 /// particular span or event. Instead, if a subscriber is only interested in a subset
@@ -208,6 +208,7 @@ feature! {
 /// [`register_callsite`]: Subscribe::register_callsite()
 /// [`enabled`]: Subscribe::enabled()
 /// [`on_enter`]: Subscribe::on_enter()
+/// [`event_enabled`]: Subscribe::on_enter()
 pub trait Subscribe<C>
 where
     C: Collect,
@@ -315,6 +316,30 @@ where
     // seems like a good future-proofing measure as it may grow other methods later...
     fn on_follows_from(&self, _span: &span::Id, _follows: &span::Id, _ctx: Context<'_, C>) {}
 
+    /// Called before `on_event`, to determine if `on_event` should be called.
+    ///
+    /// <div class="example-wrap" style="display:inline-block">
+    /// <pre class="ignore" style="white-space:normal;font:inherit;">
+    ///
+    /// **Note**: This method determines whether an event is globally enabled,
+    /// *not* whether the individual subscriber will be notified about the
+    /// event. This is intended to be used by layers that implement filtering
+    /// for the entire stack. Layers which do not wish to be notified about
+    /// certain events but do not wish to globally disable them should ignore
+    /// those events in their [on_event][Self::on_event].
+    ///
+    /// </pre></div>
+    ///
+    /// See [the trait-level documentation] for more information on filtering
+    /// with `Subscriber`s.
+    ///
+    /// [`Interest`]: tracing_core::Interest
+    /// [the trait-level documentation]: #filtering-with-subscribers
+    #[inline] // collapse this to a constant please mrs optimizer
+    fn event_enabled(&self, _event: &Event<'_>, _ctx: Context<'_, C>) -> bool {
+        true
+    }
+
     /// Notifies this subscriber that an event has occurred.
     fn on_event(&self, _event: &Event<'_>, _ctx: Context<'_, C>) {}
 
@@ -631,8 +656,10 @@ where
     }
 
     fn event(&self, event: &Event<'_>) {
-        self.inner.event(event);
-        self.subscriber.on_event(event, self.ctx());
+        if self.subscriber.event_enabled(event, self.ctx()) {
+            self.inner.event(event);
+            self.subscriber.on_event(event, self.ctx());
+        }
     }
 
     fn enter(&self, span: &span::Id) {
@@ -754,6 +781,17 @@ where
         self.subscriber.on_follows_from(span, follows, ctx);
     }
 
+    #[inline]
+    fn event_enabled(&self, event: &Event<'_>, ctx: Context<'_, C>) -> bool {
+        if self.subscriber.event_enabled(event, ctx.clone()) {
+            // if the outer subscriber enables the callsite metadata, ask the inner subscriber.
+            self.inner.event_enabled(event, ctx)
+        } else {
+            // otherwise, the callsite is disabled by this subscriber
+            false
+        }
+    }
+
     #[inline]
     fn on_event(&self, event: &Event<'_>, ctx: Context<'_, C>) {
         self.inner.on_event(event, ctx.clone());
@@ -845,6 +883,14 @@ where
         }
     }
 
+    #[inline]
+    fn event_enabled(&self, event: &Event<'_>, ctx: Context<'_, C>) -> bool {
+        match self {
+            Some(ref inner) => inner.event_enabled(event, ctx),
+            None => true,
+        }
+    }
+
     #[inline]
     fn on_event(&self, event: &Event<'_>, ctx: Context<'_, C>) {
         if let Some(ref inner) = self {
@@ -926,6 +972,11 @@ feature! {
                 self.deref().on_follows_from(span, follows, ctx)
             }
 
+            #[inline]
+            fn event_enabled(&self, event: &Event<'_>, ctx: Context<'_, C>) -> bool {
+                self.deref().event_enabled(event, ctx)
+            }
+
             #[inline]
             fn on_event(&self, event: &Event<'_>, ctx: Context<'_, C>) {
                 self.deref().on_event(event, ctx)