subscriber: implement per-subscriber filtering (#1523)

## Motivation

Currently, filtering with `Subscribe`s is always _global_. If a
`Subscribe` performs filtering, it will disable a span or event for
_all_ subscribers that compose the current subscriber. In some cases,
however, it is often desirable for individual subscribers to see
different spans or events than the rest of the `Subscribe` stack.

Issues in other projects, such as tokio-rs/console#64 and
tokio-rs/console#76, linkerd/linkerd2-proxy#601,
influxdata/influxdb_iox#1012 and influxdata/influxdb_iox#1681,
jackwaudby/spaghetti#86, etc; as well as `tracing` feature requests like
#302, #597, and #1348, all indicate that there is significant demand for
the ability to add filters to individual `Subscribe`s.

Unfortunately, doing this nicely is somewhat complicated. Although a
naive implementation that simply skips events/spans in
`Subscribe::on_event` and `Subscribe::new_span` based on some filter is
relatively simple, this wouldn't really be an ideal solution, for a
number of reasons. A proper per-subscriber filtering implementation
would satisfy the following _desiderata_:

* If a per-subscriber filter disables a span, it shouldn't be present
  _for the subscriber that filter is attached to_ when iterating over
  span contexts (such as `Context::event_scope`, `SpanRef::scope`, etc),
  or when looking up a span's parents.
* When _all_ per-subscriber filters disable a span or event, it should
  be completely disabled, rather than simply skipped by those particular
  subscribers. This means that per-subscriber filters should participate
  in `enabled`, as well as being able to skip spans and events in
  `new_span` and `on_event`.
* Per-subscriber filters shouldn't interfere with non-filtered
  `Subscribe`s. If a subscriber contains subscribers without any
  filters, as well as subscribers with per-subscriber filters, the
  non-filtered `Subscribe`s should behave exactly as they would without
  any per-subscriber filters present.
* Similarly, per-subscriber filters shouldn't interfere with global
  filtering. If a `Subscribe` in a stack is performing global filtering
  (e.g. the current filtering behavior), per-subscriber filters should
  also be effected by the global filter.
* Per-subscriber filters _should_ be able to participate in `Interest`
  caching, _but only when doing so doesn't interfere with
  non-per-subscriber-filtered subscribers_.
* Per-subscriber filters should be _tree-shaped_. If a `Subscriber`
  consists of multiple subscribers that have been `Subscribeed` together
  to form new `Subscribe`s, and some of the `Subscribeed` subscribers
  have per-subscriber filters, those per-subscriber filters should
  effect all subscribers in that subtree. Similarly, if `Subscribe`s in
  a per-subscriber filtered subtree have their _own_ per-subscriber
  filters, those subscribers should be effected by the union of their
  own filters and any per-subscriber filters that wrap them at higher
  levels in the tree.

Meeting all these requirements means that implementing per-subscriber
filtering correctly is somewhat more complex than simply skipping events
and spans in a `Subscribe`'s `on_event` and `new_span` callbacks.

## Solution

This branch has a basic working implementation of per-subscriber
filtering for `tracing-subscriber` v0.2. It should be possible to add
this in a point release without a breaking change --- in particular, the
new per-subscriber filtering feature _doesn't_ interfere with the global
filtering behavior that subscribers currently have when not using the
new per-subscriber filtering APIs, so existing configurations should all
behave exactly as they do now.

### Implementation

The new per-subscriber filtering API consists of a `Filter` trait that
defines a filtering strategy and a `Filtered` type that combines a
`Filter` with a `Subscribe`. When `enabled` is called on a `Filtered`
subscriber, it checks the metadata against the `Filter` and stores
whether that filter enabled the span/event in a thread-local cell. If
every per-subscriber filter disables a span/event, and there are no
non-per-subscriber-filtered subscribers in use, the `Registry`'s
`enabled` will return `false`. Otherwise, when the span/event is
recorded, each `Filtered` subscriber's `on_event` and `new_span` method
checks with the `Registry` to see whether _it_ enabled or disabled that
span/event, and skips passing it to the wrapped subscriber if it was
disabled. When a span is recorded, the `Registry` which filters disabled
it, so that they can skip it when looking up spans in other callbacks.

 A new `on_subscriber` method was added to the `Subscribe` trait, which
 allows running a callback when adding a `Subscribe` to a `Subscriber`.
 This allows mutating both the `Subscribe` and the `Subscriber`, since
 they are both passed by value when subscribering a `Subscribe` onto a
 `Subscriber`, and a new `register_filter` method was added to
 `LookupSpan`, which returns a `FilterId` type. When a `Filtered`
 subscriber is added to a subscriber that implements `LookupSpan`, it
 calls `register_filter` in its `on_subscriber` hook to generate a
 unique ID for its filter. `FilterId` can be used to look up whether a
 span or event was enabled by the `Filtered` subscriber.

Internally, the filter results are stored in a simple bitset to reduce
the performance overhead of adding per-subscriber filtering as much as
possible. The `FilterId` type is actually a bitmask that's used to set
the bit corresponding to a `Filtered` subscriber in that bitmap when it
disables a span or event, and check whether the bit is set when querying
if a span or event was disabled. Setting a bit in the bitset and testing
whether its set is extremely efficient --- it's just a couple of bitwise
ops. Additionally, the "filter map" that tracks which filters disabled a
span or event is just a single `u64` --- we don't need to allocate a
costly `HashSet` or similar every time we want to filter an event. This
*does* mean that we are limited to 64 unique filters per
`Registry`...however, I would be _very_ surprised if anyone _ever_
builds a subscriber with 64 subscribers, let alone 64 separate
per-subscriber filters.

Additionally, the `Context` and `SpanRef` types now also potentially
carry `FilterId`s, so that they can filter span lookups and span
traversals for `Filtered` subscribers. When `Filtered` subscribers are
nested using the `Subscribeed`, the `Context` _combines_ the filter ID
of the inner and outer subscribers so that spans can be skipped if they
were disabled by either.

Finally, an implementation of the `Filter` trait was added for
`LevelFilter`, and new `FilterFn` and `DynFilterFn` structs were added
to produce a `Filter` implementation from a function pointer or closure.

### Notes

Some other miscellaneous factors to consider when reviewing this change
include:

* I did a bit of unrelated refactoring in the `subscriber` module. Since
  we were adding more code to the `Subscribe` module (the `Filter`
  trait), the `subscriber.rs` file got significantly longer and was
  becoming somewhat hard to navigate. Therefore, I split out the
  `Context` and `Subscribeed` types into their own files. Most of the
  code in those files is the same as it was before, although some of it
  was changed in order to support per-subscriber filtering. Apologies in
  advance to reviewers for this...
* In order to test this change, I added some new test-support code in
  `tracing-subscriber`. The `tracing` test support code provides a
  `Subscriber` implementation that can be configured to expect a
  particular set of spans and events, and assert that they occurred as
  expected. In order to test per-subscriber filtering, I added a new
  `Subscribe` implementation that does the same thing, but implements
  `Subscribe` rather than `Subscriber`. This allows testing
  per-subscriber filter behavior in the same way we test global filter
  behavior. Reviewers should feel free to just skim the test the test
  support changes, or skip reviewing them entirely.

## Future Work

There is a bunch of additional stuff we can do once this change lands.
In order to limit the size of this PR, I didn't make these changes yet,
but once this merges, we will want to consider some important follow up
changes:

* [ ] Currently, _only_ `tracing-subscriber`'s `Registry` is capable of
      serving as a root subscriber for per-subscriber filters. This is
      in order to avoid stabilizing a lot of the per-subscriber
      filtering design as public API without taking the time to ensure
      there are no serious issues. In the future, we will want to add
      new APIs to allow users to implement their own root `Subscriber`s
      that can host subscribers with per-subscriber filtering.
* [ ] The `EnvFilter` type doesn't implement the `Filter` trait and thus
      cannot currently be used as a per-subscriber filter. We should add
      an implementation of `Filter` for `EnvFilter`.
* [ ] The new `FilterFn` and `DynFilterFn` types could conceivably
      _also_ be used as global filters. We could consider adding
      `Subscribe` implementations to allow them to be used for global
      filtering.
* [ ] We could consider adding new `Filter` implementations for
      performing common filtering predicates. For example, "only enable
      spans/events with a particular target or a set of targets", "only
      enable spans", "only enable events", etc. Then, we could add a
      combinator API to combine these predicates to build new filters.
* [ ] It would be nice if the `Interest` system could be used to cache
      the result of multiple individual per-subscriber filter
      evaluations, so that we don't need to always use `enabled` when
      there are several per-subscriber filters that disagree on whether
      or not they want a given span/event. There are a lot of unused
      bits in `Interest` that could potentially be used for this
      purpose. However, this would require new API changes in
      `tracing-core`, and is therefore out of scope for this PR.

Closes #302
Closes #597
Closes #1348

Signed-off-by: Eliza Weisman <[email protected]>

Author: hawkw

tracing-mock/src/collector.rs

@@ -3,7 +3,6 @@ use super::{
     event::MockEvent,
     field as mock_field,
     span::{MockSpan, NewSpan},
-    Parent,
 };
 use std::{
     collections::{HashMap, VecDeque},
@@ -22,7 +21,7 @@ use tracing::{
 };
 
 #[derive(Debug, Eq, PartialEq)]
-enum Expect {
+pub enum Expect {
     Event(MockEvent),
     Enter(MockSpan),
     Exit(MockSpan),
@@ -221,7 +220,8 @@ where
                 if let Some(name) = expected_span.name() {
                     assert_eq!(name, span.name);
                 }
-                let mut checker = expected_values.checker(format!("span {}: ", span.name));
+                let context = format!("span {}: ", span.name);
+                let mut checker = expected_values.checker(&context, &self.name);
                 values.record(&mut checker);
                 checker.finish();
             }
@@ -234,69 +234,18 @@ where
         match self.expected.lock().unwrap().pop_front() {
             None => {}
             Some(Expect::Event(mut expected)) => {
-                let spans = self.spans.lock().unwrap();
-                expected.check(event);
-                match expected.parent {
-                    Some(Parent::ExplicitRoot) => {
-                        assert!(
-                            event.is_root(),
-                            "[{}] expected {:?} to be an explicit root event",
-                            self.name,
-                            name
-                        );
-                    }
-                    Some(Parent::Explicit(expected_parent)) => {
-                        let actual_parent =
-                            event.parent().and_then(|id| spans.get(id)).map(|s| s.name);
-                        assert_eq!(
-                            Some(expected_parent.as_ref()),
-                            actual_parent,
-                            "[{}] expected {:?} to have explicit parent {:?}",
-                            self.name,
-                            name,
-                            expected_parent,
-                        );
-                    }
-                    Some(Parent::ContextualRoot) => {
-                        assert!(
-                            event.is_contextual(),
-                            "[{}] expected {:?} to have a contextual parent",
-                            self.name,
-                            name
-                        );
-                        assert!(
-                            self.current.lock().unwrap().last().is_none(),
-                            "[{}] expected {:?} to be a root, but we were inside a span",
-                            self.name,
-                            name
-                        );
-                    }
-                    Some(Parent::Contextual(expected_parent)) => {
-                        assert!(
-                            event.is_contextual(),
-                            "[{}] expected {:?} to have a contextual parent",
-                            self.name,
-                            name
-                        );
-                        let stack = self.current.lock().unwrap();
-                        let actual_parent =
-                            stack.last().and_then(|id| spans.get(id)).map(|s| s.name);
-                        assert_eq!(
-                            Some(expected_parent.as_ref()),
-                            actual_parent,
-                            "[{}] expected {:?} to have contextual parent {:?}",
-                            self.name,
-                            name,
-                            expected_parent,
-                        );
-                    }
-                    None => {}
-                }
+                let get_parent_name = || {
+                    let stack = self.current.lock().unwrap();
+                    let spans = self.spans.lock().unwrap();
+                    event
+                        .parent()
+                        .and_then(|id| spans.get(id))
+                        .or_else(|| stack.last().and_then(|id| spans.get(id)))
+                        .map(|s| s.name.to_string())
+                };
+                expected.check(event, get_parent_name, &self.name);
             }
-            Some(ex) => ex.bad(
-                &self.name,
-                format_args!("[{}] observed event {:?}", self.name, event),
-            ),
+            Some(ex) => ex.bad(&self.name, format_args!("observed event {:#?}", event)),
         }
     }
 
@@ -320,70 +269,14 @@ where
         let mut spans = self.spans.lock().unwrap();
         if was_expected {
             if let Expect::NewSpan(mut expected) = expected.pop_front().unwrap() {
-                let name = meta.name();
-                expected
-                    .span
-                    .metadata
-                    .check(meta, format_args!("span `{}`", name));
-                let mut checker = expected.fields.checker(name.to_string());
-                span.record(&mut checker);
-                checker.finish();
-                match expected.parent {
-                    Some(Parent::ExplicitRoot) => {
-                        assert!(
-                            span.is_root(),
-                            "[{}] expected {:?} to be an explicit root span",
-                            self.name,
-                            name
-                        );
-                    }
-                    Some(Parent::Explicit(expected_parent)) => {
-                        let actual_parent =
-                            span.parent().and_then(|id| spans.get(id)).map(|s| s.name);
-                        assert_eq!(
-                            Some(expected_parent.as_ref()),
-                            actual_parent,
-                            "[{}] expected {:?} to have explicit parent {:?}",
-                            self.name,
-                            name,
-                            expected_parent,
-                        );
-                    }
-                    Some(Parent::ContextualRoot) => {
-                        assert!(
-                            span.is_contextual(),
-                            "[{}] expected {:?} to have a contextual parent",
-                            self.name,
-                            name
-                        );
-                        assert!(
-                            self.current.lock().unwrap().last().is_none(),
-                            "[{}] expected {:?} to be a root, but we were inside a span",
-                            self.name,
-                            name
-                        );
-                    }
-                    Some(Parent::Contextual(expected_parent)) => {
-                        assert!(
-                            span.is_contextual(),
-                            "[{}] expected {:?} to have a contextual parent",
-                            self.name,
-                            name
-                        );
-                        let stack = self.current.lock().unwrap();
-                        let actual_parent =
-                            stack.last().and_then(|id| spans.get(id)).map(|s| s.name);
-                        assert_eq!(
-                            Some(expected_parent.as_ref()),
-                            actual_parent,
-                            "[{}] expected {:?} to have contextual parent {:?}",
-                            self.name,
-                            name,
-                            expected_parent,
-                        );
-                    }
-                    None => {}
-                }
+                let get_parent_name = || {
+                    let stack = self.current.lock().unwrap();
+                    span.parent()
+                        .and_then(|id| spans.get(id))
+                        .or_else(|| stack.last().and_then(|id| spans.get(id)))
+                        .map(|s| s.name.to_string())
+                };
+                expected.check(span, get_parent_name, &self.name);
             }
         }
         spans.insert(
@@ -537,11 +430,15 @@ where
 }
 
 impl MockHandle {
+    pub fn new(expected: Arc<Mutex<VecDeque<Expect>>>, name: String) -> Self {
+        Self(expected, name)
+    }
+
     pub fn assert_finished(&self) {
         if let Ok(ref expected) = self.0.lock() {
             assert!(
                 !expected.iter().any(|thing| thing != &Expect::Nothing),
-                "[{}] more notifications expected: {:?}",
+                "\n[{}] more notifications expected: {:#?}",
                 self.1,
                 **expected
             );
@@ -550,26 +447,44 @@ impl MockHandle {
 }
 
 impl Expect {
-    fn bad(&self, name: impl AsRef<str>, what: fmt::Arguments<'_>) {
+    pub fn bad(&self, name: impl AsRef<str>, what: fmt::Arguments<'_>) {
         let name = name.as_ref();
         match self {
-            Expect::Event(e) => panic!("[{}] expected event {}, but {} instead", name, e, what,),
-            Expect::Enter(e) => panic!("[{}] expected to enter {} but {} instead", name, e, what,),
-            Expect::Exit(e) => panic!("[{}] expected to exit {} but {} instead", name, e, what,),
+            Expect::Event(e) => panic!(
+                "\n[{}] expected event {}\n[{}] but instead {}",
+                name, e, name, what,
+            ),
+            Expect::Enter(e) => panic!(
+                "\n[{}] expected to enter {}\n[{}] but instead {}",
+                name, e, name, what,
+            ),
+            Expect::Exit(e) => panic!(
+                "\n[{}] expected to exit {}\n[{}] but instead {}",
+                name, e, name, what,
+            ),
             Expect::CloneSpan(e) => {
-                panic!("[{}] expected to clone {} but {} instead", name, e, what,)
+                panic!(
+                    "\n[{}] expected to clone {}\n[{}] but instead {}",
+                    name, e, name, what,
+                )
             }
             Expect::DropSpan(e) => {
-                panic!("[{}] expected to drop {} but {} instead", name, e, what,)
+                panic!(
+                    "\n[{}] expected to drop {}\n[{}] but instead {}",
+                    name, e, name, what,
+                )
             }
             Expect::Visit(e, fields) => panic!(
-                "[{}] expected {} to record {} but {} instead",
-                name, e, fields, what,
+                "\n[{}] expected {} to record {}\n[{}] but instead {}",
+                name, e, fields, name, what,
+            ),
+            Expect::NewSpan(e) => panic!(
+                "\n[{}] expected {}\n[{}] but instead {}",
+                name, e, name, what
             ),
-            Expect::NewSpan(e) => panic!("[{}] expected {} but {} instead", name, e, what),
             Expect::Nothing => panic!(
-                "[{}] expected nothing else to happen, but {} instead",
-                name, what,
+                "\n[{}] expected nothing else to happen\n[{}] but {} instead",
+                name, name, what,
             ),
         }
     }

tracing-mock/src/event.rs

@@ -1,16 +1,17 @@
 #![allow(missing_docs)]
-use super::{field, metadata, Parent};
+use super::{field, metadata, span, Parent};
 
 use std::fmt;
 
 /// A mock event.
 ///
 /// This is intended for use with the mock subscriber API in the
 /// `subscriber` module.
-#[derive(Debug, Default, Eq, PartialEq)]
+#[derive(Default, Eq, PartialEq)]
 pub struct MockEvent {
     pub fields: Option<field::Expect>,
     pub(crate) parent: Option<Parent>,
+    in_spans: Vec<span::MockSpan>,
     metadata: metadata::Expect,
 }
 
@@ -20,6 +21,10 @@ pub fn mock() -> MockEvent {
     }
 }
 
+pub fn msg(message: impl fmt::Display) -> MockEvent {
+    mock().with_fields(field::msg(message))
+}
+
 impl MockEvent {
     pub fn named<I>(self, name: I) -> Self
     where
@@ -78,17 +83,49 @@ impl MockEvent {
         }
     }
 
-    pub(crate) fn check(&mut self, event: &tracing::Event<'_>) {
+    pub fn check(
+        &mut self,
+        event: &tracing::Event<'_>,
+        get_parent_name: impl FnOnce() -> Option<String>,
+        collector_name: &str,
+    ) {
         let meta = event.metadata();
         let name = meta.name();
         self.metadata
-            .check(meta, format_args!("event \"{}\"", name));
-        assert!(meta.is_event(), "expected {}, but got {:?}", self, event);
+            .check(meta, format_args!("event \"{}\"", name), collector_name);
+        assert!(
+            meta.is_event(),
+            "[{}] expected {}, but got {:?}",
+            collector_name,
+            self,
+            event
+        );
         if let Some(ref mut expected_fields) = self.fields {
-            let mut checker = expected_fields.checker(name.to_string());
+            let mut checker = expected_fields.checker(name, collector_name);
             event.record(&mut checker);
             checker.finish();
         }
+
+        if let Some(ref expected_parent) = self.parent {
+            let actual_parent = get_parent_name();
+            expected_parent.check_parent_name(
+                actual_parent.as_deref(),
+                event.parent().cloned(),
+                event.metadata().name(),
+                collector_name,
+            )
+        }
+    }
+
+    pub fn in_scope(self, spans: impl IntoIterator<Item = span::MockSpan>) -> Self {
+        Self {
+            in_spans: spans.into_iter().collect(),
+            ..self
+        }
+    }
+
+    pub fn scope_mut(&mut self) -> &mut [span::MockSpan] {
+        &mut self.in_spans[..]
     }
 }
 
@@ -97,3 +134,35 @@ impl fmt::Display for MockEvent {
         write!(f, "an event{}", self.metadata)
     }
 }
+
+impl fmt::Debug for MockEvent {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let mut s = f.debug_struct("MockEvent");
+
+        if let Some(ref name) = self.metadata.name {
+            s.field("name", name);
+        }
+
+        if let Some(ref target) = self.metadata.target {
+            s.field("target", target);
+        }
+
+        if let Some(ref level) = self.metadata.level {
+            s.field("level", &format_args!("{:?}", level));
+        }
+
+        if let Some(ref fields) = self.fields {
+            s.field("fields", fields);
+        }
+
+        if let Some(ref parent) = self.parent {
+            s.field("parent", &format_args!("{:?}", parent));
+        }
+
+        if !self.in_spans.is_empty() {
+            s.field("in_spans", &self.in_spans);
+        }
+
+        s.finish()
+    }
+}