From 21ba5256567b8a4e5f8e1d64fff32cac2e42089e Mon Sep 17 00:00:00 2001 From: not-matthias Date: Tue, 14 Jul 2026 10:52:12 +0200 Subject: [PATCH 01/14] refactor: split memtrack into module --- .../{memtrack.rs => memtrack/mod.rs} | 128 ++++++++++++------ .../src/artifacts/memtrack/writer.rs | 43 ++++++ 2 files changed, 132 insertions(+), 39 deletions(-) rename crates/runner-shared/src/artifacts/{memtrack.rs => memtrack/mod.rs} (62%) create mode 100644 crates/runner-shared/src/artifacts/memtrack/writer.rs diff --git a/crates/runner-shared/src/artifacts/memtrack.rs b/crates/runner-shared/src/artifacts/memtrack/mod.rs similarity index 62% rename from crates/runner-shared/src/artifacts/memtrack.rs rename to crates/runner-shared/src/artifacts/memtrack/mod.rs index 64c90e82..ad843b28 100644 --- a/crates/runner-shared/src/artifacts/memtrack.rs +++ b/crates/runner-shared/src/artifacts/memtrack/mod.rs @@ -1,6 +1,12 @@ use libc::pid_t; use serde::{Deserialize, Serialize}; -use std::io::{BufReader, BufWriter, Read, Write}; +use std::io::{BufReader, Read, Write}; + +mod pipeline; +mod writer; + +pub use pipeline::*; +pub use writer::*; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct MemtrackArtifact { @@ -86,43 +92,6 @@ impl Iterator for MemtrackEventStream { } } -/// Streaming writer for memtrack events with compression -pub struct MemtrackWriter { - serializer: rmp_serde::Serializer>>, -} - -impl MemtrackWriter { - pub fn new(writer: W) -> anyhow::Result { - // We're dealing with a lot of events, so we want to compress as much as possible - // while not taking too much time to compress. - const COMPRESSION_LEVEL: i32 = 1; - const BUFFER_SIZE: usize = 256 * 1024 /* 256 KB */; - - let writer = BufWriter::with_capacity(BUFFER_SIZE, writer); - let encoder = zstd::Encoder::new(writer, COMPRESSION_LEVEL)?; - Ok(Self { - serializer: rmp_serde::Serializer::new(encoder), - }) - } - - /// Write a single event to the stream - pub fn write_event(&mut self, event: &MemtrackEvent) -> anyhow::Result<()> { - event.serialize(&mut self.serializer)?; - Ok(()) - } - - /// Finish writing and flush the compression stream - pub fn finish(self) -> anyhow::Result<()> { - let encoder = self.serializer.into_inner(); - let mut writer = encoder.finish()?; - - // Flush the writer to ensure all data is written to the underlying writer - writer.flush()?; - - Ok(()) - } -} - #[cfg(test)] mod tests { use crate::artifacts::ArtifactExt; @@ -162,6 +131,87 @@ mod tests { Ok(()) } + #[test] + fn manual_serialize_is_byte_identical_to_derive() { + #[derive(serde::Serialize)] + struct Shadow { + pid: libc::pid_t, + tid: libc::pid_t, + timestamp: u64, + addr: u64, + #[serde(flatten)] + kind: MemtrackEventKind, + } + + let kinds = [ + MemtrackEventKind::Malloc { size: 7 }, + MemtrackEventKind::Free, + MemtrackEventKind::Realloc { + old_addr: Some(0x1000), + size: 42, + }, + MemtrackEventKind::Realloc { + old_addr: None, + size: 42, + }, + MemtrackEventKind::Calloc { size: 9 }, + MemtrackEventKind::AlignedAlloc { size: 9 }, + MemtrackEventKind::Mmap { size: 9 }, + MemtrackEventKind::Munmap { size: 9 }, + MemtrackEventKind::Brk { size: 9 }, + ]; + + for kind in kinds { + let event = MemtrackEvent { + pid: -7, + tid: 42, + timestamp: 0xDEAD, + addr: 0xBEEF, + kind, + }; + let shadow = Shadow { + pid: -7, + tid: 42, + timestamp: 0xDEAD, + addr: 0xBEEF, + kind, + }; + + assert_eq!( + rmp_serde::to_vec(&event).unwrap(), + rmp_serde::to_vec(&shadow).unwrap() + ); + } + } + + #[test] + fn concatenated_frames_decode_in_order() -> anyhow::Result<()> { + let events: Vec<_> = (0..2500) + .map(|i| MemtrackEvent { + pid: 1, + tid: 1, + timestamp: i, + addr: i, + kind: MemtrackEventKind::Malloc { size: i }, + }) + .collect(); + + let mut file = Vec::new(); + for batch in events.chunks(1000) { + let mut writer = MemtrackWriter::new(Vec::::new())?; + for event in batch { + writer.write_event(event)?; + } + let frame = writer.finish()?; + file.extend_from_slice(&frame); + } + + let decoded: Vec<_> = MemtrackArtifact::decode_streamed(Cursor::new(file))?.collect(); + assert_eq!(decoded, events); + + Ok(()) + } + #[test] fn test_artifact_is_empty() -> anyhow::Result<()> { let artifact = MemtrackArtifact { events: vec![] }; @@ -179,7 +229,7 @@ mod tests { fn test_deserialize_realloc_compat() -> anyhow::Result<()> { // The file contains a single serialized event using the old format without `old_addr`: // MemtrackEventKind::Realloc { size: 42 } - let buf = include_bytes!("../../testdata/realloc.MemtrackArtifact.msgpack"); + let buf = include_bytes!("../../../testdata/realloc.MemtrackArtifact.msgpack"); assert_eq!( MemtrackArtifact::decode_streamed(Cursor::new(buf))?.count(), 1 diff --git a/crates/runner-shared/src/artifacts/memtrack/writer.rs b/crates/runner-shared/src/artifacts/memtrack/writer.rs new file mode 100644 index 00000000..1f6d8c05 --- /dev/null +++ b/crates/runner-shared/src/artifacts/memtrack/writer.rs @@ -0,0 +1,43 @@ +use serde::Serialize; +use std::io::{BufWriter, Write}; + +use super::MemtrackEvent; + +/// Streaming writer for memtrack events, serializing into a zstd-compressed sink. +pub struct MemtrackWriter { + serializer: rmp_serde::Serializer, +} + +impl MemtrackWriter>> { + pub fn new(writer: W) -> anyhow::Result { + // We're dealing with a lot of events, so we want to compress as much as possible + // while not taking too much time to compress. + const COMPRESSION_LEVEL: i32 = 1; + const BUFFER_SIZE: usize = 256 * 1024 /* 256 KB */; + + let writer = BufWriter::with_capacity(BUFFER_SIZE, writer); + let encoder = zstd::Encoder::new(writer, COMPRESSION_LEVEL)?; + Ok(Self { + serializer: rmp_serde::Serializer::new(encoder), + }) + } + + /// Finish writing and flush the compression stream + pub fn finish(self) -> anyhow::Result<()> { + let encoder = self.serializer.into_inner(); + let mut writer = encoder.finish()?; + + // Flush the writer to ensure all data is written to the underlying writer + writer.flush()?; + + Ok(()) + } +} + +impl MemtrackWriter { + /// Write a single event to the stream + pub fn write_event(&mut self, event: &MemtrackEvent) -> anyhow::Result<()> { + event.serialize(&mut self.serializer)?; + Ok(()) + } +} From aadcfb8aaefdca32ec812ef728143d2086a316dd Mon Sep 17 00:00:00 2001 From: not-matthias Date: Tue, 14 Jul 2026 10:54:24 +0200 Subject: [PATCH 02/14] fix: wire MemtrackWriter serializer through BufWriter to fix per-event zstd compression bottleneck --- crates/runner-shared/src/artifacts/memtrack/writer.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/runner-shared/src/artifacts/memtrack/writer.rs b/crates/runner-shared/src/artifacts/memtrack/writer.rs index 1f6d8c05..d1f4de14 100644 --- a/crates/runner-shared/src/artifacts/memtrack/writer.rs +++ b/crates/runner-shared/src/artifacts/memtrack/writer.rs @@ -15,10 +15,10 @@ impl MemtrackWriter>> { const COMPRESSION_LEVEL: i32 = 1; const BUFFER_SIZE: usize = 256 * 1024 /* 256 KB */; - let writer = BufWriter::with_capacity(BUFFER_SIZE, writer); let encoder = zstd::Encoder::new(writer, COMPRESSION_LEVEL)?; + let writer = BufWriter::with_capacity(BUFFER_SIZE, encoder); Ok(Self { - serializer: rmp_serde::Serializer::new(encoder), + serializer: rmp_serde::Serializer::new(writer), }) } From 1a1ab2911b657a73c0eb3b0bedb6dc7003cfe98c Mon Sep 17 00:00:00 2001 From: not-matthias Date: Tue, 14 Jul 2026 10:55:03 +0200 Subject: [PATCH 03/14] fix: lower zstd compression level to -5 to reduce serialization overhead --- crates/runner-shared/src/artifacts/memtrack/writer.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/runner-shared/src/artifacts/memtrack/writer.rs b/crates/runner-shared/src/artifacts/memtrack/writer.rs index d1f4de14..5fc746dd 100644 --- a/crates/runner-shared/src/artifacts/memtrack/writer.rs +++ b/crates/runner-shared/src/artifacts/memtrack/writer.rs @@ -12,7 +12,7 @@ impl MemtrackWriter>> { pub fn new(writer: W) -> anyhow::Result { // We're dealing with a lot of events, so we want to compress as much as possible // while not taking too much time to compress. - const COMPRESSION_LEVEL: i32 = 1; + const COMPRESSION_LEVEL: i32 = -5; const BUFFER_SIZE: usize = 256 * 1024 /* 256 KB */; let encoder = zstd::Encoder::new(writer, COMPRESSION_LEVEL)?; From 68b9eeb7a5b81543d05b57a78b3f212ce47f353b Mon Sep 17 00:00:00 2001 From: not-matthias Date: Tue, 14 Jul 2026 11:37:05 +0200 Subject: [PATCH 04/14] ref(memtrack): return the underlying sink from MemtrackWriter::finish Callers that encode a frame into an in-memory buffer need the buffer back after the zstd stream is finalized, so frames can be composed into a larger artifact stream. Refs COD-3071 --- .../runner-shared/src/artifacts/memtrack/writer.rs | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/crates/runner-shared/src/artifacts/memtrack/writer.rs b/crates/runner-shared/src/artifacts/memtrack/writer.rs index 5fc746dd..2f665766 100644 --- a/crates/runner-shared/src/artifacts/memtrack/writer.rs +++ b/crates/runner-shared/src/artifacts/memtrack/writer.rs @@ -22,15 +22,13 @@ impl MemtrackWriter>> { }) } - /// Finish writing and flush the compression stream - pub fn finish(self) -> anyhow::Result<()> { - let encoder = self.serializer.into_inner(); + /// Finish writing, flush the compression stream, and return the sink + pub fn finish(self) -> anyhow::Result { + let buffered = self.serializer.into_inner(); + let encoder = buffered.into_inner().map_err(|e| e.into_error())?; let mut writer = encoder.finish()?; - - // Flush the writer to ensure all data is written to the underlying writer writer.flush()?; - - Ok(()) + Ok(writer) } } From f573bf1a4d35b0aafa06d723b74e3ba66ad98189 Mon Sep 17 00:00:00 2001 From: not-matthias Date: Tue, 14 Jul 2026 11:37:51 +0200 Subject: [PATCH 05/14] perf(memtrack): add parallel zstd encode pipeline for artifact events Group events into fixed-size frames and compress them across a rayon pool, writing windows in input order. Bounds peak memory regardless of run length and removes the single-threaded compression bottleneck. Refs COD-3071 --- Cargo.lock | 2 + Cargo.toml | 2 + crates/runner-shared/Cargo.toml | 1 + .../src/artifacts/memtrack/pipeline.rs | 138 ++++++++++++++++++ 4 files changed, 143 insertions(+) create mode 100644 crates/runner-shared/src/artifacts/memtrack/pipeline.rs diff --git a/Cargo.lock b/Cargo.lock index f7816f66..c53e2087 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2378,6 +2378,7 @@ dependencies = [ "anyhow", "bindgen", "clap", + "crossbeam-channel", "env_logger", "glob", "insta", @@ -3598,6 +3599,7 @@ dependencies = [ "linux-perf-event-reader 0.10.2 (registry+https://github.com/rust-lang/crates.io-index)", "log", "rand 0.8.6", + "rayon", "rmp", "rmp-serde", "serde", diff --git a/Cargo.toml b/Cargo.toml index ab974cb6..4afc3aa8 100644 --- a/Cargo.toml +++ b/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" diff --git a/crates/runner-shared/Cargo.toml b/crates/runner-shared/Cargo.toml index d0c67f77..8b8f6ab9 100644 --- a/crates/runner-shared/Cargo.toml +++ b/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 } diff --git a/crates/runner-shared/src/artifacts/memtrack/pipeline.rs b/crates/runner-shared/src/artifacts/memtrack/pipeline.rs new file mode 100644 index 00000000..70ddfc9c --- /dev/null +++ b/crates/runner-shared/src/artifacts/memtrack/pipeline.rs @@ -0,0 +1,138 @@ +use std::io::{BufWriter, Write}; + +use itertools::Itertools; +use rayon::prelude::*; + +use super::MemtrackEvent; +use super::writer::MemtrackWriter; + +/// Events per self-contained zstd frame. Larger frames compress better; smaller +/// frames cap the work (and memory) a single worker holds while encoding. +const FRAME_EVENTS: usize = 64 * 1024; +/// Frames compressed in parallel per window. A window is encoded across the +/// worker pool and then written before the next one starts, so this bounds peak +/// memory to roughly `FRAME_EVENTS * WINDOW_FRAMES` events regardless of how long +/// the source runs. +const WINDOW_FRAMES: usize = 16; + +/// Encode a stream of events into a single compressed artifact stream, +/// compressing frames in parallel across a Rayon pool of `n_workers` threads. +/// +/// Events are grouped into fixed-size frames; each frame is one self-contained +/// zstd frame. Frames are encoded a window at a time: a window is compressed in +/// parallel, then its frames are written in input order before the next window +/// starts, so the output matches the input order and peak memory stays bounded. +/// +/// Blocks the calling thread until `events` is exhausted. Returns the total +/// number of events written. +pub fn encode_events(events: S, out: W, n_workers: usize) -> anyhow::Result +where + S: IntoIterator, + W: Write, +{ + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(n_workers.max(1)) + .build()?; + + let mut out = BufWriter::new(out); + let mut total = 0u64; + let mut wrote_any = false; + + for window in &events.into_iter().chunks(FRAME_EVENTS * WINDOW_FRAMES) { + let window: Vec = window.collect(); + total += window.len() as u64; + + let frames: Vec> = pool.install(|| { + window + .par_chunks(FRAME_EVENTS) + .map(encode_frame) + .collect::>() + })?; + + for frame in frames { + out.write_all(&frame)?; + } + wrote_any = true; + } + + // Always emit at least one (possibly empty) frame so the artifact stream is + // valid and decodable even when no events were recorded. + if !wrote_any { + out.write_all(&encode_frame(&[])?)?; + } + + out.flush()?; + Ok(total) +} + +/// Encode one batch as a single self-contained zstd frame. +fn encode_frame(batch: &[MemtrackEvent]) -> anyhow::Result> { + let mut writer = MemtrackWriter::new(Vec::new())?; + for event in batch { + writer.write_event(event)?; + } + writer.finish() +} + +#[cfg(test)] +mod tests { + use std::io::Cursor; + + use super::super::{MemtrackArtifact, MemtrackEventKind}; + use super::*; + + fn malloc_events(range: std::ops::Range) -> Vec { + range + .map(|i| MemtrackEvent { + pid: 1, + tid: 1, + timestamp: i, + addr: i, + kind: MemtrackEventKind::Malloc { size: i }, + }) + .collect() + } + + #[test] + fn preserves_order_across_parallel_frames() -> anyhow::Result<()> { + // More events than fit in one frame, so ordering has to hold across the + // frames the worker pool compresses in parallel. + let events = malloc_events(0..(FRAME_EVENTS as u64 * 3 + 7)); + + let mut out = Vec::new(); + let total = encode_events(events.clone(), &mut out, 4)?; + assert_eq!(total, events.len() as u64); + + let decoded: Vec<_> = MemtrackArtifact::decode_streamed(Cursor::new(out))?.collect(); + assert_eq!(decoded, events); + + Ok(()) + } + + #[test] + fn preserves_order_across_window_boundary() -> anyhow::Result<()> { + let events = malloc_events(0..(FRAME_EVENTS * WINDOW_FRAMES + 1) as u64); + + let mut out = Vec::new(); + let total = encode_events(events.clone(), &mut out, 4)?; + assert_eq!(total, events.len() as u64); + + let decoded: Vec<_> = MemtrackArtifact::decode_streamed(Cursor::new(out))?.collect(); + assert_eq!(decoded, events); + + Ok(()) + } + + #[test] + fn empty_source_writes_a_valid_stream() -> anyhow::Result<()> { + let events: Vec = Vec::new(); + + let mut out = Vec::new(); + let total = encode_events(events, &mut out, 4)?; + assert_eq!(total, 0); + + assert!(MemtrackArtifact::is_empty(Cursor::new(out))); + + Ok(()) + } +} From efa5d31090c292396c10edf64d67e0769a290094 Mon Sep 17 00:00:00 2001 From: not-matthias Date: Tue, 14 Jul 2026 11:38:01 +0200 Subject: [PATCH 06/14] test(memtrack): benchmark encode_events with a realistic allocation trace Generate a seeded malloc/free/realloc/mmap workload with a live-heap model instead of uniform random events, and sweep worker counts. Refs COD-3071 --- .../runner-shared/benches/memtrack_writer.rs | 95 ++++++++++++++++++- 1 file changed, 94 insertions(+), 1 deletion(-) diff --git a/crates/runner-shared/benches/memtrack_writer.rs b/crates/runner-shared/benches/memtrack_writer.rs index 482b2a7b..62aa2451 100644 --- a/crates/runner-shared/benches/memtrack_writer.rs +++ b/crates/runner-shared/benches/memtrack_writer.rs @@ -1,7 +1,7 @@ use divan::Bencher; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; -use runner_shared::artifacts::{MemtrackEvent, MemtrackEventKind, MemtrackWriter}; +use runner_shared::artifacts::{MemtrackEvent, MemtrackEventKind, MemtrackWriter, encode_events}; fn main() { divan::main(); @@ -53,3 +53,96 @@ fn write_events(bencher: Bencher, n: usize) { writer.finish().unwrap(); }); } + +fn generate_realistic_events(n: usize) -> Vec { + const SIZES: [u64; 8] = [16, 24, 32, 48, 64, 96, 128, 512]; + let mut rng = StdRng::seed_from_u64(42); + let mut events = Vec::with_capacity(n); + let mut live_heap: Vec = Vec::new(); + let mut live_mmap: Vec<(u64, u64)> = Vec::new(); + let mut free_list: Vec = Vec::new(); + let mut next_addr: u64 = 0x5555_5555_0000; + let mut ts: u64 = 1_700_000_000_000_000_000; + let pid = 4242; + let tids = [4242, 4243, 4244, 4245]; + + while events.len() < n { + ts += rng.gen_range(50..2_000); + let tid = tids[rng.gen_range(0..tids.len())]; + let roll = rng.gen_range(0..100); + let (addr, kind) = if roll < 48 || live_heap.is_empty() { + let size = if rng.gen_range(0..100) < 95 { + SIZES[rng.gen_range(0..SIZES.len())] + } else { + rng.gen_range(4096..1 << 20) + }; + let addr = free_list.pop().unwrap_or_else(|| { + let addr = next_addr; + next_addr += (size + 15) & !15; + addr + }); + let kind = match rng.gen_range(0..20) { + 0 => MemtrackEventKind::Calloc { size }, + 1 => MemtrackEventKind::Mmap { size }, + _ => MemtrackEventKind::Malloc { size }, + }; + if let MemtrackEventKind::Mmap { size } = kind { + live_mmap.push((addr, size)); + } else { + live_heap.push(addr); + } + (addr, kind) + } else if roll < 90 { + let idx = rng.gen_range(0..live_heap.len() + live_mmap.len()); + if idx < live_heap.len() { + let addr = live_heap.swap_remove(idx); + free_list.push(addr); + (addr, MemtrackEventKind::Free) + } else { + let (addr, size) = live_mmap.swap_remove(idx - live_heap.len()); + free_list.push(addr); + (addr, MemtrackEventKind::Munmap { size }) + } + } else { + let idx = rng.gen_range(0..live_heap.len()); + let old_addr = live_heap[idx]; + let size = SIZES[rng.gen_range(0..SIZES.len())] * 2; + let new_addr = if rng.r#gen() { + old_addr + } else { + free_list.push(old_addr); + free_list.swap_remove(rng.gen_range(0..free_list.len())) + }; + live_heap[idx] = new_addr; + ( + new_addr, + MemtrackEventKind::Realloc { + old_addr: Some(old_addr), + size, + }, + ) + }; + events.push(MemtrackEvent { + pid, + tid, + timestamp: ts, + addr, + kind, + }); + } + + events +} + +const REALISTIC_EVENTS: usize = 1_000_000; + +#[divan::bench(args = [16, 8, 4], max_time = 10.0)] +fn encode_events_realistic(bencher: Bencher, n_workers: usize) { + let events = generate_realistic_events(REALISTIC_EVENTS); + + bencher.bench_local(|| { + let mut output = Vec::new(); + encode_events(events.iter().copied(), &mut output, n_workers).unwrap(); + output + }); +} From 09ae1d96955563b41efd5f6169a3affa95947959 Mon Sep 17 00:00:00 2001 From: not-matthias Date: Tue, 14 Jul 2026 11:38:09 +0200 Subject: [PATCH 07/14] ref(memtrack): let the tracker own the poller and drain on shutdown Forward ring buffer events over a crossbeam channel directly instead of through an extra keep-alive forwarding thread. On shutdown, consume the ring buffer once more so events emitted after the last poll are not lost, then close the channel. Refs COD-3071 --- crates/memtrack/Cargo.toml | 1 + crates/memtrack/src/ebpf/memtrack.rs | 4 +- crates/memtrack/src/ebpf/poller.rs | 72 +++++++++++----------------- crates/memtrack/src/ebpf/tracker.rs | 41 +++++++--------- 4 files changed, 49 insertions(+), 69 deletions(-) diff --git a/crates/memtrack/Cargo.toml b/crates/memtrack/Cargo.toml index ea6bee60..7b76ad1a 100644 --- a/crates/memtrack/Cargo.toml +++ b/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 } diff --git a/crates/memtrack/src/ebpf/memtrack.rs b/crates/memtrack/src/ebpf/memtrack.rs index 400ab669..649d1797 100644 --- a/crates/memtrack/src/ebpf/memtrack.rs +++ b/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, + crossbeam_channel::Receiver, )> { // Use the syscalls skeleton's ring buffer (both programs share the same one) RingBufferPoller::with_channel(&self.skel.maps.events, poll_timeout_ms) diff --git a/crates/memtrack/src/ebpf/poller.rs b/crates/memtrack/src/ebpf/poller.rs index 1f4ef2b4..ab482dc6 100644 --- a/crates/memtrack/src/ebpf/poller.rs +++ b/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; - -/// 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, poll_thread: Option>, } 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( + /// Poll `rb_map` and forward each parsed event over the returned channel. + pub fn with_channel( rb_map: &M, - handler: EventHandler, poll_timeout_ms: u64, - ) -> Result { + ) -> Result<(Self, Receiver)> { + let (tx, rx) = unbounded::(); + 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( - rb_map: &M, - poll_timeout_ms: u64, - ) -> Result<(Self, mpsc::Receiver)> { - 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() { diff --git a/crates/memtrack/src/ebpf/tracker.rs b/crates/memtrack/src/ebpf/tracker.rs index 91064ade..eb53e6e7 100644 --- a/crates/memtrack/src/ebpf/tracker.rs +++ b/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, } 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> { - // 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 From 1c47c8010e534219e402ef542a86ce6a86d8a78a Mon Sep 17 00:00:00 2001 From: not-matthias Date: Tue, 14 Jul 2026 11:38:16 +0200 Subject: [PATCH 08/14] perf(memtrack): replace drain/writer threads with the encode pipeline Feed the poller's event channel straight into encode_events on one thread, sized to available parallelism. Shutdown is now deterministic: stop the poller, which drains stragglers and closes the channel, then join the pipeline. Fixes COD-3071 --- crates/memtrack/src/main.rs | 86 ++++++++----------------------------- 1 file changed, 17 insertions(+), 69 deletions(-) diff --git a/crates/memtrack/src/main.rs b/crates/memtrack/src/main.rs index e9e06685..57dad028 100644 --- a/crates/memtrack/src/main.rs +++ b/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::(); - - // 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 From 78d5ae241f4f06700ab002eb3c86a23640f8d3a7 Mon Sep 17 00:00:00 2001 From: not-matthias Date: Tue, 14 Jul 2026 11:57:54 +0200 Subject: [PATCH 09/14] perf: pre-size and reuse the memtrack encode window buffer itertools::chunks erases the iterator size hint, so collecting each window grew by doubling (~29% of encode driver time in reallocation). A single pre-sized buffer reused across windows removes the growth and the chunk-adapter bookkeeping: encode_events_realistic median drops 21%/17%/9% at 16/8/4 workers. --- .../runner-shared/src/artifacts/memtrack/pipeline.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/crates/runner-shared/src/artifacts/memtrack/pipeline.rs b/crates/runner-shared/src/artifacts/memtrack/pipeline.rs index 70ddfc9c..c47b3aed 100644 --- a/crates/runner-shared/src/artifacts/memtrack/pipeline.rs +++ b/crates/runner-shared/src/artifacts/memtrack/pipeline.rs @@ -1,6 +1,5 @@ use std::io::{BufWriter, Write}; -use itertools::Itertools; use rayon::prelude::*; use super::MemtrackEvent; @@ -38,8 +37,15 @@ where let mut total = 0u64; let mut wrote_any = false; - for window in &events.into_iter().chunks(FRAME_EVENTS * WINDOW_FRAMES) { - let window: Vec = window.collect(); + let cap = FRAME_EVENTS * WINDOW_FRAMES; + let mut events = events.into_iter(); + let mut window: Vec = Vec::with_capacity(cap); + loop { + window.clear(); + window.extend(events.by_ref().take(cap)); + if window.is_empty() { + break; + } total += window.len() as u64; let frames: Vec> = pool.install(|| { From 88a6136d3852d46a322b0f57ea0dd8caaaaa6d3d Mon Sep 17 00:00:00 2001 From: not-matthias Date: Wed, 15 Jul 2026 19:24:50 +0200 Subject: [PATCH 10/14] ci: also run runner-shared benchmarks in walltime mode --- .github/workflows/ci.yml | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f0c77f7c..49543b75 100644 --- a/.github/workflows/ci.yml +++ b/.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: From d06d42e396918e7c15f46ed835c46ba1a99ad830 Mon Sep 17 00:00:00 2001 From: not-matthias Date: Wed, 15 Jul 2026 19:35:27 +0200 Subject: [PATCH 11/14] ci: install rust toolchain by explicit channel for older rustup --- .github/actions/install-rust/action.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/actions/install-rust/action.yml b/.github/actions/install-rust/action.yml index 3773347e..517f27fb 100644 --- a/.github/actions/install-rust/action.yml +++ b/.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 From 880233383d4847280651d1a7a691e0ea537cb49b Mon Sep 17 00:00:00 2001 From: not-matthias Date: Wed, 15 Jul 2026 20:36:38 +0200 Subject: [PATCH 12/14] perf(memtrack): wake the ringbuf consumer only past a data watermark Submitting with default flags wakes the consumer on nearly every event, and the wakeup dominates submission cost at high event rates. Submit with BPF_RB_NO_WAKEUP until 64KB of unconsumed data has accumulated, then force a wakeup. Kernel ringbuf benchmarks show up to ~20x higher sustained throughput with sampled notification. The poller's 10ms poll timeout flushes the tail that never reaches the watermark. Refs COD-3071 --- crates/memtrack/src/ebpf/c/memtrack.bpf.c | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/crates/memtrack/src/ebpf/c/memtrack.bpf.c b/crates/memtrack/src/ebpf/c/memtrack.bpf.c index 06601e27..18c5761c 100644 --- a/crates/memtrack/src/ebpf/c/memtrack.bpf.c +++ b/crates/memtrack/src/ebpf/c/memtrack.bpf.c @@ -94,6 +94,18 @@ int tracepoint_sched_fork(struct trace_event_raw_sched_process_fork* ctx) { /* == Helper functions for the allocation tracking == */ +/* Wake the consumer only once this much unconsumed data has accumulated. + * Per-event wakeups dominate submission cost at high event rates; batching + * them behind a data watermark amortizes the wakeup to ~1 per thousand + * events. The userspace poller's poll timeout flushes the tail that never + * reaches the watermark. */ +#define WAKEUP_DATA_SIZE (64 * 1024) + +static __always_inline long wake_flags(void) { + long avail = bpf_ringbuf_query(&events, BPF_RB_AVAIL_DATA); + return avail >= WAKEUP_DATA_SIZE ? BPF_RB_FORCE_WAKEUP : BPF_RB_NO_WAKEUP; +} + /* Helper to check if tracking is currently enabled */ static __always_inline int is_enabled(void) { __u32 key = 0; @@ -157,7 +169,7 @@ static __always_inline __u64* take_param(void* map) { \ fill_data; \ \ - bpf_ringbuf_submit(e, 0); \ + bpf_ringbuf_submit(e, wake_flags()); \ return 0; \ } From 4da30e951c733d6a7f98fa7d3888ce7f74fb9d8e Mon Sep 17 00:00:00 2001 From: not-matthias Date: Wed, 15 Jul 2026 20:37:03 +0200 Subject: [PATCH 13/14] perf(memtrack): grow the event ring buffer to 256 MiB 16 MiB holds only ~300K records (~60ms of headroom at multi-M events/s bursts), so a short consumer scheduling stall overflows the buffer and drops events. 256 MiB gives ~1s of worst-case burst headroom. The buffer is memcg-charged and vmap'd, so the only cost is the memory itself while tracking runs. Refs COD-3071 --- crates/memtrack/src/ebpf/c/memtrack.bpf.c | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/memtrack/src/ebpf/c/memtrack.bpf.c b/crates/memtrack/src/ebpf/c/memtrack.bpf.c index 18c5761c..ad64d73b 100644 --- a/crates/memtrack/src/ebpf/c/memtrack.bpf.c +++ b/crates/memtrack/src/ebpf/c/memtrack.bpf.c @@ -40,7 +40,7 @@ BPF_HASH_MAP(tracked_pids, __u32, __u8, 10000); /* Map to store parent-child relationships to detect hierarchy */ BPF_HASH_MAP(pids_ppid, __u32, __u32, 10000); /* Ring buffer for sending events to userspace */ -BPF_RINGBUF(events, 16 * 1024 * 1024); +BPF_RINGBUF(events, 256 * 1024 * 1024); /* Map to control whether tracking is enabled (0 = disabled, 1 = enabled) */ BPF_ARRAY_MAP(tracking_enabled, __u8, 1); /* Counter for events that couldn't be added to the ring buffer */ From 5c27231b94437287b71ab505faba533d06e0cb24 Mon Sep 17 00:00:00 2001 From: not-matthias Date: Wed, 15 Jul 2026 20:37:54 +0200 Subject: [PATCH 14/14] perf(memtrack): drain ringbuf bursts eagerly and reserve poller cores Two consumer-side fixes for event drops under allocation bursts: - After each poll() the poller calls consume() once, which drains the ring buffer to empty without paying an epoll wakeup round-trip per poll cycle. - The encode pipeline no longer claims every core: rayon workers on all cores starve the poll thread (and the tracked command), which is what let the ring buffer fill up. Reserve two cores. Refs COD-3071 --- crates/memtrack/src/ebpf/poller.rs | 5 +++++ crates/memtrack/src/main.rs | 5 ++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/crates/memtrack/src/ebpf/poller.rs b/crates/memtrack/src/ebpf/poller.rs index ab482dc6..6d169afe 100644 --- a/crates/memtrack/src/ebpf/poller.rs +++ b/crates/memtrack/src/ebpf/poller.rs @@ -39,6 +39,11 @@ impl RingBufferPoller { move || { while !shutdown.load(Ordering::Relaxed) { let _ = ringbuf.poll(timeout); + + // consume() drains the buffer to empty in a single call, so + // records produced while poll() was draining are picked up + // here without paying another epoll_wait round-trip. + let _ = ringbuf.consume(); } // Events may still be sitting in the ring buffer after the last diff --git a/crates/memtrack/src/main.rs b/crates/memtrack/src/main.rs index 57dad028..24fc75af 100644 --- a/crates/memtrack/src/main.rs +++ b/crates/memtrack/src/main.rs @@ -126,8 +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))?; + // Leave headroom for the ring buffer poll thread and the tracked + // command: encode workers on every core starve the poller during + // allocation bursts, which overflows the kernel ring buffer. let n_workers = thread::available_parallelism() - .map(|n| n.get()) + .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));