turbo_trace_server/
bottom_up.rs1use std::{env, sync::Arc};
2
3use hashbrown::HashMap;
4
5use crate::{
6 span::{SpanBottomUp, SpanIndex},
7 span_ref::SpanRef,
8 string_tuple_ref::StringTupleRef,
9};
10
11pub struct SpanBottomUpBuilder {
12 pub self_spans: Vec<SpanIndex>,
14 pub children: HashMap<(String, String), SpanBottomUpBuilder>,
15 pub example_span: SpanIndex,
16}
17
18impl SpanBottomUpBuilder {
19 pub fn new(example_span: SpanIndex) -> Self {
20 Self {
21 self_spans: vec![],
22 children: HashMap::default(),
23 example_span,
24 }
25 }
26
27 pub fn build(self) -> SpanBottomUp {
28 SpanBottomUp::new(
29 self.self_spans,
30 self.example_span,
31 self.children
32 .into_values()
33 .map(|child| Arc::new(child.build()))
34 .collect(),
35 )
36 }
37}
38
39pub fn build_bottom_up_graph<'a>(
40 spans: impl Iterator<Item = SpanRef<'a>>,
41) -> Vec<Arc<SpanBottomUp>> {
42 let max_depth = env::var("BOTTOM_UP_DEPTH")
43 .ok()
44 .and_then(|s| s.parse().ok())
45 .unwrap_or(usize::MAX);
46 let mut roots: HashMap<(String, String), SpanBottomUpBuilder> = HashMap::default();
47
48 let mut current_iterators: Vec<Box<dyn Iterator<Item = SpanRef<'_>>>> =
54 vec![Box::new(spans.flat_map(|span| span.children()))];
55
56 let mut current_path: Vec<((&'_ str, &'_ str), SpanIndex)> = vec![];
57 while let Some(mut iter) = current_iterators.pop() {
58 if let Some(child) = iter.next() {
59 current_iterators.push(iter);
60
61 let (category, name) = child.group_name();
62 let (_, mut bottom_up) = roots
63 .raw_entry_mut()
64 .from_key(&StringTupleRef(category, name))
65 .or_insert_with(|| {
66 (
67 (category.to_string(), name.to_string()),
68 SpanBottomUpBuilder::new(child.index()),
69 )
70 });
71 bottom_up.self_spans.push(child.index());
72 let mut prev = None;
73 for &((category, title), example_span) in current_path.iter().rev().take(max_depth) {
74 if prev == Some((category, title)) {
75 continue;
76 }
77 let (_, child_bottom_up) = bottom_up
78 .children
79 .raw_entry_mut()
80 .from_key(&StringTupleRef(category, title))
81 .or_insert_with(|| {
82 (
83 (category.to_string(), title.to_string()),
84 SpanBottomUpBuilder::new(example_span),
85 )
86 });
87 child_bottom_up.self_spans.push(child.index());
88 bottom_up = child_bottom_up;
89 prev = Some((category, title));
90 }
91
92 current_path.push((child.group_name(), child.index()));
93 current_iterators.push(Box::new(child.children()));
94 } else {
95 current_path.pop();
96 }
97 }
98 roots.into_values().map(|b| Arc::new(b.build())).collect()
99}