Repository: vercel/next.js
Pull request: [turbopack] Simplify snapshotting logic

## Summary

- **Use `modified` flags instead of `restored` flags** to decide which task data categories to encode during snapshots. Previously, `process` used `restored` flags (encoding any category that had been loaded into memory), while `process_snapshot` used `modified` flags. The `restored` check was overly conservative — it would re-serialize categories that were loaded from disk but never actually changed. Using `modified` for both paths only encodes categories that were actually dirtied, avoiding redundant re-persistence of clean data.

- **Unify `process` and `process_snapshot` into a single callback** on `take_snapshot`. Since both closures now use the same logic (`modified` flags), the two-callback API was unnecessary. `take_snapshot` now takes a single `Fn(TaskId, &TaskStorage, &mut TurboBincodeBuffer) -> SnapshotItem` — for snapshot copies the `Box<TaskStorage>` is simply dereferenced before calling.

- **Change transient task handling from silent empty return to `unreachable!`**, since transient tasks should never appear in the persistence path.

## Test plan

- Existing snapshot/persistence tests should continue to pass
- The `modified` flags are always set for any task that ends up in the modified map (via `track_modification_internal`), so no categories are missed

Changed files:
filename: turbopack/crates/turbo-tasks-backend/src/backend/mod.rs
patch: @@ -1111,107 +1111,81 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         let task_cache_stats: Mutex<FxHashMap<_, TaskCacheStats>> =
             Mutex::new(FxHashMap::default());
 
