diff --git a/Cargo.lock b/Cargo.lock index e94bda5e..8b13c34d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2378,6 +2378,7 @@ dependencies = [ "anyhow", "bindgen", "clap", + "crossbeam-channel", "env_logger", "glob", "insta", @@ -3592,6 +3593,7 @@ dependencies = [ "anyhow", "bincode", "codspeed-divan-compat", + "crossbeam-channel", "itertools 0.14.0", "libc", "linux-perf-event-reader 0.10.2 (registry+https://github.com/rust-lang/crates.io-index)", diff --git a/Cargo.toml b/Cargo.toml index 48aced24..0e0a62f1 100644 --- a/Cargo.toml +++ b/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" diff --git a/crates/memtrack/Cargo.toml b/crates/memtrack/Cargo.toml index ae481941..d5bb32c5 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 3a1bfa26..f16e885e 100644 --- a/crates/memtrack/src/ebpf/memtrack.rs +++ b/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, + 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) + RingBufferPoller::with_batch_channel(&self.skel.maps.events, poll_timeout_ms, batch_size) } } diff --git a/crates/memtrack/src/ebpf/poller.rs b/crates/memtrack/src/ebpf/poller.rs index 1f4ef2b4..bdf30f53 100644 --- a/crates/memtrack/src/ebpf/poller.rs +++ b/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; +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> = 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, 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 deliver events in ordered batches of at most `batch_size`. + pub fn with_batch_channel( rb_map: &M, - handler: EventHandler, poll_timeout_ms: u64, - ) -> Result { + batch_size: usize, + ) -> 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); + 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( - 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 9f1af587..68d90cb2 100644 --- a/crates/memtrack/src/ebpf/tracker.rs +++ b/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, } 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> { - // 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>> { 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 diff --git a/crates/memtrack/src/main.rs b/crates/memtrack/src/main.rs index b14487eb..80fd9633 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_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::(); - - // 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 diff --git a/crates/memtrack/tests/shared.rs b/crates/memtrack/tests/shared.rs index 054b3b8e..f5653bcc 100644 --- a/crates/memtrack/tests/shared.rs +++ b/crates/memtrack/tests/shared.rs @@ -149,8 +149,8 @@ pub fn track_command( let rx = tracker.track(root_pid)?; let mut events = Vec::new(); - while let Ok(event) = rx.recv_timeout(Duration::from_secs(10)) { - events.push(event); + while let Ok(batch) = rx.recv_timeout(Duration::from_secs(10)) { + events.extend(batch); } // Drop the tracker in a new thread to not block the test diff --git a/crates/runner-shared/Cargo.toml b/crates/runner-shared/Cargo.toml index d0c67f77..09006a44 100644 --- a/crates/runner-shared/Cargo.toml +++ b/crates/runner-shared/Cargo.toml @@ -10,6 +10,7 @@ serde = { workspace = true } serde_json = { workspace = true } # Pinned to 1.x: 2.0 changes the wire format and serde integration bincode = "1.3" +crossbeam-channel = { workspace = true } itertools = { workspace = true } linux-perf-event-reader = { workspace = true } log = { workspace = true } diff --git a/crates/runner-shared/benches/memtrack_writer.rs b/crates/runner-shared/benches/memtrack_writer.rs index 482b2a7b..a82021d1 100644 --- a/crates/runner-shared/benches/memtrack_writer.rs +++ b/crates/runner-shared/benches/memtrack_writer.rs @@ -1,7 +1,10 @@ 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_batches_with_level, +}; +use std::io; fn main() { divan::main(); @@ -53,3 +56,31 @@ fn write_events(bencher: Bencher, n: usize) { writer.finish().unwrap(); }); } + +fn run_encode_pipeline(bencher: Bencher, n: usize, level: i32) { + const BATCH_EVENTS: usize = 64 * 1024; + let events = generate_events(n); + + bencher + .with_inputs(|| { + events + .chunks(BATCH_EVENTS) + .map(<[MemtrackEvent]>::to_vec) + .collect::>() + }) + .bench_values(|batches| { + encode_batches_with_level(batches, io::sink(), 4, level).unwrap(); + }); +} + +/// Level-1 zstd. +#[divan::bench(args = [100_000, 1_000_000])] +fn encode_pipeline(bencher: Bencher, n: usize) { + run_encode_pipeline(bencher, n, 1); +} + +/// Fast zstd (production default): negative level selects the "fast" super-strategy. +#[divan::bench(args = [100_000, 1_000_000])] +fn encode_pipeline_zstd_fast(bencher: Bencher, n: usize) { + run_encode_pipeline(bencher, n, -5); +} diff --git a/crates/runner-shared/src/artifacts/memtrack.rs b/crates/runner-shared/src/artifacts/memtrack.rs deleted file mode 100644 index 64c90e82..00000000 --- a/crates/runner-shared/src/artifacts/memtrack.rs +++ /dev/null @@ -1,201 +0,0 @@ -use libc::pid_t; -use serde::{Deserialize, Serialize}; -use std::io::{BufReader, BufWriter, Read, Write}; - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct MemtrackArtifact { - pub events: Vec, -} -impl super::ArtifactExt for MemtrackArtifact { - fn encode_to_writer(&self, writer: W) -> anyhow::Result<()> { - let mut writer = MemtrackWriter::new(writer)?; - for event in &self.events { - writer.write_event(event)?; - } - writer.finish()?; - Ok(()) - } -} - -impl MemtrackArtifact { - pub fn decode_streamed( - reader: R, - ) -> anyhow::Result>>> { - let decoder = zstd::Decoder::new(reader)?; - Ok(MemtrackEventStream { - deserializer: rmp_serde::Deserializer::new(decoder), - }) - } - - pub fn is_empty(reader: R) -> bool { - let Ok(mut stream) = MemtrackArtifact::decode_streamed(BufReader::new(reader)) else { - return true; - }; - stream.next().is_none() - } -} - -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] -pub struct MemtrackEvent { - pub pid: pid_t, - pub tid: pid_t, - pub timestamp: u64, - pub addr: u64, - #[serde(flatten)] - pub kind: MemtrackEventKind, -} - -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] -#[serde(tag = "type")] -pub enum MemtrackEventKind { - Malloc { - size: u64, - }, - Free, - Realloc { - #[serde(default, skip_serializing_if = "Option::is_none")] - old_addr: Option, - size: u64, - }, - Calloc { - size: u64, - }, - AlignedAlloc { - size: u64, - }, - Mmap { - size: u64, - }, - Munmap { - size: u64, - }, - Brk { - size: u64, - }, -} - -pub struct MemtrackEventStream { - deserializer: rmp_serde::Deserializer>, -} - -impl Iterator for MemtrackEventStream { - type Item = MemtrackEvent; - - fn next(&mut self) -> Option { - MemtrackEvent::deserialize(&mut self.deserializer).ok() - } -} - -/// 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; - - use super::*; - use std::io::Cursor; - - #[test] - fn test_decode_streamed() -> anyhow::Result<()> { - let events = vec![ - MemtrackEvent { - pid: 1, - tid: 11, - timestamp: 100, - addr: 0x10, - kind: MemtrackEventKind::Malloc { size: 64 }, - }, - MemtrackEvent { - pid: 1, - tid: 12, - timestamp: 200, - addr: 0x20, - kind: MemtrackEventKind::Free, - }, - ]; - - let artifact = MemtrackArtifact { - events: events.clone(), - }; - let mut buf = Vec::new(); - artifact.encode_to_writer(&mut buf)?; - - let stream = MemtrackArtifact::decode_streamed(Cursor::new(buf))?; - let collected: Vec<_> = stream.collect(); - assert_eq!(collected, events); - - Ok(()) - } - - #[test] - fn test_artifact_is_empty() -> anyhow::Result<()> { - let artifact = MemtrackArtifact { events: vec![] }; - - let mut buf = Vec::new(); - artifact.encode_to_writer(&mut buf)?; - - let reader = Cursor::new(buf); - assert!(MemtrackArtifact::is_empty(reader)); - - Ok(()) - } - - #[test] - 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"); - assert_eq!( - MemtrackArtifact::decode_streamed(Cursor::new(buf))?.count(), - 1 - ); - - let event = MemtrackArtifact::decode_streamed(Cursor::new(buf))? - .next() - .unwrap(); - assert!(matches!( - event.kind, - MemtrackEventKind::Realloc { - old_addr: None, - size: 42 - } - )); - - Ok(()) - } -} diff --git a/crates/runner-shared/src/artifacts/memtrack/mod.rs b/crates/runner-shared/src/artifacts/memtrack/mod.rs new file mode 100644 index 00000000..95a70afc --- /dev/null +++ b/crates/runner-shared/src/artifacts/memtrack/mod.rs @@ -0,0 +1,283 @@ +use libc::pid_t; +use serde::{Deserialize, Serialize}; +use std::io::{BufReader, Read, Write}; + +mod pipeline; +mod writer; + +pub use pipeline::*; +pub use writer::*; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MemtrackArtifact { + pub events: Vec, +} +impl super::ArtifactExt for MemtrackArtifact { + fn encode_to_writer(&self, writer: W) -> anyhow::Result<()> { + let mut writer = MemtrackWriter::new(writer)?; + for event in &self.events { + writer.write_event(event)?; + } + writer.finish()?; + Ok(()) + } +} + +impl MemtrackArtifact { + pub fn decode_streamed( + reader: R, + ) -> anyhow::Result>>> { + let decoder = zstd::Decoder::new(reader)?; + Ok(MemtrackEventStream { + deserializer: rmp_serde::Deserializer::new(decoder), + }) + } + + pub fn is_empty(reader: R) -> bool { + let Ok(mut stream) = MemtrackArtifact::decode_streamed(BufReader::new(reader)) else { + return true; + }; + stream.next().is_none() + } +} + +#[derive(Debug, Clone, Copy, Deserialize, PartialEq, Eq)] +pub struct MemtrackEvent { + pub pid: pid_t, + pub tid: pid_t, + pub timestamp: u64, + pub addr: u64, + #[serde(flatten)] + pub kind: MemtrackEventKind, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(tag = "type")] +pub enum MemtrackEventKind { + Malloc { + size: u64, + }, + Free, + Realloc { + #[serde(default, skip_serializing_if = "Option::is_none")] + old_addr: Option, + size: u64, + }, + Calloc { + size: u64, + }, + AlignedAlloc { + size: u64, + }, + Mmap { + size: u64, + }, + Munmap { + size: u64, + }, + Brk { + size: u64, + }, +} + +impl serde::Serialize for MemtrackEvent { + fn serialize(&self, serializer: S) -> Result { + use serde::ser::SerializeMap; + + let (type_name, size, old_addr): (&str, Option, Option) = match self.kind { + MemtrackEventKind::Malloc { size } => ("Malloc", Some(size), None), + MemtrackEventKind::Free => ("Free", None, None), + MemtrackEventKind::Realloc { old_addr, size } => ("Realloc", Some(size), old_addr), + MemtrackEventKind::Calloc { size } => ("Calloc", Some(size), None), + MemtrackEventKind::AlignedAlloc { size } => ("AlignedAlloc", Some(size), None), + MemtrackEventKind::Mmap { size } => ("Mmap", Some(size), None), + MemtrackEventKind::Munmap { size } => ("Munmap", Some(size), None), + MemtrackEventKind::Brk { size } => ("Brk", Some(size), None), + }; + + let len = 5 + usize::from(old_addr.is_some()) + usize::from(size.is_some()); + let mut map = serializer.serialize_map(Some(len))?; + map.serialize_entry("pid", &self.pid)?; + map.serialize_entry("tid", &self.tid)?; + map.serialize_entry("timestamp", &self.timestamp)?; + map.serialize_entry("addr", &self.addr)?; + map.serialize_entry("type", type_name)?; + if let Some(old_addr) = old_addr { + map.serialize_entry("old_addr", &old_addr)?; + } + if let Some(size) = size { + map.serialize_entry("size", &size)?; + } + map.end() + } +} + +pub struct MemtrackEventStream { + deserializer: rmp_serde::Deserializer>, +} + +impl Iterator for MemtrackEventStream { + type Item = MemtrackEvent; + + fn next(&mut self) -> Option { + MemtrackEvent::deserialize(&mut self.deserializer).ok() + } +} + +#[cfg(test)] +mod tests { + use crate::artifacts::ArtifactExt; + + use super::*; + use std::io::Cursor; + + #[test] + fn test_decode_streamed() -> anyhow::Result<()> { + let events = vec![ + MemtrackEvent { + pid: 1, + tid: 11, + timestamp: 100, + addr: 0x10, + kind: MemtrackEventKind::Malloc { size: 64 }, + }, + MemtrackEvent { + pid: 1, + tid: 12, + timestamp: 200, + addr: 0x20, + kind: MemtrackEventKind::Free, + }, + ]; + + let artifact = MemtrackArtifact { + events: events.clone(), + }; + let mut buf = Vec::new(); + artifact.encode_to_writer(&mut buf)?; + + let stream = MemtrackArtifact::decode_streamed(Cursor::new(buf))?; + let collected: Vec<_> = stream.collect(); + assert_eq!(collected, events); + + 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![] }; + + let mut buf = Vec::new(); + artifact.encode_to_writer(&mut buf)?; + + let reader = Cursor::new(buf); + assert!(MemtrackArtifact::is_empty(reader)); + + Ok(()) + } + + #[test] + 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"); + assert_eq!( + MemtrackArtifact::decode_streamed(Cursor::new(buf))?.count(), + 1 + ); + + let event = MemtrackArtifact::decode_streamed(Cursor::new(buf))? + .next() + .unwrap(); + assert!(matches!( + event.kind, + MemtrackEventKind::Realloc { + old_addr: None, + size: 42 + } + )); + + Ok(()) + } +} 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..d84f1f2f --- /dev/null +++ b/crates/runner-shared/src/artifacts/memtrack/pipeline.rs @@ -0,0 +1,213 @@ +use std::collections::BTreeMap; +use std::io::{BufWriter, Write}; +use std::thread; + +use serde::Serialize; + +use super::MemtrackEvent; +use super::writer::COMPRESSION_LEVEL; + +/// Encode ordered event batches into a single compressed artifact stream, using +/// `n_workers` threads to serialize and compress batches in parallel. +/// +/// Each batch becomes one self-contained frame. Frames are reordered by sequence +/// number before being written, so the output matches the input order regardless +/// of which worker finishes first. Returns the total number of events written. +/// +/// The caller's thread drives the dispatch loop and blocks until `batches` is +/// exhausted, so a channel receiver source keeps the pipeline alive until the +/// channel is closed. +pub fn encode_batches(batches: S, out: W, n_workers: usize) -> anyhow::Result +where + S: IntoIterator>, + W: Write + Send, +{ + encode_batches_with_level(batches, out, n_workers, COMPRESSION_LEVEL) +} + +/// Same as [`encode_batches`] but with an explicit per-frame zstd `level` +/// (negative levels select the "fast" strategy). Lets benchmarks probe faster +/// levels without changing the production default. +pub fn encode_batches_with_level( + batches: S, + out: W, + n_workers: usize, + level: i32, +) -> anyhow::Result +where + S: IntoIterator>, + W: Write + Send, +{ + let n_workers = n_workers.max(1); + let (work_tx, work_rx) = crossbeam_channel::bounded::<(u64, Vec)>(2 * n_workers); + let (frame_tx, frame_rx) = crossbeam_channel::unbounded::<(u64, Vec)>(); + // Recycle frame buffers from the writer back to the workers, so steady-state + // compression reuses allocations instead of allocating a fresh output Vec per + // frame. Best-effort: an empty pool just means a worker allocates a new buffer. + let (pool_tx, pool_rx) = crossbeam_channel::unbounded::>(); + + thread::scope(|scope| { + let workers: Vec<_> = (0..n_workers) + .map(|_| { + let work_rx = work_rx.clone(); + let frame_tx = frame_tx.clone(); + let pool_rx = pool_rx.clone(); + scope.spawn(move || -> anyhow::Result<()> { + // One compressor and one scratch buffer per worker, reused across + // frames so we amortize the zstd CCtx allocation and avoid a fresh + // msgpack Vec on every batch. + let mut compressor = zstd::bulk::Compressor::new(level)?; + let mut scratch = Vec::new(); + while let Ok((seq, batch)) = work_rx.recv() { + let mut frame = pool_rx.try_recv().unwrap_or_default(); + encode_frame(&mut compressor, &mut scratch, &batch, &mut frame)?; + if frame_tx.send((seq, frame)).is_err() { + break; + } + } + Ok(()) + }) + }) + .collect(); + drop(work_rx); + drop(frame_tx); + drop(pool_rx); + + let writer_thread = scope.spawn(move || -> anyhow::Result<()> { + let mut out = BufWriter::new(out); + let mut next_seq = 0; + let mut pending = BTreeMap::new(); + + while let Ok((seq, frame)) = frame_rx.recv() { + pending.insert(seq, frame); + while let Some(mut frame) = pending.remove(&next_seq) { + out.write_all(&frame)?; + next_seq += 1; + // Return the buffer for a worker to reuse (best-effort). + frame.clear(); + let _ = pool_tx.send(frame); + } + } + + if !pending.is_empty() { + anyhow::bail!("memtrack writer stopped with missing frame {next_seq}"); + } + + out.flush()?; + Ok(()) + }); + + // Tag batches with a sequence number in arrival order and hand them to the + // workers. The bounded work channel applies backpressure when the workers + // fall behind. + // + // A send failure means a worker (and likely the writer) has already died. + // Record it but don't return yet: joining below lets the writer's own error + // (e.g. an IO failure) surface as the root cause instead of this "workers + // stopped early" symptom. + let mut seq = 0; + let mut total = 0; + let mut workers_alive = true; + for batch in batches { + total += batch.len() as u64; + if work_tx.send((seq, batch)).is_err() { + workers_alive = false; + break; + } + seq += 1; + } + + // Always emit at least one (possibly empty) frame so the artifact stream is + // valid and decodable even when no events were recorded. + if workers_alive && seq == 0 && work_tx.send((0, Vec::new())).is_err() { + workers_alive = false; + } + drop(work_tx); + + // Join workers, then the writer. A worker error or the writer's IO error is + // surfaced here (via `?`) ahead of the `workers_alive` symptom below. + for worker in workers { + worker + .join() + .map_err(|_| anyhow::anyhow!("memtrack encode worker panicked"))??; + } + writer_thread + .join() + .map_err(|_| anyhow::anyhow!("memtrack writer thread panicked"))??; + + if !workers_alive { + anyhow::bail!("memtrack encode workers stopped early"); + } + Ok(total) + }) +} + +/// Serialize one batch into `scratch` as concatenated msgpack, then compress it +/// into `frame` as a single self-contained zstd frame with the worker's reused +/// `compressor`. `frame` is overwritten from the start, so a recycled buffer can +/// be passed in to avoid an allocation. +fn encode_frame( + compressor: &mut zstd::bulk::Compressor, + scratch: &mut Vec, + batch: &[MemtrackEvent], + frame: &mut Vec, +) -> anyhow::Result<()> { + scratch.clear(); + let mut serializer = rmp_serde::Serializer::new(&mut *scratch); + for event in batch { + event.serialize(&mut serializer)?; + } + + frame.clear(); + frame.reserve(zstd::zstd_safe::compress_bound(scratch.len())); + compressor.compress_to_buffer(scratch.as_slice(), frame)?; + Ok(()) +} + +#[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_workers() -> anyhow::Result<()> { + let events = malloc_events(0..10_000); + let batches: Vec<_> = events.chunks(1000).map(<[_]>::to_vec).collect(); + + let mut out = Vec::new(); + let total = encode_batches(batches, &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 batches: Vec> = Vec::new(); + + let mut out = Vec::new(); + let total = encode_batches(batches, &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..78f71b94 --- /dev/null +++ b/crates/runner-shared/src/artifacts/memtrack/writer.rs @@ -0,0 +1,49 @@ +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, +} + +// Negative levels select zstd's "fast" super-strategy. Events are streamed to +// disk in parallel while the tracked process runs, so throughput (keeping the +// encode ahead of the allocation rate) matters more than ratio; the repeated +// map keys that dominate the raw bytes compress away even at a fast level. +pub(super) const COMPRESSION_LEVEL: i32 = -5; +const BUFFER_SIZE: usize = 256 * 1024 /* 256 KB */; + +impl MemtrackWriter>> { + pub fn new(writer: W) -> anyhow::Result { + Self::new_with_level(writer, COMPRESSION_LEVEL) + } + + pub fn new_with_level(writer: W, level: i32) -> anyhow::Result { + // Serializing an event emits ~13 tiny writes (map header + per-field key/value). + // The BufWriter must wrap the encoder so those writes are coalesced before reaching + // zstd's streaming compressor, whose per-call overhead dwarfs the buffer copy. + let encoder = zstd::Encoder::new(writer, level)?; + let writer = BufWriter::with_capacity(BUFFER_SIZE, encoder); + Ok(Self { + serializer: rmp_serde::Serializer::new(writer), + }) + } + + /// Finish writing and flush the compression stream + pub fn finish(self) -> anyhow::Result { + let writer = self.serializer.into_inner(); + let encoder = writer.into_inner().map_err(|e| anyhow::anyhow!("{e}"))?; + let inner = encoder.finish()?; + Ok(inner) + } +} + +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(()) + } +}