perf(memtrack): batch ring-buffer events over the channel

The poll thread sent one item per channel message, and std's mpsc allocates a
31-slot block per 31 messages, so a run that captured 704k events also
allocated 22.7k blocks purely to hand them over.

The callback now fills a shared 1024-item buffer and sends it whole. Partial
batches are flushed after every poll, after the drain-path consume before its
ack, and on the shutdown consume, so `drain()` keeps promising that all
pending entries sit in the channel once it returns.

Consumers take `Vec<T>` and flatten: the encode pipeline keeps its
`IntoIterator<Item = MemtrackEvent>` contract unchanged.

Author: not-matthias

.github/workflows/ci.yml

@@ -159,8 +159,8 @@ jobs:
       fail-fast: false
       matrix:
         mode: [walltime, memory]
-        # One config file per workload so each job stays short; both physical
-        # variants of a workload live in the same file to be measured back to back.
+        # Both physical variants of a workload live in the same file, so they
+        # are measured back to back.
         workload: [ls, dd, tar]
     env:
       CODSPEED_REV: ${{ github.event.pull_request.head.sha || github.sha }}

crates/memtrack/src/ebpf/attach_worker.rs

@@ -121,7 +121,7 @@ impl Drop for AttachWorker {
 
 struct Worker {
     poller: RingBufferPoller,
-    rx: mpsc::Receiver<AttachRequest>,
+    rx: mpsc::Receiver<Vec<AttachRequest>>,
     bpf: Arc<Mutex<MemtrackBpf>>,
     shutdown: Arc<AtomicBool>,
     fatal: Arc<Mutex<Option<String>>>,
@@ -139,7 +139,7 @@ impl Worker {
                     self.record_fatal(e);
                     return;
                 }
-                let mut batch: Vec<AttachRequest> = self.rx.try_iter().collect();
+                let mut batch: Vec<AttachRequest> = self.rx.try_iter().flatten().collect();
                 if batch.is_empty() {
                     break;
                 }
@@ -151,12 +151,14 @@ impl Worker {
             }
 
             let first = match self.rx.recv_timeout(RECV_TIMEOUT) {
-                Ok(req) => req,
+                Ok(reqs) => reqs,
                 Err(RecvTimeoutError::Timeout) => continue,
                 Err(RecvTimeoutError::Disconnected) => return,
             };
-            let mut batch: Vec<AttachRequest> =
-                std::iter::once(first).chain(self.rx.try_iter()).collect();
+            let mut batch: Vec<AttachRequest> = first
+                .into_iter()
+                .chain(self.rx.try_iter().flatten())
+                .collect();
 
             if let Err(e) = self.process_batch(&mut batch, &mut known) {
                 self.record_fatal(e);
@@ -195,7 +197,7 @@ impl Worker {
 
             // Every producer is stopped, so a synchronous drain is complete.
             self.poller.drain()?;
-            batch.extend(self.rx.try_iter());
+            batch.extend(self.rx.try_iter().flatten());
         }
 
         let mut seen: HashSet<(u64, u64)> = HashSet::new();

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

@@ -218,7 +218,7 @@ impl MemtrackBpf {
     pub fn poll_events_with_channel(
         &self,
         poll_interval_ms: u64,
-        tx: std::sync::mpsc::Sender<runner_shared::artifacts::MemtrackEvent>,
+        tx: std::sync::mpsc::Sender<Vec<runner_shared::artifacts::MemtrackEvent>>,
     ) -> Result<RingBufferPoller> {
         with_skel!(self, skel => RingBufferPoller::new(
             &skel.maps.events,
@@ -233,7 +233,7 @@ impl MemtrackBpf {
     pub(crate) fn poll_attach_with_channel(
         &self,
         poll_interval_ms: u64,
-        tx: std::sync::mpsc::Sender<crate::ebpf::events::AttachRequest>,
+        tx: std::sync::mpsc::Sender<Vec<crate::ebpf::events::AttachRequest>>,
     ) -> Result<RingBufferPoller> {
         with_skel!(self, skel => RingBufferPoller::new(
             &skel.maps.attach_requests,

crates/memtrack/src/ebpf/poller.rs

@@ -1,11 +1,58 @@
 use anyhow::{Context, Result};
 use libbpf_rs::{MapCore, RingBufferBuilder};
+use parking_lot::Mutex;
+use std::sync::Arc;
 use std::sync::mpsc::{self, RecvTimeoutError, Sender};
 use std::thread::JoinHandle;
 use std::time::Duration;
 
+/// Items buffered before a channel send. `std::sync::mpsc` allocates a block
+/// every 31 messages, so sending one item at a time makes that allocation
+/// dominate the pipeline; batching amortizes it over a whole batch.
+const BATCH_ITEMS: usize = 1024;
+
+/// The lock is released before the send so a slow consumer never blocks the
+/// ring-buffer callback.
+fn flush_batch<T>(batch: &Mutex<Vec<T>>, tx: &Sender<Vec<T>>) {
+    let mut buf = batch.lock();
+    if buf.is_empty() {
+        return;
+    }
+    let items = std::mem::replace(&mut *buf, Vec::with_capacity(BATCH_ITEMS));
+    drop(buf);
+    let _ = tx.send(items);
+}
+
+fn poll_iteration<T>(
+    control: std::result::Result<Sender<()>, RecvTimeoutError>,
+    consume: impl FnOnce(),
+    poll: impl FnOnce(),
+    batch: &Mutex<Vec<T>>,
+    tx: &Sender<Vec<T>>,
+) -> bool {
+    match control {
+        Ok(ack) => {
+            consume();
+            // `drain` promises pending entries are in the channel before returning.
+            flush_batch(batch, tx);
+            let _ = ack.send(());
+            true
+        }
+        Err(RecvTimeoutError::Timeout) => {
+            poll();
+            flush_batch(batch, tx);
+            true
+        }
+        Err(RecvTimeoutError::Disconnected) => {
+            consume();
+            flush_batch(batch, tx);
+            false
+        }
+    }
+}
+
 /// Polls a BPF ring buffer in a background thread, parsing raw entries with a
-/// user-supplied closure and forwarding them to an mpsc channel.
+/// user-supplied closure and forwarding them to an mpsc channel in batches.
 ///
 /// The poll thread runs until the poller is dropped, doing a final full
 /// `consume()` on shutdown so no buffered entries are lost.
@@ -15,17 +62,36 @@ pub struct RingBufferPoller {
 }
 
 impl RingBufferPoller {
-    pub fn new<M, T, F>(rb_map: &M, parse: F, tx: Sender<T>, poll_interval_ms: u64) -> Result<Self>
+    pub fn new<M, T, F>(
+        rb_map: &M,
+        parse: F,
+        tx: Sender<Vec<T>>,
+        poll_interval_ms: u64,
+    ) -> Result<Self>
     where
         M: MapCore,
         T: Send + 'static,
         F: Fn(&[u8]) -> Option<T> + Send + 'static,
     {
+        // `Arc<Mutex<_>>` rather than `Rc<RefCell<_>>`: the built `RingBuffer` moves
+        // into the poll thread, so the callback must be `Send`.
+        let batch = Arc::new(Mutex::new(Vec::with_capacity(BATCH_ITEMS)));
+        let cb_batch = Arc::clone(&batch);
+        let cb_tx = tx.clone();
+
         let mut builder = RingBufferBuilder::new();
         builder.add(rb_map, move |data| {
-            if let Some(item) = parse(data) {
-                let _ = tx.send(item);
+            let Some(item) = parse(data) else {
+                return 0;
+            };
+            let mut buf = cb_batch.lock();
+            buf.push(item);
+            if buf.len() < BATCH_ITEMS {
+                return 0;
             }
+            let items = std::mem::replace(&mut *buf, Vec::with_capacity(BATCH_ITEMS));
+            drop(buf);
+            let _ = cb_tx.send(items);
             0
         })?;
         let ringbuf = builder.build()?;
@@ -35,21 +101,17 @@ impl RingBufferPoller {
         // poll tick, and disconnection is the shutdown signal.
         let (ctl, ctl_rx) = mpsc::channel::<Sender<()>>();
         let poll_thread = std::thread::spawn(move || {
-            loop {
-                match ctl_rx.recv_timeout(Duration::from_millis(poll_interval_ms)) {
-                    Ok(ack) => {
-                        let _ = ringbuf.consume();
-                        let _ = ack.send(());
-                    }
-                    Err(RecvTimeoutError::Timeout) => {
-                        let _ = ringbuf.poll(Duration::ZERO);
-                    }
-                    Err(RecvTimeoutError::Disconnected) => {
-                        let _ = ringbuf.consume();
-                        break;
-                    }
-                }
-            }
+            while poll_iteration(
+                ctl_rx.recv_timeout(Duration::from_millis(poll_interval_ms)),
+                || {
+                    let _ = ringbuf.consume();
+                },
+                || {
+                    let _ = ringbuf.poll(Duration::ZERO);
+                },
+                &batch,
+                &tx,
+            ) {}
         });
 
         Ok(Self {
@@ -77,3 +139,72 @@ impl Drop for RingBufferPoller {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::cell::Cell;
+
+    fn partial_batch() -> Vec<u8> {
+        vec![42; BATCH_ITEMS - 1]
+    }
+
+    #[test]
+    fn timeout_flushes_partial_batch() {
+        let expected = partial_batch();
+        let batch = Mutex::new(expected.clone());
+        let (tx, rx) = mpsc::channel();
+        let polled = Cell::new(false);
+
+        assert!(poll_iteration(
+            Err(RecvTimeoutError::Timeout),
+            || unreachable!(),
+            || polled.set(true),
+            &batch,
+            &tx,
+        ));
+
+        assert!(polled.get());
+        assert_eq!(rx.recv().unwrap(), expected);
+    }
+
+    #[test]
+    fn drain_flushes_partial_batch() {
+        let expected = partial_batch();
+        let batch = Mutex::new(expected.clone());
+        let (tx, rx) = mpsc::channel();
+        let (ack_tx, ack_rx) = mpsc::channel();
+        let consumed = Cell::new(false);
+
+        assert!(poll_iteration(
+            Ok(ack_tx),
+            || consumed.set(true),
+            || unreachable!(),
+            &batch,
+            &tx,
+        ));
+
+        assert!(consumed.get());
+        assert_eq!(rx.recv().unwrap(), expected);
+        ack_rx.recv().unwrap();
+    }
+
+    #[test]
+    fn shutdown_flushes_partial_batch() {
+        let expected = partial_batch();
+        let batch = Mutex::new(expected.clone());
+        let (tx, rx) = mpsc::channel();
+        let consumed = Cell::new(false);
+
+        assert!(!poll_iteration(
+            Err(RecvTimeoutError::Disconnected),
+            || consumed.set(true),
+            || unreachable!(),
+            &batch,
+            &tx,
+        ));
+
+        assert!(consumed.get());
+        assert_eq!(rx.recv().unwrap(), expected);
+    }
+}

crates/memtrack/src/main.rs

@@ -130,7 +130,8 @@ fn track_command(
         .map(|n| n.get().saturating_sub(2).max(1))
         .unwrap_or(4);
 
-    let pipeline_thread = thread::spawn(move || encode_events(event_rx, out_file, n_workers));
+    let pipeline_thread =
+        thread::spawn(move || encode_events(event_rx.into_iter().flatten(), out_file, n_workers));
 
     // Wait for the command to complete
     let status = session.wait().context("Failed to wait for command")?;

crates/memtrack/src/session.rs

@@ -8,14 +8,14 @@ use std::sync::mpsc::Receiver;
 /// stays alive as long as the session does; dropping it stops event delivery.
 pub struct Session {
     child: Child,
-    events: Option<Receiver<MemtrackEvent>>,
+    events: Option<Receiver<Vec<MemtrackEvent>>>,
     _poller: RingBufferPoller,
 }
 
 impl Session {
     pub(crate) fn new(
         child: Child,
-        events: Receiver<MemtrackEvent>,
+        events: Receiver<Vec<MemtrackEvent>>,
         poller: RingBufferPoller,
     ) -> Self {
         Self {
@@ -30,7 +30,7 @@ impl Session {
     }
 
     /// Take ownership of the event receiver. Can only be taken once.
-    pub fn take_events(&mut self) -> Result<Receiver<MemtrackEvent>> {
+    pub fn take_events(&mut self) -> Result<Receiver<Vec<MemtrackEvent>>> {
         self.events.take().context("events already taken")
     }
 

crates/memtrack/tests/shared.rs

@@ -361,7 +361,7 @@ fn run_tracked<T>(
     // Dropping the session does a final ring buffer drain and closes the
     // channel, so collecting terminates without a silence timeout.
     drop(session);
-    let events: Vec<Event> = rx.iter().collect();
+    let events: Vec<Event> = rx.into_iter().flatten().collect();
 
     tracker.finish()?;