perf(memtrack): reduce serialization bottleneck

crates/memtrack/Cargo.toml

@@ -21,6 +21,7 @@ ebpf = ["dep:libbpf-rs", "dep:libbpf-cargo", "dep:vmlinux"]
 [dependencies]
 anyhow = { workspace = true }
 clap = { workspace = true }
+crossbeam-channel = "0.5.15"
 libc = { workspace = true }
 log = { workspace = true }
 env_logger = { workspace = true }

crates/memtrack/src/ebpf/memtrack.rs

@@ -453,16 +453,17 @@ impl MemtrackBpf {
         Ok(())
     }
 
-    /// Start polling with an mpsc channel for events
-    pub fn start_polling_with_channel(
+    /// Start polling, delivering events in ordered batches of at most `batch_size`.
+    pub fn start_polling_with_batches(
         &self,
         poll_timeout_ms: u64,
+        batch_size: usize,
     ) -> Result<(
         RingBufferPoller,
-        std::sync::mpsc::Receiver<runner_shared::artifacts::MemtrackEvent>,
+        crossbeam_channel::Receiver<Vec<runner_shared::artifacts::MemtrackEvent>>,
     )> {
         // Use the syscalls skeleton's ring buffer (both programs share the same one)
-        RingBufferPoller::with_channel(&self.skel.maps.events, poll_timeout_ms)
+        RingBufferPoller::with_batch_channel(&self.skel.maps.events, poll_timeout_ms, batch_size)
     }
 }
 

crates/memtrack/src/ebpf/poller.rs

@@ -1,78 +1,100 @@
 use anyhow::Result;
+use crossbeam_channel::{Receiver, Sender, unbounded};
 use libbpf_rs::{MapCore, RingBufferBuilder};
 use runner_shared::artifacts::MemtrackEvent as Event;
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::{Arc, mpsc};
+use std::sync::{Arc, Mutex};
 use std::thread::JoinHandle;
 use std::time::Duration;
 
 use super::events::parse_event;
 
-/// A handler function for processing ring buffer events
-pub type EventHandler = Box<dyn Fn(Event) + Send>;
+/// Coalesces parsed events into fixed-size batches on the poll thread, so the
+/// per-event cost stays a `Vec` push and events cross the channel one batch at a
+/// time instead of one at a time.
+struct BatchAccumulator {
+    batch: Vec<Event>,
+    batch_size: usize,
+    tx: Sender<Vec<Event>>,
+}
 
-/// RingBufferPoller manages polling a BPF ring buffer in a background thread
-/// and sending events to handlers
+impl BatchAccumulator {
+    fn push(&mut self, event: Event) {
+        self.batch.push(event);
+        if self.batch.len() >= self.batch_size {
+            self.flush();
+        }
+    }
+
+    fn flush(&mut self) {
+        if self.batch.is_empty() {
+            return;
+        }
+        let full = std::mem::replace(&mut self.batch, Vec::with_capacity(self.batch_size));
+        let _ = self.tx.send(full);
+    }
+}
+
+/// Polls a BPF ring buffer on a background thread and hands off batches of events.
+///
+/// The ring buffer is single-consumer, so one poll thread drains it. Batching in
+/// the poll callback keeps that thread cheap enough to keep up with the kernel.
 pub struct RingBufferPoller {
     shutdown: Arc<AtomicBool>,
     poll_thread: Option<JoinHandle<()>>,
 }
 
 impl RingBufferPoller {
-    /// Create a new RingBufferPoller for the given ring buffer map
-    ///
-    /// # Arguments
-    /// * `rb_map` - The BPF ring buffer map to poll
-    /// * `handler` - Callback function to handle each event
-    /// * `poll_timeout_ms` - How long to wait for events in each poll iteration
-    pub fn new<M: MapCore + 'static>(
+    /// Poll `rb_map` and deliver events in ordered batches of at most `batch_size`.
+    pub fn with_batch_channel<M: MapCore + 'static>(
         rb_map: &M,
-        handler: EventHandler,
         poll_timeout_ms: u64,
-    ) -> Result<Self> {
+        batch_size: usize,
+    ) -> Result<(Self, Receiver<Vec<Event>>)> {
+        let (tx, rx) = unbounded();
+        let accum = Arc::new(Mutex::new(BatchAccumulator {
+            batch: Vec::with_capacity(batch_size),
+            batch_size,
+            tx,
+        }));
+
+        let cb_accum = accum.clone();
         let mut builder = RingBufferBuilder::new();
         builder.add(rb_map, move |data| {
             if let Some(event) = parse_event(data) {
-                handler(event);
+                cb_accum.lock().unwrap().push(event);
             }
             0
         })?;
-
         let ringbuf = builder.build()?;
+
         let shutdown = Arc::new(AtomicBool::new(false));
         let shutdown_clone = shutdown.clone();
-
         let poll_thread = std::thread::spawn(move || {
+            let timeout = Duration::from_millis(poll_timeout_ms);
             while !shutdown_clone.load(Ordering::Relaxed) {
-                let _ = ringbuf.poll(Duration::from_millis(poll_timeout_ms));
+                let _ = ringbuf.poll(timeout);
             }
-        });
 
-        Ok(Self {
-            shutdown,
-            poll_thread: Some(poll_thread),
-        })
-    }
+            // The poll loop stopped, but events emitted just before the tracked
+            // process exited may still be in the ring buffer. Drain what's there,
+            // wait briefly for stragglers, drain again, then flush the tail batch.
+            let _ = ringbuf.consume();
+            std::thread::sleep(Duration::from_millis(50));
+            let _ = ringbuf.consume();
+            accum.lock().unwrap().flush();
+        });
 
