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