-        // Helper to encode a TaskStorage into a SnapshotItem
-        // encode_meta/encode_data control whether to encode each category
-        let encode_snapshot_item =
-            |task_id: TaskId,
-             inner: &TaskStorage,
-             encode_meta: bool,
-             encode_data: bool,
-             buffer: &mut TurboBincodeBuffer| {
-                let encode_category = |task_id: TaskId,
-                                       data: &TaskStorage,
-                                       category: SpecificTaskDataCategory,
-                                       buffer: &mut TurboBincodeBuffer|
-                 -> Option<TurboBincodeBuffer> {
-                    match encode_task_data(task_id, data, category, buffer) {
-                        Ok(encoded) => {
-                            #[cfg(feature = "print_cache_item_size")]
-                            {
-                                let mut stats = task_cache_stats.lock();
-                                let entry = stats
-                                    .entry(self.get_task_name(task_id, turbo_tasks))
-                                    .or_default();
-                                match category {
-                                    SpecificTaskDataCategory::Meta => entry.add_meta(&encoded),
-                                    SpecificTaskDataCategory::Data => entry.add_data(&encoded),
-                                }
+        // Encode each task's modified categories. We only encode categories with `modified` set,
+        // meaning the category was actually dirtied. Categories restored from disk but never
+        // modified don't need re-persisting since the on-disk version is still valid.
+        // For tasks accessed during snapshot mode, a frozen copy was made and its `modified`
+        // flags were copied from the live task at snapshot creation time, reflecting which
+        // categories were dirtied before the snapshot was taken.
+        let process = |task_id: TaskId, inner: &TaskStorage, buffer: &mut TurboBincodeBuffer| {
+            let encode_category = |task_id: TaskId,
+                                   data: &TaskStorage,
+                                   category: SpecificTaskDataCategory,
+                                   buffer: &mut TurboBincodeBuffer|
+             -> Option<TurboBincodeBuffer> {
+                match encode_task_data(task_id, data, category, buffer) {
+                    Ok(encoded) => {
+                        #[cfg(feature = "print_cache_item_size")]
+                        {
+                            let mut stats = task_cache_stats.lock();
+                            let entry = stats
+                                .entry(self.get_task_name(task_id, turbo_tasks))
+                                .or_default();
+                            match category {
+                                SpecificTaskDataCategory::Meta => entry.add_meta(&encoded),
+                                SpecificTaskDataCategory::Data => entry.add_data(&encoded),
                             }
-                            Some(encoded)
-                        }
-                        Err(err) => {
-                            eprintln!(
-                                "Serializing task {} failed ({:?}): {:?}",
-                                self.debug_get_task_description(task_id),
-                                category,
-                                err
-                            );
-                            None
                         }
+                        Some(encoded)
+                    }
+                    Err(err) => {
+                        eprintln!(
+                            "Serializing task {} failed ({:?}): {:?}",
+                            self.debug_get_task_description(task_id),
+                            category,
+                            err
+                        );
+                        None
                     }
-                };
-                if task_id.is_transient() {
-                    return SnapshotItem {
-                        task_id,
-                        data: None,
-                        meta: None,
-                    };
                 }
+            };
+            if task_id.is_transient() {
+                unreachable!("transient task_ids should never be enqueued to be persisted");
+            }
 
-                #[cfg(feature = "print_cache_item_size")]
-                if encode_meta {
-                    task_cache_stats
-                        .lock()
-                        .entry(self.get_task_name(task_id, turbo_tasks))
-                        .or_default()
-                        .add_counts(inner);
-                }
+            let encode_meta = inner.flags.meta_modified();
+            let encode_data = inner.flags.data_modified();
 
-                let meta = if encode_meta {
-                    encode_category(task_id, inner, SpecificTaskDataCategory::Meta, buffer)
-                } else {
-                    None
-                };
+            #[cfg(feature = "print_cache_item_size")]
+            if encode_meta {
+                task_cache_stats
+                    .lock()
+                    .entry(self.get_task_name(task_id, turbo_tasks))
+                    .or_default()
+                    .add_counts(inner);
+            }
 
-                let data = if encode_data {
-                    encode_category(task_id, inner, SpecificTaskDataCategory::Data, buffer)
-                } else {
-                    None
-                };
+            let meta = if encode_meta {
+                encode_category(task_id, inner, SpecificTaskDataCategory::Meta, buffer)
+            } else {
+                None
+            };
 
-                SnapshotItem {
-                    task_id,
-                    meta,
-                    data,
-                }
+            let data = if encode_data {
+                encode_category(task_id, inner, SpecificTaskDataCategory::Data, buffer)
+            } else {
+                None
             };
 
-        // Process tasks from the main storage map (uses restored flags)
-        let process = |task_id: TaskId, inner: &TaskStorage, buffer: &mut TurboBincodeBuffer| {
-            encode_snapshot_item(
+            SnapshotItem {
                 task_id,
-                inner,
-                inner.flags.meta_restored(),
-                inner.flags.data_restored(),
-                buffer,
-            )
+                meta,
+                data,
+            }
         };
 
-        // Process tasks that were accessed during snapshot mode (uses modified flags)
-        let process_snapshot =
-            |task_id: TaskId, inner: Box<TaskStorage>, buffer: &mut TurboBincodeBuffer| {
-                encode_snapshot_item(
-                    task_id,
-                    &inner,
-                    inner.flags.meta_modified(),
-                    inner.flags.data_modified(),
-                    buffer,
-                )
-            };
-
-        // 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);
+        // take_snapshot already filters empty items and empty shards in parallel
+        let task_snapshots = self.storage.take_snapshot(&process);
 
         swap_retain(&mut persisted_task_cache_log, |shard| !shard.is_empty());
 

filename: turbopack/crates/turbo-tasks-backend/src/backend/storage.rs
patch: @@ -106,7 +106,7 @@ impl Storage {
         }
     }
 
-    /// Processes every modified item (resp. a snapshot of it) with the given functions and returns
+    /// Processes every modified item (resp. a snapshot of it) with the given function and returns
     /// the results. Ends snapshot mode afterwards.
     /// process is called while holding a read lock on the task storage, so it can access
     /// the TaskStorage directly without cloning.
@@ -120,12 +120,10 @@ impl Storage {
     pub fn take_snapshot<
         'l,
         P: for<'a> Fn(TaskId, &'a TaskStorage, &mut TurboBincodeBuffer) -> SnapshotItem + Sync,
-        PS: Fn(TaskId, Box<TaskStorage>, &mut TurboBincodeBuffer) -> SnapshotItem + Sync,
     >(
         &'l self,
         process: &'l P,
-        process_snapshot: &'l PS,
-    ) -> Vec<SnapshotShard<'l, P, PS>> {
+    ) -> Vec<SnapshotShard<'l, P>> {
         if !self.snapshot_mode() {
             self.start_snapshot();
         }
@@ -174,7 +172,6 @@ impl Storage {
                 modified,
                 storage: self,
                 process,
-                process_snapshot,
                 _guard: guard.clone(),
             })
         })
@@ -237,12 +234,12 @@ impl Storage {
         // And update the flags
         for key in removed_snapshots {
             if let Some(mut inner) = self.map.get_mut(&key) {
-                if inner.flags.meta_snapshot() {
-                    inner.flags.set_meta_snapshot(false);
+                if inner.flags.meta_modified_during_snapshot() {
+                    inner.flags.set_meta_modified_during_snapshot(false);
                     inner.flags.set_meta_modified(true);
                 }
-                if inner.flags.data_snapshot() {
-                    inner.flags.set_data_snapshot(false);
+                if inner.flags.data_modified_during_snapshot() {
+                    inner.flags.set_data_modified_during_snapshot(false);
                     inner.flags.set_data_modified(true);
                 }
             }
@@ -318,16 +315,17 @@ impl StorageWriteGuard<'_> {
         #[cfg(feature = "trace_task_modification")] name: &str,
     ) {
         let flags = &self.inner.flags;
-        if flags.is_snapshot(category) {
+        if flags.is_modified_during_snapshot(category) {
             return;
         }
         let modified = flags.is_modified(category);
+        let already_modified_during_snapshot = flags.any_modified_during_snapshot();
         #[cfg(feature = "trace_task_modification")]
         let _span = (!modified).then(|| tracing::trace_span!("mark_modified", name).entered());
         match (self.storage.snapshot_mode(), modified) {
             (false, false) => {
                 // Not in snapshot mode and item is unmodified
-                if !flags.any_snapshot() && !flags.any_modified() {
+                if !already_modified_during_snapshot && !flags.any_modified() {
                     self.storage
                         .modified
                         .insert(*self.inner.key(), ModifiedState::Modified);
@@ -340,18 +338,22 @@ impl StorageWriteGuard<'_> {
             }
             (true, false) => {
                 // In snapshot mode and item is unmodified (so it's not part of the snapshot)
-                if !flags.any_snapshot() {
+                // Mark it so it gets re-added as Modified after this snapshot completes
+                if !already_modified_during_snapshot {
                     self.storage
                         .modified
                         .insert(*self.inner.key(), ModifiedState::Snapshot(None));
                 }
-                self.inner.flags.set_snapshot(category, true);
+                self.inner
+                    .flags
+                    .set_modified_during_snapshot(category, true);
             }
             (true, true) => {
                 // In snapshot mode and item is modified (so it's part of the snapshot)
                 // We need to store the original version that is part of the snapshot
-                if !flags.any_snapshot() {
-                    // Snapshot all non-transient fields but keep the modified bits.
+                if !already_modified_during_snapshot {
+                    // Snapshot all non-transient fields but keep the modified bits since
+                    // save_snapshot relies on them
                     let mut snapshot = self.inner.clone_snapshot();
                     snapshot.flags.set_data_modified(flags.data_modified());
                     snapshot.flags.set_meta_modified(flags.meta_modified());
@@ -360,7 +362,9 @@ impl StorageWriteGuard<'_> {
                         ModifiedState::Snapshot(Some(Box::new(snapshot))),
                     );
                 }
-                self.inner.flags.set_snapshot(category, true);
+                self.inner
+                    .flags
+                    .set_modified_during_snapshot(category, true);
             }
         }
     }
@@ -451,23 +455,21 @@ impl Drop for SnapshotGuard<'_> {
     }
 }
 
-pub struct SnapshotShard<'l, P, PS> {
+pub struct SnapshotShard<'l, P> {
     direct_snapshots: Vec<(TaskId, Box<TaskStorage>)>,
     modified: SmallVec<[TaskId; 4]>,
     storage: &'l Storage,
     process: &'l P,
-    process_snapshot: &'l PS,
     /// 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>
+impl<'l, P> IntoIterator for SnapshotShard<'l, P>
 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>;
+    type IntoIter = SnapshotShardIter<'l, P>;
 
     fn into_iter(self) -> Self::IntoIter {
         let buffer = self._guard.take_scratch_buffer();
@@ -480,28 +482,27 @@ where
 
 /// 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>,
+pub struct SnapshotShardIter<'l, P> {
+    shard: SnapshotShard<'l, P>,
     buffer: TurboBincodeBuffer,
 }
 
-impl<'l, P, PS> Iterator for SnapshotShardIter<'l, P, PS>
+impl<'l, P> Iterator for SnapshotShardIter<'l, P>
 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((task_id, snapshot)) = self.shard.direct_snapshots.pop() {
-            let item = (self.shard.process_snapshot)(task_id, snapshot, &mut self.buffer);
+            let item = (self.shard.process)(task_id, &snapshot, &mut self.buffer);
             if !item.is_empty() {
                 return Some(item);
             }
         }
         while let Some(task_id) = self.shard.modified.pop() {
             let inner = self.shard.storage.map.get(&task_id).unwrap();
-            if !inner.flags.any_snapshot() {
+            if !inner.flags.any_modified_during_snapshot() {
                 let item = (self.shard.process)(task_id, &inner, &mut self.buffer);
                 if !item.is_empty() {
                     return Some(item);
@@ -516,7 +517,7 @@ where
                     snapshot.take()
                 };
                 if let Some(snapshot) = maybe_snapshot {
-                    let item = (self.shard.process_snapshot)(task_id, snapshot, &mut self.buffer);
+                    let item = (self.shard.process)(task_id, &snapshot, &mut self.buffer);
                     if !item.is_empty() {
                         return Some(item);
                     }
@@ -527,7 +528,7 @@ where
     }
 }
 
-impl<P, PS> Drop for SnapshotShardIter<'_, P, PS> {
+impl<P> Drop for SnapshotShardIter<'_, P> {
     fn drop(&mut self) {
         self.shard
             ._guard

filename: turbopack/crates/turbo-tasks-backend/src/backend/storage_schema.rs
patch: @@ -171,11 +171,11 @@ struct TaskStorageSchema {
 
     /// Whether meta was modified after snapshot mode was entered (snapshot taken).
     #[field(storage = "flag", category = "transient")]
-    meta_snapshot: bool,
+    meta_modified_during_snapshot: bool,
 
     /// Whether data was modified after snapshot mode was entered (snapshot taken).
     #[field(storage = "flag", category = "transient")]
-    data_snapshot: bool,
+    data_modified_during_snapshot: bool,
 
     /// Whether dependencies have been prefetched.
     #[field(storage = "flag", category = "transient")]
@@ -334,8 +334,8 @@ impl TaskFlags {
     }
 
     /// Check if any snapshot flag is set
-    pub fn any_snapshot(&self) -> bool {
-        self.meta_snapshot() || self.data_snapshot()
+    pub fn any_modified_during_snapshot(&self) -> bool {
+        self.meta_modified_during_snapshot() || self.data_modified_during_snapshot()
     }
 
     /// Check if any modified flag is set
@@ -360,18 +360,22 @@ impl TaskFlags {
     }
 
     /// Check if the specified category has a snapshot
-    pub fn is_snapshot(&self, category: SpecificTaskDataCategory) -> bool {
+    pub fn is_modified_during_snapshot(&self, category: SpecificTaskDataCategory) -> bool {
         match category {
-            SpecificTaskDataCategory::Meta => self.meta_snapshot(),
-            SpecificTaskDataCategory::Data => self.data_snapshot(),
+            SpecificTaskDataCategory::Meta => self.meta_modified_during_snapshot(),
+            SpecificTaskDataCategory::Data => self.data_modified_during_snapshot(),
         }
     }
 
     /// Set the snapshot flag for the specified category
-    pub fn set_snapshot(&mut self, category: SpecificTaskDataCategory, value: bool) {
+    pub fn set_modified_during_snapshot(
+        &mut self,
+        category: SpecificTaskDataCategory,
+        value: bool,
+    ) {
         match category {
-            SpecificTaskDataCategory::Meta => self.set_meta_snapshot(value),
-            SpecificTaskDataCategory::Data => self.set_data_snapshot(value),
+            SpecificTaskDataCategory::Meta => self.set_meta_modified_during_snapshot(value),
+            SpecificTaskDataCategory::Data => self.set_data_modified_during_snapshot(value),
         }
     }
 }
@@ -739,8 +743,8 @@ mod tests {
         assert!(!storage.flags.data_restored());
         assert!(!storage.flags.meta_modified());
         assert!(!storage.flags.data_modified());
-        assert!(!storage.flags.meta_snapshot());
-        assert!(!storage.flags.data_snapshot());
+        assert!(!storage.flags.meta_modified_during_snapshot());
+        assert!(!storage.flags.data_modified_during_snapshot());
         assert!(!storage.flags.prefetched());
 
         // Test setting restored flags
@@ -756,10 +760,10 @@ mod tests {
         assert!(storage.flags.data_modified());
 
         // Test setting snapshot flags
-        storage.flags.set_meta_snapshot(true);
-        storage.flags.set_data_snapshot(true);
-        assert!(storage.flags.meta_snapshot());
-        assert!(storage.flags.data_snapshot());
+        storage.flags.set_meta_modified_during_snapshot(true);
+        storage.flags.set_data_modified_during_snapshot(true);
+        assert!(storage.flags.meta_modified_during_snapshot());
+        assert!(storage.flags.data_modified_during_snapshot());
 
         // Test prefetched flag
         storage.flags.set_prefetched(true);