perf(memtrack): reduce serialization bottleneck

Cargo.toml

@@ -101,6 +101,7 @@ exclude = ["crates/samply-codspeed"]
 [workspace.dependencies]
 anyhow = "1.0"
 clap = { version = "4.6", features = ["derive", "env"] }
+crossbeam-channel = "0.5.15"
 libc = "0.2"
 log = "0.4.28"
 serde_json = "1.0"

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 = { workspace = true }
 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,112 @@
 use anyhow::Result;
+use crossbeam_channel::{Receiver, unbounded};
 use libbpf_rs::{MapCore, RingBufferBuilder};
 use runner_shared::artifacts::MemtrackEvent as Event;
+use std::cell::RefCell;
+use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::{Arc, mpsc};
 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>;
+thread_local! {
+    /// Events staged by the ring-buffer callback and drained by the poll loop.
+    /// libbpf invokes the callback synchronously on the poll thread, so staging
+    /// through a thread-local cell keeps the per-event cost a plain `Vec` push
+    /// (no lock) while still letting the poll loop reach the events to flush them.
+    static STAGED: RefCell<Vec<Event>> = const { RefCell::new(Vec::new()) };
+}
+
+/// After the poll loop stops, keep consuming stragglers until this many rounds
+/// pass with no new events, then give up.
+const STRAGGLER_IDLE_ROUNDS: u32 = 3;
+/// Delay between straggler-drain rounds once the ring buffer looks empty.
+const STRAGGLER_POLL_INTERVAL: Duration = Duration::from_millis(5);
 
-/// RingBufferPoller manages polling a BPF ring buffer in a background thread
-/// and sending events to handlers
+/// 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::<Vec<Event>>();
+
         let mut builder = RingBufferBuilder::new();
         builder.add(rb_map, move |data| {
             if let Some(event) = parse_event(data) {
-                handler(event);
+                STAGED.with_borrow_mut(|staged| staged.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);
+
+            // Move staged events into the channel in batches of at most `batch_size`.
+            // High rates emit full batches; the trailing partial batch is streamed
+            // each poll cycle so events never wait for shutdown to be delivered.
+            let drain = || {
+                STAGED.with_borrow_mut(|staged| {
+                    while staged.len() > batch_size {
+                        let rest = staged.split_off(batch_size);
+                        let batch = std::mem::replace(staged, rest);
+                        let _ = tx.send(batch);
+                    }
+                    if !staged.is_empty() {
+                        let _ = tx.send(std::mem::take(staged));
+                    }
+                });
+            };
+
             while !shutdown_clone.load(Ordering::Relaxed) {
-                let _ = ringbuf.poll(Duration::from_millis(poll_timeout_ms));
+                let _ = ringbuf.poll(timeout);
+                drain();
             }
-        });
 
-        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, and a few stragglers
+            // can still arrive as it tears down. Consume repeatedly, sleeping only
+            // between empty rounds: keep going as long as new events show up, and
+            // stop after a few consecutive idle rounds. This drains as fast as the
+            // events arrive instead of always paying a fixed worst-case delay.
+            let mut idle_rounds = 0;
+            while idle_rounds < STRAGGLER_IDLE_ROUNDS {
+                let before = STAGED.with_borrow(Vec::len);
+                let _ = ringbuf.consume();
+                if STAGED.with_borrow(Vec::len) == before {
+                    idle_rounds += 1;
+                    std::thread::sleep(STRAGGLER_POLL_INTERVAL);
+                } else {
+                    idle_rounds = 0;
+                }
+                drain();
+            }
+        });
 
-    /// 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,27 @@ 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,32 @@ 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);
 
-        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