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: loop {
187 let Some(line_end) = buffer.iter().position(|b| *b == b'\n') else {
188 break;
189 };
190 let full_line = &buffer[..line_end];
191 buffer = &buffer[line_end + 1..];
192 bytes_read += full_line.len() + 1;
193
194 if full_line.is_empty() {
195 continue;
196 }
197 let ty = full_line[0];
198 let mut line = &full_line[2..];
199
200 match ty {
202 b'v' => {
203 let _ = read_hex(&mut line)?;
204 self.version = read_hex(&mut line)? as u32;
205 if self.version != 2 && self.version != 3 {
206 bail!("Unsupported version: {} (expected 2 or 3)", self.version);
207 }
208 }
209 b's' => {
210 let string = if self.version == 2 {
211 String::from_utf8(line.to_vec())?
212 } else {
213 read_sized_string(&mut line)?
214 };
215 self.strings.push(demangle(&string).to_string());
216 }
217 b't' => {
218 let TraceNode {
219 ip_index,
220 parent_index,
221 } = TraceNode::read(&mut line)?;
222 let ip_index = *self
223 .trace_instruction_pointers
224 .get(ip_index)
225 .context("ip not found")?;
226 let (ip, ip_info) = self
227 .instruction_pointers
228 .get_index(ip_index)
229 .context("ip not found")?;
230 if parent_index == 0
232 && let Some(trace_index) = ip_info.first_trace_of_ip
233 {
234 let trace = self.traces.get(trace_index).context("trace not found")?;
235 self.traces.push(*trace);
236 continue;
237 }
238 let parent = if parent_index > 0 {
240 let parent = *self.traces.get(parent_index).context("parent not found")?;
241 if let Some(trace_index) =
243 self.ip_parent_map.get(&(ip_index, parent.span_index))
244 {
245 let trace = self.traces.get(*trace_index).context("trace not found")?;
246 self.traces.push(*trace);
247 continue;
248 }
249 if parent.ip_index == ip_index {
251 self.traces.push(parent);
252 continue;
253 }
254 if !self.expand_recursion {
255 let mut current = parent.parent_trace_index;
257 while current > 0 {
258 let current_parent =
259 self.traces.get(current).context("parent not found")?;
260 current = current_parent.parent_trace_index;
261 if current_parent.ip_index == ip_index {
262 if parent.ip_index == RECURSION_IP {
263 self.traces.push(parent);
265 } else if let Some(trace_index) =
266 self.ip_parent_map.get(&(RECURSION_IP, parent.span_index))
267 {
268 let trace = self
270 .traces
271 .get(*trace_index)
272 .context("trace not found")?;
273 self.traces.push(*trace);
274 } else {
275 let span_index = store.add_span(
277 Some(parent.span_index),
278 self.last_timestamp,
279 "".to_string(),
280 "recursion".to_string(),
281 Vec::new(),
282 &mut outdated_spans,
283 );
284 store.complete_span(span_index);
285 let index = self.traces.len();
286 self.traces.push(TraceData {
287 ip_index: RECURSION_IP,
288 parent_trace_index: parent_index,
289 span_index,
290 });
291 self.ip_parent_map
292 .insert((RECURSION_IP, parent.span_index), index);
293 }
294 continue 'outer;
295 }
296 }
297 }
298 Some(parent.span_index)
299 } else {
300 None
301 };
302 let InstructionPointer {
303 module_index,
304 frames,
305 custom_name,
306 } = ip;
307 let module = self
308 .strings
309 .get(*module_index)
310 .context("module not found")?;
311 let name = if let Some(name) = custom_name.as_ref() {
312 name.to_string()
313 } else if let Some(first_frame) = frames.first() {
314 let file = self
315 .strings
316 .get(first_frame.file_index)
317 .context("file not found")?;
318 let function = self
319 .strings
320 .get(first_frame.function_index)
321 .context("function not found")?;
322 format!("{} @ {file}:{}", function, first_frame.line)
323 } else {
324 "unknown".to_string()
325 };
326 let mut args = Vec::new();
327 for Frame {
328 function_index,
329 file_index,
330 line,
331 } in frames.iter()
332 {
333 let file = self.strings.get(*file_index).context("file not found")?;
334 let function = self
335 .strings
336 .get(*function_index)
337 .context("function not found")?;
338 args.push((
339 "location".to_string(),
340 format!("{function} @ {file}:{line}"),
341 ));
342 }
343
344 let span_index = store.add_span(
345 parent,
346 self.last_timestamp,
347 module.to_string(),
348 name,
349 args,
350 &mut outdated_spans,
351 );
352 store.complete_span(span_index);
353 self.spans += 1;
354 let index = self.traces.len();
355 self.traces.push(TraceData {
356 span_index,
357 ip_index,
358 parent_trace_index: parent_index,
359 });
360 self.instruction_pointers
361 .get_index_mut(ip_index)
362 .unwrap()
363 .1
364 .first_trace_of_ip
365 .get_or_insert(index);
366 if let Some(parent) = parent {
367 self.ip_parent_map.insert((ip_index, parent), index);
368 }
369 }
370 b'i' => {
371 let mut ip = InstructionPointer::read(&mut line)?;
372 if let Some(frame) = ip.frames.first()
373 && let Some(function) = self.strings.get(frame.function_index)
374 {
375 let crate_name = function
376 .strip_prefix('<')
377 .unwrap_or(function)
378 .split("::")
379 .next()
380 .unwrap()
381 .split('[')
382 .next()
383 .unwrap();
384 if self.collapse_crates.contains(crate_name)
385 || !self.expand_crates.is_empty()
386 && !self.expand_crates.contains(crate_name)
387 {
388 ip.frames.clear();
389 ip.custom_name = Some(crate_name.to_string());
390 }
391 }
392 match self.instruction_pointers.entry(ip) {
393 Entry::Occupied(e) => {
394 self.trace_instruction_pointers.push(e.index());
395 }
396 Entry::Vacant(e) => {
397 self.trace_instruction_pointers.push(e.index());
398 e.insert(InstructionPointerExtraInfo {
399 first_trace_of_ip: None,
400 });
401 }
402 }
403 }
404 b'#' => {
405 }
407 b'X' => {
408 let line = from_utf8(line)?;
409 println!("Debuggee: {line}");
410 }
411 b'c' => {
412 let timestamp = read_hex(&mut line)?;
414 self.last_timestamp = Timestamp::from_micros(timestamp);
415 }
416 b'a' => {
417 let info = AllocationInfo::read(&mut line)?;
419 self.allocations.push(info);
420 }
421 b'+' => {
422 let index = read_hex_index(&mut line)?;
424 let AllocationInfo { size, trace_index } = self
425 .allocations
426 .get(index)
427 .context("allocation not found")?;
428 if *trace_index > 0 {
429 let TraceData { span_index, .. } =
430 self.traces.get(*trace_index).context("trace not found")?;
431 store.add_allocation(*span_index, *size, 1, &mut outdated_spans);
432 self.allocated_memory += *size;
433 }
434 }
435 b'-' => {
436 let index = read_hex_index(&mut line)?;
438 let AllocationInfo { size, trace_index } = self
439 .allocations
440 .get(index)
441 .context("allocation not found")?;
442 if *trace_index > 0 {
443 let TraceData { span_index, .. } =
444 self.traces.get(*trace_index).context("trace not found")?;
445 store.add_deallocation(*span_index, *size, 1, &mut outdated_spans);
446 self.allocated_memory -= *size;
447 self.temp_allocated_memory += *size;
448 }
449 }
450 b'R' => {
451 }
453 b'A' => {
454 }
457 b'S' => {
458 }
461 b'I' => {
462 }
465 _ => {
466 let line = from_utf8(line)?;
467 println!("{} {line}", ty as char)
468 }
469 }
470 }
471 store.invalidate_outdated_spans(&outdated_spans);
472 Ok(bytes_read)
473 }
474}
475
476fn read_hex_index(s: &mut &[u8]) -> anyhow::Result<usize> {
477 Ok(read_hex(s)? as usize)
478}
479
480fn read_hex(s: &mut &[u8]) -> anyhow::Result<u64> {
481 let mut n: u64 = 0;
482 loop {
483 if let Some(c) = s.first() {
484 match c {
485 b'0'..=b'9' => {
486 n *= 16;
487 n += (*c - b'0') as u64;
488 }
489 b'a'..=b'f' => {
490 n *= 16;
491 n += (*c - b'a' + 10) as u64;
492 }
493 b' ' => {
494 *s = &s[1..];
495 return Ok(n);
496 }
497 _ => {
498 bail!("Expected hex char");
499 }
500 }
501 *s = &s[1..];
502 } else {
503 return Ok(n);
504 }
505 }
506}
507
508fn read_sized_string(s: &mut &[u8]) -> anyhow::Result<String> {
509 let size = read_hex(s)? as usize;
510 let str = &s[..size];
511 *s = &s[size..];
512 Ok(String::from_utf8(str.to_vec())?)
513}
514
515fn read_all<T>(
516 s: &mut &[u8],
517 f: impl Fn(&mut &[u8]) -> anyhow::Result<T>,
518) -> anyhow::Result<Vec<T>> {
519 let mut res = Vec::new();
520 while !s.is_empty() {
521 res.push(f(s)?);
522 }
523 Ok(res)
524}