Protect against HashMap rehash spikes in UnresolvedSamples

While adding support for in-kernel unwinding with sample rates as high
as 4KHz, sample dropping became an issue. Looking closer using a
`perf record` of compiler running on 16 CPU, 4KHz samples, rehashing was
taking over 260ms in the worst case (see "Stack Chart" and highlighting
rehash): https://share.firefox.dev/3Y32wsT

This changes so that UnresolvedStacks uses a BTreeMap for the interning
storage, which has a much more bounded O(log n) worst case insertion.
An LRU hash map serves on top as an O(1) fast path.

The new implementation is roughly on par and perhaps slightly faster,
and processing time spikes are almost gone:
https://share.firefox.dev/4pMM9wy

Footnote, the `stacks` array itself also suffers from O(n) realloc
behavior, but that doesn't seem to be problematic *yet*, so I'll leave
this for when it's found to be a problem in the future.

Author: ishitatsuyuki

Cargo.lock

@@ -43,6 +43,7 @@ serde = "1.0.204"
 wholesym = { version = "0.8.1", path = "../wholesym", features = ["api"]}
 platform-dirs = "0.3"
 rustc-hash = "2"
+schnellru = "0.2"
 mio = { version = "1", features = ["os-ext", "os-poll"] }
 ctrlc = "3.4.4"
 log = "0.4.21"

samply/Cargo.toml

@@ -7,7 +7,7 @@ use linux_perf_event_reader::CpuMode;
 
 pub type FastHashMap<K, V> = rustc_hash::FxHashMap<K, V>;
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub enum StackMode {
     User,
     Kernel,
@@ -37,7 +37,7 @@ impl From<CpuMode> for StackMode {
     }
 }
 
-#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub enum StackFrame {
     InstructionPointer(u64, StackMode),
     ReturnAddress(u64, StackMode),

samply/src/shared/types.rs

@@ -1,6 +1,9 @@
 use std::collections::hash_map::Entry;
+use std::collections::BTreeMap;
 
 use fxprof_processed_profile::{CpuDelta, FrameHandle, MarkerHandle, ThreadHandle, Timestamp};
+use rustc_hash::FxBuildHasher;
+use schnellru::{ByLength, LruMap};
 
 use super::types::{FastHashMap, StackFrame, StackMode};
 
@@ -154,18 +157,43 @@ pub struct SampleData {
     pub weight: i32,
 }
 
-#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
+#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord)]
 pub struct UnresolvedStackHandle(u32);
 
 impl UnresolvedStackHandle {
     /// Represents the empty stack / the root stack node
     pub const EMPTY: Self = Self(u32::MAX);
 }
 
-#[derive(Debug, Clone, Default)]
+/// An interner that recursively interns stacks in the form of (prefix, frame).
+///
+/// The interner avoids the worst case O(n) runtime behavior of hash maps due to rehashes. Under the
+/// hood, BTreeMap is used to store the full mapping with worst case O(log n) inserts, while an LRU
+/// hash map maintains a O(1) fast path for the majority of accesses.
 pub struct UnresolvedStacks {
     pub stacks: Vec<(UnresolvedStackHandle, StackFrame)>, // (prefix, frame)
-    pub stack_lookup: FastHashMap<(UnresolvedStackHandle, StackFrame), UnresolvedStackHandle>, // (prefix, frame) -> stack index
+    stack_lookup: BTreeMap<(UnresolvedStackHandle, StackFrame), UnresolvedStackHandle>, // (prefix, frame) -> stack index
+    stack_cache: LruMap<(UnresolvedStackHandle, StackFrame), UnresolvedStackHandle, ByLength, FxBuildHasher>,
+}
+
+impl UnresolvedStacks {
+    // Two things to consider for the size of table:
+    // - The cache should be small enough to fit in CPU cache
+    // - LruMap uses SwissTable under the hood where erases use tombstones and rehashing is required
+    //   for compaction. The cache should be small enough to bound rehashing cost
+    //
+    // Assuming 64B per entry, a 256KB L2 can fit 4K entries, and rehashing time seems to be under 1ms.
+    const CACHE_SIZE: u32 = 4096;
+}
+
+impl Default for UnresolvedStacks {
+    fn default() -> Self {
+        Self {
+            stacks: Vec::new(),
+            stack_cache: LruMap::with_hasher(ByLength::new(Self::CACHE_SIZE), FxBuildHasher),
+            stack_lookup: BTreeMap::new(),
+        }
+    }
 }
 
 impl UnresolvedStacks {
@@ -182,11 +210,17 @@ impl UnresolvedStacks {
     ) -> UnresolvedStackHandle {
         for frame in frames {
             let x = (prefix, frame);
+            if let Some(node) = self.stack_cache.get(&x).copied() {
+                prefix = node;
+                continue;
+            }
+
             let node = *self.stack_lookup.entry(x).or_insert_with(|| {
                 let new_index = self.stacks.len() as u32;
                 self.stacks.push(x);
                 UnresolvedStackHandle(new_index)
             });
+            self.stack_cache.insert(x, node);
             prefix = node;
         }
         prefix
@@ -198,17 +232,10 @@ impl UnresolvedStacks {
         &mut self,
         frames: impl Iterator<Item = StackFrame>,
     ) -> UnresolvedStackHandle {
-        let mut prefix = UnresolvedStackHandle::EMPTY;
-        for frame in frames.filter(|f| f.stack_mode() != Some(StackMode::Kernel)) {
-            let x = (prefix, frame);
-            let node = *self.stack_lookup.entry(x).or_insert_with(|| {
-                let new_index = self.stacks.len() as u32;
-                self.stacks.push(x);
-                UnresolvedStackHandle(new_index)
-            });
-            prefix = node;
-        }
-        prefix
+        self.convert_with_prefix(
+            UnresolvedStackHandle::EMPTY,
+            frames.filter(|f| f.stack_mode() != Some(StackMode::Kernel)),
+        )
     }
 
     // Appends the stack to `buf`, starting with the callee-most frame.