diff --git a/crates/runner-shared/benches/memtrack_reader.rs b/crates/runner-shared/benches/memtrack_reader.rs index c04f2c6b..2215eb64 100644 --- a/crates/runner-shared/benches/memtrack_reader.rs +++ b/crates/runner-shared/benches/memtrack_reader.rs @@ -19,8 +19,9 @@ fn main() { divan::main(); } -/// Artifact sizes in events. -const WALLTIME_MODE_SIZES: &[usize] = &[1_000_000]; +/// Artifact sizes in events, up to a CI-sized artifact with enough frames to +/// show how the parallel frame scan scales. +const WALLTIME_MODE_SIZES: &[usize] = &[1_000_000, 10_000_000, 100_000_000]; /// Under the simulation and memory instruments, generating and searching /// larger artifacts outlasts the CI job, so these modes only search the @@ -65,9 +66,9 @@ fn bench(bencher: Bencher, n: usize, payload: Payload) { let artifact = unsafe { memmap2::Mmap::map(&file) }.unwrap(); let find = || { - MemtrackArtifact::decode_module_events(&artifact[..]) + MemtrackArtifact::decode_module_events(&artifact) .unwrap() - .count() + .len() }; assert_eq!(find(), module_events); diff --git a/crates/runner-shared/src/artifacts/memtrack/mod.rs b/crates/runner-shared/src/artifacts/memtrack/mod.rs index 05339fc7..96209507 100644 --- a/crates/runner-shared/src/artifacts/memtrack/mod.rs +++ b/crates/runner-shared/src/artifacts/memtrack/mod.rs @@ -36,20 +36,27 @@ impl MemtrackArtifact { }) } - /// Stream only the events that place modules in processes: executable - /// mappings, and the forks and execs that inherit or drop them. Offline - /// stack attribution needs these few events out of the whole artifact. - pub fn decode_module_events( - reader: R, - ) -> anyhow::Result> { - Ok(Self::decode_streamed(reader)?.filter(|event| { - matches!( - event.kind, - MemtrackEventKind::Mapping { .. } - | MemtrackEventKind::Fork { .. } - | MemtrackEventKind::Exec - ) - })) + /// Find the events that place modules in processes: executable mappings, + /// and the forks and execs that inherit or drop them. Offline stack + /// attribution needs these few events out of the whole artifact. + /// + /// Events are returned in artifact order. A truncated artifact yields the + /// events before the cut. + pub fn decode_module_events(artifact: &[u8]) -> anyhow::Result> { + use rayon::prelude::*; + + // The encoder writes self-contained zstd frames: split on them and + // decode the frames in parallel. Collecting in frame order keeps the + // artifact order. + let (frames, tail) = split_zstd_frames(artifact); + let mut events = frames + .par_iter() + .map(|frame| module_events_in_frame(frame)) + .collect::>>()? + .concat(); + // A truncated last frame cannot be split off; stream it on its own. + events.extend(Self::decode_streamed(tail)?.filter(|event| event.kind.is_module_event())); + Ok(events) } pub fn is_empty(reader: R) -> bool { @@ -60,6 +67,27 @@ impl MemtrackArtifact { } } +/// Split an artifact into its complete zstd frames, plus the unsplittable rest. +fn split_zstd_frames(mut artifact: &[u8]) -> (Vec<&[u8]>, &[u8]) { + let mut frames = Vec::new(); + while let Ok(len) = zstd::zstd_safe::find_frame_compressed_size(artifact) { + let (frame, rest) = artifact.split_at(len); + frames.push(frame); + artifact = rest; + } + (frames, artifact) +} + +/// Stream one frame and return its module events, holding one decoded event +/// at a time however large the frame decompresses. Like +/// [`MemtrackEventStream`], reading stops at the first event that fails to +/// decode. +fn module_events_in_frame(frame: &[u8]) -> anyhow::Result> { + Ok(MemtrackArtifact::decode_streamed(frame)? + .filter(|event| event.kind.is_module_event()) + .collect()) +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct MemtrackEvent { pub pid: pid_t, @@ -141,6 +169,14 @@ pub enum MemtrackEventKind { }, } +impl MemtrackEventKind { + /// Whether the event places modules in processes: an executable mapping, + /// or a fork or exec that inherits or drops them. + pub fn is_module_event(&self) -> bool { + matches!(self, Self::Mapping { .. } | Self::Fork { .. } | Self::Exec) + } +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct StackRecord { pub hash: u64, @@ -351,6 +387,81 @@ mod tests { Ok(()) } + /// Every event kind, with stack payloads, across several frames and a + /// last frame cut short: the parallel frame decoder must find exactly what + /// a full streamed decode filtered to module events finds. + #[test] + fn module_events_match_a_filtered_full_decode() -> anyhow::Result<()> { + let kinds = [ + MemtrackEventKind::Malloc { + size: 64, + stack_hash: 7, + }, + MemtrackEventKind::Free, + MemtrackEventKind::Realloc { + old_addr: Some(0x10), + size: 128, + stack_hash: 0, + }, + MemtrackEventKind::Stack { + record: Box::new(StackRecord { + hash: 7, + sp: 0x7fff_0000, + regs: vec![1; 33], + bytes: vec![0xab; 4096], + fp_chain: vec![0x5555_0000, 0x5555_0010], + truncated: true, + }), + }, + MemtrackEventKind::Fork { parent_pid: 1 }, + MemtrackEventKind::Exec, + MemtrackEventKind::Rss { + member: 1, + size: 4096, + }, + MemtrackEventKind::Mapping { + path: "/usr/lib/libexample.so".into(), + dev: 0x0800_0001, + ino: 42, + file_offset: 0x1000, + len: 0x2000, + }, + MemtrackEventKind::Exit, + MemtrackEventKind::Rmap { + member: 1, + delta: -70_000, + }, + ]; + let events: Vec<_> = (0..3000u64) + .map(|i| MemtrackEvent { + pid: 2, + tid: 3, + timestamp: i, + addr: i * 16, + kind: kinds[i as usize % kinds.len()].clone(), + }) + .collect(); + + let mut artifact = Vec::new(); + for batch in events.chunks(1000) { + let mut writer = MemtrackWriter::new(Vec::::new())?; + for event in batch { + writer.write_event(event)?; + } + artifact.extend_from_slice(&writer.finish()?); + } + let truncated = &artifact[..artifact.len() - 64]; + + for artifact in [&artifact[..], truncated] { + let expected: Vec<_> = MemtrackArtifact::decode_streamed(artifact)? + .filter(|event| event.kind.is_module_event()) + .collect(); + assert!(!expected.is_empty()); + assert_eq!(MemtrackArtifact::decode_module_events(artifact)?, expected); + } + Ok(()) + } + #[test] fn test_artifact_is_empty() -> anyhow::Result<()> { let artifact = MemtrackArtifact { events: vec![] }; diff --git a/src/executor/memory/module_artifacts.rs b/src/executor/memory/module_artifacts.rs index 0e80d183..9aaff76c 100644 --- a/src/executor/memory/module_artifacts.rs +++ b/src/executor/memory/module_artifacts.rs @@ -90,8 +90,11 @@ fn read_mappings(results_folder: &Path) -> Result> { } let file = std::fs::File::open(entry.path())?; + // SAFETY: read-only mapping of an artifact memtrack finished writing + // before it exited; nothing writes to it while the runner reads it. + let artifact = unsafe { memmap2::Mmap::map(&file)? }; mappings.extend( - read_mappings_from_artifact(file) + read_mappings_from_artifact(&artifact) .with_context(|| format!("Failed to decode {:?}", entry.path()))?, ); } @@ -102,8 +105,8 @@ fn read_mappings(results_folder: &Path) -> Result> { /// Reconstruct mappings across forks because inherited perf events do not /// synthesize mappings that already existed when a child was forked. -fn read_mappings_from_artifact(reader: R) -> Result> { - let mut timeline = MemtrackArtifact::decode_module_events(reader)?.collect::>(); +fn read_mappings_from_artifact(artifact: &[u8]) -> Result> { + let mut timeline = MemtrackArtifact::decode_module_events(artifact)?; // Ties break so exec purges before mapping, while fork inherits that mapping. timeline.sort_by_key(|event| { @@ -611,7 +614,7 @@ mod tests { artifact.encode_to_writer(&mut encoded).unwrap(); assert_eq!( - read_mappings_from_artifact(std::io::Cursor::new(encoded)).unwrap(), + read_mappings_from_artifact(&encoded).unwrap(), vec![ ProcessMapping { pid: 7, @@ -858,7 +861,7 @@ mod tests { let artifact = MemtrackArtifact { events }; let mut encoded = Vec::new(); artifact.encode_to_writer(&mut encoded).unwrap(); - read_mappings_from_artifact(std::io::Cursor::new(encoded)).unwrap() + read_mappings_from_artifact(&encoded).unwrap() } fn mapping_summary(mappings: &[ProcessMapping]) -> Vec<(pid_t, &str, u64)> {