trace-server: performance improvement (#8147)

### Description

- avoid invalidating self time tree spans when they don't have been
accessed
- used multiple core for some computations 

### Testing Instructions

<!--
  Give a quick description of steps to test your changes.
-->

crates/turbopack-trace-server/Cargo.toml

@@ -18,6 +18,7 @@ flate2 = { version = "1.0.28" }
 indexmap = { workspace = true, features = ["serde"] }
 itertools = { workspace = true }
 postcard = { workspace = true }
+rayon = "1"
 rustc-demangle = "0.1"
 serde = { workspace = true }
 serde_json = { workspace = true }

crates/turbopack-trace-server/src/server.rs

@@ -171,7 +171,7 @@ fn handle_connection(
         Ok(())
     }
     loop {
-        match websocket.read().unwrap() {
+        match websocket.read()? {
             Message::Frame(_frame) => {}
             Message::Text(text) => {
                 let message: ClientToServerMessage = serde_json::from_str(&text)?;

crates/turbopack-trace-server/src/span_ref.rs

@@ -264,6 +264,7 @@ impl<'a> SpanRef<'a> {
                     if duration == 0 {
                         continue;
                     }
+                    store.set_max_self_time_lookup(*end);
                     let concurrent_time = store.self_time_tree.lookup_range_count(*start, *end);
                     self_time += duration * duration / concurrent_time;
                 }

crates/turbopack-trace-server/src/store.rs

@@ -3,7 +3,7 @@ use std::{
     collections::HashSet,
     mem::replace,
     num::NonZeroUsize,
-    sync::OnceLock,
+    sync::{atomic::AtomicU64, OnceLock},
 };
 
 use crate::{
@@ -19,6 +19,7 @@ const CUT_OFF_DEPTH: u32 = 150;
 pub struct Store {
     pub(crate) spans: Vec<Span>,
     pub(crate) self_time_tree: SelfTimeTree<SpanIndex>,
+    max_self_time_lookup_time: AtomicU64,
 }
 
 fn new_root_span() -> Span {
@@ -52,13 +53,15 @@ impl Store {
         Self {
             spans: vec![new_root_span()],
             self_time_tree: SelfTimeTree::new(),
+            max_self_time_lookup_time: AtomicU64::new(0),
         }
     }
 
     pub fn reset(&mut self) {
         self.spans.truncate(1);
         self.spans[0] = new_root_span();
         self.self_time_tree = SelfTimeTree::new();
+        *self.max_self_time_lookup_time.get_mut() = 0;
     }
 
     pub fn has_time_info(&self) -> bool {
@@ -125,17 +128,36 @@ impl Store {
         outdated_spans.insert(span_index);
     }
 
+    pub fn set_max_self_time_lookup(&self, time: u64) {
+        let mut old = self
+            .max_self_time_lookup_time
+            .load(std::sync::atomic::Ordering::Relaxed);
+        while old < time {
+            match self.max_self_time_lookup_time.compare_exchange(
+                old,
+                time,
+                std::sync::atomic::Ordering::Relaxed,
+                std::sync::atomic::Ordering::Relaxed,
+            ) {
+                Ok(_) => break,
+                Err(real_old) => old = real_old,
+            }
+        }
+    }
+
     fn insert_self_time(
         &mut self,
         start: u64,
         end: u64,
         span_index: SpanIndex,
         outdated_spans: &mut HashSet<SpanIndex>,
     ) {
-        self.self_time_tree
-            .for_each_in_range(start, end, |_, _, span| {
-                outdated_spans.insert(*span);
-            });
+        if *self.max_self_time_lookup_time.get_mut() >= start {
+            self.self_time_tree
+                .for_each_in_range(start, end, |_, _, span| {
+                    outdated_spans.insert(*span);
+                });
+        }
         self.self_time_tree.insert(start, end, span_index);
     }
 
@@ -295,6 +317,7 @@ impl Store {
             span.total_deallocations.take();
             span.total_persistent_allocations.take();
             span.total_allocation_count.take();
+            span.total_span_count.take();
             span.extra.take();
         }
 

crates/turbopack-trace-server/src/viewer.rs

@@ -5,6 +5,7 @@ use std::{
 
 use either::Either;
 use itertools::Itertools;
+use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 use serde::{Deserialize, Serialize};
 
 use crate::{
@@ -361,6 +362,10 @@ impl Viewer {
             .iter()
             .min_by_key(|span| span.start())
             .map_or(0, |span| span.start());
+        root_spans.par_iter().for_each(|span| {
+            span.max_depth();
+            QueueItem::Span(*span).value(value_mode);
+        });
         for span in root_spans {
             if matches!(value_mode, ValueMode::Duration) {
                 // Move current to start if needed.
@@ -396,6 +401,46 @@ impl Viewer {
             }
         }
         enqueue_children(children, &mut queue);
+        queue.par_iter().for_each(|item| {
+            let QueueItem::Span(span) = item.item else {
+                return;
+            };
+            let view_mode = if span.is_complete() {
+                item.view_mode
+            } else {
+                item.view_mode.as_spans()
+            };
+
+            match (view_mode.bottom_up(), view_mode.aggregate_children()) {
+                (false, false) => {}
+                (false, true) => {
+                    span.graph()
+                        .collect::<Vec<_>>()
+                        .par_iter()
+                        .for_each(|event| {
+                            value_mode.value_from_graph_event(event);
+                        });
+                }
+                (true, false) => {
+                    span.bottom_up()
+                        .collect::<Vec<_>>()
+                        .par_iter()
+                        .for_each(|bu| {
+                            bu.spans().collect::<Vec<_>>().par_iter().for_each(|span| {
+                                value_mode.value_from_bottom_up_span(span);
+                            });
+                        });
+                }
+                (true, true) => {
+                    span.bottom_up()
+                        .collect::<Vec<_>>()
+                        .par_iter()
+                        .for_each(|bu| {
+                            value_mode.value_from_bottom_up(bu);
+                        });
+                }
+            }
+        });
 
         let mut lines: Vec<Vec<LineEntry<'_>>> = vec![];