allow to manually mark tasks as root

vercel · Jan 30, 2025 · ede728b · ede728b
1 parent 0ba2230
commit ede728b
Show file tree

Hide file tree

Showing 5 changed files with 59 additions and 4 deletions.
diff --git a/turbopack/crates/turbo-tasks-backend/src/backend/mod.rs b/turbopack/crates/turbo-tasks-backend/src/backend/mod.rs
@@ -1723,6 +1723,23 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         }
     }
 
+    fn set_own_task_aggregation_number(
+        &self,
+        task: TaskId,
+        aggregation_number: u32,
+        turbo_tasks: &dyn TurboTasksBackendApi<TurboTasksBackend<B>>,
+    ) {
+        let mut ctx = self.execute_context(turbo_tasks);
+        AggregationUpdateQueue::run(
+            AggregationUpdateJob::UpdateAggregationNumber {
+                task_id: task,
+                base_aggregation_number: aggregation_number,
+                distance: None,
+            },
+            &mut ctx,
+        );
+    }
+
     fn connect_task(
         &self,
         task: TaskId,
@@ -2014,6 +2031,16 @@ impl<B: BackingStorage> Backend for TurboTasksBackend<B> {
         self.0.mark_own_task_as_finished(task_id, turbo_tasks);
     }
 
+    fn set_own_task_aggregation_number(
+        &self,
+        task: TaskId,
+        aggregation_number: u32,
+        turbo_tasks: &dyn TurboTasksBackendApi<Self>,
+    ) {
+        self.0
+            .set_own_task_aggregation_number(task, aggregation_number, turbo_tasks);
+    }
+
     fn mark_own_task_as_session_dependent(
         &self,
         task: TaskId,

diff --git a/turbopack/crates/turbo-tasks-testing/src/lib.rs b/turbopack/crates/turbo-tasks-testing/src/lib.rs
@@ -288,6 +288,10 @@ impl TurboTasksApi for VcStorage {
         // no-op
     }
 
+    fn set_own_task_aggregation_number(&self, _task: TaskId, _aggregation_number: u32) {
+        // no-op
+    }
+
     fn detached_for_testing(
         &self,
         _f: std::pin::Pin<Box<dyn Future<Output = Result<()>> + Send + 'static>>,

diff --git a/turbopack/crates/turbo-tasks/src/backend.rs b/turbopack/crates/turbo-tasks/src/backend.rs
@@ -589,6 +589,15 @@ pub trait Backend: Sync + Send {
         // Do nothing by default
     }
 
+    fn set_own_task_aggregation_number(
+        &self,
+        _task: TaskId,
+        _aggregation_number: u32,
+        _turbo_tasks: &dyn TurboTasksBackendApi<Self>,
+    ) {
+        // Do nothing by default
+    }
+
     fn mark_own_task_as_session_dependent(
         &self,
         _task: TaskId,

diff --git a/turbopack/crates/turbo-tasks/src/lib.rs b/turbopack/crates/turbo-tasks/src/lib.rs
@@ -102,10 +102,11 @@ pub use join_iter_ext::{JoinIterExt, TryFlatJoinIterExt, TryJoinIterExt};
 pub use key_value_pair::KeyValuePair;
 pub use magic_any::MagicAny;
 pub use manager::{
-    dynamic_call, emit, mark_finished, mark_session_dependent, mark_stateful, prevent_gc, run_once,
-    run_once_with_reason, spawn_blocking, spawn_thread, trait_call, turbo_tasks, turbo_tasks_scope,
-    CurrentCellRef, ReadConsistency, TaskPersistence, TurboTasks, TurboTasksApi,
-    TurboTasksBackendApi, TurboTasksBackendApiExt, TurboTasksCallApi, Unused, UpdateInfo,
+    dynamic_call, emit, mark_finished, mark_root, mark_session_dependent, mark_stateful,
+    prevent_gc, run_once, run_once_with_reason, spawn_blocking, spawn_thread, trait_call,
+    turbo_tasks, turbo_tasks_scope, CurrentCellRef, ReadConsistency, TaskPersistence, TurboTasks,
+    TurboTasksApi, TurboTasksBackendApi, TurboTasksBackendApiExt, TurboTasksCallApi, Unused,
+    UpdateInfo,
 };
 pub use output::OutputContent;
 pub use raw_vc::{CellId, RawVc, ReadRawVcFuture, ResolveTypeError};

diff --git a/turbopack/crates/turbo-tasks/src/manager.rs b/turbopack/crates/turbo-tasks/src/manager.rs
@@ -164,6 +164,7 @@ pub trait TurboTasksApi: TurboTasksCallApi + Sync + Send {
     fn read_own_task_cell(&self, task: TaskId, index: CellId) -> Result<TypedCellContent>;
     fn update_own_task_cell(&self, task: TaskId, index: CellId, content: CellContent);
     fn mark_own_task_as_finished(&self, task: TaskId);
+    fn set_own_task_aggregation_number(&self, task: TaskId, aggregation_number: u32);
     fn mark_own_task_as_session_dependent(&self, task: TaskId);
 
     fn connect_task(&self, task: TaskId);
@@ -1393,6 +1394,11 @@ impl<B: Backend + 'static> TurboTasksApi for TurboTasks<B> {
         self.backend.mark_own_task_as_finished(task, self);
     }
 
+    fn set_own_task_aggregation_number(&self, task: TaskId, aggregation_number: u32) {
+        self.backend
+            .set_own_task_aggregation_number(task, aggregation_number, self);
+    }
+
     fn mark_own_task_as_session_dependent(&self, task: TaskId) {
         self.backend.mark_own_task_as_session_dependent(task, self);
     }
@@ -1696,6 +1702,14 @@ pub fn mark_session_dependent() {
     });
 }
 
+/// Marks the current task as finished. This excludes it from waiting for
+/// strongly consistency.
+pub fn mark_root() {
+    with_turbo_tasks(|tt| {
+        tt.set_own_task_aggregation_number(current_task("turbo_tasks::mark_root()"), u32::MAX)
+    });
+}
+
 /// Marks the current task as finished. This excludes it from waiting for
 /// strongly consistency.
 pub fn mark_finished() {