Turbopack: parallel search_index in trace server (#78357)

### What?

Improve search performance in trace server by using parallelization.

Author: sokra

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

@@ -6,7 +6,7 @@ use std::{
 };
 
 use hashbrown::HashMap;
-use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
+use rayon::iter::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator};
 use rustc_hash::FxHashSet;
 
 use crate::{
@@ -417,10 +417,28 @@ impl<'a> SpanRef<'a> {
 
     fn search_index(&self) -> &HashMap<String, Vec<SpanIndex>> {
         self.extra().search_index.get_or_init(|| {
-            let mut index: HashMap<String, Vec<SpanIndex>> = HashMap::default();
-            let mut queue = VecDeque::with_capacity(8);
-            queue.push_back(*self);
-            while let Some(span) = queue.pop_front() {
+            let mut all_spans = Vec::new();
+            all_spans.push(self.index);
+            let mut i = 0;
+            while i < all_spans.len() {
+                let index = all_spans[i];
+                let span = SpanRef {
+                    span: &self.store.spans[index],
+                    store: self.store,
+                    index,
+                };
+                for child in span.children() {
+                    all_spans.push(child.index);
+                }
+                i += 1;
+            }
+
+            enum SpanOrMap<'a> {
+                Span(SpanRef<'a>),
+                Map(HashMap<String, Vec<SpanIndex>>),
+            }
+
+            fn add_span_to_map<'a>(index: &mut HashMap<String, Vec<SpanIndex>>, span: SpanRef<'a>) {
                 if !span.is_root() {
                     let (cat, name) = span.nice_name();
                     if !cat.is_empty() {
@@ -453,11 +471,47 @@ impl<'a> SpanRef<'a> {
                             .or_insert_with(|| (name.to_string(), vec![span.index()]));
                     }
                 }
-                for child in span.children() {
-                    queue.push_back(child);
+            }
+
+            let result = all_spans
+                .into_par_iter()
+                .map(|index| {
+                    SpanOrMap::Span(SpanRef {
+                        span: &self.store.spans[index],
+                        store: self.store,
+                        index,
+                    })
+                })
+                .reduce(
+                    || SpanOrMap::Map(HashMap::default()),
+                    |a, b| {
+                        let mut map = match a {
+                            SpanOrMap::Span(span) => {
+                                let mut map = HashMap::default();
+                                add_span_to_map(&mut map, span);
+                                map
+                            }
+                            SpanOrMap::Map(map) => map,
+                        };
+                        match b {
+                            SpanOrMap::Span(span) => {
+                                add_span_to_map(&mut map, span);
+                            }
+                            SpanOrMap::Map(other_map) => {
+                                map.extend(other_map);
+                            }
+                        }
+                        SpanOrMap::Map(map)
+                    },
+                );
+            match result {
+                SpanOrMap::Span(span) => {
+                    let mut map = HashMap::default();
+                    add_span_to_map(&mut map, span);
+                    map
                 }
+                SpanOrMap::Map(map) => map,
             }
-            index
         })
     }
 }