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