perf(memtrack): reduce serialization bottleneck

.github/actions/install-rust/action.yml

@@ -14,7 +14,7 @@ runs:
   steps:
     - name: Install rust-toolchain.toml
       shell: bash
-      run: rustup toolchain install
+      run: rustup toolchain install "$(sed -n 's/.*channel = "\(.*\)".*/\1/p' rust-toolchain.toml)"
     - name: Install additional rustup components
       if: inputs.components != ''
       shell: bash

.github/workflows/ci.yml

@@ -120,7 +120,11 @@ jobs:
         run: sudo chown -R $USER:$USER . ~/.cargo
 
   benchmarks:
-    runs-on: ubuntu-latest
+    runs-on: ${{ matrix.mode == 'walltime' && 'codspeed-macro' || 'ubuntu-latest' }}
+    strategy:
+      fail-fast: false
+      matrix:
+        mode: [simulation, walltime]
     steps:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
         with:
@@ -133,11 +137,11 @@ jobs:
           tool: cargo-codspeed
 
       - name: Build benchmarks
-        run: cargo codspeed build -p runner-shared
+        run: cargo codspeed build -p runner-shared -m ${{ matrix.mode }}
       - name: Run benchmarks
         uses: CodSpeedHQ/action@3194d9a39c4d46684cb44bf7207fc56626aad8fd # v4
         with:
-          mode: simulation
+          mode: ${{ matrix.mode }}
           run: cargo codspeed run -p runner-shared
 
   check:

Cargo.toml

@@ -101,12 +101,14 @@ 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"
 serde = { version = "1.0.228", features = ["derive"] }
 ipc-channel = "0.20"
 itertools = "0.14.0"
+rayon = "1.12"
 linux-perf-event-reader = "0.10.2"                                                       # matches the version linux-perf-data resolves to
 env_logger = "0.11.10"
 tempfile = "3.27.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

@@ -575,13 +575,13 @@ impl MemtrackBpf {
         Ok(())
     }
 
-    /// Start polling with an mpsc channel for events
+    /// Start polling, forwarding each event over the returned channel.
     pub fn start_polling_with_channel(
         &self,
         poll_timeout_ms: u64,
     ) -> Result<(
         RingBufferPoller,
-        std::sync::mpsc::Receiver<runner_shared::artifacts::MemtrackEvent>,
+        crossbeam_channel::Receiver<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)

crates/memtrack/src/ebpf/poller.rs

@@ -1,78 +1,62 @@
 use anyhow::Result;
+use crossbeam_channel::{Receiver, unbounded};
 use libbpf_rs::{MapCore, RingBufferBuilder};
 use runner_shared::artifacts::MemtrackEvent as Event;
+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>;
-
-/// 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 forwards events over a channel.
 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 forward each parsed event over the returned channel.
+    pub fn with_channel<M: MapCore + 'static>(
         rb_map: &M,
-        handler: EventHandler,
         poll_timeout_ms: u64,
-    ) -> Result<Self> {
+    ) -> Result<(Self, Receiver<Event>)> {
+        let (tx, rx) = unbounded::<Event>();
+
         let mut builder = RingBufferBuilder::new();
         builder.add(rb_map, move |data| {
             if let Some(event) = parse_event(data) {
-                handler(event);
+                let _ = tx.send(event);
             }
             0
         })?;
-
         let ringbuf = builder.build()?;
+
         let shutdown = Arc::new(AtomicBool::new(false));
-        let shutdown_clone = shutdown.clone();
+        let poll_thread = std::thread::spawn({
+            let shutdown = shutdown.clone();
+            let timeout = Duration::from_millis(poll_timeout_ms);
+            move || {
+                while !shutdown.load(Ordering::Relaxed) {
+                    let _ = ringbuf.poll(timeout);
+                }
 
-        let poll_thread = std::thread::spawn(move || {
-            while !shutdown_clone.load(Ordering::Relaxed) {
-                let _ = ringbuf.poll(Duration::from_millis(poll_timeout_ms));
+                // Events may still be sitting in the ring buffer after the last
+                // poll; consume once more so they reach the channel before it closes.
+                let _ = ringbuf.consume();
             }
         });
 
-        Ok(Self {
-            shutdown,
-            poll_thread: Some(poll_thread),
-        })
-    }
-
-    /// 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,12 @@
+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};
 
 pub struct Tracker {
     bpf: MemtrackBpf,
+    poller: Option<RingBufferPoller>,
 }
 
 impl Tracker {
@@ -32,7 +34,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<()> {
@@ -43,32 +45,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.
+    /// Returns a receiver of allocation events. 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<Event>> {
-        // Add the PID to track
         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, event_rx) = self.bpf.start_polling_with_channel(10)?;
+        self.poller = Some(poller);
+
+        Ok(event_rx)
+    }
+
+    /// Stop the poll thread, draining ring-buffer stragglers. This closes the
+    /// event 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_events};
 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,61 +126,11 @@ 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;
-                    }
-                }
-            }
-        }
-    });
+    let n_workers = thread::available_parallelism()
+        .map(|n| n.get())
+        .unwrap_or(4);
 
-    // 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()?;
-
-        info!("Wrote {i} memtrack events to disk");
-
-        Ok(())
-    });
+    let pipeline_thread = thread::spawn(move || encode_events(event_rx, out_file, n_workers));
 
     // Wait for the command to complete
     let status = child.wait().context("Failed to wait for command")?;
@@ -199,19 +146,20 @@ fn track_command(
         warn!("Failed to disable tracking: {e:#}");
     }
 
-    // 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"))?;
+    // Stopping the poller closes the event channel; without this the pipeline
+    // join below would block forever.
+    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");
 
     // Detach probes explicitly: the IPC thread still holds an Arc clone, so the
     // tracker would otherwise never be dropped before process::exit and the

crates/runner-shared/Cargo.toml

@@ -13,6 +13,7 @@ bincode = "1.3"
 itertools = { workspace = true }
 linux-perf-event-reader = { workspace = true }
 log = { workspace = true }
+rayon = { workspace = true }
 rmp = "0.8.15"
 rmp-serde = "1.3.1"
 libc = { workspace = true }