Skip to main content

memtrace_utils/
parser.rs

1use indexmap::map::Entry;
2use indexmap::IndexMap;
3use std::fs::OpenOptions;
4use std::io;
5use std::io::BufRead;
6use std::path::Path;
7use std::time::Duration;
8use thiserror::Error;
9
10#[derive(Error, Debug)]
11pub enum Error {
12    #[error(transparent)]
13    Io(#[from] io::Error),
14    #[error("Invalid format: line: {0}; msg: {1}")]
15    InvalidFormat(usize, &'static str),
16    #[error("Internal {0}")]
17    Internal(String),
18}
19
20#[derive(Debug)]
21pub struct Trace {
22    pub ip_idx: u64,
23    pub parent_idx: u64,
24}
25
26#[derive(Debug)]
27pub struct InstructionPointer {
28    pub ip: u64,
29    pub module_idx: usize,
30    pub frame: Option<Frame>,
31    pub inlined: Vec<Frame>,
32}
33
34#[derive(Debug)]
35pub enum Frame {
36    Single {
37        function_idx: usize,
38    },
39    Multiple {
40        function_idx: usize,
41        file_idx: usize,
42        line_number: u32,
43    },
44}
45
46#[derive(Debug, Default)]
47pub struct AllocationData {
48    pub allocations: u64,
49    pub temporary: u64,
50    pub leaked: u64,
51    pub peak: u64,
52}
53
54#[derive(Debug)]
55pub struct AllocationInfo {
56    pub allocation_idx: u64,
57    pub size: u64,
58}
59
60impl AllocationInfo {
61    pub fn new(allocation_idx: u64, size: u64) -> Self {
62        Self {
63            allocation_idx,
64            size,
65        }
66    }
67}
68
69#[derive(Debug)]
70pub struct Allocation {
71    pub trace_idx: u64,
72    pub data: AllocationData,
73}
74
75impl Allocation {
76    pub fn new(trace_idx: u64) -> Self {
77        Self {
78            trace_idx,
79            data: Default::default(),
80        }
81    }
82}
83
84#[derive(Debug)]
85pub struct AccumulatedData {
86    pub strings: Vec<String>,
87    pub traces: Vec<Trace>,
88    pub instruction_pointers: Vec<InstructionPointer>,
89    pub allocation_indices: IndexMap<u64, u64>,
90    pub allocation_infos: Vec<AllocationInfo>,
91    pub allocations: Vec<Allocation>,
92    pub total: AllocationData,
93    pub duration: Duration,
94    pub peak_rss: u64,
95    pub page_size: u64,
96    pub pages: u64,
97}
98
99impl AccumulatedData {
100    pub fn new() -> Self {
101        Self {
102            strings: Vec::with_capacity(4096),
103            traces: Vec::with_capacity(65536),
104            instruction_pointers: Vec::with_capacity(16384),
105            allocation_indices: IndexMap::with_capacity(16384),
106            allocations: Vec::with_capacity(16384),
107            allocation_infos: Vec::with_capacity(16384),
108            total: AllocationData::default(),
109            duration: Duration::default(),
110            peak_rss: 0,
111            page_size: 0,
112            pages: 0,
113        }
114    }
115}
116
117pub struct Parser {
118    data: AccumulatedData,
119    last_ptr: u64,
120}
121
122impl Parser {
123    pub fn new() -> Self {
124        Self {
125            data: AccumulatedData::new(),
126            last_ptr: 0,
127        }
128    }
129
130    pub fn parse_file(mut self, file_path: impl AsRef<Path>) -> Result<AccumulatedData, Error> {
131        let file = OpenOptions::new().read(true).open(file_path)?;
132        let reader = io::BufReader::new(file);
133
134        for (n, line) in reader.lines().enumerate() {
135            self.parse_line(&line?).map_err(|e| Error::InvalidFormat(n+1, e))?;
136        }
137
138        Ok(self.data)
139    }
140
141    fn parse_line(&mut self, line: &str) -> Result<(), &'static str> {
142        let mut split = line.split_whitespace();
143
144        let Some(first) = split.next() else {
145            return Ok(());
146        };
147
148        match first {
149            "s" => {
150                let str_len = usize::from_str_radix(split.next().ok_or("failed to find str_len")?, 16)
151                    .map_err(|_| "failed to parse str_len")?;
152                self.data
153                    .strings
154                    .push(line[line.len() - str_len..].to_string());
155            }
156            "t" => {
157                let ip_idx = u64::from_str_radix(split.next().ok_or("failed to find ip_idx")?, 16)
158                    .map_err(|_| "failed to parse ip_idx")?;
159                let parent_idx = u64::from_str_radix(split.next().ok_or("failed to find parent_idx")?, 16)
160                    .map_err(|_| "failed to parse parent_idx")?;
161
162                self.data.traces.push(Trace { ip_idx, parent_idx })
163            }
164            "i" => {
165                let ip = u64::from_str_radix(split.next().ok_or("failed to find ip")?, 16)
166                    .map_err(|_| "failed to parse ip")?;
167                let module_idx =
168                    usize::from_str_radix(split.next().ok_or("failed to find module_idx")?, 16)
169                        .map_err(|_| "failed to parse module_idx")?;
170
171                let frame = Self::parse_frame(&mut split)?;
172
173                let mut inlined = Vec::new();
174                while let Some(frame) = Self::parse_frame(&mut split)? {
175                    inlined.push(frame);
176                }
177
178                self.data.instruction_pointers.push(InstructionPointer {
179                    ip,
180                    module_idx,
181                    frame,
182                    inlined,
183                })
184            }
185            "a" => {
186                let size = u64::from_str_radix(split.next().ok_or("failed to find size")?, 16)
187                    .map_err(|_| "failed to parse size")?;
188                let trace_idx = u64::from_str_radix(split.next().ok_or("failed to find trace_idx")?, 16)
189                    .map_err(|_| "failed to parse trace_idx")?;
190
191                let allocation_idx = self.add_allocation(trace_idx);
192                self.data
193                    .allocation_infos
194                    .push(AllocationInfo::new(allocation_idx, size));
195            }
196            "+" => {
197                let allocation_info_idx =
198                    u64::from_str_radix(split.next().ok_or("failed to find allocation_info_idx")?, 16)
199                        .map_err(|_| "failed to parse allocation_info_idx")?;
200
201                let info = &mut self.data.allocation_infos[allocation_info_idx as usize];
202
203                let allocation = self
204                    .data
205                    .allocations
206                    .get_mut(info.allocation_idx as usize)
207                    .ok_or_else(|| "allocation not found")?;
208
209                self.last_ptr = info.allocation_idx;
210
211                allocation.data.leaked += info.size;
212                if allocation.data.leaked > allocation.data.peak {
213                    allocation.data.peak = allocation.data.leaked;
214                }
215                allocation.data.allocations += 1;
216
217                self.data.total.leaked += info.size;
218                self.data.total.allocations += 1;
219
220                if self.data.total.leaked > self.data.total.peak {
221                    self.data.total.peak = self.data.total.leaked;
222                }
223            }
224            "-" => {
225                let allocation_info_idx =
226                    u64::from_str_radix(split.next().ok_or("failed to find allocation_info_idx")?, 16)
227                        .map_err(|_| "failed to parse allocation_info_idx")?;
228
229                let info = &mut self.data.allocation_infos[allocation_info_idx as usize];
230
231                let allocation = self
232                    .data
233                    .allocations
234                    .get_mut(info.allocation_idx as usize)
235                    .ok_or_else(|| "allocation not found")?;
236
237                self.data.total.leaked -= info.size;
238
239                let temporary = self.last_ptr == info.allocation_idx;
240                self.last_ptr = 0;
241
242                if temporary {
243                    self.data.total.temporary += 1;
244                }
245
246                allocation.data.leaked -= info.size;
247                if temporary {
248                    allocation.data.temporary += 1;
249                }
250            }
251            "c" => {
252                let timestamp = u64::from_str_radix(split.next().ok_or("failed to find timestamp")?, 16)
253                    .map_err(|_| "failed to parse timestamp")?;
254                self.data.duration = Duration::from_millis(timestamp);
255            }
256            "R" => {
257                let rss = u64::from_str_radix(split.next().ok_or("failed to find rss")?, 16)
258                    .map_err(|_| "failed to parse rss")?;
259                if rss > self.data.peak_rss {
260                    self.data.peak_rss = rss;
261                }
262            }
263            "I" => {
264                self.data.page_size =
265                    u64::from_str_radix(split.next().ok_or("failed to find page_size")?, 16)
266                        .map_err(|_| "failed to parse page_size")?;
267                self.data.pages =
268                    u64::from_str_radix(split.next().ok_or("failed to find pages")?, 16)
269                        .map_err(|_| "failed to parse pages")?;
270            }
271            "#" => {
272                // comment
273            }
274            _ => {}
275        }
276        Ok(())
277    }
278
279    fn add_allocation(&mut self, trace_idx: u64) -> u64 {
280        match self.data.allocation_indices.entry(trace_idx) {
281            Entry::Occupied(e) => *e.get(),
282            Entry::Vacant(e) => {
283                let idx = self.data.allocations.len() as u64;
284                e.insert(idx);
285                let allocation = Allocation::new(trace_idx);
286                self.data.allocations.push(allocation);
287                idx
288            }
289        }
290    }
291
292    fn parse_frame<'a>(mut iter: impl Iterator<Item = &'a str>) -> Result<Option<Frame>, &'static str> {
293        let Some(first) = iter.next() else {
294            return Ok(None);
295        };
296
297        let function_idx = usize::from_str_radix(first, 16).map_err(|_| "failed to parse function_idx")?;
298
299        let Some(file_val) = iter.next() else {
300            return Ok(Some(Frame::Single { function_idx }));
301        };
302
303        let file_idx = usize::from_str_radix(file_val, 16).map_err(|_| "failed to parse file_idx")?;
304        let line_number = u32::from_str_radix(iter.next().ok_or("failed to find line number")?, 16)
305            .map_err(|_| "failed to parse line_number")?;
306
307        Ok(Some(Frame::Multiple {
308            function_idx,
309            file_idx,
310            line_number,
311        }))
312    }
313}
314
315#[cfg(test)]
316mod tests {
317    use crate::parser::Parser;
318
319    #[test]
320    fn test_read_trace_file() {
321        let file = "/tmp/pipe.out";
322        let data = Parser::new().parse_file(file).unwrap();
323
324        println!("{:#?}", data);
325    }
326}