perf(memtrack): resolve stack fp chains off the ring poll thread

Each stack record needs its frame-pointer chain looked up in the
stack_traces map, which is a syscall per record. Doing that inside the
ring-buffer parse callback made the poll thread pay it, so a burst of
stacks could push it behind the producer and records were dropped.

ResolvingPoller wraps a RingBufferPoller with a dedicated resolver
thread: the poll thread only parses (event, stackid) and hands it over an
internal channel, and the resolver does the map lookup and forwards the
completed event. Drop order keeps the existing shutdown contract, the
ring is dropped first so its poll thread joins and closes the internal
sender, which lets the resolver drain what it already has before its
join returns.

Author: not-matthias

crates/memtrack/src/ebpf/memtrack/mod.rs

@@ -6,7 +6,7 @@ use std::collections::HashMap;
 use std::mem::MaybeUninit;
 use std::path::Path;
 
-use crate::ebpf::poller::RingBufferPoller;
+use crate::ebpf::poller::{RingBufferPoller, ThreadedRingBufferPoller};
 use crate::ebpf::tracker::TrackerOptions;
 
 mod token {
@@ -248,33 +248,37 @@ impl MemtrackBpf {
         ))
     }
 
-    /// Poll the stack-record ring buffer into `tx`.
+    /// Poll the stack-record ring buffer into `tx`. The map lookup for each
+    /// stack's frame-pointer chain is a syscall, so it runs on a resolver
+    /// thread rather than the ring poll thread: see [`ThreadedRingBufferPoller`].
     pub(crate) fn poll_stacks(
         &self,
         poll_interval_ms: u64,
         tx: std::sync::mpsc::Sender<runner_shared::artifacts::MemtrackEvent>,
-    ) -> Result<RingBufferPoller> {
+    ) -> Result<ThreadedRingBufferPoller> {
         use crate::ebpf::events;
         use runner_shared::artifacts::MemtrackEventKind;
 
-        // The poller outlives this borrow of the skeleton, so the chain lookup
-        // needs an owned handle rather than a reference to the skeleton map.
+        // The resolver thread outlives this borrow of the skeleton, so the
+        // chain lookup needs an owned handle rather than a reference to the
+        // skeleton map.
         let stack_traces = with_skel!(self, skel => {
             libbpf_rs::MapHandle::try_from(&skel.maps.stack_traces)
                 .context("Failed to create handle for stack_traces map")?
         });
 
-        let parse = move |data: &[u8]| {
-            let (mut event, stackid) = events::parse_stack(data)?;
-            if let MemtrackEventKind::Stack { record } = &mut event.kind {
-                record.fp_chain = events::fp_chain(&stack_traces, stackid);
-            }
-            Some(event)
-        };
+        let resolve =
+            move |(mut event, stackid): (runner_shared::artifacts::MemtrackEvent, i64)| {
+                if let MemtrackEventKind::Stack { record } = &mut event.kind {
+                    record.fp_chain = events::fp_chain(&stack_traces, stackid);
+                }
+                event
+            };
 
-        with_skel!(self, skel => RingBufferPoller::new(
+        with_skel!(self, skel => ThreadedRingBufferPoller::new(
             &skel.maps.stacks,
-            parse,
+            events::parse_stack,
+            resolve,
             tx,
             poll_interval_ms,
         ))

crates/memtrack/src/ebpf/poller.rs

@@ -77,3 +77,63 @@ impl Drop for RingBufferPoller {
         }
     }
 }
+
+/// A [`RingBufferPoller`] whose parsed items need a further, potentially
+/// expensive step (e.g. a BPF map lookup, which is a syscall) before they are
+/// forwarded on `tx`. That step runs on a dedicated resolver thread instead
+/// of the poll thread, so a slow per-record resolve can't make the poll
+/// thread fall behind the ring and drop records.
+pub struct ThreadedRingBufferPoller {
+    // Declaration order is the shutdown order: dropping `ring` first
+    // disconnects its control channel and joins its poll thread, which drops
+    // the internal sender the resolver reads from. Only then can the
+    // resolver's `recv` loop observe disconnection, finish forwarding
+    // whatever it already has, and let `resolver`'s join below return. This
+    // gives callers the same "disconnect, fully drain, then join" contract
+    // as a plain `RingBufferPoller`.
+    ring: Option<RingBufferPoller>,
+    resolver: Option<JoinHandle<()>>,
+}
+
+impl ThreadedRingBufferPoller {
+    /// Poll `rb_map` with `parse` like [`RingBufferPoller::new`], but run
+    /// `resolve` on a separate thread: `parse` results are forwarded over an
+    /// internal channel, and `resolve` turns each one into the value sent on
+    /// `tx`.
+    pub fn new<M, T, U, F, R>(
+        rb_map: &M,
+        parse: F,
+        resolve: R,
+        tx: Sender<U>,
+        poll_interval_ms: u64,
+    ) -> Result<Self>
+    where
+        M: MapCore,
+        T: Send + 'static,
+        U: Send + 'static,
+        F: Fn(&[u8]) -> Option<T> + Send + 'static,
+        R: Fn(T) -> U + Send + 'static,
+    {
+        let (parsed_tx, parsed_rx) = mpsc::channel::<T>();
+        let ring = RingBufferPoller::new(rb_map, parse, parsed_tx, poll_interval_ms)?;
+        let resolver = std::thread::spawn(move || {
+            for item in parsed_rx {
+                let _ = tx.send(resolve(item));
+            }
+        });
+
+        Ok(Self {
+            ring: Some(ring),
+            resolver: Some(resolver),
+        })
+    }
+}
+
+impl Drop for ThreadedRingBufferPoller {
+    fn drop(&mut self) {
+        drop(self.ring.take());
+        if let Some(resolver) = self.resolver.take() {
+            let _ = resolver.join();
+        }
+    }
+}

crates/memtrack/src/session.rs

@@ -1,4 +1,4 @@
-use crate::ebpf::poller::RingBufferPoller;
+use crate::ebpf::poller::{RingBufferPoller, ThreadedRingBufferPoller};
 use crate::perf_mappings::PerfMappingPoller;
 use crate::prelude::*;
 use runner_shared::artifacts::MemtrackEvent;
@@ -17,7 +17,7 @@ pub struct Session {
     // join their poll threads before PerfMappingPoller drops and emits its
     // buffered Mapping records as the terminal stream suffix.
     _poller: RingBufferPoller,
-    _stack_poller: Option<RingBufferPoller>,
+    _stack_poller: Option<ThreadedRingBufferPoller>,
     _perf_mapping_poller: Option<PerfMappingPoller>,
 }
 
@@ -26,7 +26,7 @@ impl Session {
         child: Child,
         events: Receiver<MemtrackEvent>,
         poller: RingBufferPoller,
-        stack_poller: Option<RingBufferPoller>,
+        stack_poller: Option<ThreadedRingBufferPoller>,
         perf_mapping_poller: Option<PerfMappingPoller>,
     ) -> Self {
         Self {