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 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: 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/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/c/memtrack.bpf.c b/crates/memtrack/src/ebpf/c/memtrack.bpf.c index 06601e27..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 */ @@ -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; \ } 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..6d169afe 100644 --- a/crates/memtrack/src/ebpf/poller.rs +++ b/crates/memtrack/src/ebpf/poller.rs @@ -1,78 +1,67 @@ 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); + + // 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(); + } - 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 diff --git a/crates/memtrack/src/main.rs b/crates/memtrack/src/main.rs index e9e06685..24fc75af 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,14 @@ 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; - } - } - } - } - }); + // 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().saturating_sub(2).max(1)) + .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 +149,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 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/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 + }); +} 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/pipeline.rs b/crates/runner-shared/src/artifacts/memtrack/pipeline.rs new file mode 100644 index 00000000..c47b3aed --- /dev/null +++ b/crates/runner-shared/src/artifacts/memtrack/pipeline.rs @@ -0,0 +1,144 @@ +use std::io::{BufWriter, Write}; + +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; + + 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(|| { + 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(()) + } +} 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..2f665766 --- /dev/null +++ b/crates/runner-shared/src/artifacts/memtrack/writer.rs @@ -0,0 +1,41 @@ +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 = -5; + const BUFFER_SIZE: usize = 256 * 1024 /* 256 KB */; + + let encoder = zstd::Encoder::new(writer, COMPRESSION_LEVEL)?; + let writer = BufWriter::with_capacity(BUFFER_SIZE, encoder); + Ok(Self { + serializer: rmp_serde::Serializer::new(writer), + }) + } + + /// 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()?; + writer.flush()?; + Ok(writer) + } +} + +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(()) + } +}