Turbopack: optimize SelfTimeTree performance, memory usage, and fix range query boundaries (#92931)

### What?

Optimize the `SelfTimeTree` data structure in the Turbopack trace server
for better performance and lower memory usage when loading large trace
files, and fix a correctness bug in range queries.

### Why?

Two issues were identified:

1. **Performance**: The tree was splitting and rebalancing on every
insert during bulk loading, and each `distribute_entries` call
recursively triggered further split checks per entry, creating
significant overhead.

2. **Correctness**: Range query boundary conditions used exclusive
comparisons (`<`/`>`), incorrectly excluding entries that touched the
exact start or end timestamp of the queried range.

### How?

**Performance optimizations:**

1. **Batch split checks**: Add `insert_without_check` for inserting
entries without triggering a split check per entry. In
`distribute_entries`, entries are moved to children using
`insert_without_check`, and `check_for_split` is called once per child
after all entries are distributed (rather than once per entry).

2. **Lazy optimization with memory reclamation**: Add an `optimize()`
method and a `for_each_in_range_optimize()` traversal method that
distribute and rebalance nodes lazily:
- `store.optimize()` is called when the reader reaches end-of-data
(before waiting for more), doing a bulk pass to properly structure the
tree after initial load. After distributing entries to children, each
node calls `shrink_to_fit()` on its entries vec to release excess
capacity.
- `for_each_in_range_optimize()` distributes and rebalances each node it
visits, so query paths are progressively optimized during normal use.

3. **Fix `rebalance()` bug**: `check_for_split()` was incorrectly called
after merging subtrees during rebalance — replaced with
`distribute_entries()` since the node already has children and needs
redistribution, not a fresh split.

**Correctness fix:**

All range boundary comparisons in `for_each_in_range`,
`for_each_in_range_optimize`, and `lookup_range_count` updated from
exclusive (`<`/`>`) to inclusive (`<=`/`>=`), ensuring entries with
timestamps exactly equal to the range endpoints are included in results.

<!-- NEXT_JS_LLM_PR -->

---------

Co-authored-by: Tobias Koppers <sokra@users.noreply.github.com>
Co-authored-by: Claude <noreply@anthropic.com>

Author: sokra

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

@@ -210,6 +210,7 @@ impl TraceReader {
             match file.read(&mut chunk) {
                 Ok(bytes_read) => {
                     if bytes_read == 0 {
+                        self.store.write().optimize();
                         if let Some(value) = self.wait_for_more_data(
                             &mut file,
                             &mut initial_read,
@@ -308,6 +309,7 @@ impl TraceReader {
                     if err.kind() == io::ErrorKind::UnexpectedEof
                         || err.kind() == io::ErrorKind::InvalidInput
                     {
+                        self.store.write().optimize();
                         if let Some(value) = self.wait_for_more_data(
                             &mut file,
                             &mut initial_read,

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

@@ -53,6 +53,11 @@ impl<T> SelfTimeTree<T> {
         self.check_for_split();
     }
 
+    fn insert_without_check(&mut self, start: Timestamp, end: Timestamp, item: T) {
+        self.count += 1;
+        self.entries.push(SelfTimeEntry { start, end, item });
+    }
+
     fn check_for_split(&mut self) {
         if self.entries.len() >= SPLIT_COUNT {
             let spanning_entries = if let Some(children) = &mut self.children {
@@ -66,6 +71,17 @@ impl<T> SelfTimeTree<T> {
         }
     }
 
+    pub fn optimize(&mut self) {
+        if self.children.is_some() {
+            self.distribute_entries();
+            self.rebalance();
+            let children = self.children.as_mut().unwrap();
+            children.left.optimize();
+            children.right.optimize();
+        }
+        self.entries.shrink_to_fit();
+    }
+
     fn split(&mut self) {
         debug_assert!(!self.entries.is_empty());
         self.distribute_entries();
@@ -100,16 +116,18 @@ impl<T> SelfTimeTree<T> {
             let SelfTimeEntry { start, end, .. } = self.entries[i];
             if end <= children.split_point {
                 let SelfTimeEntry { start, end, item } = self.entries.swap_remove(i);
-                children.left.insert(start, end, item);
+                children.left.insert_without_check(start, end, item);
             } else if start >= children.split_point {
                 let SelfTimeEntry { start, end, item } = self.entries.swap_remove(i);
-                children.right.insert(start, end, item);
+                children.right.insert_without_check(start, end, item);
             } else {
                 self.entries.swap(i, children.spanning_entries);
                 children.spanning_entries += 1;
                 i += 1;
             }
         }
+        children.left.check_for_split();
+        children.right.check_for_split();
     }
 
     fn rebalance(&mut self) {
@@ -188,7 +206,7 @@ impl<T> SelfTimeTree<T> {
                     self.entries.append(right_entries);
                     *right = take(right_right);
                     *spanning_entries = 0;
-                    self.check_for_split();
+                    self.distribute_entries();
                 }
             }
         }
@@ -198,18 +216,18 @@ impl<T> SelfTimeTree<T> {
     pub fn lookup_range_count(&self, start: Timestamp, end: Timestamp) -> Timestamp {
         let mut total_count = Timestamp::ZERO;
         for entry in &self.entries {
-            if entry.start < end && entry.end > start {
+            if entry.start <= end && entry.end >= start {
                 let start = std::cmp::max(entry.start, start);
                 let end = std::cmp::min(entry.end, end);
                 let span = end - start;
                 total_count += span;
             }
         }
         if let Some(children) = &self.children {
-            if start < children.split_point {
+            if start <= children.split_point {
                 total_count += children.left.lookup_range_count(start, end);
             }
-            if end > children.split_point {
+            if end >= children.split_point {
                 total_count += children.right.lookup_range_count(start, end);
             }
         }
@@ -225,7 +243,7 @@ impl<T> SelfTimeTree<T> {
         }
         let mut current_count = 0;
         let mut changes = Vec::new();
-        self.for_each_in_range(start, end, |s, e, _| {
+        self.for_each_in_range(start, end, &mut |s, e, _| {
             if s <= start {
                 current_count += 1;
             } else {
@@ -260,28 +278,44 @@ impl<T> SelfTimeTree<T> {
         &self,
         start: Timestamp,
         end: Timestamp,
-        mut f: impl FnMut(Timestamp, Timestamp, &T),
+        f: &mut impl FnMut(Timestamp, Timestamp, &T),
     ) {
-        self.for_each_in_range_ref(start, end, &mut f);
+        for entry in &self.entries {
+            if entry.start <= end && entry.end >= start {
+                f(entry.start, entry.end, &entry.item);
+            }
+        }
+        if let Some(children) = &self.children {
+            if start <= children.split_point {
+                children.left.for_each_in_range(start, end, f);
+            }
+            if end >= children.split_point {
+                children.right.for_each_in_range(start, end, f);
+            }
+        }
     }
 
-    fn for_each_in_range_ref(
-        &self,
+    pub fn for_each_in_range_optimize(
+        &mut self,
         start: Timestamp,
         end: Timestamp,
         f: &mut impl FnMut(Timestamp, Timestamp, &T),
     ) {
+        if self.children.is_some() {
+            self.distribute_entries();
+            self.rebalance();
+        }
         for entry in &self.entries {
-            if entry.start < end && entry.end > start {
+            if entry.start <= end && entry.end >= start {
                 f(entry.start, entry.end, &entry.item);
             }
         }
-        if let Some(children) = &self.children {
-            if start < children.split_point {
-                children.left.for_each_in_range_ref(start, end, f);
+        if let Some(children) = &mut self.children {
+            if start <= children.split_point {
+                children.left.for_each_in_range_optimize(start, end, f);
             }
-            if end > children.split_point {
-                children.right.for_each_in_range_ref(start, end, f);
+            if end >= children.split_point {
+                children.right.for_each_in_range_optimize(start, end, f);
             }
         }
     }

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

@@ -91,6 +91,12 @@ impl Store {
         self.memory_samples.clear();
     }
 
+    pub fn optimize(&mut self) {
+        if let Some(tree) = self.self_time_tree.as_mut() {
+            tree.optimize();
+        }
+    }
+
     pub fn has_time_info(&self) -> bool {
         self.self_time_tree
             .as_ref()
@@ -192,7 +198,7 @@ impl Store {
     ) {
         if let Some(tree) = self.self_time_tree.as_mut() {
             if Timestamp::from_value(*self.max_self_time_lookup_time.get_mut()) >= start {
-                tree.for_each_in_range(start, end, |_, _, span| {
+                tree.for_each_in_range_optimize(start, end, &mut |_, _, span| {
                     outdated_spans.insert(*span);
                 });
             }