From 72d3048152dcaccfa9b7f907f7e95de9851e2a92 Mon Sep 17 00:00:00 2001
From: Tobias Koppers <tobias.koppers@googlemail.com>
Date: Wed, 9 Oct 2024 10:55:00 +0200
Subject: [PATCH] fix race condition when scheduling dirty tasks

---
 .../turbo-tasks-backend/src/backend/mod.rs    | 10 ++------
 .../backend/operation/aggregation_update.rs   | 24 ++++++++++++-------
 .../crates/turbo-tasks-backend/src/data.rs    |  6 +++++
 3 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/turbopack/crates/turbo-tasks-backend/src/backend/mod.rs b/turbopack/crates/turbo-tasks-backend/src/backend/mod.rs
index 5814a996ef65b..4de4664f0c021 100644
--- a/turbopack/crates/turbo-tasks-backend/src/backend/mod.rs
+++ b/turbopack/crates/turbo-tasks-backend/src/backend/mod.rs
@@ -353,14 +353,8 @@ impl TurboTasksBackendInner {
                 task = ctx.task(task_id, TaskDataCategory::All);
             }
 
-            let is_dirty = get!(task, Dirty)
-                .map(|dirty_state| {
-                    dirty_state
-                        .clean_in_session
-                        .map(|clean_in_session| clean_in_session != self.session_id)
-                        .unwrap_or(true)
-                })
-                .unwrap_or(false);
+            let is_dirty =
+                get!(task, Dirty).map_or(false, |dirty_state| dirty_state.get(self.session_id));
 
             // Check the dirty count of the root node
             let dirty_tasks = get!(task, AggregatedDirtyContainerCount)
diff --git a/turbopack/crates/turbo-tasks-backend/src/backend/operation/aggregation_update.rs b/turbopack/crates/turbo-tasks-backend/src/backend/operation/aggregation_update.rs
index 3b30770466ffe..bfb91b6dae1b6 100644
--- a/turbopack/crates/turbo-tasks-backend/src/backend/operation/aggregation_update.rs
+++ b/turbopack/crates/turbo-tasks-backend/src/backend/operation/aggregation_update.rs
@@ -183,6 +183,15 @@ impl AggregatedDataUpdate {
         } = self;
         let mut result = Self::default();
         if let Some((dirty_container_id, count)) = dirty_container_update {
+            // When a dirty container count is increased and the task is considered as active
+            // `AggregateRoot` we need to schedule the dirty tasks in the new dirty container
+            let current_session_update = count.get(session_id);
+            if current_session_update > 0 && task.has_key(&CachedDataItemKey::AggregateRoot {}) {
+                queue.push(AggregationUpdateJob::FindAndScheduleDirty {
+                    task_ids: vec![*dirty_container_id],
+                })
+            }
+
             let mut aggregated_update = Default::default();
             update!(
                 task,
@@ -195,12 +204,7 @@ impl AggregatedDataUpdate {
                     (!new.is_default()).then_some(new)
                 }
             );
-            let current_session_update = aggregated_update.get(session_id);
-            if current_session_update > 0 && task.has_key(&CachedDataItemKey::AggregateRoot {}) {
-                queue.push(AggregationUpdateJob::FindAndScheduleDirty {
-                    task_ids: vec![*dirty_container_id],
-                })
-            }
+
             let dirty_state = get!(task, Dirty).copied();
             let task_id = task.id();
             update!(task, AggregatedDirtyContainerCount, |old: Option<
@@ -563,20 +567,22 @@ impl AggregationUpdateQueue {
         }
         if let Some(task_id) = popped {
             let mut task = ctx.task(task_id, TaskDataCategory::Meta);
-            #[allow(clippy::collapsible_if, reason = "readablility")]
-            if task.has_key(&CachedDataItemKey::Dirty {}) {
+            let session_id = ctx.session_id();
+            let dirty = get!(task, Dirty).map_or(false, |d| d.get(session_id));
+            if dirty {
                 let description = ctx.backend.get_task_desc_fn(task_id);
                 if task.add(CachedDataItem::new_scheduled(description)) {
                     ctx.turbo_tasks.schedule(task_id);
                 }
             }
             if is_aggregating_node(get_aggregation_number(&task)) {
+                // TODO if it has an `AggregateRoot` we can skip visiting the nested nodes since
+                // this would already be scheduled by the `AggregateRoot`
                 if !task.has_key(&CachedDataItemKey::AggregateRoot {}) {
                     task.insert(CachedDataItem::AggregateRoot {
                         value: RootState::new(ActiveType::CachedActiveUntilClean, task_id),
                     });
                 }
-                let session_id = ctx.session_id();
                 let dirty_containers: Vec<_> = get_many!(task, AggregatedDirtyContainer { task } count if count.get(session_id) > 0 => task);
                 if !dirty_containers.is_empty() {
                     self.push(AggregationUpdateJob::FindAndScheduleDirty {
diff --git a/turbopack/crates/turbo-tasks-backend/src/data.rs b/turbopack/crates/turbo-tasks-backend/src/data.rs
index 2224f81ebf3f2..5ff4aa12241ec 100644
--- a/turbopack/crates/turbo-tasks-backend/src/data.rs
+++ b/turbopack/crates/turbo-tasks-backend/src/data.rs
@@ -90,6 +90,12 @@ pub struct DirtyState {
     pub clean_in_session: Option<SessionId>,
 }
 
+impl DirtyState {
+    pub fn get(&self, session: SessionId) -> bool {
+        self.clean_in_session != Some(session)
+    }
+}
+
 fn add_with_diff(v: &mut i32, u: i32) -> i32 {
     let old = *v;
     *v += u;