[turbopack] Share scratch buffer across shards using thread local (#90167)

Instead of allocating a scratch buffer for every shard, buffer them in a thread local.

Author: lukesandberg

Cargo.lock

@@ -57,6 +57,7 @@ turbo-persistence = { workspace = true }
 turbo-rcstr = { workspace = true }
 turbo-tasks = { workspace = true }
 turbo-tasks-hash = { workspace = true }
+thread_local = { workspace = true }
 
 [dev-dependencies]
 criterion = { workspace = true, features = ["async_tokio"] }

turbopack/crates/turbo-tasks-backend/Cargo.toml

@@ -1209,7 +1209,8 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
                 )
             };
 
-        // take_snapshot already filters empty items and empty shards in parallel
+        // take_snapshot filters empty shards (no modified/snapshot entries) in parallel.
+        // Individual empty SnapshotItems are filtered by the iterator.
         let task_snapshots = self.storage.take_snapshot(&process, &process_snapshot);
 
         swap_retain(&mut persisted_task_cache_log, |shard| !shard.is_empty());

turbopack/crates/turbo-tasks-backend/src/backend/mod.rs

@@ -1,11 +1,12 @@
 use std::{
+    cell::Cell,
     hash::Hash,
-    iter::Peekable,
     ops::{Deref, DerefMut},
     sync::{Arc, atomic::AtomicBool},
 };
 
 use smallvec::SmallVec;
+use thread_local::ThreadLocal;
 use turbo_bincode::TurboBincodeBuffer;
 use turbo_tasks::{FxDashMap, TaskId, parallel};
 
@@ -113,7 +114,9 @@ impl Storage {
     /// receives an owned Box<TaskStorage> snapshot.
     /// Both callbacks receive a mutable scratch buffer that can be reused across iterations
     /// to avoid repeated allocations.
-    /// The returned iterators are guaranteed to be non-empty and only yield non-empty items.
+    /// The returned shards implement `IntoIterator`. Empty shards (no modified or snapshot
+    /// entries) are filtered out, but shards may still yield no items if all entries produce
+    /// empty `SnapshotItem`s (this is rare and only happens under error conditions).
     pub fn take_snapshot<
         'l,
         P: for<'a> Fn(TaskId, &'a TaskStorage, &mut TurboBincodeBuffer) -> SnapshotItem + Sync,
@@ -122,12 +125,15 @@ impl Storage {
         &'l self,
         process: &'l P,
         process_snapshot: &'l PS,
-    ) -> Vec<Peekable<SnapshotShard<'l, P, PS>>> {
+    ) -> Vec<SnapshotShard<'l, P, PS>> {
         if !self.snapshot_mode() {
             self.start_snapshot();
         }
 
-        let guard = Arc::new(SnapshotGuard { storage: self });
+        let guard = Arc::new(SnapshotGuard {
+            storage: self,
+            scratch_buffers: ThreadLocal::new(),
+        });
 
         // The number of shards is much larger than the number of threads, so the effect of the
         // locks held is negligible.
@@ -163,22 +169,14 @@ impl Storage {
                 return None;
             }
 
-            /// How big of a buffer to allocate initially.  Based on metrics from a large
-            /// application this should cover about 98% of values with no resizes
-            const SCRATCH_BUFFER_SIZE: usize = 4096;
-            let shard = SnapshotShard {
+            Some(SnapshotShard {
                 direct_snapshots,
                 modified,
                 storage: self,
-                guard: Some(guard.clone()),
                 process,
                 process_snapshot,
-                scratch_buffer: TurboBincodeBuffer::with_capacity(SCRATCH_BUFFER_SIZE),
-            };
-
-            // Peek to filter out shards that only produce empty items
-            let mut iter = shard.peekable();
-            iter.peek().is_some().then_some(iter)
+                _guard: guard.clone(),
+            })
         })
         .into_iter()
         .flatten()
@@ -382,8 +380,69 @@ impl DerefMut for StorageWriteGuard<'_> {
     }
 }
 
+/// How big of a buffer to allocate initially. Based on metrics from a large
+/// application this should cover about 98% of values with no resizes.
+const SCRATCH_BUFFER_INITIAL_SIZE: usize = 4096;
+
+/// State machine for a per-thread scratch buffer slot.
+///
+/// Transitions:
+/// - `Uninit` → `Taken` (first take)
+/// - `Available` → `Taken` (subsequent takes)
+/// - `Taken` → `Available` (return)
+///
+/// Any other transition is a bug (e.g. double-take or double-return).
+#[derive(Default)]
+enum ScratchBufferSlot {
+    /// No buffer has been allocated on this thread yet.
+    #[default]
+    Uninit,
+    /// The buffer is currently checked out.
+    Taken,
+    /// The buffer is available for reuse.
+    Available(TurboBincodeBuffer),
+}
+
 pub struct SnapshotGuard<'l> {
     storage: &'l Storage,
