subscriber: make Registry::enter/exit much faster (#1058) (#1059)

This backports #1058 to v0.1.x. This was already approved on master.

## Motivation

We've tried very hard to make sure that entering and exiting spans is
lightweight...in the `tracing-core` core machinery. Unfortunately, we
haven't done any benchmarking of how subscriber implementations actually
handle enter/exit events. It turns out that in `tracing-subscriber`'s
`Registry`, there's actually significant overhead for entering a span:
calling `span.enter()` may take as long as 69 ns (on my machine).

## Solution

I've written some microbenchmarks for entering and exiting enabled spans
using `tracing-subscriber::fmt`, comparing them with the overhead of
calling `enter` on an enabled span. Based on this, I've made some
performance improvements. These optimizations include:

- Removing the `HashSet` that's used for tracking previously entered
  span IDs, in favor of linear search. Span stacks probably never deep
  enough for a hashmap to be faster than iterating over a couple of
  vec indices.
- Preallocating the vec used for the span stack to hold at least 64
  elements. This means we'll never have a lag spike from reallocating,
  as I think it'll almost never be deeper than 64 IDs.
- Only cloning/dropping an ID's ref count for the _first_ ID in the stack.

This makes entering and exiting enabled spans significantly faster:
![image](https://user-images.githubusercontent.com/2796466/96798681-3fc85000-13b6-11eb-9e85-7602d918ee09.png)

It would be nice to continue optimizing this, but this might be about
the best it gets given the other requirements that we're now making
assertions about.

Signed-off-by: Eliza Weisman <eliza@buoyant.io>

tracing-subscriber/Cargo.toml

@@ -84,3 +84,7 @@ harness = false
 [[bench]]
 name = "fmt"
 harness = false
+
+[[bench]]
+name = "enter"
+harness = false

tracing-subscriber/benches/enter.rs

@@ -0,0 +1,64 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+use tracing_subscriber::prelude::*;
+
+fn enter(c: &mut Criterion) {
+    let mut group = c.benchmark_group("enter");
+    let _subscriber = tracing_subscriber::fmt()
+        .with_max_level(tracing::Level::INFO)
+        .finish()
+        .set_default();
+    group.bench_function("enabled", |b| {
+        let span = tracing::info_span!("foo");
+        b.iter_with_large_drop(|| span.enter())
+    });
+    group.bench_function("disabled", |b| {
+        let span = tracing::debug_span!("foo");
+        b.iter_with_large_drop(|| span.enter())
+    });
+}
+
+fn enter_exit(c: &mut Criterion) {
+    let mut group = c.benchmark_group("enter_exit");
+    let _subscriber = tracing_subscriber::fmt()
+        .with_max_level(tracing::Level::INFO)
+        .finish()
+        .set_default();
+    group.bench_function("enabled", |b| {
+        let span = tracing::info_span!("foo");
+        b.iter(|| span.enter())
+    });
+    group.bench_function("disabled", |b| {
+        let span = tracing::debug_span!("foo");
+        b.iter(|| span.enter())
+    });
+}
+
+fn enter_many(c: &mut Criterion) {
+    let mut group = c.benchmark_group("enter_many");
+    let _subscriber = tracing_subscriber::fmt()
+        .with_max_level(tracing::Level::INFO)
+        .finish()
+        .set_default();
+    group.bench_function("enabled", |b| {
+        let span1 = tracing::info_span!("span1");
+        let _e1 = span1.enter();
+        let span2 = tracing::info_span!("span2");
+        let _e2 = span2.enter();
+        let span3 = tracing::info_span!("span3");
+        let _e3 = span3.enter();
+        let span = tracing::info_span!("foo");
+        b.iter_with_large_drop(|| span.enter())
+    });
+    group.bench_function("disabled", |b| {
+        let span1 = tracing::info_span!("span1");
+        let _e1 = span1.enter();
+        let span2 = tracing::info_span!("span2");
+        let _e2 = span2.enter();
+        let span3 = tracing::info_span!("span3");
+        let _e3 = span3.enter();
+        let span = tracing::debug_span!("foo");
+        b.iter_with_large_drop(|| span.enter())
+    });
+}
+criterion_group!(benches, enter, enter_exit, enter_many);
+criterion_main!(benches);

tracing-subscriber/src/registry/sharded.rs

@@ -202,19 +202,21 @@ impl Subscriber for Registry {
     fn event(&self, _: &Event<'_>) {}
 
     fn enter(&self, id: &span::Id) {
-        self.current_spans
+        if self
+            .current_spans
             .get_or_default()
             .borrow_mut()
-            .push(self.clone_span(id));
+            .push(id.clone())
+        {
+            self.clone_span(id);
+        }
     }
 
     fn exit(&self, id: &span::Id) {
-        if let Some(id) = self
-            .current_spans
-            .get()
-            .and_then(|spans| spans.borrow_mut().pop(id))
-        {
-            dispatcher::get_default(|dispatch| dispatch.try_close(id.clone()));
+        if let Some(spans) = self.current_spans.get() {
+            if spans.borrow_mut().pop(id) {
+                dispatcher::get_default(|dispatch| dispatch.try_close(id.clone()));
+            }
         }
     }
 

tracing-subscriber/src/registry/stack.rs

@@ -1,5 +1,3 @@
-use std::collections::HashSet;
-
 pub(crate) use tracing_core::span::Id;
 
 #[derive(Debug)]
@@ -15,34 +13,29 @@ struct ContextId {
 #[derive(Debug, Default)]
 pub(crate) struct SpanStack {
     stack: Vec<ContextId>,
-    ids: HashSet<Id>,
 }
 
 impl SpanStack {
-    pub(crate) fn push(&mut self, id: Id) {
-        let duplicate = self.ids.contains(&id);
-        if !duplicate {
-            self.ids.insert(id.clone());
-        }
-        self.stack.push(ContextId { id, duplicate })
+    #[inline]
+    pub(crate) fn push(&mut self, id: Id) -> bool {
+        let duplicate = self.stack.iter().any(|i| i.id == id);
+        self.stack.push(ContextId { id, duplicate });
+        !duplicate
     }
 
-    pub(crate) fn pop(&mut self, expected_id: &Id) -> Option<Id> {
+    #[inline]
+    pub(crate) fn pop(&mut self, expected_id: &Id) -> bool {
         if let Some((idx, _)) = self
             .stack
             .iter()
             .enumerate()
             .rev()
             .find(|(_, ctx_id)| ctx_id.id == *expected_id)
         {
-            let ContextId { id, duplicate } = self.stack.remove(idx);
-            if !duplicate {
-                self.ids.remove(&id);
-            }
-            Some(id)
-        } else {
-            None
+            let ContextId { id: _, duplicate } = self.stack.remove(idx);
+            return !duplicate;
         }
+        false
     }
 
     #[inline]
@@ -65,7 +58,7 @@ mod tests {
         let id = Id::from_u64(1);
         stack.push(id.clone());
 
-        assert_eq!(Some(id.clone()), stack.pop(&id));
+        assert!(stack.pop(&id));
     }
 
     #[test]
@@ -75,6 +68,6 @@ mod tests {
         stack.push(Id::from_u64(2));
 
         let id = Id::from_u64(1);
-        assert_eq!(Some(id.clone()), stack.pop(&id));
+        assert!(stack.pop(&id));
     }
 }