@@ -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 ) ]
64103pub 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 ) ]
145192pub 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 ! [ ] } ;
0 commit comments