Repository: vercel/next.js
Pull request: turbo-tasks-backend: batch find_and_schedule_dirty using for_each_task_meta

### What?

Batch-process `find_and_schedule_dirty` in `aggregation_update.rs` by collecting all queued jobs (up to `FIND_AND_SCHEDULE_BATCH_SIZE` = 10 000) into a `SmallVec` and pre-fetching their task metadata with a batched `ctx.for_each_task_meta(...)` call.

### Why?

`find_and_schedule` can accumulate thousands of tasks during invalidation cascades. The previous implementation issued one `ctx.task(...)` call per task inside `process()`, serializing backing-storage fetches one at a time.

`ctx.for_each_task_meta` triggers a batched fetch of all task metadata from the backing store: keys are sorted by hash for cache-friendly sequential access to the storage layer. The callback is invoked per-task once data is ready, with the task guard handed directly to the callback — no second lock acquisition needed.

The batch limit is set to 10 000, because find-and-schedule jobs are cheap (metadata read + optional schedule) compared to aggregation-update jobs, so yielding less often is safe and beneficial.

### How?

**`fn process` — `find_and_schedule` branch:**
- Replace the one-at-a-time loop with a single `drain(..FIND_AND_SCHEDULE_BATCH_SIZE)` that collects up to 10 000 `FindAndScheduleJob` structs into a `SmallVec`, then calls `find_and_schedule_dirty` once with the whole batch.

**`fn find_and_schedule_dirty`:**
- Change parameter from `task_id: TaskId` to `jobs: SmallVec<[FindAndScheduleJob; 4]>`.
- Call `ctx.for_each_task_meta(...)` to batch-prefetch all task metadata (sorted by hash for cache-friendly access) and process each task in the callback.

**`prepare_tasks_with_callback` fix:**
- Release `TaskLockCounter` *before* calling `prepared_task_callback` instead of after. This ensures the counter is 0 when the callback runs, so callbacks that drop their task guard and then call `ctx.task()` (like `find_and_schedule_dirty_internal` → `ctx.schedule()`) no longer trigger the "Concurrent task lock acquisition detected" panic in debug builds.
- `for_each_task` now uses `acquire()` instead of `reacquire()` since counter is guaranteed 0 at callback entry. `reacquire()` is removed.

**Cleanup:**
- Use `FxHashMap` (already imported) instead of `std::collections::HashMap`.
- Combine consecutive `#[cfg(trace_find_and_schedule)]` let bindings for clarity.
- Add `FIND_AND_SCHEDULE_BATCH_SIZE = 10_000` as a self-contained constant (not derived from `MAX_COUNT_BEFORE_YIELD`).

Changed files:
filename: turbopack/crates/turbo-tasks-backend/src/backend/operation/aggregation_update.rs
patch: @@ -43,6 +43,10 @@ type FxRingSet<T> = RingSet<T, FxBuildHasher>;
 
 pub const LEAF_NUMBER: u32 = 16;
 const MAX_COUNT_BEFORE_YIELD: usize = 1000;
+/// Batch size for find_and_schedule processing. These jobs are cheaper than aggregation
+/// updates (they only read task metadata and optionally schedule a task), so we can process more
+/// of them per `process()` call before yielding.
+const FIND_AND_SCHEDULE_BATCH_SIZE: usize = 10000;
 const MAX_UPPERS_FOLLOWER_PRODUCT: usize = 31;
 
 type TaskIdVec = SmallVec<[TaskId; 4]>;
@@ -1237,22 +1241,13 @@ impl AggregationUpdateQueue {
             self.optimize_task(ctx, task_id);
             false
         } else if !self.find_and_schedule.is_empty() {
-            let mut remaining = MAX_COUNT_BEFORE_YIELD;
-            while remaining > 0 {
-                if let Some(FindAndScheduleJob {
-                    task_id,
-                    #[cfg(feature = "trace_find_and_schedule")]
-                    span,
-                }) = self.find_and_schedule.pop_front()
-                {
-                    #[cfg(feature = "trace_find_and_schedule")]
-                    let _guard = span.map(|s| s.entered());
-                    self.find_and_schedule_dirty(task_id, ctx);
-                    remaining -= 1;
-                } else {
-                    break;
-                }
-            }
+            let count = self
+                .find_and_schedule
+                .len()
+                .min(FIND_AND_SCHEDULE_BATCH_SIZE);
+            let jobs: SmallVec<[FindAndScheduleJob; 4]> =
+                self.find_and_schedule.drain(..count).collect();
+            self.find_and_schedule_dirty(jobs, ctx);
             false
         } else {
             true
@@ -1422,23 +1417,32 @@ impl AggregationUpdateQueue {
         }
     }
 