-    /// Create a new RingBufferPoller with an mpsc channel for events
-    ///
-    /// Returns the RingBufferPoller and the receiver end of the channel
-    pub fn with_channel<M: MapCore + 'static>(
-        rb_map: &M,
-        poll_timeout_ms: u64,
-    ) -> Result<(Self, mpsc::Receiver<Event>)> {
-        let (tx, rx) = mpsc::channel();
-        let poller = Self::new(
-            rb_map,
-            Box::new(move |event| {
-                let _ = tx.send(event);
-            }),
-            poll_timeout_ms,
-        )?;
-        Ok((poller, rx))
+        Ok((
+            Self {
+                shutdown,
+                poll_thread: Some(poll_thread),
+            },
+            rx,
+        ))
     }
 
-    /// Stop the polling thread and wait for it to finish
+    /// Stop the polling thread and wait for it to finish draining.
     pub fn shutdown(&mut self) {
         self.shutdown.store(true, Ordering::Relaxed);
         if let Some(thread) = self.poll_thread.take() {

crates/memtrack/src/ebpf/tracker.rs

@@ -1,10 +1,15 @@
+use crate::ebpf::poller::RingBufferPoller;
 use crate::prelude::*;
 use crate::{AllocatorLib, ebpf::MemtrackBpf};
+use crossbeam_channel::Receiver;
 use runner_shared::artifacts::MemtrackEvent as Event;
-use std::sync::mpsc::{self, Receiver};
+
+/// Events are drained into batches of this size before crossing the channel.
+const DRAIN_BATCH_EVENTS: usize = 64 * 1024;
 
 pub struct Tracker {
     bpf: MemtrackBpf,
+    poller: Option<RingBufferPoller>,
 }
 
 impl Tracker {
@@ -32,7 +37,7 @@ impl Tracker {
         let mut bpf = MemtrackBpf::new()?;
         bpf.attach_tracepoints()?;
 
-        Ok(Self { bpf })
+        Ok(Self { bpf, poller: None })
     }
 
     pub fn attach_allocators(&mut self, libs: &[AllocatorLib]) -> Result<()> {
@@ -48,32 +53,25 @@ impl Tracker {
         self.bpf.attach_allocator_probes(lib.kind, &lib.path)
     }
 
-    /// Start tracking allocations for a specific PID
+    /// Start tracking allocations for a specific PID.
     ///
-    /// Returns a receiver channel that will receive allocation events.
-    /// The receiver will continue to produce events until the tracker is dropped.
-    pub fn track(&mut self, pid: i32) -> Result<Receiver<Event>> {
-        // Add the PID to track
+    /// Returns a receiver of ordered event batches. The poller is owned by the
+    /// tracker and keeps running until [`Tracker::stop_polling`] is called or the
+    /// tracker is dropped.
+    pub fn track(&mut self, pid: i32) -> Result<Receiver<Vec<Event>>> {
         self.bpf.add_tracked_pid(pid)?;
         debug!("Tracking PID {pid}");
 
-        // Start polling with channel
-        let (_poller, event_rx) = self.bpf.start_polling_with_channel(10)?;
-
-        // Keep the poller alive by moving it into the channel
-        // When the receiver is dropped, the poller will also be dropped
-        let (tx, rx) = mpsc::channel();
-        std::thread::spawn(move || {
-            // Keep poller alive
-            let _p = _poller;
-            while let Ok(event) = event_rx.recv() {
-                if tx.send(event).is_err() {
-                    break;
-                }
-            }
-        });
-
-        Ok(rx)
+        let (poller, batch_rx) = self.bpf.start_polling_with_batches(10, DRAIN_BATCH_EVENTS)?;
+        self.poller = Some(poller);
+
+        Ok(batch_rx)
+    }
+
+    /// Stop the poll thread, draining ring-buffer stragglers and flushing the
+    /// final partial batch. This closes the batch channel returned by [`track`].
+    pub fn stop_polling(&mut self) {
+        self.poller.take();
     }
 
     /// Bump RLIMIT_MEMLOCK for kernels older than 5.11. Newer kernels account BPF

crates/memtrack/src/main.rs

@@ -2,15 +2,12 @@ use clap::Parser;
 use ipc_channel::ipc;
 use memtrack::prelude::*;
 use memtrack::{MemtrackIpcMessage, Tracker, handle_ipc_message};
-use runner_shared::artifacts::{ArtifactExt, MemtrackArtifact, MemtrackEvent, MemtrackWriter};
+use runner_shared::artifacts::{ArtifactExt, MemtrackArtifact, encode_batches};
 use std::os::unix::process::CommandExt;
 use std::path::{Path, PathBuf};
 use std::process::Command;
-use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::mpsc::channel;
 use std::sync::{Arc, Mutex};
 use std::thread;
-use std::time::Duration;
 
 #[derive(Parser)]
 #[command(name = "memtrack")]
@@ -129,79 +126,34 @@ fn track_command(
     let file_name = MemtrackArtifact::file_name(Some(root_pid));
     let out_file = std::fs::File::create(out_dir.join(file_name))?;
 
-    let (write_tx, write_rx) = channel::<MemtrackEvent>();
-
-    // Stage A: Fast drain thread - This is required so that we immediately clear the ring buffer
-    // because it only has a limited size.
-    static DRAIN_EVENTS: AtomicBool = AtomicBool::new(true);
-    let write_tx_clone = write_tx.clone();
-    let drain_thread = thread::spawn(move || {
-        // Regular draining loop
-        while DRAIN_EVENTS.load(Ordering::Relaxed) {
-            let Ok(event) = event_rx.recv_timeout(Duration::from_millis(100)) else {
-                continue;
-            };
-            let _ = write_tx_clone.send(event);
-        }
-
-        // Final aggressive drain - keep trying until truly empty
-        loop {
-            match event_rx.try_recv() {
-                Ok(event) => {
-                    let _ = write_tx_clone.send(event);
-                }
-                Err(_) => {
-                    // Sleep briefly and try once more to catch late arrivals
-                    thread::sleep(Duration::from_millis(50));
-                    if let Ok(event) = event_rx.try_recv() {
-                        let _ = write_tx_clone.send(event);
-                    } else {
-                        break;
-                    }
-                }
-            }
-        }
-    });
-
-    // Stage B: Writer thread - Immediately writes the events to disk
-    let writer_thread = thread::spawn(move || -> anyhow::Result<()> {
-        let mut writer = MemtrackWriter::new(out_file)?;
-
-        let mut i = 0;
-        while let Ok(first) = write_rx.recv() {
-            writer.write_event(&first)?;
-            i += 1;
-
-            // Drain any backlog in a tight loop (batching)
-            while let Ok(ev) = write_rx.try_recv() {
-                writer.write_event(&ev)?;
-                i += 1;
-            }
-        }
-        writer.finish()?;
+    let n_workers = std::thread::available_parallelism()
+        .map(|n| n.get())
+        .unwrap_or(4)
+        .min(8);
 
-        info!("Wrote {i} memtrack events to disk");
-
-        Ok(())
-    });
+    // Encode batches into the artifact on a dedicated thread. It drains the batch
+    // channel until the poller closes it, so it runs alongside the tracked command.
+    let pipeline_thread =
+        thread::spawn(move || encode_batches(event_rx, out_file, n_workers));
 
     // Wait for the command to complete
     let status = child.wait().context("Failed to wait for command")?;
     debug!("Command exited with status: {status}");
 
-    // Wait for drain thread to finish
-    debug!("Waiting for the drain thread to finish");
-    DRAIN_EVENTS.store(false, Ordering::Relaxed);
-    drain_thread
-        .join()
-        .map_err(|_| anyhow::anyhow!("Failed to join drain thread"))?;
+    // Stop the poll thread: drains ring-buffer stragglers, flushes the final batch,
+    // and closes the batch channel so the encode pipeline drains to completion.
+    debug!("Stopping the ring buffer poller");
+    tracker_arc
+        .lock()
+        .map_err(|_| anyhow!("tracker mutex poisoned"))?
+        .stop_polling();
 
-    // Wait for writer thread to finish and propagate errors
-    debug!("Waiting for the writer thread to finish");
-    drop(write_tx);
-    writer_thread
+    debug!("Waiting for the encode pipeline to finish");
+    let total = pipeline_thread
         .join()
-        .map_err(|_| anyhow::anyhow!("Failed to join writer thread"))??;
+        .map_err(|_| anyhow::anyhow!("Failed to join memtrack encode pipeline"))??;
+
+    info!("Wrote {total} memtrack events to disk");
 
     // Read the eBPF dropped-event counter after the run is complete.
     // A non-zero value means the ring buffer overflowed and the trace is

crates/memtrack/tests/shared.rs

@@ -149,8 +149,8 @@ pub fn track_command(
     let rx = tracker.track(root_pid)?;
 
     let mut events = Vec::new();
-    while let Ok(event) = rx.recv_timeout(Duration::from_secs(10)) {
-        events.push(event);
+    while let Ok(batch) = rx.recv_timeout(Duration::from_secs(10)) {
+        events.extend(batch);
     }
 
     // Drop the tracker in a new thread to not block the test

crates/runner-shared/Cargo.toml

@@ -10,6 +10,7 @@ serde = { workspace = true }
 serde_json = { workspace = true }
 # Pinned to 1.x: 2.0 changes the wire format and serde integration
 bincode = "1.3"
+crossbeam-channel = "0.5.15"
 itertools = { workspace = true }
 linux-perf-event-reader = { workspace = true }
 log = { workspace = true }