Repository: vercel/next.js
Pull request: Turbopack: parallel connect children

### What?

parallize connecting children when there are a lot of children. This avoid the long single threaded work.

Closes PACK-5579

Changed files:
filename: turbopack/crates/turbo-tasks-backend/src/backend/mod.rs
patch: @@ -1752,13 +1752,20 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         // suspend in `CleanupOldEdgesOperation`), but that's ok as the task is still dirty and
         // would be executed again.
 
-        if self.task_execution_completed_connect(&mut ctx, task_id, new_children, is_now_immutable)
+        let has_new_children = !new_children.is_empty();
+        if has_new_children
+            && self.task_execution_completed_connect(&mut ctx, task_id, new_children)
         {
             // Task was stale and has been rescheduled
             return true;
         }
 
-        if self.task_execution_completed_finish(&mut ctx, task_id, &mut removed_data) {
+        if self.task_execution_completed_finish(
+            &mut ctx,
+            task_id,
+            &mut removed_data,
+            is_now_immutable,
+        ) {
             // Task was stale and has been rescheduled
             return true;
         }
@@ -1997,8 +2004,9 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         ctx: &mut impl ExecuteContext<'_>,
         task_id: TaskId,
         new_children: FxHashSet<TaskId>,
-        is_now_immutable: bool,
     ) -> bool {
+        debug_assert!(!new_children.is_empty());
+
         let mut task = ctx.task(task_id, TaskDataCategory::All);
         let Some(in_progress) = get!(task, InProgress) else {
             panic!("Task execution completed, but task is not in progress: {task:#?}");
@@ -2039,35 +2047,16 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
             return true;
         }
 
-        // If the task is not stateful and has no mutable children, it does not have a way to be
-        // invalidated and we can mark it as immutable.
-        if is_now_immutable {
-            let _ = task.add(CachedDataItem::Immutable { value: () });
-        }
-
-        let mut queue = AggregationUpdateQueue::new();
-
-        let has_new_children = !new_children.is_empty();
-        if has_new_children {
-            let has_active_count = ctx.should_track_activeness()
-                && get!(task, Activeness).map_or(false, |activeness| activeness.active_counter > 0);
-            connect_children(
-                task_id,
-                &mut task,
-                new_children,
-                &mut queue,
-                has_active_count,
-                ctx.should_track_activeness(),
-            );
-        }
-
-        drop(task);
-
-        if has_new_children {
-            #[cfg(feature = "trace_task_completion")]
-            let _span = tracing::trace_span!("connect new children").entered();
-            queue.execute(ctx);
-        }
+        let has_active_count = ctx.should_track_activeness()
+            && get!(task, Activeness).map_or(false, |activeness| activeness.active_counter > 0);
+        connect_children(
+            ctx,
+            task_id,
+            task,
+            new_children,
+            has_active_count,
+            ctx.should_track_activeness(),
+        );
 
         false
     }
@@ -2077,6 +2066,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         ctx: &mut impl ExecuteContext<'_>,
         task_id: TaskId,
         removed_data: &mut Vec<CachedDataItem>,
