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 <[email protected]>
Co-authored-by: Claude <[email protected]>
This commit is contained in:
Tobias Koppers
2026-04-17 14:36:27 +02:00
committed by GitHub
co-authored by Tobias Koppers Claude
parent e2e8ba9aca
commit 3d3def23a1
3 changed files with 64 additions and 22 deletions
@@ -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,
@@ -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,7 +216,7 @@ 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;
@@ -206,10 +224,10 @@ impl<T> SelfTimeTree<T> {
}
}
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 {
@@ -257,31 +275,47 @@ impl<T> SelfTimeTree<T> {
}
pub fn for_each_in_range(
&self,
start: Timestamp,
end: Timestamp,
mut f: impl FnMut(Timestamp, Timestamp, &T),
) {
self.for_each_in_range_ref(start, end, &mut f);
}
fn for_each_in_range_ref(
&self,
start: Timestamp,
end: Timestamp,
f: &mut impl FnMut(Timestamp, Timestamp, &T),
) {
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 start <= children.split_point {
children.left.for_each_in_range(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(start, end, f);
}
}
}
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 {
f(entry.start, entry.end, &entry.item);
}
}
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_optimize(start, end, f);
}
}
}
@@ -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);
});
}