Skip to content

Commit 1663e23

Browse files
committed
perf(runner-shared): decode memtrack frames for module events in parallel
Finding the Mapping, Fork and Exec events decoded the whole artifact through one streaming deserializer on one thread. The streaming reader copies every key, string and stack payload into a fresh allocation before serde sees it. Every frame the encoder writes is a self-contained zstd frame, so decode_module_events now splits the artifact into its frames and decodes them on the rayon pool. Each worker decompresses a frame into a reused buffer and decodes events from that slice, so values are read in place instead of copied out of a stream. Results are concatenated in artifact order. A truncated last frame cannot be split off and is still streamed, keeping the events before the cut. The runner maps the artifact instead of streaming it from the file. The Mapping/Fork/Exec filter moves to MemtrackEventKind::is_module_event, and a test checks the new decoder against a filtered full decode over every event kind, several frames and a truncated last frame.
1 parent 7044982 commit 1663e23

3 files changed

Lines changed: 146 additions & 21 deletions

File tree

‎crates/runner-shared/benches/memtrack_reader.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,9 +49,9 @@ fn bench(bencher: Bencher, n: usize, payload: Payload) {
4949
let artifact = unsafe { memmap2::Mmap::map(&file) }.unwrap();
5050

5151
let find = || {
52-
MemtrackArtifact::decode_module_events(&artifact[..])
52+
MemtrackArtifact::decode_module_events(&artifact)
5353
.unwrap()
54-
.count()
54+
.len()
5555
};
5656
assert_eq!(find(), module_events);
5757

‎crates/runner-shared/src/artifacts/memtrack/mod.rs‎

Lines changed: 136 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -36,20 +36,27 @@ impl MemtrackArtifact {
3636
})
3737
}
3838

