perf(runner-shared): stream memtrack encoder frames instead of windows

.github/workflows/ci.yml

@@ -133,7 +133,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        mode: [simulation, walltime]
+        mode: [simulation, walltime, memory]
     steps:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
         with:

.pre-commit-config.yaml

@@ -27,6 +27,9 @@ repos:
       - id: clang-format
         files: ^crates/memtrack/src/ebpf/c/.*\.(c|h|bpf\.c)$
         args: [--style=file, -i, --version=22.1.2]
+        # The hook pip-installs clang-format on first use; parallel batches
+        # race on that install and fail with a PermissionError on a cold cache.
+        require_serial: true
   - repo: local
     hooks:
       - id: init-worktree

crates/runner-shared/src/artifacts/memtrack/pipeline.rs

@@ -1,26 +1,32 @@
+use std::cell::RefCell;
+use std::collections::VecDeque;
 use std::io::{BufWriter, Write};
+use std::sync::mpsc::{self, Receiver, TryRecvError};
 
-use rayon::prelude::*;
+use serde::Serialize;
 
 use super::MemtrackEvent;
-use super::writer::MemtrackWriter;
+use super::writer::COMPRESSION_LEVEL;
 
 /// 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;
+/// Frames allowed in flight per worker. At the cap the reader waits for the
+/// oldest frame, bounding memory to about `cap + 1` frames.
+const MAX_IN_FLIGHT_PER_WORKER: usize = 2;
+
+/// An encoded frame, plus the event buffer it was built from so the reader can
+/// reuse it for the next frame instead of allocating a new one.
+type EncodedFrame = (anyhow::Result<Vec<u8>>, Vec<MemtrackEvent>);
 
 /// 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.
+/// zstd frame. Full frames are submitted as soon as they fill, and completed
+/// frames are written in input order. At most
+/// `MAX_IN_FLIGHT_PER_WORKER * n_workers` frames are in flight; at the cap the
+/// reader waits for the oldest frame before pulling more events.
 ///
 /// Blocks the calling thread until `events` is exhausted. Returns the total
 /// number of events written.
@@ -32,33 +38,64 @@ where
     let pool = rayon::ThreadPoolBuilder::new()
         .num_threads(n_workers.max(1))
         .build()?;
+    let max_in_flight = MAX_IN_FLIGHT_PER_WORKER * n_workers.max(1);
 
     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<MemtrackEvent> = Vec::with_capacity(cap);