-    /// Schedules the task if it's dirty.
+    /// Schedules the tasks if they're dirty.
     ///
     /// Only used when activeness is tracked.
-    fn find_and_schedule_dirty(&mut self, task_id: TaskId, ctx: &mut impl ExecuteContext) {
+    fn find_and_schedule_dirty(
+        &mut self,
+        jobs: SmallVec<[FindAndScheduleJob; 4]>,
+        ctx: &mut impl ExecuteContext,
+    ) {
         #[cfg(feature = "trace_find_and_schedule")]
-        let _span = trace_span!(
-            "find and schedule",
-            %task_id,
-            name = ctx.get_task_description(task_id)
-        )
-        .entered();
-        let task = ctx.task(
-            task_id,
-            // For performance reasons this should stay `Meta` and not `All`
-            TaskDataCategory::Meta,
-        );
-        self.find_and_schedule_dirty_internal(task_id, task, ctx);
+        let mut spans: FxHashMap<TaskId, Option<Span>> = jobs
+            .iter()
+            .map(|job| (job.task_id, job.span.clone()))
+            .collect();
+        // For performance reasons this should stay `Meta` and not `All`
+        ctx.for_each_task_meta(jobs.into_iter().map(|job| job.task_id), |task, ctx| {
+            let task_id = task.id();
+            // Enter the enqueue-time span and create a per-task child span with the
+            // task description. Both guards must live until the end of the closure.
+            #[cfg(feature = "trace_find_and_schedule")]
+            let _trace = (
+                spans.remove(&task_id).flatten().map(|s| s.entered()),
+                trace_span!("find and schedule", %task_id, name = task.get_task_description())
+                    .entered(),
+            );
+            self.find_and_schedule_dirty_internal(task_id, task, ctx);
+        });
     }
 
     fn find_and_schedule_dirty_internal(

filename: turbopack/crates/turbo-tasks-backend/src/backend/operation/mod.rs
patch: @@ -137,13 +137,6 @@ impl TaskLockCounter {
         }
     }
 
-    /// Increment the count by 1 without checking for concurrent access.
-    /// Used when the caller knows another guard already validated exclusive access.
-    fn reacquire(&self) {
-        #[cfg(debug_assertions)]
-        self.0.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
-    }
-
     /// Decrement the count by 1.
     fn release(&self) {
         #[cfg(debug_assertions)]
@@ -297,10 +290,10 @@ impl<'e, B: BackingStorage> ExecuteContextImpl<'e, B> {
                 tasks_to_restore_for_meta_indicies.push(i);
                 ready = false;
             }
+            self.task_lock_counter.release();
             if ready {
                 prepared_task_callback(self, task_id, category, task);
             }
-            self.task_lock_counter.release();
         }
         if tasks_to_restore_for_meta.is_empty() && tasks_to_restore_for_data.is_empty() {
             return;
@@ -378,8 +371,8 @@ impl<'e, B: BackingStorage> ExecuteContextImpl<'e, B> {
                 task.restore_from(storage, TaskDataCategory::Meta);
                 task.flags.set_restored(TaskDataCategory::Meta);
             }
-            prepared_task_callback(self, task_id, category, task);
             self.task_lock_counter.release();
+            prepared_task_callback(self, task_id, category, task);
             if let Some(task_type) = task_type {
                 // Insert into the task cache to avoid future lookups
                 self.backend.task_cache.entry(task_type).or_insert(task_id);
@@ -463,9 +456,9 @@ impl<'e, B: BackingStorage> ExecuteContext<'e> for ExecuteContextImpl<'e, B> {
     ) {
         let task_lock_counter = self.task_lock_counter.clone();
         self.prepare_tasks_with_callback(task_ids, true, |this, task_id, _category, task| {
-            // The prepare_tasks_with_callback already checked for concurrent access
-            // but will also decrement, so we re-increment here since Drop will decrement.
-            task_lock_counter.reacquire();
+            // prepare_tasks_with_callback releases the counter before calling this callback,
+            // so the counter is 0 here. Acquire for the TaskGuardImpl that will release on Drop.
+            task_lock_counter.acquire();
 
             let guard = TaskGuardImpl {
                 task,