turbo_trace_server/reader/
heaptrack.rs1use std::{env, str::from_utf8, sync::Arc};
2
3use anyhow::{Context, Result, bail};
4use indexmap::map::Entry;
5use rustc_demangle::demangle;
6use rustc_hash::{FxHashMap, FxHashSet};
7
8use super::TraceFormat;
9use crate::{FxIndexMap, span::SpanIndex, store_container::StoreContainer, timestamp::Timestamp};
10
11#[derive(Debug, Clone, Copy)]
12struct TraceNode {
13 ip_index: usize,
14 parent_index: usize,
15}
16
17impl TraceNode {
18 pub fn read(s: &mut &[u8]) -> Result<Self> {
19 Ok(Self {
20 ip_index: read_hex_index(s)?,
21 parent_index: read_hex_index(s)?,
22 })
23 }
24}
25
26#[derive(Debug, Hash, PartialEq, Eq)]
27struct InstructionPointer {
28 module_index: usize,
29 frames: Vec<Frame>,
30 custom_name: Option<String>,
31}
32
33impl InstructionPointer {
34 pub fn read(s: &mut &[u8]) -> Result<Self> {
35 let _ip = read_hex(s)?;
36 Ok(Self {
37 module_index: read_hex_index(s)?,
38 frames: read_all(s, Frame::read)?,
39 custom_name: None,
40 })
41 }
42}
43
44#[derive(Debug, Hash, PartialEq, Eq)]
45struct Frame {
46 function_index: usize,
47 file_index: usize,
48 line: u64,
49}
50
51impl Frame {
52 pub fn read(s: &mut &[u8]) -> Result<Self> {
53 Ok(Self {
54 function_index: read_hex_index(s)?,
55 file_index: read_hex_index(s)?,
56 line: read_hex(s)?,
57 })
58 }
59}
60
61#[derive(Debug)]
62struct AllocationInfo {
63 size: u64,
64 trace_index: usize,
65}
66
67impl AllocationInfo {
68 pub fn read(s: &mut &[u8]) -> Result<Self> {
69 Ok(Self {
70 size: read_hex(s)?,
71 trace_index: read_hex_index(s)?,
72 })
73 }
74}
75
76struct InstructionPointerExtraInfo {
77 first_trace_of_ip: Option<usize>,
78}
79
80#[derive(Clone, Copy)]
81struct TraceData {
82 span_index: SpanIndex,
83 ip_index: usize,
84 parent_trace_index: usize,
85}
86
87pub struct HeaptrackFormat {
88 store: Arc<StoreContainer>,
89 version: u32,
90 last_timestamp: Timestamp,
91 strings: Vec<String>,
92 traces: Vec<TraceData>,
93 ip_parent_map: FxHashMap<(usize, SpanIndex), usize>,
94 trace_instruction_pointers: Vec<usize>,
95 instruction_pointers: FxIndexMap<InstructionPointer, InstructionPointerExtraInfo>,
96 allocations: Vec<AllocationInfo>,
97 spans: usize,
98 collapse_crates: FxHashSet<String>,
99 expand_crates: FxHashSet<String>,
100 expand_recursion: bool,
101 allocated_memory: u64,
102 temp_allocated_memory: u64,
103}
104
105const RECURSION_IP: usize = 1;
106
107impl HeaptrackFormat {
108 pub fn new(store: Arc<StoreContainer>) -> Self {
109 Self {
110 store,
111 version: 0,
112 last_timestamp: Timestamp::ZERO,
113 strings: vec!["".to_string()],
114 traces: vec![TraceData {
115 span_index: SpanIndex::new(usize::MAX).unwrap(),
116 ip_index: 0,
117 parent_trace_index: 0,
118 }],
119 ip_parent_map: FxHashMap::default(),
120 instruction_pointers: {
121 let mut map = FxIndexMap::with_capacity_and_hasher(2, Default::default());
122 map.insert(
123 InstructionPointer {
124 module_index: 0,
125 frames: Vec::new(),
126 custom_name: Some("root".to_string()),
127 },
128 InstructionPointerExtraInfo {
129 first_trace_of_ip: None,
130 },
131 );
132 map.insert(
133 InstructionPointer {
134 module_index: 0,
135 frames: Vec::new(),
136 custom_name: Some("recursion".to_string()),
137 },
138 InstructionPointerExtraInfo {
139 first_trace_of_ip: None,
140 },
141 );
142 map
143 },
144 trace_instruction_pointers: vec![0],
145 allocations: vec![],
146 spans: 0,
147 collapse_crates: env::var("COLLAPSE_CRATES")
148 .unwrap_or_default()
149 .split(',')
150 .map(|s| s.to_string())
151 .collect(),
152 expand_crates: env::var("EXPAND_CRATES")
153 .unwrap_or_default()
154 .split(',')
155 .filter(|s| !s.is_empty())
156 .map(|s| s.to_string())
157 .collect(),
158 expand_recursion: env::var("EXPAND_RECURSION").is_ok(),
159 allocated_memory: 0,
160 temp_allocated_memory: 0,
161 }
162 }
163}
164
165impl TraceFormat for HeaptrackFormat {
166 fn stats(&self) -> String {
167 format!(
168 "{} spans, {} strings, {} ips, {} traces, {} allocations, {:.2} GB allocated, {:.2} \
169 GB temporarily allocated",
170 self.spans,
171 self.strings.len() - 1,
172 self.trace_instruction_pointers.len() - 1,
173 self.traces.len() - 1,
174 self.allocations.len() - 1,
175 (self.allocated_memory / 1024 / 1024) as f32 / 1024.0,
176 (self.temp_allocated_memory / 1024 / 1024) as f32 / 1024.0,
177 )
178 }
179
180 type Reused = ();
181
182 fn read(&mut self, mut buffer: &[u8], _reuse: &mut Self::Reused) -> anyhow::Result<usize> {
183 let mut bytes_read = 0;
184 let mut outdated_spans = FxHashSet::default();
185 let mut store = self.store.write();
186 'outer: while let Some(line_end) = buffer.iter().position(|b| *b == b'\n') {
187 let full_line = &buffer[..line_end];
188 buffer = &buffer[line_end + 1..];
189 bytes_read += full_line.len() + 1;
190
191 if full_line.is_empty() {
192 continue;
193 }
194 let ty = full_line[0];
195 let mut line = &full_line[2..];
196
197 match ty {
199 b'v' => {
200 let _ = read_hex(&mut line)?;
201 self.version = read_hex(&mut line)? as u32;
202 if self.version != 2 && self.version != 3 {
203 bail!("Unsupported version: {} (expected 2 or 3)", self.version);
204 }
205 }
206 b's' => {
207 let string = if self.version == 2 {
208 String::from_utf8(line.to_vec())?
209 } else {
210 read_sized_string(&mut line)?
211 };
212 self.strings.push(demangle(&string).to_string());
213 }
214 b't' => {
215 let TraceNode {
216 ip_index,
217 parent_index,
218 } = TraceNode::read(&mut line)?;
219 let ip_index = *self
220 .trace_instruction_pointers
221 .get(ip_index)
222 .context("ip not found")?;
223 let (ip, ip_info) = self
224 .instruction_pointers
225 .get_index(ip_index)
226 .context("ip not found")?;
227 if parent_index == 0
229 && let Some(trace_index) = ip_info.first_trace_of_ip
230 {
231 let trace = self.traces.get(trace_index).context("trace not found")?;
232 self.traces.push(*trace);
233 continue;
234 }
235 let parent = if parent_index > 0 {
237 let parent = *self.traces.get(parent_index).context("parent not found")?;
238 if let Some(trace_index) =
240 self.ip_parent_map.get(&(ip_index, parent.span_index))
241 {
242 let trace = self.traces.get(*trace_index).context("trace not found")?;
243 self.traces.push(*trace);
244 continue;
245 }
246 if parent.ip_index == ip_index {
248 self.traces.push(parent);
249 continue;
250 }
251 if !self.expand_recursion {
252 let mut current = parent.parent_trace_index;
254 while current > 0 {
255 let current_parent =
256 self.traces.get(current).context("parent not found")?;
257 current = current_parent.parent_trace_index;
258 if current_parent.ip_index == ip_index {
259 if parent.ip_index == RECURSION_IP {
260 self.traces.push(parent);
262 } else if let Some(trace_index) =
263 self.ip_parent_map.get(&(RECURSION_IP, parent.span_index))
264 {
265 let trace = self
267 .traces
268 .get(*trace_index)
269 .context("trace not found")?;
270 self.traces.push(*trace);
271 } else {
272 let span_index = store.add_span(
274 Some(parent.span_index),
275 self.last_timestamp,
276 "".to_string(),
277 "recursion".to_string(),
278 Vec::new(),
279 &mut outdated_spans,
280 );
281 store.complete_span(span_index);
282 let index = self.traces.len();
283 self.traces.push(TraceData {
284 ip_index: RECURSION_IP,
285 parent_trace_index: parent_index,
286 span_index,
287 });
288 self.ip_parent_map
289 .insert((RECURSION_IP, parent.span_index), index);
290 }
291 continue 'outer;
292 }
293 }
294 }
295 Some(parent.span_index)
296 } else {
297 None
298 };
299 let InstructionPointer {
300 module_index,
301 frames,
302 custom_name,
303 } = ip;
304 let module = self
305 .strings
306 .get(*module_index)
307 .context("module not found")?;
308 let name = if let Some(name) = custom_name.as_ref() {
309 name.to_string()
310 } else if let Some(first_frame) = frames.first() {
311 let file = self
312 .strings
313 .get(first_frame.file_index)
314 .context("file not found")?;
315 let function = self
316 .strings
317 .get(first_frame.function_index)
318 .context("function not found")?;
319 format!("{} @ {file}:{}", function, first_frame.line)
320 } else {
321 "unknown".to_string()
322 };
323 let mut args = Vec::new();
324 for Frame {
325 function_index,
326 file_index,
327 line,
328 } in frames.iter()
329 {
330 let file = self.strings.get(*file_index).context("file not found")?;
331 let function = self
332 .strings
333 .get(*function_index)
334 .context("function not found")?;
335 args.push((
336 "location".to_string(),
337 format!("{function} @ {file}:{line}"),
338 ));
339 }
340
341 let span_index = store.add_span(
342 parent,
343 self.last_timestamp,
344 module.to_string(),
345 name,
346 args,
347 &mut outdated_spans,
348 );
349 store.complete_span(span_index);
350 self.spans += 1;
351 let index = self.traces.len();
352 self.traces.push(TraceData {
353 span_index,
354 ip_index,
355 parent_trace_index: parent_index,
356 });
357 self.instruction_pointers
358 .get_index_mut(ip_index)
359 .unwrap()
360 .1
361 .first_trace_of_ip
362 .get_or_insert(index);
363 if let Some(parent) = parent {
364 self.ip_parent_map.insert((ip_index, parent), index);
365 }
366 }
367 b'i' => {
368 let mut ip = InstructionPointer::read(&mut line)?;
369 if let Some(frame) = ip.frames.first()
370 && let Some(function) = self.strings.get(frame.function_index)
371 {
372 let crate_name = function
373 .strip_prefix('<')
374 .unwrap_or(function)
375 .split("::")
376 .next()
377 .unwrap()
378 .split('[')
379 .next()
380 .unwrap();
381 if self.collapse_crates.contains(crate_name)
382 || !self.expand_crates.is_empty()
383 && !self.expand_crates.contains(crate_name)
384 {
385 ip.frames.clear();
386 ip.custom_name = Some(crate_name.to_string());
387 }
388 }
389 match self.instruction_pointers.entry(ip) {
390 Entry::Occupied(e) => {
391 self.trace_instruction_pointers.push(e.index());
392 }
393 Entry::Vacant(e) => {
394 self.trace_instruction_pointers.push(e.index());
395 e.insert(InstructionPointerExtraInfo {
396 first_trace_of_ip: None,
397 });
398 }
399 }
400 }
401 b'#' => {
402 }
404 b'X' => {
405 let line = from_utf8(line)?;
406 println!("Debuggee: {line}");
407 }
408 b'c' => {
409 let timestamp = read_hex(&mut line)?;
411 self.last_timestamp = Timestamp::from_micros(timestamp);
412 }
413 b'a' => {
414 let info = AllocationInfo::read(&mut line)?;
416 self.allocations.push(info);
417 }
418 b'+' => {
419 let index = read_hex_index(&mut line)?;
421 let AllocationInfo { size, trace_index } = self
422 .allocations
423 .get(index)
424 .context("allocation not found")?;
425 if *trace_index > 0 {
426 let TraceData { span_index, .. } =
427 self.traces.get(*trace_index).context("trace not found")?;
428 store.add_allocation(*span_index, *size, 1, &mut outdated_spans);
429 self.allocated_memory += *size;
430 }
431 }
432 b'-' => {
433 let index = read_hex_index(&mut line)?;
435 let AllocationInfo { size, trace_index } = self
436 .allocations
437 .get(index)
438 .context("allocation not found")?;
439 if *trace_index > 0 {
440 let TraceData { span_index, .. } =
441 self.traces.get(*trace_index).context("trace not found")?;
442 store.add_deallocation(*span_index, *size, 1, &mut outdated_spans);
443 self.allocated_memory -= *size;
444 self.temp_allocated_memory += *size;
445 }
446 }
447 b'R' => {
448 }
450 b'A' => {
451 }
454 b'S' => {
455 }
458 b'I' => {
459 }
462 _ => {
463 let line = from_utf8(line)?;
464 println!("{} {line}", ty as char)
465 }
466 }
467 }
468 store.invalidate_outdated_spans(&outdated_spans);
469 Ok(bytes_read)
470 }
471}
472
473fn read_hex_index(s: &mut &[u8]) -> anyhow::Result<usize> {
474 Ok(read_hex(s)? as usize)
475}
476
477fn read_hex(s: &mut &[u8]) -> anyhow::Result<u64> {
478 let mut n: u64 = 0;
479 loop {
480 if let Some(c) = s.first() {
481 match c {
482 b'0'..=b'9' => {
483 n *= 16;
484 n += (*c - b'0') as u64;
485 }
486 b'a'..=b'f' => {
487 n *= 16;
488 n += (*c - b'a' + 10) as u64;
489 }
490 b' ' => {
491 *s = &s[1..];
492 return Ok(n);
493 }
494 _ => {
495 bail!("Expected hex char");
496 }
497 }
498 *s = &s[1..];
499 } else {
500 return Ok(n);
501 }
502 }
503}
504
505fn read_sized_string(s: &mut &[u8]) -> anyhow::Result<String> {
506 let size = read_hex(s)? as usize;
507 let str = &s[..size];
508 *s = &s[size..];
509 Ok(String::from_utf8(str.to_vec())?)
510}
511
512fn read_all<T>(
513 s: &mut &[u8],
514 f: impl Fn(&mut &[u8]) -> anyhow::Result<T>,
515) -> anyhow::Result<Vec<T>> {
516 let mut res = Vec::new();
517 while !s.is_empty() {
518 res.push(f(s)?);
519 }
520 Ok(res)
521}