perf(runner-shared): decode memtrack frames for module events in parallel

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);
 

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<R: std::io::Read>(
-        reader: R,
-    ) -> anyhow::Result<impl Iterator<Item = MemtrackEvent>> {
-        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<Vec<MemtrackEvent>> {
+        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::<anyhow::Result<Vec<_>>>()?
+            .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<R: std::io::Read>(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<Vec<MemtrackEvent>> {
+    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::<u8>::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![] };

src/executor/memory/module_artifacts.rs

@@ -90,8 +90,11 @@ fn read_mappings(results_folder: &Path) -> Result<Vec<ProcessMapping>> {
         }
 
         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<Vec<ProcessMapping>> {
 
 /// 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<R: std::io::Read>(reader: R) -> Result<Vec<ProcessMapping>> {
-    let mut timeline = MemtrackArtifact::decode_module_events(reader)?.collect::<Vec<_>>();
+fn read_mappings_from_artifact(artifact: &[u8]) -> Result<Vec<ProcessMapping>> {
+    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)> {