-    loop {
-        window.clear();
-        window.extend(events.by_ref().take(cap));
-        if window.is_empty() {
-            break;
+    let mut in_flight: VecDeque<Receiver<EncodedFrame>> = VecDeque::with_capacity(max_in_flight);
+    // Event buffers handed back by workers, reused for the next frames.
+    let mut spare: Vec<Vec<MemtrackEvent>> = Vec::new();
+
+    let submit = |frame: Vec<MemtrackEvent>| {
+        let (tx, rx) = mpsc::sync_channel(1);
+        pool.spawn(move || {
+            let encoded = encode_frame(&frame);
+            let _ = tx.send((encoded, frame));
+        });
+        rx
+    };
+    let mut collect = |(encoded, mut frame): EncodedFrame, spare: &mut Vec<Vec<MemtrackEvent>>| {
+        out.write_all(&encoded?)?;
+        wrote_any = true;
+        frame.clear();
+        spare.push(frame);
+        anyhow::Ok(())
+    };
+
+    let mut frame: Vec<MemtrackEvent> = Vec::with_capacity(FRAME_EVENTS);
+    for event in events {
+        frame.push(event);
+        if frame.len() < FRAME_EVENTS {
+            continue;
         }
-        total += window.len() as u64;
-
-        let frames: Vec<Vec<u8>> = pool.install(|| {
-            window
-                .par_chunks(FRAME_EVENTS)
-                .map(encode_frame)
-                .collect::<anyhow::Result<_>>()
-        })?;
-
-        for frame in frames {
-            out.write_all(&frame)?;
+        total += frame.len() as u64;
+        let next = spare
+            .pop()
+            .unwrap_or_else(|| Vec::with_capacity(FRAME_EVENTS));
+        in_flight.push_back(submit(std::mem::replace(&mut frame, next)));
+
+        // Write finished frames in input order. Wait for the oldest frame only
+        // while the queue is at the cap.
+        while let Some(rx) = in_flight.front() {
+            let at_cap = in_flight.len() >= max_in_flight;
+            let result = match rx.try_recv() {
+                Ok(result) => result,
+                Err(TryRecvError::Empty) if !at_cap => break,
+                // Blocks at the cap; fails at once if the worker is gone.
+                Err(_) => recv_frame(rx)?,
+            };
+            in_flight.pop_front();
+            collect(result, &mut spare)?;
         }
-        wrote_any = true;
+    }
+
+    if !frame.is_empty() {
+        total += frame.len() as u64;
+        in_flight.push_back(submit(frame));
+    }
+    for rx in in_flight.drain(..) {
+        collect(recv_frame(&rx)?, &mut spare)?;
     }
 
     // Always emit at least one (possibly empty) frame so the artifact stream is
@@ -71,20 +108,70 @@ where
     Ok(total)
 }
 
+/// Block until the frame behind `rx` is encoded.
+fn recv_frame(rx: &Receiver<EncodedFrame>) -> anyhow::Result<EncodedFrame> {
+    rx.recv()
+        .map_err(|_| anyhow::anyhow!("frame encoder worker exited without a result"))
+}
+
+/// Upper estimate of the msgpack size of one event, used to size the frame
+/// buffer up front. Growing it by doubling instead would leave each worker
+/// pinning about twice a frame's msgpack size.
+const MSGPACK_BYTES_PER_EVENT: usize = 80;
+
+/// Per-thread scratch state reused across frames: the msgpack buffer and the
+/// zstd compression context.
+///
+/// Each worker keeps its buffer (about one frame's msgpack, ~5 MiB) until the
+/// pool exits, even when idle. Peak usage needs it anyway, since every busy
+/// worker holds one, and allocating it per frame measured about 9% slower with
+/// a single worker.
+struct FrameEncoder {
+    msgpack: Vec<u8>,
+    compressor: zstd::bulk::Compressor<'static>,
+}
+
+thread_local! {
+    static FRAME_ENCODER: RefCell<Option<FrameEncoder>> = const { RefCell::new(None) };
+}
+
 /// Encode one batch as a single self-contained zstd frame.
+///
+/// The batch is serialized with `rmp_serde` into a reused buffer, then
+/// compressed in one shot with a reused zstd context. The decoded bytes are the
+/// same msgpack stream `MemtrackWriter` produces.
 fn encode_frame(batch: &[MemtrackEvent]) -> anyhow::Result<Vec<u8>> {
-    let mut writer = MemtrackWriter::new(Vec::new())?;
-    for event in batch {
-        writer.write_event(event)?;
-    }
-    writer.finish()
+    FRAME_ENCODER.with(|cell| {
+        let mut slot = cell.borrow_mut();
+        let enc = match slot.as_mut() {
+            Some(enc) => enc,
+            None => slot.insert(FrameEncoder {
+                msgpack: Vec::new(),
+                compressor: zstd::bulk::Compressor::new(COMPRESSION_LEVEL)?,
+            }),
+        };
+
+        enc.msgpack.clear();
+        enc.msgpack
+            .reserve_exact(batch.len() * MSGPACK_BYTES_PER_EVENT);
+        let mut serializer = rmp_serde::Serializer::new(&mut enc.msgpack);
+        for event in batch {
+            event.serialize(&mut serializer)?;
+        }
+
+        let mut compressed = enc.compressor.compress(&enc.msgpack)?;
+        // Trim the worst-case compression bound so in-flight frames only hold
+        // their actual size.
+        compressed.shrink_to_fit();
+        Ok(compressed)
+    })
 }
 
 #[cfg(test)]
 mod tests {
     use std::io::Cursor;
 
-    use super::super::{MemtrackArtifact, MemtrackEventKind};
+    use super::super::{MemtrackArtifact, MemtrackEventKind, MemtrackWriter};
     use super::*;
 
     fn malloc_events(range: std::ops::Range<u64>) -> Vec<MemtrackEvent> {
@@ -116,11 +203,11 @@ mod tests {
     }
 
     #[test]
-    fn preserves_order_across_window_boundary() -> anyhow::Result<()> {
-        let events = malloc_events(0..(FRAME_EVENTS * WINDOW_FRAMES + 1) as u64);
+    fn preserves_order_beyond_in_flight_cap() -> anyhow::Result<()> {
+        let events = malloc_events(0..(FRAME_EVENTS as u64 * 5 + 3));
 
         let mut out = Vec::new();
-        let total = encode_events(events.clone(), &mut out, 4)?;
+        let total = encode_events(events.clone(), &mut out, 1)?;
         assert_eq!(total, events.len() as u64);
 
         let decoded: Vec<_> = MemtrackArtifact::decode_streamed(Cursor::new(out))?.collect();
@@ -129,6 +216,25 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn frame_payload_matches_memtrack_writer() -> anyhow::Result<()> {
+        let events = malloc_events(0..10_000);
+
+        let mut reference = MemtrackWriter::new(Vec::new())?;
+        for event in &events {
+            reference.write_event(event)?;
+        }
+        let reference = zstd::decode_all(Cursor::new(reference.finish()?))?;
+
+        // Encode twice to also exercise the reused per-thread buffers.
+        for _ in 0..2 {
+            let frame = zstd::decode_all(Cursor::new(encode_frame(&events)?))?;
+            assert_eq!(frame, reference);
+        }
+
+        Ok(())
+    }
+
     #[test]
     fn empty_source_writes_a_valid_stream() -> anyhow::Result<()> {
         let events: Vec<MemtrackEvent> = Vec::new();

crates/runner-shared/src/artifacts/memtrack/writer.rs

@@ -3,16 +3,17 @@ use std::io::{BufWriter, Write};
 
 use super::MemtrackEvent;
 
+/// zstd level used for memtrack artifacts. We're dealing with a lot of events, so
+/// we want to compress as much as possible while not taking too much time.
+pub(crate) const COMPRESSION_LEVEL: i32 = -5;
+
 /// Streaming writer for memtrack events, serializing into a zstd-compressed sink.
 pub struct MemtrackWriter<B: Write> {
     serializer: rmp_serde::Serializer<B>,
 }
 
 impl<W: Write> MemtrackWriter<BufWriter<zstd::Encoder<'static, W>>> {
     pub fn new(writer: W) -> anyhow::Result<Self> {
-        // 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)?;