turbo_trace_server/
store.rs

1use std::{
2    cmp::{max, min},
3    env,
4    num::NonZeroUsize,
5    sync::{OnceLock, atomic::AtomicU64},
6};
7
8use rustc_hash::FxHashSet;
9use turbo_rcstr::{RcStr, rcstr};
10
11use crate::{
12    chunked_vec::ChunkedVec,
13    self_time_tree::SelfTimeTree,
14    span::{Span, SpanArgs, SpanEvent, SpanIndex, SpanTimeData},
15    span_ref::SpanRef,
16    timestamp::Timestamp,
17};
18
19pub type SpanId = NonZeroUsize;
20
21/// This max depth is used to avoid deep recursion in the span tree,
22/// which can lead to stack overflows and performance issues.
23/// Spans deeper than this depth will be re-parented to an ancestor
24/// at the cut-off depth (Flattening).
25const CUT_OFF_DEPTH: u32 = 80;
26
27/// A single memory usage sample: (timestamp, memory_bytes, memory_pressure).
28/// Sorted by timestamp. `memory_pressure` is an OS-reported pressure value in
29/// the range `0..=100`; `0` is used when the reporter platform did not expose
30/// a pressure signal.
31type MemorySample = (Timestamp, u64, u8);
32
33/// Maximum number of memory samples returned in a query result.
34const MAX_MEMORY_SAMPLES: usize = 200;
35
36pub struct Store {
37    pub(crate) spans: ChunkedVec<Span>,
38    pub(crate) self_time_tree: Option<SelfTimeTree<SpanIndex>>,
39    max_self_time_lookup_time: AtomicU64,
40    /// Global sorted list of memory samples (timestamp, memory_bytes).
41    memory_samples: Vec<MemorySample>,
42}
43
44fn new_root_span() -> Span {
45    Span {
46        parent: None,
47        depth: 0,
48        start: Timestamp::MAX,
49        category: RcStr::default(),
50        name: rcstr!("(root)"),
51        args: SpanArgs::new(),
52        events: Default::default(),
53        is_complete: true,
54        self_allocations: 0,
55        self_allocation_count: 0,
56        self_deallocations: 0,
57        self_deallocation_count: 0,
58        totals: OnceLock::new(),
59        time_data: SpanTimeData {
60            self_end: Timestamp::MAX,
61            ..Default::default()
62        },
63        extra: OnceLock::new(),
64        names: OnceLock::new(),
65    }
66}
67
68impl Store {
69    pub fn new() -> Self {
70        let mut spans = ChunkedVec::new();
71        spans.push(new_root_span());
72        Self {
73            spans,
74            self_time_tree: env::var("NO_CORRECTED_TIME")
75                .ok()
76                .is_none()
77                .then(SelfTimeTree::new),
78            max_self_time_lookup_time: AtomicU64::new(0),
79            memory_samples: Vec::new(),
80        }
81    }
82
83    pub fn reset(&mut self) {
84        self.spans = ChunkedVec::new();
85        self.spans.push(new_root_span());
86        if let Some(tree) = self.self_time_tree.as_mut() {
87            *tree = SelfTimeTree::new();
88        }
89        *self.max_self_time_lookup_time.get_mut() = 0;
90        self.memory_samples.clear();
91    }
92
93    pub fn optimize(&mut self) {
94        if let Some(tree) = self.self_time_tree.as_mut() {
95            tree.optimize();
96        }
97    }
98
99    pub fn has_time_info(&self) -> bool {
100        self.self_time_tree
101            .as_ref()
102            .is_none_or(|tree| tree.len() > 0)
103    }
104
105    pub fn add_span(
106        &mut self,
107        parent: Option<SpanIndex>,
108        start: Timestamp,
109        category: RcStr,
110        name: RcStr,
111        args: SpanArgs,
112        outdated_spans: &mut FxHashSet<SpanIndex>,
113    ) -> SpanIndex {
114        let id = SpanIndex::new(self.spans.len()).unwrap();
115        let ignore_self_time = &name == "thread" || &name == "blocking";
116        self.spans.push(Span {
117            parent,
118            depth: 0,
119            start,
120            category,
121            name,
122            args,
123            events: Default::default(),
124            is_complete: false,
125            self_allocations: 0,
126            self_allocation_count: 0,
127            self_deallocations: 0,
128            self_deallocation_count: 0,
129            totals: OnceLock::new(),
130            time_data: SpanTimeData {
131                self_end: start,
132                ignore_self_time,
133                ..Default::default()
134            },
135            extra: OnceLock::new(),
136            names: OnceLock::new(),
137        });
138        let mut parent = if let Some(parent) = parent {
139            outdated_spans.insert(parent);
140            &mut self.spans[parent.get()]
141        } else {
142            &mut self.spans[0]
143        };
144        let mut depth = parent.depth + 1;
145        if depth >= CUT_OFF_DEPTH
146            && let Some(parent_of_parent) = parent.parent
147        {
148            outdated_spans.insert(parent_of_parent);
149            self.spans[id.get()].parent = Some(parent_of_parent);
150            parent = &mut self.spans[parent_of_parent.get()];
151            depth = CUT_OFF_DEPTH - 1;
152        }
153        if depth < CUT_OFF_DEPTH {
154            parent.events.push(SpanEvent::Child { start, index: id });
155        }
156        parent.start = min(parent.start, start);
157        let span = &mut self.spans[id.get()];
158        span.depth = depth;
159        id
160    }
161
162    pub fn add_args(
163        &mut self,
164        span_index: SpanIndex,
165        args: SpanArgs,
166        outdated_spans: &mut FxHashSet<SpanIndex>,
167    ) {
168        let span = &mut self.spans[span_index.get()];
169        span.args.extend(args);
170        outdated_spans.insert(span_index);
171    }
172
173    pub fn set_max_self_time_lookup(&self, time: Timestamp) {
174        let time = *time;
175        let mut old = self
176            .max_self_time_lookup_time
177            .load(std::sync::atomic::Ordering::Relaxed);
178        while old < time {
179            match self.max_self_time_lookup_time.compare_exchange(
180                old,
181                time,
182                std::sync::atomic::Ordering::Relaxed,
183                std::sync::atomic::Ordering::Relaxed,
184            ) {
185                Ok(_) => break,
186                Err(real_old) => old = real_old,
187            }
188        }
189    }
190
191    fn insert_self_time(
192        &mut self,
193        start: Timestamp,
194        end: Timestamp,
195        span_index: SpanIndex,
196        outdated_spans: &mut FxHashSet<SpanIndex>,
197    ) {
198        if let Some(tree) = self.self_time_tree.as_mut() {
199            if Timestamp::from_value(*self.max_self_time_lookup_time.get_mut()) >= start {
200                tree.for_each_in_range_optimize(start, end, &mut |_, _, span| {
201                    outdated_spans.insert(*span);
202                });
203            }
204            tree.insert(start, end, span_index);
205        }
206    }
207
208    pub fn add_self_time(
209        &mut self,
210        span_index: SpanIndex,
211        start: Timestamp,
212        end: Timestamp,
213        outdated_spans: &mut FxHashSet<SpanIndex>,
214    ) {
215        let event = SpanEvent::self_time(start, end);
216        let span = &mut self.spans[span_index.get()];
217        let time_data = &mut span.time_data;
218        if time_data.ignore_self_time {
219            return;
220        }
221        outdated_spans.insert(span_index);
222        time_data.self_time += end - start;
223        time_data.self_end = max(time_data.self_end, end);
224        if let Some(event) = event {
225            span.events.push(event);
226            self.insert_self_time(start, end, span_index, outdated_spans);
227        }
228    }
229
230    pub fn set_total_time(
231        &mut self,
232        span_index: SpanIndex,
233        start_time: Timestamp,
234        total_time: Timestamp,
235        outdated_spans: &mut FxHashSet<SpanIndex>,
236    ) {
237        let span = SpanRef {
238            span: &self.spans[span_index.get()],
239            store: self,
240            index: span_index.get(),
241        };
242        let mut children = span
243            .children()
244            .map(|c| (c.span.start, c.span.time_data.self_end, c.index()))
245            .collect::<Vec<_>>();
246        children.sort();
247        let self_end = start_time + total_time;
248        let mut self_time = Timestamp::ZERO;
249        let mut current = start_time;
250        let mut events = Vec::new();
251        for (start, end, index) in children {
252            if start > current {
253                if start > self_end {
254                    if let Some(event) = SpanEvent::self_time(current, self_end) {
255                        events.push(event);
256                        self.insert_self_time(current, self_end, span_index, outdated_spans);
257                        self_time += self_end - current;
258                    }
259                    break;
260                }
261                if let Some(event) = SpanEvent::self_time(current, start) {
262                    events.push(event);
263                    self.insert_self_time(current, start, span_index, outdated_spans);
264                    self_time += start - current;
265                }
266            }
267            events.push(SpanEvent::Child { start, index });
268            current = max(current, end);
269        }
270        current -= start_time;
271        if current < total_time {
272            self_time += total_time - current;
273            let st = current + start_time;
274            let en = start_time + total_time;
275            if let Some(event) = SpanEvent::self_time(st, en) {
276                events.push(event);
277                self.insert_self_time(st, en, span_index, outdated_spans);
278            }
279        }
280        let span = &mut self.spans[span_index.get()];
281        outdated_spans.insert(span_index);
282        let time_data = &mut span.time_data;
283        time_data.self_time = self_time;
284        time_data.self_end = self_end;
285        span.events = events.into();
286        span.start = start_time;
287    }
288
289    pub fn set_parent(
290        &mut self,
291        span_index: SpanIndex,
292        parent: SpanIndex,
293        outdated_spans: &mut FxHashSet<SpanIndex>,
294    ) {
295        outdated_spans.insert(span_index);
296        let span = &mut self.spans[span_index.get()];
297        let span_start = span.start;
298
299        let old_parent = span.parent.replace(parent);
300        let old_parent = if let Some(parent) = old_parent {
301            outdated_spans.insert(parent);
302            &mut self.spans[parent.get()]
303        } else {
304            &mut self.spans[0]
305        };
306        old_parent.events.retain_unordered(
307            |event: &SpanEvent| !matches!(event, SpanEvent::Child { index, .. } if *index == span_index),
308        );
309
310        outdated_spans.insert(parent);
311        let parent = &mut self.spans[parent.get()];
312        parent.events.push(SpanEvent::Child {
313            start: span_start,
314            index: span_index,
315        });
316    }
317
318    pub fn add_allocation(
319        &mut self,
320        span_index: SpanIndex,
321        allocation: u64,
322        count: u64,
323        outdated_spans: &mut FxHashSet<SpanIndex>,
324    ) {
325        let span = &mut self.spans[span_index.get()];
326        outdated_spans.insert(span_index);
327        span.self_allocations += allocation;
328        span.self_allocation_count += count;
329    }
330
331    pub fn add_deallocation(
332        &mut self,
333        span_index: SpanIndex,
334        deallocation: u64,
335        count: u64,
336        outdated_spans: &mut FxHashSet<SpanIndex>,
337    ) {
338        let span = &mut self.spans[span_index.get()];
339        outdated_spans.insert(span_index);
340        span.self_deallocations += deallocation;
341        span.self_deallocation_count += count;
342    }
343
344    pub fn add_memory_sample(&mut self, ts: Timestamp, memory: u64, memory_pressure: u8) {
345        // Samples arrive nearly sorted (roughly chronological from the trace
346        // writer), so an insertion-sort step is efficient: push to the end
347        // then swap backward until the timestamp ordering is restored.
348        self.memory_samples.push((ts, memory, memory_pressure));
349        let mut i = self.memory_samples.len() - 1;
350        while i > 0 && self.memory_samples[i - 1].0 > ts {
351            self.memory_samples.swap(i, i - 1);
352            i -= 1;
353        }
354    }
355
356    /// Returns up to `MAX_MEMORY_SAMPLES` memory samples in the range
357    /// `[start, end]`. When more samples exist, groups of N consecutive
358    /// samples are merged by taking the maximum memory value in each group.
359    pub fn memory_samples_for_range(&self, start: Timestamp, end: Timestamp) -> Vec<u64> {
360        self.memory_samples_for_range_with_ts(start, end)
361            .into_iter()
362            .map(|(_, mem, _)| mem)
363            .collect()
364    }
365
366    /// Like `memory_samples_for_range` but keeps the timestamps and the
367    /// memory-pressure byte. Timestamps are absolute store timestamps (same
368    /// reference frame as span start/end). When the raw slice exceeds
369    /// `MAX_MEMORY_SAMPLES`, each merged group is represented by the sample
370    /// whose memory value was the group's max (its timestamp and pressure
371    /// byte are kept alongside it).
372    pub fn memory_samples_for_range_with_ts(
373        &self,
374        start: Timestamp,
375        end: Timestamp,
376    ) -> Vec<MemorySample> {
377        let slice = self.memory_samples_slice(start, end);
378        let count = slice.len();
379        if count == 0 {
380            return Vec::new();
381        }
382
383        if count <= MAX_MEMORY_SAMPLES {
384            return slice.to_vec();
385        }
386
387        // Merge groups of N samples, taking the max memory in each group and
388        // keeping the timestamp and pressure of that max sample.
389        let n = count.div_ceil(MAX_MEMORY_SAMPLES);
390        slice
391            .chunks(n)
392            .map(|chunk| *chunk.iter().max_by_key(|(_, mem, _)| *mem).unwrap())
393            .collect()
394    }
395
396    /// Returns up to `MAX_MEMORY_SAMPLES` memory pressure values in the range
397    /// `[start, end]`. The returned slice has the same length and group
398    /// boundaries as [`Self::memory_samples_for_range`] so that the two
399    /// results can be rendered in parallel. Each group is downsampled by
400    /// taking the maximum pressure value.
401    pub fn memory_pressure_samples_for_range(&self, start: Timestamp, end: Timestamp) -> Vec<u8> {
402        let slice = self.memory_samples_slice(start, end);
403        let count = slice.len();
404        if count == 0 {
405            return Vec::new();
406        }
407
408        if count <= MAX_MEMORY_SAMPLES {
409            return slice.iter().map(|(_, _, p)| *p).collect();
410        }
411
412        let n = count.div_ceil(MAX_MEMORY_SAMPLES);
413        slice
414            .chunks(n)
415            .map(|chunk| chunk.iter().map(|(_, _, p)| *p).max().unwrap())
416            .collect()
417    }
418
419    fn memory_samples_slice(&self, start: Timestamp, end: Timestamp) -> &[MemorySample] {
420        // Binary search for the first sample >= start
421        let lo = self
422            .memory_samples
423            .partition_point(|(ts, _, _)| *ts < start);
424        // Binary search for the first sample > end
425        let hi = self.memory_samples.partition_point(|(ts, _, _)| *ts <= end);
426        &self.memory_samples[lo..hi]
427    }
428
429    pub fn complete_span(&mut self, span_index: SpanIndex) {
430        let span = &mut self.spans[span_index.get()];
431        span.is_complete = true;
432    }
433
434    pub fn invalidate_outdated_spans(&mut self, outdated_spans: &FxHashSet<SpanId>) {
435        fn invalidate_span(span: &mut Span) {
436            span.time_data.end.take();
437            span.time_data.total_time.take();
438            span.time_data.corrected_self_time.take();
439            span.time_data.corrected_total_time.take();
440            for event in span.events.iter_mut_unordered() {
441                if let SpanEvent::SelfTime(self_time) = event {
442                    self_time.corrected_self_time.take();
443                }
444            }
445            span.totals.take();
446            span.extra.take();
447        }
448
449        for id in outdated_spans.iter() {
450            let mut span = &mut self.spans[id.get()];
451            loop {
452                invalidate_span(span);
453                let Some(parent) = span.parent else {
454                    break;
455                };
456                if outdated_spans.contains(&parent) {
457                    break;
458                }
459                span = &mut self.spans[parent.get()];
460            }
461        }
462
463        invalidate_span(&mut self.spans[0]);
464    }
465
466    pub fn root_spans(&self) -> impl Iterator<Item = SpanRef<'_>> {
467        self.spans[0].events.iter().filter_map(|event| match event {
468            &SpanEvent::Child { index: id, .. } => Some(SpanRef {
469                span: &self.spans[id.get()],
470                store: self,
471                index: id.get(),
472            }),
473            _ => None,
474        })
475    }
476
477    pub fn root_span(&self) -> SpanRef<'_> {
478        SpanRef {
479            span: &self.spans[0],
480            store: self,
481            index: 0,
482        }
483    }
484
485    pub fn span(&self, id: SpanId) -> Option<(SpanRef<'_>, bool)> {
486        let id = id.get();
487        let is_graph = id & 1 == 1;
488        let index = id >> 1;
489        self.spans.get(index).map(|span| {
490            (
491                SpanRef {
492                    span,
493                    store: self,
494                    index,
495                },
496                is_graph,
497            )
498        })
499    }
500}
turbo_trace_server/store.rs

turbo_trace_server/
store.rs
