Repository: vercel/next.js
Pull request: Reland "[turbopack] Fix some modification tracking bugs in TaskStorage"

Reverts vercel/next.js#90096

* always `track_mutation()` before the mutation
* only `track_mutation()` if we are actually going to mutate

The previous attempt was reverted because we would call `track_modification` too often which lead to a test failure in the filesystem-cache-test because the DB grew too much.

Changed files:
filename: test/e2e/filesystem-cache/filesystem-cache.test.ts
patch: @@ -97,6 +97,17 @@ for (const cacheEnabled of [false, true]) {
       )
     }
 
+    // Ensure each test starts with a fresh server and no test leaks a
+    // running server into the next one.
+    beforeEach(async () => {
+      // stop() is a no-op if already stopped; start() rebuilds + starts.
+      await stop()
+      await start()
+    })
+    afterEach(async () => {
+      await stop()
+    })
+
     // Very flakey with Webpack enabled
     ;(process.env.IS_TURBOPACK_TEST ? it : it.skip)(
       'should cache or not cache loaders',

filename: turbopack/crates/turbo-tasks-backend/src/backend/storage_schema.rs
patch: @@ -393,13 +393,44 @@ impl TaskStorage {
     ///
     /// The `extract` closure should return `Some(&mut T)` for the matching variant,
     /// or `None` for non-matching variants.
-    pub fn find_lazy_mut<T>(
+    fn find_lazy_mut<T>(
         &mut self,
         extract: impl Fn(&mut LazyField) -> Option<&mut T>,
     ) -> Option<&mut T> {
         self.lazy.iter_mut().find_map(extract)
     }
 
+    /// Find and extract a lazy field, returning its index and a reference to the inner value.
+    ///
+    /// Combines index lookup and extraction into a single scan. The returned index
+    /// can be used with `lazy_at_mut` for subsequent mutation without re-scanning.
+    fn find_lazy_ref<T>(&self, extract: impl Fn(&LazyField) -> Option<&T>) -> Option<(usize, &T)> {
+        self.lazy
+            .iter()
+            .enumerate()
+            .find_map(|(idx, field)| extract(field).map(|val| (idx, val)))
+    }
+
+    /// Access a lazy field by index (mutable), extracting the inner value.
+    ///
+    /// # Panics
+    /// Panics if `idx` is out of bounds or the extractor returns `None`.
+    fn lazy_at_mut<T>(
+        &mut self,
+        idx: usize,
+        extract: impl FnOnce(&mut LazyField) -> Option<&mut T>,
+    ) -> &mut T {
+        extract(&mut self.lazy[idx]).unwrap()
+    }
+
+    /// Take a lazy field by known index, removing it from the Vec via swap_remove.
+    ///
+    /// # Panics
+    /// Panics if `idx` is out of bounds.
+    fn lazy_take_at<T>(&mut self, idx: usize, extract: impl FnOnce(LazyField) -> T) -> T {
+        extract(self.lazy.swap_remove(idx))
+    }
+
     /// Get or create a lazy field, returning a mutable reference.
     ///
     /// Uses a single `extract` closure that serves as both the matcher (by returning Some/None)

filename: turbopack/crates/turbo-tasks-macros/src/derive/task_storage_macro.rs
patch: @@ -1476,8 +1476,8 @@ fn generate_trait_accessor_methods(field: &FieldInfo) -> TokenStream {
                     #check_access
                     let current = self.typed().flags.#field_name();
                     if current != value {
-                        self.typed_mut().flags.#set_name(value);
                         #track_modification
+                        self.typed_mut().flags.#set_name(value);
                     }
                 }
             }
@@ -1533,8 +1533,6 @@ fn generate_direct_accessors(field: &FieldInfo) -> TokenStream {
                     // For fields with default semantics, always return Some(&mut self.field)
                     quote! {
                         #[doc = "Get a mutable reference to the field value."]
-                        #[doc = ""]
-                        #[doc = "Tracks modification pessimistically - assumes caller will mutate."]
                         fn #get_mut_name(&mut self) -> &mut #value_type {
                             #check_access
                             #track_modification
@@ -1571,6 +1569,67 @@ fn generate_direct_accessors(field: &FieldInfo) -> TokenStream {
         }
     };
 
+    let set_body = if field.is_transient() {
+        quote! {
+            #set_expr(value)
+        }
+    } else if field.lazy {
+        // For lazy fields, combine equality check and set into one operation to avoid
+        // double-scanning the lazy vec (get_expr scans via find_lazy, then set_expr
+        // scans again via set_lazy).
+        let extractor = field.lazy_extractor_closure();
+        let unwraper = field.lazy_unwrap_closure();
+        let constructor = field.lazy_constructor(quote! { value });
+        quote! {
+            if let Some((idx, old_ref)) = self.typed().find_lazy_ref(#extractor) {
+                if old_ref == &value {
+                    return None;
+                }
+                #track_modification
+                let old = std::mem::replace(&mut self.typed_mut().lazy[idx], #constructor);
+                Some((#unwraper)(old))
+            } else {
+                #track_modification
+                self.typed_mut().lazy.push(#constructor);
+                None
+            }
+        }
+    } else {
+        quote! {
+            if #get_expr.is_some_and(|old| old == &value) {
+                return None;
+            }
+            #track_modification
+            #set_expr(value)
+        }
+    };
+
+    let take_body = if field.is_transient() {
+        quote! {
+            #take_expr
+        }
+    } else if field.lazy {
+        // For lazy fields, combine existence check and take into one operation to avoid
+        // double-scanning the lazy vec (get_expr scans via find_lazy, then take_expr
+        // scans again via take_lazy).
+        let extractor = field.lazy_extractor_closure();
+        let unwraper = field.lazy_unwrap_closure();
+        quote! {
+            let (idx, _) = self.typed().find_lazy_ref(#extractor)?;
+            #track_modification
+            Some(self.typed_mut().lazy_take_at(idx, #unwraper))
+        }
+    } else {
+        quote! {
+            if #get_expr.is_some() {
+                #track_modification
+                #take_expr
+            } else {
+                None
+            }
+        }
+    };
+
     quote! {
         #[doc = "Get a reference to the field value (if present)"]
         fn #get_name(&self) -> Option<&#value_type> {
@@ -1584,23 +1643,16 @@ fn generate_direct_accessors(field: &FieldInfo) -> TokenStream {
             #get_expr.is_some()
         }
 
-        #[doc = "Set the field value, returning the old value if present"]
+        #[doc = "Set the field value, returning the old value if present."]
         fn #set_name(&mut self, value: #value_type) -> Option<#value_type> {
             #check_access
-            #track_modification
-            #set_expr(value)
+            #set_body
         }
 
         #[doc = "Take the field value, clearing it"]
-        #[doc = ""]
-        #[doc = "Only tracks modification if there was a value to take."]
         fn #take_name(&mut self) -> Option<#value_type> {
             #check_access
-            let value = #take_expr;
-            if value.is_some() {
-                #track_modification
-            }
-            value
+            #take_body
         }
 
         #get_mut_accessor
@@ -1664,45 +1716,170 @@ fn generate_autoset_ops(field: &FieldInfo) -> TokenStream {
         quote! { #ref_expr.is_empty() }
     };
 
-    // Remove uses find_lazy_mut for lazy to avoid allocation.
-    let remove_body = if is_option {
-        let extractor = field.lazy_extractor_closure();
+    // For transient fields, track_modification is a no-op so skip all guards.
+    let remove_body;
+    let add_body;
+    let set_body;
+    let extend_body;
 
-        quote! {
-            let Some(set) = self.typed_mut().find_lazy_mut(#extractor) else {
-                return false;
-            };
-            let removed = set.remove(item);
-            if removed {
+    if field.is_transient() {
+        remove_body = quote! {
+            #mut_expr.remove(item)
+        };
+        add_body = quote! {
+            #mut_expr.insert(item)
+        };
+        set_body = if is_option {
+            let unwraper = field.lazy_unwrap_closure();
+            let matches = field.lazy_matches_closure();
+            let ctor = field.lazy_constructor(quote! {set});
+            quote! {
+                self.typed_mut().set_lazy(#matches, #unwraper, #ctor)
+            }
+        } else {
+            quote! {
+                let old = #take_expr;
+                *#mut_expr = set;
+                Some(old)
+            }
+        };
+        extend_body = quote! {
+            #mut_expr.extend(items);
+        };
+    } else {
+        // Remove: only track modification if the item exists.
+        // For lazy fields, find the index once and reuse it to avoid double-scanning.
+        remove_body = if is_option {
+            let extractor = field.lazy_extractor_closure();
+            quote! {
+                let Some((idx, val)) = self.typed().find_lazy_ref(#extractor) else {
+                    return false;
+                };
+                if !val.contains(item) {
+                    return false;
+                }
+                #track_modification
+                self.typed_mut().lazy_at_mut(idx, #extractor).remove(item)
+            }
+        } else {
+            quote! {
+                if !#ref_expr.contains(item) {
+                    return false;
+                }
                 #track_modification
+                #mut_expr.remove(item)
             }
-            return removed;
+        };
 
-        }
-    } else {
-        quote! {
-            let removed = #mut_expr.remove(item);
-            if removed {
+        // Add: only track modification if the item is actually new.
+        // For lazy fields, use find_lazy_ref + lazy_at_mut to avoid double-scanning.
+        add_body = if is_option {
+            let extractor = field.lazy_extractor_closure();
+            let ctor = field.lazy_constructor(quote! { set });
+            quote! {
+                if let Some((idx, existing)) = self.typed().find_lazy_ref(#extractor) {
+                    if existing.contains(&item) {
+                        return false;
+                    }
+                    #track_modification
+                    self.typed_mut().lazy_at_mut(idx, #extractor).insert(item)
+                } else {
+                    #track_modification
+                    let mut set = <#field_type as Default>::default();
+                    set.insert(item);
+                    self.typed_mut().lazy.push(#ctor);
+                    true
+                }
+            }
+        } else {
+            quote! {
+                if #ref_expr.contains(&item) {
+                    return false;
+                }
                 #track_modification
+                #mut_expr.insert(item)
             }
-            removed
-        }
-    };
+        };
 
-    let set_body = if is_option {
-        let unwraper = field.lazy_unwrap_closure();
-        let matches = field.lazy_matches_closure();
-        let ctor = field.lazy_constructor(quote! {set});
-        quote! {
-             self.typed_mut().set_lazy(#matches, #unwraper, #ctor)
-        }
-    } else {
-        quote! {
-            let old = #take_expr;
-            *#mut_expr = set;
-            Some(old)
+        if is_option {
+            // For lazy fields, combine guard and set into one operation to avoid
+            // double-scanning the lazy vec (set_guard would scan via find_lazy,
+            // then set_lazy would scan again via position).
+            let extractor = field.lazy_extractor_closure();
+            let unwraper = field.lazy_unwrap_closure();
+            let ctor = field.lazy_constructor(quote! {set});
+            set_body = quote! {
+                if let Some((idx, old_ref)) = self.typed().find_lazy_ref(#extractor) {
+                    if old_ref == &set {
+                        return None;
+                    }
+                    #track_modification
+                    let old = std::mem::replace(&mut self.typed_mut().lazy[idx], #ctor);
+                    Some((#unwraper)(old))
+                } else {
+                    #track_modification
+                    self.typed_mut().lazy.push(#ctor);
+                    None
+                }
+            };
+        } else {
+            set_body = quote! {
+                if #ref_expr == &set {
+                    return None;
+                }
+                #track_modification
+                let old = #take_expr;
+                *#mut_expr = set;
+                Some(old)
+            };
         }
-    };
+
+        // Extend: use peekable iterator to avoid Vec allocation.
+        // For lazy fields, look up the set once via find_lazy_ref to avoid repeated scans.
+        extend_body = if is_option {
+            let extractor = field.lazy_extractor_closure();
+            let ctor = field.lazy_constructor(quote! { set });
+            quote! {
+                let mut iter = items.into_iter().peekable();
+                if let Some((idx, existing)) = self.typed().find_lazy_ref(#extractor) {
+                    // Skip items already in the set
+                    loop {
+                        match iter.peek() {
+                            None => return,
+                            Some(item) if existing.contains(item) => { iter.next(); }
+                            Some(_) => break,
+                        }
+                    }
+                    // Found a new item - track and extend using the known index
+                    #track_modification
+                    self.typed_mut().lazy_at_mut(idx, #extractor).extend(iter);
+                } else {
+                    // Set doesn't exist yet - if iterator is empty, nothing to do
+                    if iter.peek().is_none() {
+                        return;
+                    }
+                    #track_modification
+                    let set: #field_type = iter.collect();
+                    self.typed_mut().lazy.push(#ctor);
+                }
+            }
+        } else {
+            quote! {
+                let mut iter = items.into_iter().peekable();
+                // Skip items already in the set until we find a new one
+                loop {
+                    match iter.peek() {
+                        None => return,
+                        Some(item) if #ref_expr.contains(item) => { iter.next(); }
+                        Some(_) => break,
+                    }
+                }
+                // Found a new item - track modification and insert remaining items
+                #track_modification
+                #mut_expr.extend(iter);
+            }
+        };
+    }
 
     quote! {
         #[doc = "Check if the set contains an item"]
@@ -1713,30 +1890,18 @@ fn generate_autoset_ops(field: &FieldInfo) -> TokenStream {
 
         #[doc = "Add an item to the set."]
         #[doc = "Returns true if the item was newly added, false if it already existed."]
+        #[doc = "Only tracks modification if the item is actually added."]
         #[must_use]
         fn #add_name(&mut self, item: #element_type) -> bool {
             #check_access
-            let added = #mut_expr.insert(item);
-            if added {
-                #track_modification
-            }
-            added
+            #add_body
         }
 
         #[doc = "Add multiple items to the set from an iterator."]
         #[doc = "Only tracks modification if at least one item is actually added."]
         fn #extend_name(&mut self, items: impl IntoIterator<Item = #element_type>) {
             #check_access
-            let set = #mut_expr;
-            let mut any_added = false;
-            for item in items {
-                if set.insert(item) {
-                    any_added = true;
-                }
-            }
-            if any_added {
-                #track_modification
-            }
+            #extend_body
         }
 
         #[doc = "Remove an item from the set."]
@@ -1746,12 +1911,11 @@ fn generate_autoset_ops(field: &FieldInfo) -> TokenStream {
             #remove_body
         }
 
-        #[doc = "Remove multiple items from the set."]
-        #[doc = "Only tracks modification if at least one item is actually removed."]
+        #[doc = "Replace the entire set, returning the old set if present."]
+        #[doc = "Only tracks modification if the set actually changes."]
         fn #set_name(&mut self, set: #field_type) -> Option<#field_type>
         {
             #check_access
-            #track_modification
             #set_body
         }
 
@@ -1844,25 +2008,26 @@ fn generate_countermap_ops(field: &FieldInfo) -> TokenStream {
 
     // Generate remove body - for lazy fields, we need to check if the map exists first
     // without allocating it. For inline fields, we can use the mut_expr directly.
-    let remove_body = if is_option {
-        // Lazy: use find_lazy_mut to avoid allocating, only track modification if something was
+    // Only track modification if the key exists (check before mutating).
+    // For transient fields, skip guards since track_modification is a no-op.
+    let remove_body = if field.is_transient() {
+        quote! {
+            #track_modification
+            #mut_expr.remove(key)
+        }
+    } else if is_option {
         let extractor = field.lazy_extractor_closure();
         quote! {
-            let map = self.typed_mut().find_lazy_mut(#extractor)?;
-            let result = map.remove(key);
-            if result.is_some() {
-                #track_modification
-            }
-            result
+            let (idx, val) = self.typed().find_lazy_ref(#extractor)?;
+            val.get(key)?;
+            #track_modification
+            self.typed_mut().lazy_at_mut(idx, #extractor).remove(key)
         }
     } else {
-        // Inline: direct access, only track modification if something was removed
         quote! {
-            let result = #mut_expr.remove(key);
-            if result.is_some() {
-                #track_modification
-            }
-            result
+            self.#get_name(key)?;
+            #track_modification
+            #mut_expr.remove(key)
         }
     };
 
@@ -1890,21 +2055,130 @@ fn generate_countermap_ops(field: &FieldInfo) -> TokenStream {
     // Generate signed-type-specific methods only for i32
     let signed_methods = if is_signed {
         let update_positive_crossing_name = field.infixed_ident("update", "positive_crossing");
+        let update_positive_crossing_body = if field.is_transient() {
+            quote! {
+                #mut_expr.update_positive_crossing(key, delta)
+            }
+        } else {
+            quote! {
+                if delta == 0 {
+                    return false;
+                }
+                #track_modification
+                #mut_expr.update_positive_crossing(key, delta)
+            }
+        };
 
         quote! {
             #[doc = "Update a signed counter by the given delta."]
             #[doc = "Returns true if the count crossed the positive boundary (became positive or non-positive)."]
             #[must_use]
             fn #update_positive_crossing_name(&mut self, key: #key_type, delta: #value_type) -> bool {
                 #check_access
-                #track_modification
-                #mut_expr.update_positive_crossing(key, delta)
+                #update_positive_crossing_body
             }
         }
     } else {
         quote! {}
     };
 
+    let update_count_body = if field.is_transient() {
+        quote! {
+            #mut_expr.update_count(key, delta)
+        }
+    } else {
+        quote! {
+            if delta == 0 {
+                return false;
+            }
+            #track_modification
+            #mut_expr.update_count(key, delta)
+        }
+    };
+
+    let update_counts_body = if field.is_transient() {
+        quote! {
+            let map = #mut_expr;
+            for key in keys {
+                map.update_count(key, delta);
+            }
+        }
+    } else {
+        quote! {
+            if delta == 0 {
+                return;
+            }
+            #track_modification
+            let map = #mut_expr;
+            for key in keys {
+                map.update_count(key, delta);
+            }
+        }
+    };
+
+    let update_and_get_body = if field.is_transient() {
+        quote! {
+            #mut_expr.update_and_get(key, delta)
+        }
+    } else {
+        quote! {
+            if delta == 0 {
+                return self.#get_name(&key).copied().unwrap_or_default();
+            }
+            #track_modification
+            #mut_expr.update_and_get(key, delta)
+        }
+    };
+
+    let update_with_body = if field.is_transient() {
+        quote! {
+            #mut_expr.update_with(key, f)
+        }
+    } else if is_option {
+        let extractor = field.lazy_extractor_closure();
+        let constructor = field.lazy_constructor(quote! {new_map});
+        quote! {
+            let (position, old_value) = if let Some((index, map)) = self.typed().find_lazy_ref(#extractor) {
+                // This copy is very cheap
+                (Some(index), map.get(&key).copied())
+            } else {
+                (None, None)
+            };
+            let new_value = f(old_value);
+            if old_value != new_value {
+                #track_modification
+                match new_value {
+                    Some(value) => {
+                        if let Some(position) = position {
+                            self.typed_mut().lazy_at_mut(position, #extractor).insert(key, value);
+                        } else {
+                            let mut new_map = CounterMap::default();
+                            new_map.insert(key, value);
+                            self.typed_mut().lazy.push(#constructor);
+                        }
+                    }
+                    None => {
+                        // the position must be available, otherwise `f` would have mapped None to None and thus the != check above would have failed.
+                        self.typed_mut().lazy_at_mut(position.unwrap(), #extractor).remove(&key);
+
+                    }
+                }
+            }
+        }
+    } else {
+        quote! {
+            let old = self.#get_name(&key).copied();
+            let new = f(old);
+            if old != new {
+                #track_modification
+                match new {
+                    Some(value) => { #mut_expr.insert(key, value); }
+                    None => { #mut_expr.remove(&key); }
+                }
+            }
+        }
+    };
+
     quote! {
         #[doc = "Get a single entry from the counter map"]
         fn #get_name(&self, key: &#key_type) -> Option<&#value_type> {
@@ -1917,26 +2191,20 @@ fn generate_countermap_ops(field: &FieldInfo) -> TokenStream {
         #[must_use]
         fn #update_count_name(&mut self, key: #key_type, delta: #value_type) -> bool {
             #check_access
-            #track_modification
-            #mut_expr.update_count(key, delta)
+            #update_count_body
         }
 
         #[doc = "Update multiple counters by the given delta."]
         #[doc = "More efficient than calling update_count in a loop."]
         fn #update_counts_name(&mut self, keys: impl Iterator<Item = #key_type>, delta: #value_type) {
             #check_access
-            #track_modification
-            let map = #mut_expr;
-            for key in keys {
-                map.update_count(key, delta);
-            }
+            #update_counts_body
         }
 
         #[doc = "Update a counter by the given delta and return the new value."]
         fn #update_and_get_name(&mut self, key: #key_type, delta: #value_type) -> #value_type {
             #check_access
-            #track_modification
-            #mut_expr.update_and_get(key, delta)
+            #update_and_get_body
         }
 
         #[doc = "Update a counter using a closure that receives the current value"]
@@ -1946,8 +2214,7 @@ fn generate_countermap_ops(field: &FieldInfo) -> TokenStream {
             F: FnOnce(Option<#value_type>) -> Option<#value_type>,
         {
             #check_access
-            #track_modification
-            #mut_expr.update_with(key, f)
+            #update_with_body
         }
 
         #[doc = "Add a new entry, panicking if the entry already exists."]
@@ -2055,28 +2322,58 @@ fn generate_automap_ops(field: &FieldInfo) -> TokenStream {
 
     let take_expression = {
         let take_name = field.take_ident();
-        quote! {self.typed_mut().#take_name();}
+        quote! {self.typed_mut().#take_name()}
     };
 
     // Generate remove body - for lazy fields, avoid allocation if map doesn't exist.
-    // Using ? operator to early-return None if map doesn't exist.
-    let remove_body = if is_option {
+    // Only track modification if the key exists (check before mutating).
+    // For transient fields, skip guards since track_modification is a no-op.
+    let remove_body = if field.is_transient() {
+        quote! {
+            #mut_expr.remove(key)
+        }
+    } else if is_option {
         let extractor = field.lazy_extractor_closure();
         quote! {
-            let map = self.typed_mut().find_lazy_mut(#extractor)?;
-            let result = map.remove(key);
-            if result.is_some() {
-                #track_modification
+            let (idx, val) = self.typed().find_lazy_ref(#extractor)?;
+            val.get(key)?;
+            #track_modification
+            self.typed_mut().lazy_at_mut(idx, #extractor).remove(key)
+        }
+    } else {
+        quote! {
+            if !self.#has_entry_name(key) {
+                return None;
             }
-            result
+            #track_modification
+            #mut_expr.remove(key)
+        }
+    };
+
+    let take_body = if field.is_transient() {
+        quote! {
+            #take_expression
+        }
+    } else if is_option {
+        // For lazy fields, use find_lazy_ref to check existence and emptiness in one scan,
+        // then lazy_take_at to take by known index without re-scanning.
+        let extractor = field.lazy_extractor_closure();
+        let unwraper = field.lazy_unwrap_closure();
+        quote! {
+            let (idx, val) = self.typed().find_lazy_ref(#extractor)?;
+            if val.is_empty() {
+                return None;
+            }
+            #track_modification
+            Some(self.typed_mut().lazy_take_at(idx, #unwraper))
         }
     } else {
         quote! {
-            let result = #mut_expr.remove(key);
-            if result.is_some() {
-                #track_modification
+            if self.#is_empty_name() {
+                return None;
             }
-            result
+            #track_modification
+            #take_expression
         }
     };
 
@@ -2110,14 +2407,10 @@ fn generate_automap_ops(field: &FieldInfo) -> TokenStream {
 
 
         #[doc = "Remove the full map and return it"]
-        #[doc = "Only tracks modification if an entry was actually removed."]
+        #[doc = "Only tracks modification if the map is non-empty."]
         fn #take_name(&mut self) -> Option<#field_type> {
             #check_access
-            let value = #take_expression;
-            if value.is_some() {
-                #track_modification
-            }
-            value
+            #take_body
         }
 
         #[doc = "Iterate over all key-value pairs in the map"]