+        is_now_immutable: bool,
     ) -> bool {
         let mut task = ctx.task(task_id, TaskDataCategory::All);
         let Some(in_progress) = remove!(task, InProgress) else {
@@ -2107,6 +2097,12 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
             return true;
         }
 
+        // If the task is not stateful and has no mutable children, it does not have a way to be
+        // invalidated and we can mark it as immutable.
+        if is_now_immutable {
+            let _ = task.add(CachedDataItem::Immutable { value: () });
+        }
+
         // Notify in progress cells
         removed_data.extend(task.extract_if(
             CachedDataItemType::InProgressCell,

filename: turbopack/crates/turbo-tasks-backend/src/backend/operation/connect_children.rs
patch: @@ -1,29 +1,31 @@
 use rustc_hash::FxHashSet;
 use smallvec::SmallVec;
-use turbo_tasks::TaskId;
+use turbo_tasks::{
+    TaskId,
+    scope::scope_and_block,
+    util::{good_chunk_size, into_chunks},
+};
 
 use crate::{
     backend::operation::{
-        AggregationUpdateJob, AggregationUpdateQueue, TaskGuard,
-        aggregation_update::InnerOfUppersHasNewFollowersJob, get_aggregation_number, get_uppers,
-        is_aggregating_node,
+        AggregationUpdateJob, AggregationUpdateQueue, ChildExecuteContext, ExecuteContext,
+        Operation, TaskGuard, aggregation_update::InnerOfUppersHasNewFollowersJob,
+        get_aggregation_number, get_uppers, is_aggregating_node,
     },
     data::CachedDataItem,
 };
 
 pub fn connect_children(
+    ctx: &mut impl ExecuteContext<'_>,
     parent_task_id: TaskId,
-    parent_task: &mut impl TaskGuard,
+    mut parent_task: impl TaskGuard,
     new_children: FxHashSet<TaskId>,
-    queue: &mut AggregationUpdateQueue,
     has_active_count: bool,
     should_track_activeness: bool,
 ) {
-    if new_children.is_empty() {
-        return;
-    }
+    debug_assert!(!new_children.is_empty());
 
-    let parent_aggregation = get_aggregation_number(parent_task);
+    let parent_aggregation = get_aggregation_number(&parent_task);
 
     for &new_child in new_children.iter() {
         parent_task.add_new(CachedDataItem::Child {
@@ -35,39 +37,117 @@ pub fn connect_children(
     let new_follower_ids: SmallVec<_> = new_children.into_iter().collect();
 
     let aggregating_node = is_aggregating_node(parent_aggregation);
-    let upper_ids = (!aggregating_node).then(|| get_uppers(&*parent_task));
+    let upper_ids = (!aggregating_node).then(|| get_uppers(&parent_task));
 
-    if let Some(upper_ids) = upper_ids {
-        // Parent is a leaf node, the children are followers of it now.
-        if !upper_ids.is_empty() {
-            queue.push(
-                InnerOfUppersHasNewFollowersJob {
-                    upper_ids,
-                    new_follower_ids: new_follower_ids.clone(),
-                }
-                .into(),
-            );
-        }
-        // We need to decrease the active count because we temporarily increased it during
-        // connect_child. We need to increase the active count when the parent has active
-        // count, because it's added as follower.
-        if should_track_activeness && !has_active_count {
-            queue.push(AggregationUpdateJob::DecreaseActiveCounts {
-                task_ids: new_follower_ids,
-            })
-        }
-    } else {
-        // Parent is an aggregating node. We run the normal code to connect the children.
-        queue.push(AggregationUpdateJob::InnerOfUpperHasNewFollowers {
-            upper_id: parent_task_id,
-            new_follower_ids: new_follower_ids.clone(),
-        });
-        // We need to decrease the active count because we temporarily increased it during
-        // connect_child.
-        if should_track_activeness {
+    drop(parent_task);
+
+    fn process_new_children(
+        ctx: &mut impl ExecuteContext<'_>,
+        new_follower_ids: SmallVec<[TaskId; 4]>,
+        upper_ids: Option<SmallVec<[TaskId; 4]>>,
+        parent_task_id: TaskId,
+        has_active_count: bool,
+        should_track_activeness: bool,
+    ) {
+        debug_assert!(!new_follower_ids.is_empty());
+
+        let mut queue = AggregationUpdateQueue::new();
+
+        if let Some(upper_ids) = upper_ids {
+            // We need to add new followers when there are upper ids as the parent is a leaf node
+            // and new children are new followers.
+            let add_followers = !upper_ids.is_empty();
+
+            // And we need decrease the active count when parent doesn't have an active count as the
+            // active count was temporarily increased during connect_child. We need to
+            // increase the active count when the parent has active count, because it's
+            // added as follower.
+            let decrease_active_count = should_track_activeness && !has_active_count;
+
+            // We special case the situation when we need to do both operations to avoid
+            // cloning the new follower ids unnecessarily.
+            if decrease_active_count && add_followers {
+                queue.push(
+                    InnerOfUppersHasNewFollowersJob {
+                        upper_ids,
+                        new_follower_ids: new_follower_ids.clone(),
+                    }
+                    .into(),
+                );
+                queue.push(AggregationUpdateJob::DecreaseActiveCounts {
+                    task_ids: new_follower_ids,
+                })
+            } else if decrease_active_count {
+                queue.push(AggregationUpdateJob::DecreaseActiveCounts {
+                    task_ids: new_follower_ids,
+                })
+            } else if add_followers {
+                queue.push(
+                    InnerOfUppersHasNewFollowersJob {
+                        upper_ids,
+                        new_follower_ids,
+                    }
+                    .into(),
+                );
+            }
+        } else if should_track_activeness {
+            // Parent is an aggregating node. We run the normal code to connect the children.
+            queue.push(AggregationUpdateJob::InnerOfUpperHasNewFollowers {
+                upper_id: parent_task_id,
+                new_follower_ids: new_follower_ids.clone(),
+            });
+            // We need to decrease the active count because we temporarily increased it during
+            // connect_child.
             queue.push(AggregationUpdateJob::DecreaseActiveCounts {
                 task_ids: new_follower_ids,
             });
+        } else {
+            // Parent is an aggregating node. We run the normal code to connect the children.
+            queue.push(AggregationUpdateJob::InnerOfUpperHasNewFollowers {
+                upper_id: parent_task_id,
+                new_follower_ids,
+            });
+        }
+
+        {
+            #[cfg(feature = "trace_task_completion")]
+            let _span = tracing::trace_span!("connect new children").entered();
+            queue.execute(ctx);
         }
     }
+
+    const MIN_CHILDREN_FOR_PARALLEL: usize = 10000;
+
+    let len = new_follower_ids.len();
+    if len >= MIN_CHILDREN_FOR_PARALLEL {
+        let new_follower_ids = new_follower_ids.into_vec();
+        let chunk_size = good_chunk_size(len);
+        let _ = scope_and_block(len.div_ceil(chunk_size), |scope| {
+            for chunk in into_chunks(new_follower_ids, chunk_size) {
+                let upper_ids = &upper_ids;
+                let child_ctx = ctx.child_context();
+                scope.spawn(move || {
+                    let mut ctx = child_ctx.create();
+                    let new_follower_ids = chunk.collect::<SmallVec<[_; 4]>>();
+                    process_new_children(
+                        &mut ctx,
+                        new_follower_ids,
+                        upper_ids.clone(),
+                        parent_task_id,
+                        has_active_count,
+                        should_track_activeness,
+                    );
+                });
+            }
+        });
+    } else {
+        process_new_children(
+            ctx,
+            new_follower_ids,
+            upper_ids,
+            parent_task_id,
+            has_active_count,
+            should_track_activeness,
+        );
+    }
 }

filename: turbopack/crates/turbo-tasks-backend/src/backend/operation/mod.rs
patch: @@ -48,6 +48,9 @@ enum TransactionState<'a, 'tx, B: BackingStorage> {
 }
 
 pub trait ExecuteContext<'e>: Sized {
+    fn child_context<'l, 'r>(&'r self) -> impl ChildExecuteContext<'l> + use<'e, 'l, Self>
+    where
+        'e: 'l;
     fn session_id(&self) -> SessionId;
     fn task(&mut self, task_id: TaskId, category: TaskDataCategory) -> impl TaskGuard + 'e;
     fn is_once_task(&self, task_id: TaskId) -> bool;
@@ -69,6 +72,10 @@ pub trait ExecuteContext<'e>: Sized {
     fn should_track_activeness(&self) -> bool;
 }
 
+pub trait ChildExecuteContext<'e>: Send + Sized {
+    fn create(self) -> impl ExecuteContext<'e>;
+}
+
 pub struct ExecuteContextImpl<'e, 'tx, B: BackingStorage>
 where
     Self: 'e,
@@ -156,6 +163,16 @@ impl<'e, 'tx, B: BackingStorage> ExecuteContext<'e> for ExecuteContextImpl<'e, '
 where
     'tx: 'e,
 {
+    fn child_context<'l, 'r>(&'r self) -> impl ChildExecuteContext<'l> + use<'e, 'tx, 'l, B>
+    where
+        'e: 'l,
+    {
+        ChildExecuteContextImpl {
+            backend: self.backend,
+            turbo_tasks: self.turbo_tasks,
+        }
+    }
+
     fn session_id(&self) -> SessionId {
         self.backend.session_id()
     }
@@ -297,6 +314,22 @@ where
     }
 }
 
+struct ChildExecuteContextImpl<'e, B: BackingStorage> {
+    backend: &'e TurboTasksBackendInner<B>,
+    turbo_tasks: &'e dyn TurboTasksBackendApi<TurboTasksBackend<B>>,
+}
+
+impl<'e, B: BackingStorage> ChildExecuteContext<'e> for ChildExecuteContextImpl<'e, B> {
+    fn create(self) -> impl ExecuteContext<'e> {
+        ExecuteContextImpl {
+            backend: self.backend,
+            turbo_tasks: self.turbo_tasks,
+            _operation_guard: None,
+            transaction: TransactionState::None,
+        }
+    }
+}
+
 pub trait TaskGuard: Debug {
     fn id(&self) -> TaskId;
     #[must_use]