test(runner-shared): benchmark finding module events in a memtrack artifact

Since 5.4.0 the memory executor decodes every event of the memtrack
artifact after the run, only to keep the few Mapping, Fork and Exec events
it needs for module artifacts. On large artifacts this takes minutes
between memtrack's last log line and the upload.

The benchmark builds an artifact laid out like memtrack writes it:
allocation events with forks and execs spread through the stream and the
mapping suffix at the end. It runs without stacks and with an 8 KiB stack
record every 400 events, from 262k events up to 100M, closer to the tens of
millions a CI memory run writes, and checks the expected number of module
events is found.

Events are generated as the encoder consumes them and the artifact is
written to a file under the target directory and mapped, so even the
100M-event artifact with stacks (about 5 GB) never has to fit in memory.

Author: not-matthias

Cargo.lock

@@ -27,8 +27,14 @@ zstd = "0.13.3"
 
 [dev-dependencies]
 divan = { version = "4.2.0", package = "codspeed-divan-compat" }
+memmap2 = "0.9.10"
 rand = "0.8"
+tempfile = { workspace = true }
 
 [[bench]]
 name = "memtrack_writer"
 harness = false
+
+[[bench]]
+name = "memtrack_reader"
+harness = false

crates/runner-shared/Cargo.toml

@@ -0,0 +1,138 @@
+//! Finding the module events (mappings, forks, execs) the runner needs for
+//! offline stack attribution inside a full memtrack artifact. These are a
+//! handful of records among millions of allocation events, so the cost is
+//! dominated by everything the search has to read past.
+
+use divan::Bencher;
+use divan::counter::{BytesCount, ItemsCount};
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
+use runner_shared::artifacts::{
+    MemtrackArtifact, MemtrackEvent, MemtrackEventKind, StackRecord, encode_events,
+};
+use std::io::Write;
+
+fn main() {
+    divan::main();
+}
+
+/// Artifact sizes in events: from four encoder frames, small enough for the
+/// simulation instrument, up to a CI-sized artifact with enough frames to show
+/// how the parallel frame scan scales.
+const SIZES: [usize; 4] = [4 * 64 * 1024, 1_000_000, 10_000_000, 100_000_000];
+
+#[derive(Clone, Copy)]
+enum Payload {
+    Allocations,
+    AllocationsAndStacks,
+}
+
+/// An artifact without stack capture.
+#[divan::bench(args = SIZES, max_time = 10.0)]
+fn find_module_events(bencher: Bencher, n: usize) {
+    bench(bencher, n, Payload::Allocations);
+}
+
+/// An artifact with captured stacks, which make up most of its bytes.
+#[divan::bench(args = SIZES, max_time = 10.0)]
+fn find_module_events_with_stacks(bencher: Bencher, n: usize) {
+    bench(bencher, n, Payload::AllocationsAndStacks);
+}
+
+fn bench(bencher: Bencher, n: usize, payload: Payload) {
+    // The largest artifacts do not fit in memory, so they are written to disk
+    // and mapped: mapped pages stay reclaimable. `/tmp` may be a tmpfs, hence
+    // the target directory.
+    let file = tempfile::tempfile_in(env!("CARGO_TARGET_TMPDIR")).unwrap();
+    let module_events = write_artifact(n, payload, &file);
+    // SAFETY: the file is unlinked and only this process holds it.
+    let artifact = unsafe { memmap2::Mmap::map(&file) }.unwrap();
+
+    let find = || {
+        MemtrackArtifact::decode_module_events(&artifact[..])
+            .unwrap()
+            .count()
+    };
+    assert_eq!(find(), module_events);
+
+    bencher
+        .counter(ItemsCount::new(n))
+        .counter(BytesCount::new(artifact.len()))
+        .bench_local(find);
+}
+
+/// Write an encoded artifact of `n` events laid out like memtrack writes it,
+/// and return the number of module events in it: allocations with a fork or
+/// exec every so often, then the mapping suffix. Events are generated as the
+/// encoder consumes them, so they are never all in memory at once.
+fn write_artifact(n: usize, payload: Payload, out: impl Write) -> usize {
+    // Executable mappings memtrack appends after the event rings drain.
+    const MAPPINGS: usize = 64;
+    // Forks and execs come from BPF tracepoints, interleaved with allocations.
+    const LIFECYCLE_INTERVAL: usize = 16 * 1024;
+    // Keeps the largest artifact with stacks at about 5 GB on disk.
+    const STACK_INTERVAL: usize = 400;
+    const ENCODE_WORKERS: usize = 4;
+    const PID: i32 = 4242;
+
+    let first_mapping = n - MAPPINGS;
+    let mut rng = StdRng::seed_from_u64(42);
+    let events = (0..n).map(|i| {
+        let kind = if i >= first_mapping {
+            MemtrackEventKind::Mapping {
+                path: format!("/usr/lib/libmodule{i}.so"),
+                dev: 0x0800_0001,
+                ino: i as u64,
+                file_offset: 0x1000,
+                len: 0x8_0000,
+            }
+        } else if i % LIFECYCLE_INTERVAL == 0 {
+            if (i / LIFECYCLE_INTERVAL) % 2 == 0 {
+                MemtrackEventKind::Fork { parent_pid: PID }
+            } else {
+                MemtrackEventKind::Exec
+            }
+        } else if matches!(payload, Payload::AllocationsAndStacks) && i % STACK_INTERVAL == 0 {
+            stack_event(&mut rng)
+        } else if i % 2 == 0 {
+            MemtrackEventKind::Malloc {
+                size: rng.gen_range(8..8192),
+                stack_hash: rng.r#gen(),
+            }
+        } else {
+            MemtrackEventKind::Free
+        };
+
+        MemtrackEvent {
+            pid: PID,
+            tid: PID,
+            timestamp: i as u64,
+            addr: rng.r#gen(),
+            kind,
+        }
+    });
+
+    encode_events(events, out, ENCODE_WORKERS).unwrap();
+    MAPPINGS + first_mapping.div_ceil(LIFECYCLE_INTERVAL)
+}
+
+/// A captured stack sized like memtrack's: its default stack copy budget of
+/// incompressible bytes, plus the x86_64 register set.
+fn stack_event(rng: &mut StdRng) -> MemtrackEventKind {
+    const STACK_BYTES: usize = 8192;
+    const REGS: usize = 33;
+    const FP_FRAMES: usize = 16;
+
+    let mut bytes = vec![0; STACK_BYTES];
+    rng.fill(&mut bytes[..]);
+    MemtrackEventKind::Stack {
+        record: Box::new(StackRecord {
+            hash: rng.r#gen(),
+            sp: rng.r#gen(),
+            regs: (0..REGS).map(|_| rng.r#gen()).collect(),
+            bytes,
+            fp_chain: (0..FP_FRAMES).map(|_| rng.r#gen()).collect(),
+            truncated: true,
+        }),
+    }
+}