fix scheduling after persistent caching

vercel · Jan 31, 2025 · 6e384eb · 6e384eb
1 parent bd02063
commit 6e384eb
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 11 deletions.
diff --git a/turbopack/crates/turbo-tasks-backend/src/backend/mod.rs b/turbopack/crates/turbo-tasks-backend/src/backend/mod.rs
@@ -462,14 +462,18 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
                 .unwrap_or_default()
                 .get(self.session_id);
             if dirty_tasks > 0 || is_dirty {
-                let root = get!(task, Activeness);
+                let root = get_mut!(task, Activeness);
                 let mut task_ids_to_schedule: Vec<_> = Vec::new();
                 // When there are dirty task, subscribe to the all_clean_event
                 let root = if let Some(root) = root {
+                    // This makes sure all tasks stay active and this task won't stale.
+                    // active_until_clean is automatically removed when this
+                    // task is clean.
+                    root.set_active_until_clean();
                     root
                 } else {
                     // If we don't have a root state, add one. This also makes sure all tasks stay
-                    // active and this task won't stale. CachedActiveUntilClean
+                    // active and this task won't stale. active_until_clean
                     // is automatically removed when this task is clean.
                     get_mut_or_insert_with!(task, Activeness, || ActivenessState::new(task_id))
                         .set_active_until_clean();
@@ -482,9 +486,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
                             task
                         }
                     );
-                    if is_dirty {
-                        task_ids_to_schedule.push(task_id);
-                    }
+                    task_ids_to_schedule.push(task_id);
                     get!(task, Activeness).unwrap()
                 };
                 let listener = root.all_clean_event.listen_with_note(move || {

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
@@ -1177,17 +1177,18 @@ impl AggregationUpdateQueue {
                 ctx.schedule(task_id);
             }
         }
-        if is_aggregating_node(get_aggregation_number(&task)) {
+        let aggregation_number = get_aggregation_number(&task);
+        if is_aggregating_node(aggregation_number) {
             // if it has `Activeness` we can skip visiting the nested nodes since
             // this would already be scheduled by the `Activeness`
-            if !task.has_key(&CachedDataItemKey::Activeness {}) {
+            let is_active_until_clean =
+                get!(task, Activeness).is_some_and(|a| a.active_until_clean);
+            if !is_active_until_clean {
                 let dirty_containers: Vec<_> = get_many!(task, AggregatedDirtyContainer { task } count if count.get(session_id) > 0 => task);
                 if !dirty_containers.is_empty() || dirty {
-                    let mut activeness_state = ActivenessState::new(task_id);
+                    let activeness_state =
+                        get_mut_or_insert_with!(task, Activeness, || ActivenessState::new(task_id));
                     activeness_state.set_active_until_clean();
-                    task.insert(CachedDataItem::Activeness {
-                        value: activeness_state,
-                    });
                     drop(task);
 
                     self.extend_find_and_schedule_dirty(dirty_containers);