+    /// Per-thread scratch buffers for encoding task data. Buffers are taken
+    /// by `SnapshotShardIter` on creation and returned on drop, allowing reuse
+    /// across multiple shards processed by the same thread. When the guard is
+    /// dropped (after all iterators are done), the `ThreadLocal` drops too,
+    /// freeing all buffers.
+    scratch_buffers: ThreadLocal<Cell<ScratchBufferSlot>>,
+}
+
+impl SnapshotGuard<'_> {
+    fn take_scratch_buffer(&self) -> TurboBincodeBuffer {
+        let cell = self.scratch_buffers.get_or_default();
+        match cell.take() {
+            ScratchBufferSlot::Available(buf) => {
+                cell.set(ScratchBufferSlot::Taken);
+                buf
+            }
+            ScratchBufferSlot::Uninit => {
+                cell.set(ScratchBufferSlot::Taken);
+                TurboBincodeBuffer::with_capacity(SCRATCH_BUFFER_INITIAL_SIZE)
+            }
+            ScratchBufferSlot::Taken => {
+                panic!("scratch buffer taken twice without being returned");
+            }
+        }
+    }
+
+    fn return_scratch_buffer(&self, buffer: TurboBincodeBuffer) {
+        let cell = self.scratch_buffers.get_or_default();
+        match cell.take() {
+            ScratchBufferSlot::Taken => cell.set(ScratchBufferSlot::Available(buffer)),
+            ScratchBufferSlot::Available(_) => {
+                panic!("scratch buffer returned without being taken (already available)");
+            }
+            ScratchBufferSlot::Uninit => {
+                panic!("scratch buffer returned without being taken (uninit)");
+            }
+        }
+    }
 }
 
 impl Drop for SnapshotGuard<'_> {
@@ -396,66 +455,82 @@ pub struct SnapshotShard<'l, P, PS> {
     direct_snapshots: Vec<(TaskId, Box<TaskStorage>)>,
     modified: SmallVec<[TaskId; 4]>,
     storage: &'l Storage,
-    guard: Option<Arc<SnapshotGuard<'l>>>,
     process: &'l P,
     process_snapshot: &'l PS,
-    /// Scratch buffer for encoding task data, reused across iterations to avoid allocations
-    scratch_buffer: TurboBincodeBuffer,
+    /// Held for its `Drop` impl — ensures snapshot mode ends when all shards are done.
+    _guard: Arc<SnapshotGuard<'l>>,
+}
+
+impl<'l, P, PS> IntoIterator for SnapshotShard<'l, P, PS>
+where
+    P: Fn(TaskId, &TaskStorage, &mut TurboBincodeBuffer) -> SnapshotItem + Sync,
+    PS: Fn(TaskId, Box<TaskStorage>, &mut TurboBincodeBuffer) -> SnapshotItem + Sync,
+{
+    type Item = SnapshotItem;
+    type IntoIter = SnapshotShardIter<'l, P, PS>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        let buffer = self._guard.take_scratch_buffer();
+        SnapshotShardIter {
+            shard: self,
+            buffer,
+        }
+    }
+}
+
+/// Iterator over a single shard's snapshot items. Holds a thread-local scratch
+/// buffer for the duration of iteration and returns it on drop.
+pub struct SnapshotShardIter<'l, P, PS> {
+    shard: SnapshotShard<'l, P, PS>,
+    buffer: TurboBincodeBuffer,
 }
 