39-
/// Stream only the events that place modules in processes: executable
40-
/// mappings, and the forks and execs that inherit or drop them. Offline
41-
/// stack attribution needs these few events out of the whole artifact.
42-
pub fn decode_module_events<R: std::io::Read>(
43-
reader: R,
44-
) -> anyhow::Result<impl Iterator<Item = MemtrackEvent>> {
45-
Ok(Self::decode_streamed(reader)?.filter(|event| {
46-
matches!(
47-
event.kind,
48-
MemtrackEventKind::Mapping { .. }
49-
| MemtrackEventKind::Fork { .. }
50-
| MemtrackEventKind::Exec
51-
)
52-
}))
39+
/// Find the events that place modules in processes: executable mappings,
40+
/// and the forks and execs that inherit or drop them. Offline stack
41+
/// attribution needs these few events out of the whole artifact.
42+
///
43+
/// Every frame the encoder writes is a self-contained zstd frame, so frames
44+
/// are decoded in parallel; events keep their artifact order. A truncated
45+
/// last frame cannot be split off and is streamed instead, so the events
46+
/// before the cut are still found.
47+
pub fn decode_module_events(artifact: &[u8]) -> anyhow::Result<Vec<MemtrackEvent>> {
48+
use rayon::prelude::*;
49+
50+
let (frames, tail) = split_zstd_frames(artifact);
51+
let mut events = frames
52+
.par_iter()
53+
.map_init(Vec::new, |msgpack, frame| {
54+
module_events_in_frame(frame, msgpack)
55+
})
56+
.collect::<anyhow::Result<Vec<_>>>()?
57+
.concat();
58+
events.extend(Self::decode_streamed(tail)?.filter(|event| event.kind.is_module_event()));
59+
Ok(events)
5360
}
5461

5562
pub fn is_empty<R: std::io::Read>(reader: R) -> bool {
@@ -60,6 +67,38 @@ impl MemtrackArtifact {
6067
}
6168
}
6269

70+
/// Split an artifact into its complete zstd frames, plus the unsplittable rest.
71+
fn split_zstd_frames(mut artifact: &[u8]) -> (Vec<&[u8]>, &[u8]) {
72+
let mut frames = Vec::new();
73+
while let Ok(len) = zstd::zstd_safe::find_frame_compressed_size(artifact) {
74+
let (frame, rest) = artifact.split_at(len);
75+
frames.push(frame);
76+
artifact = rest;
77+
}
78+
(frames, artifact)
79+
}
80+
81+
/// Decompress one frame into `msgpack` and return its module events. Decoding
82+
/// from the buffer lets strings and stack payloads be read in place instead of
83+
/// copied out of a stream first. Like [`MemtrackEventStream`], reading stops at
84+
/// the first event that fails to decode.
85+
fn module_events_in_frame(
86+
frame: &[u8],
87+
msgpack: &mut Vec<u8>,
88+
) -> anyhow::Result<Vec<MemtrackEvent>> {
89+
msgpack.clear();
90+
zstd::stream::copy_decode(frame, &mut *msgpack)?;
91+
92+
let mut deserializer = rmp_serde::Deserializer::from_read_ref(msgpack.as_slice());
93+
let mut events = Vec::new();
94+
while let Ok(event) = MemtrackEvent::deserialize(&mut deserializer) {
95+
if event.kind.is_module_event() {
96+
events.push(event);
97+
}
98+
}
99+
Ok(events)
100+
}
101+
63102
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
64103
pub struct MemtrackEvent {
65104
pub pid: pid_t,
@@ -141,6 +180,14 @@ pub enum MemtrackEventKind {
141180
},
142181
}
143182

183+
impl MemtrackEventKind {
184+
/// Whether the event places modules in processes: an executable mapping,
185+
/// or a fork or exec that inherits or drops them.
186+
pub fn is_module_event(&self) -> bool {
187+
matches!(self, Self::Mapping { .. } | Self::Fork { .. } | Self::Exec)
188+
}
189+
}
190+
144191
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
145192
pub struct StackRecord {
146193
pub hash: u64,
@@ -351,6 +398,81 @@ mod tests {
351398
Ok(())
352399
}
353400

401+
/// Every event kind, with stack payloads, across several frames and a
402+
/// last frame cut short: the parallel frame decoder must find exactly what
403+
/// a full streamed decode filtered to module events finds.
404+
#[test]
405+
fn module_events_match_a_filtered_full_decode() -> anyhow::Result<()> {
406+
let kinds = [
407+
MemtrackEventKind::Malloc {
408+
size: 64,
409+
stack_hash: 7,
410+
},
411+
MemtrackEventKind::Free,
412+
MemtrackEventKind::Realloc {
413+
old_addr: Some(0x10),
414+
size: 128,
415+
stack_hash: 0,
416+
},
417+
MemtrackEventKind::Stack {
418+
record: Box::new(StackRecord {
419+
hash: 7,
420+
sp: 0x7fff_0000,
421+
regs: vec![1; 33],
422+
bytes: vec![0xab; 4096],
423+
fp_chain: vec![0x5555_0000, 0x5555_0010],
424+
truncated: true,
425+
}),
426+
},
427+
MemtrackEventKind::Fork { parent_pid: 1 },
428+
MemtrackEventKind::Exec,
429+
MemtrackEventKind::Rss {
430+
member: 1,
431+
size: 4096,
432+
},
433+
MemtrackEventKind::Mapping {
434+
path: "/usr/lib/libexample.so".into(),
435+
dev: 0x0800_0001,
436+
ino: 42,
437+
file_offset: 0x1000,
438+
len: 0x2000,
439+
},
440+
MemtrackEventKind::Exit,
441+
MemtrackEventKind::Rmap {
442+
member: 1,
443+
delta: -70_000,
444+
},
445+
];
446+
let events: Vec<_> = (0..3000u64)
447+
.map(|i| MemtrackEvent {
448+
pid: 2,
449+
tid: 3,
450+
timestamp: i,
451+
addr: i * 16,
452+
kind: kinds[i as usize % kinds.len()].clone(),
453+
})
454+
.collect();
455+
456+
let mut artifact = Vec::new();
457+
for batch in events.chunks(1000) {
458+
let mut writer = MemtrackWriter::new(Vec::<u8>::new())?;
459+
for event in batch {
460+
writer.write_event(event)?;
461+
}
462+
artifact.extend_from_slice(&writer.finish()?);
463+
}
464+
let truncated = &artifact[..artifact.len() - 64];
465+
466+
for artifact in [&artifact[..], truncated] {
467+
let expected: Vec<_> = MemtrackArtifact::decode_streamed(artifact)?
468+
.filter(|event| event.kind.is_module_event())
469+
.collect();
470+
assert!(!expected.is_empty());
471+
assert_eq!(MemtrackArtifact::decode_module_events(artifact)?, expected);
472+
}
473+
Ok(())
474+
}
475+
354476
#[test]
355477
fn test_artifact_is_empty() -> anyhow::Result<()> {
356478
let artifact = MemtrackArtifact { events: vec![] };

‎src/executor/memory/module_artifacts.rs‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -90,8 +90,11 @@ fn read_mappings(results_folder: &Path) -> Result<Vec<ProcessMapping>> {
9090
}
9191

9292
let file = std::fs::File::open(entry.path())?;
93+
// SAFETY: read-only mapping of an artifact memtrack finished writing
94+
// before it exited; nothing writes to it while the runner reads it.
95+
let artifact = unsafe { memmap2::Mmap::map(&file)? };
9396
mappings.extend(
94-
read_mappings_from_artifact(file)
97+
read_mappings_from_artifact(&artifact)
9598
.with_context(|| format!("Failed to decode {:?}", entry.path()))?,
9699
);
97100
}
@@ -102,8 +105,8 @@ fn read_mappings(results_folder: &Path) -> Result<Vec<ProcessMapping>> {
102105

103106
/// Reconstruct mappings across forks because inherited perf events do not
104107
/// synthesize mappings that already existed when a child was forked.
105-
fn read_mappings_from_artifact<R: std::io::Read>(reader: R) -> Result<Vec<ProcessMapping>> {
106-
let mut timeline = MemtrackArtifact::decode_module_events(reader)?.collect::<Vec<_>>();
108+
fn read_mappings_from_artifact(artifact: &[u8]) -> Result<Vec<ProcessMapping>> {
109+
let mut timeline = MemtrackArtifact::decode_module_events(artifact)?;
107110

108111
// Ties break so exec purges before mapping, while fork inherits that mapping.
109112
timeline.sort_by_key(|event| {
@@ -611,7 +614,7 @@ mod tests {
611614
artifact.encode_to_writer(&mut encoded).unwrap();
612615

613616
assert_eq!(
614-
read_mappings_from_artifact(std::io::Cursor::new(encoded)).unwrap(),
617+
read_mappings_from_artifact(&encoded).unwrap(),
615618
vec![
616619
ProcessMapping {
617620
pid: 7,
@@ -858,7 +861,7 @@ mod tests {
858861
let artifact = MemtrackArtifact { events };
859862
let mut encoded = Vec::new();
860863
artifact.encode_to_writer(&mut encoded).unwrap();
861-
read_mappings_from_artifact(std::io::Cursor::new(encoded)).unwrap()
864+
read_mappings_from_artifact(&encoded).unwrap()
862865
}
863866

864867
fn mapping_summary(mappings: &[ProcessMapping]) -> Vec<(pid_t, &str, u64)> {

0 commit comments

Comments
 (0)