feat(subscriber): spill callsites into hash set (#97)

This changes the `Callsites` behavior when the array of callsites is
full. Currently, we panic in this case. This means that we are quite
generous with array sizes, for cases where, e.g., multiple async
runtimes are in use. However, being more generous with the array's
length makes the linear search performance worse.

This branch replaces the panicking behavior with a spillover behavior.
Once the array of callsites is full, we will now store any additional
callsites in a `HashSet`, rather than panicking. This means we can make
the arrays a bit shorter, and (perhaps more importantly) it means we
will no longer panic in the (rare) case where an app contains a big pile
of interesting callsites.

The spillover `HashSet` is protected by a `RwLock`, which is kind of
a bummer, since it may be locked when checking if a span/event is
in the set of callsites we care about (but only if we have spilled over).
However, it should be _contended_ only very rarely, since writes only
occur when registering a new callsite. 

I added an optional `parking_lot` feature to use `parking_lot`'s
`RwLock` implementation, which likely offers better performance than
`std`'s lock (especially when uncontended, which this lock often is).
The feature is disabled by default, for users who don't want the
additional dependency.

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

console-subscriber/Cargo.toml

@@ -5,6 +5,9 @@ authors = ["Eliza Weisman <eliza@buoyant.io>"]
 edition = "2018"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+[features]
+default = []
+parking_lot = ["parking_lot_crate", "tracing-subscriber/parking_lot"]
 
 [dependencies]
 
@@ -19,6 +22,9 @@ futures = { version = "0.3", default-features = false }
 hdrhistogram = { version = "7.3.0", default-features = false, features = ["serialization"] }
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
+# The parking_lot dependency is renamed, because we want our `parking_lot`
+# feature to also enable `tracing-subscriber`'s parking_lot feature flag.
+parking_lot_crate = { package = "parking_lot", version = "0.11", optional = true }
 
 [dev-dependencies]
 

console-subscriber/src/callsites.rs

@@ -1,12 +1,15 @@
+use crate::sync::RwLock;
 use std::{
+    collections::HashSet,
     fmt, ptr,
     sync::atomic::{AtomicPtr, AtomicUsize, Ordering},
 };
-use tracing_core::Metadata;
+use tracing_core::{callsite, Metadata};
 
 pub(crate) struct Callsites<const MAX_CALLSITES: usize> {
     ptrs: [AtomicPtr<Metadata<'static>>; MAX_CALLSITES],
     len: AtomicUsize,
+    spill: RwLock<HashSet<callsite::Identifier>>,
 }
 
 impl<const MAX_CALLSITES: usize> Callsites<MAX_CALLSITES> {
@@ -20,31 +23,58 @@ impl<const MAX_CALLSITES: usize> Callsites<MAX_CALLSITES> {
         }
 
         let idx = self.len.fetch_add(1, Ordering::AcqRel);
-        assert!(
-            idx < MAX_CALLSITES,
-            "you tried to store more than {} callsites, \
-            time to make the callsite sets bigger i guess \
-            (please open an issue for this)",
-            MAX_CALLSITES,
-        );
-        self.ptrs[idx]
-            .compare_exchange(
-                ptr::null_mut(),
-                callsite as *const _ as *mut _,
-                Ordering::AcqRel,
-                Ordering::Acquire,
-            )
-            .expect("a callsite would have been clobbered by `insert` (this is a bug)");
+        if idx <= MAX_CALLSITES {
+            // If there's still room in the callsites array, stick the address
+            // in there.
+            self.ptrs[idx]
+                .compare_exchange(
+                    ptr::null_mut(),
+                    callsite as *const _ as *mut _,
+                    Ordering::AcqRel,
+                    Ordering::Acquire,
+                )
+                .expect("a callsite would have been clobbered by `insert` (this is a bug)");
+        } else {
+            // Otherwise, we've filled the callsite array (sad!). Spill over
+            // into a hash set.
+            self.spill.write().insert(callsite.callsite());
+        }
     }
 
     pub(crate) fn contains(&self, callsite: &'static Metadata<'static>) -> bool {
-        let len = self.len.load(Ordering::Acquire);
-        for cs in &self.ptrs[..len] {
-            if ptr::eq(cs.load(Ordering::Acquire), callsite) {
-                return true;
+        let mut start = 0;
+        let mut len = self.len.load(Ordering::Acquire);
+        loop {
+            for cs in &self.ptrs[start..len] {
+                if ptr::eq(cs.load(Ordering::Acquire), callsite) {
+                    return true;
+                }
             }
+
+            // Did the length change while we were iterating over the callsite array?
+            let new_len = self.len.load(Ordering::Acquire);
+            if new_len > len {
+                // If so, check again to see if the callsite is contained in any
+                // callsites that were pushed since the last time we loaded `self.len`.
+                start = len;
+                len = new_len;
+                continue;
+            }
+
+            // If the callsite array is not full, we have checked everything.
+            if len <= MAX_CALLSITES {
+                return false;
+            }
+
+            // Otherwise, we may have spilled over to the slower fallback hash
+            // set. Check that.
+            return self.check_spill(callsite);
         }
-        false
+    }
+
+    #[cold]
+    fn check_spill(&self, callsite: &'static Metadata<'static>) -> bool {
+        self.spill.read().contains(&callsite.callsite())
     }
 }
 
@@ -69,6 +99,7 @@ impl<const MAX_CALLSITES: usize> Default for Callsites<MAX_CALLSITES> {
         Self {
             ptrs: [NULLPTR; MAX_CALLSITES],
             len: AtomicUsize::new(0),
+            spill: Default::default(),
         }
     }
 }
@@ -80,6 +111,7 @@ impl<const MAX_CALLSITES: usize> fmt::Debug for Callsites<MAX_CALLSITES> {
             .field("ptrs", &&self.ptrs[..len])
             .field("len", &len)
             .field("max_callsites", &MAX_CALLSITES)
+            .field("spill", &self.spill)
             .finish()
     }
 }

console-subscriber/src/lib.rs

@@ -21,6 +21,7 @@ mod builder;
 mod callsites;
 mod init;
 mod record;
+pub(crate) mod sync;
 
 use aggregator::Aggregator;
 pub use builder::Builder;
@@ -40,15 +41,16 @@ pub struct TasksLayer {
     /// Set of callsites for spans representing spawned tasks.
     ///
     /// For task spans, each runtime these will have like, 1-5 callsites in it, max, so
-    /// 16 is probably fine. For async operations, we may need a bigger callsites array.
-    spawn_callsites: Callsites<16>,
+    /// 8 should be plenty. If several runtimes are in use, we may have to spill
+    /// over into the backup hashmap, but it's unlikely.
+    spawn_callsites: Callsites<8>,
 
     /// Set of callsites for events representing waker operations.
     ///
-    /// 32 is probably a reasonable number of waker ops; it's a bit generous if
+    /// 16 is probably a reasonable number of waker ops; it's a bit generous if
     /// there's only one async runtime library in use, but if there are multiple,
     /// they might all have their own sets of waker ops.
-    waker_callsites: Callsites<32>,
+    waker_callsites: Callsites<16>,
 }
 
 pub struct Server {

console-subscriber/src/sync.rs

@@ -0,0 +1,49 @@
+// Some of these methods and re-exports may not be used currently.
+#![allow(dead_code, unused_imports)]
+
+#[cfg(feature = "parking_lot")]
+pub(crate) use parking_lot_crate::{RwLock, RwLockReadGuard, RwLockWriteGuard};
+
+#[cfg(not(feature = "parking_lot"))]
+pub(crate) use self::std_impl::*;
+
+#[cfg(not(feature = "parking_lot"))]
+mod std_impl {
+    use std::sync::{self, PoisonError, TryLockError};
+    pub use std::sync::{RwLockReadGuard, RwLockWriteGuard};
+
+    #[derive(Debug, Default)]
+    pub(crate) struct RwLock<T: ?Sized>(sync::RwLock<T>);
+
+    impl<T> RwLock<T> {
+        pub(crate) fn new(data: T) -> Self {
+            Self(sync::RwLock::new(data))
+        }
+    }
+
+    impl<T: ?Sized> RwLock<T> {
+        pub(crate) fn read(&self) -> RwLockReadGuard<'_, T> {
+            self.0.read().unwrap_or_else(PoisonError::into_inner)
+        }
+
+        pub(crate) fn try_read(&self) -> Option<RwLockReadGuard<'_, T>> {
+            match self.0.try_read() {
+                Ok(guard) => Some(guard),
+                Err(TryLockError::Poisoned(p)) => Some(p.into_inner()),
+                Err(TryLockError::WouldBlock) => None,
+            }
+        }
+
+        pub(crate) fn write(&self) -> RwLockWriteGuard<'_, T> {
+            self.0.write().unwrap_or_else(PoisonError::into_inner)
+        }
+
+        pub(crate) fn try_write(&self) -> Option<RwLockWriteGuard<'_, T>> {
+            match self.0.try_write() {
+                Ok(guard) => Some(guard),
+                Err(TryLockError::Poisoned(p)) => Some(p.into_inner()),
+                Err(TryLockError::WouldBlock) => None,
+            }
+        }
+    }
+}