-impl<'l, P, PS> SnapshotShard<'l, P, PS>
+impl<'l, P, PS> Iterator for SnapshotShardIter<'l, P, PS>
 where
     P: Fn(TaskId, &TaskStorage, &mut TurboBincodeBuffer) -> SnapshotItem + Sync,
     PS: Fn(TaskId, Box<TaskStorage>, &mut TurboBincodeBuffer) -> SnapshotItem + Sync,
 {
-    fn next_item(&mut self) -> Option<SnapshotItem> {
-        if let Some((task_id, snapshot)) = self.direct_snapshots.pop() {
-            return Some((self.process_snapshot)(
-                task_id,
-                snapshot,
-                &mut self.scratch_buffer,
-            ));
+    type Item = SnapshotItem;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        while let Some((task_id, snapshot)) = self.shard.direct_snapshots.pop() {
+            let item = (self.shard.process_snapshot)(task_id, snapshot, &mut self.buffer);
+            if !item.is_empty() {
+                return Some(item);
+            }
         }
-        while let Some(task_id) = self.modified.pop() {
-            let inner = self.storage.map.get(&task_id).unwrap();
+        while let Some(task_id) = self.shard.modified.pop() {
+            let inner = self.shard.storage.map.get(&task_id).unwrap();
             if !inner.flags.any_snapshot() {
-                return Some((self.process)(task_id, &inner, &mut self.scratch_buffer));
+                let item = (self.shard.process)(task_id, &inner, &mut self.buffer);
+                if !item.is_empty() {
+                    return Some(item);
+                }
             } else {
                 drop(inner);
                 let maybe_snapshot = {
-                    let mut modified_state = self.storage.modified.get_mut(&task_id).unwrap();
+                    let mut modified_state = self.shard.storage.modified.get_mut(&task_id).unwrap();
                     let ModifiedState::Snapshot(snapshot) = &mut *modified_state else {
                         unreachable!("The snapshot bit was set, so it must be in Snapshot state");
                     };
                     snapshot.take()
                 };
                 if let Some(snapshot) = maybe_snapshot {
-                    return Some((self.process_snapshot)(
-                        task_id,
-                        snapshot,
-                        &mut self.scratch_buffer,
-                    ));
+                    let item = (self.shard.process_snapshot)(task_id, snapshot, &mut self.buffer);
+                    if !item.is_empty() {
+                        return Some(item);
+                    }
                 }
             }
         }
         None
     }
 }
 
-impl<'l, P, PS> Iterator for SnapshotShard<'l, P, PS>
-where
-    P: Fn(TaskId, &TaskStorage, &mut TurboBincodeBuffer) -> SnapshotItem + Sync,
-    PS: Fn(TaskId, Box<TaskStorage>, &mut TurboBincodeBuffer) -> SnapshotItem + Sync,
-{
-    type Item = SnapshotItem;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        while let Some(item) = self.next_item() {
-            if !item.is_empty() {
-                return Some(item);
-            }
-        }
-        self.guard = None;
-        None
+impl<P, PS> Drop for SnapshotShardIter<'_, P, PS> {
+    fn drop(&mut self) {
+        self.shard
+            ._guard
+            .return_scratch_buffer(std::mem::take(&mut self.buffer));
     }
 }

turbopack/crates/turbo-tasks-backend/src/backend/storage.rs

@@ -62,7 +62,7 @@ pub trait BackingStorageSealed: 'static + Send + Sync {
         snapshots: Vec<I>,
     ) -> Result<()>
     where
-        I: Iterator<Item = SnapshotItem> + Send + Sync;
+        I: IntoIterator<Item = SnapshotItem> + Send + Sync;
     /// Returns all task IDs that match the given task type (hash collision candidates).
     ///
     /// Since TaskCache uses hash-based keys, multiple task types may (rarely) hash to the same key.
@@ -125,7 +125,7 @@ where
         snapshots: Vec<I>,
     ) -> Result<()>
     where
-        I: Iterator<Item = SnapshotItem> + Send + Sync,
+        I: IntoIterator<Item = SnapshotItem> + Send + Sync,
     {
         either::for_both!(self, this => this.save_snapshot(
             operations,

turbopack/crates/turbo-tasks-backend/src/backing_storage.rs

@@ -231,7 +231,7 @@ impl<T: KeyValueDatabase + Send + Sync + 'static> BackingStorageSealed
         snapshots: Vec<I>,
     ) -> Result<()>
     where
-        I: Iterator<Item = SnapshotItem> + Send + Sync,
+        I: IntoIterator<Item = SnapshotItem> + Send + Sync,
     {
         let _span = tracing::info_span!("save snapshot", operations = operations.len()).entered();
         let batch = self.inner.database.write_batch()?;
@@ -451,7 +451,7 @@ fn process_task_data<'a, B: ConcurrentWriteBatch<'a> + Send + Sync, I>(
     batch: &B,
 ) -> Result<()>
 where
-    I: Iterator<Item = SnapshotItem> + Send + Sync,
+    I: IntoIterator<Item = SnapshotItem> + Send + Sync,
 {
     parallel::try_for_each_owned(tasks, |tasks| {
         for SnapshotItem {