simplify session dependent tasks (#91729)

Author: lukesandberg

turbopack/crates/turbo-tasks-backend/src/backend/mod.rs

@@ -363,6 +363,37 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         self.options.dependency_tracking
     }
 
+    /// Sets the initial aggregation number for a newly created task. Root tasks get `u32::MAX`
+    /// to stay at the top. Session-dependent tasks get a high (but not max) aggregation number
+    /// because they change on every session restore, behaving like dirty leaf nodes — keeping
+    /// them near the leaves prevents long dirty-propagation chains through intermediate
+    /// aggregated nodes.
+    fn set_initial_aggregation_number(
+        &self,
+        task_id: TaskId,
+        is_root: bool,
+        is_session_dependent: bool,
+        ctx: &mut impl ExecuteContext<'_>,
+    ) {
+        let base_aggregation_number = if is_root {
+            u32::MAX
+        } else if is_session_dependent && self.should_track_dependencies() {
+            const SESSION_DEPENDENT_AGGREGATION_NUMBER: u32 = u32::MAX >> 2;
+            SESSION_DEPENDENT_AGGREGATION_NUMBER
+        } else {
+            return;
+        };
+
+        AggregationUpdateQueue::run(
+            AggregationUpdateJob::UpdateAggregationNumber {
+                task_id,
+                base_aggregation_number,
+                distance: None,
+            },
+            ctx,
+        );
+    }
+
     fn should_track_activeness(&self) -> bool {
         self.options.active_tracking
     }
@@ -470,21 +501,6 @@ struct TaskExecutionCompletePrepareResult {
 
 // Operations
 impl<B: BackingStorage> TurboTasksBackendInner<B> {
-    fn connect_child(
-        &self,
-        parent_task: Option<TaskId>,
-        child_task: TaskId,
-        task_type: Option<ArcOrOwned<CachedTaskType>>,
-        turbo_tasks: &dyn TurboTasksBackendApi<TurboTasksBackend<B>>,
-    ) {
-        operation::ConnectChildOperation::run(
-            parent_task,
-            child_task,
-            task_type,
-            self.execute_context(turbo_tasks),
-        );
-    }
-
     fn try_read_task_output(
         self: &Arc<Self>,
         task_id: TaskId,
@@ -1513,24 +1529,23 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         turbo_tasks: &dyn TurboTasksBackendApi<TurboTasksBackend<B>>,
     ) -> TaskId {
         let is_root = task_type.native_fn.is_root;
-
+        let is_session_dependent = task_type.native_fn.is_session_dependent;
+        // Create a single ExecuteContext for both lookup and connect_child
+        let mut ctx = self.execute_context(turbo_tasks);
         // First check if the task exists in the cache which only uses a read lock
         // .map(|r| *r) copies the TaskId and drops the DashMap Ref (releasing the read lock)
         // before ConnectChildOperation::run, which may re-enter task_cache with a write lock.
         if let Some(task_id) = self.task_cache.get(&task_type).map(|r| *r) {
             self.track_cache_hit(&task_type);
-            self.connect_child(
+            operation::ConnectChildOperation::run(
                 parent_task,
                 task_id,
                 Some(ArcOrOwned::Owned(task_type)),
-                turbo_tasks,
+                ctx,
             );
             return task_id;
         }
 
-        // Create a single ExecuteContext for both lookup and connect_child
-        let mut ctx = self.execute_context(turbo_tasks);
-
         let mut is_new = false;
         let (task_id, task_type) = if let Some(task_id) = ctx.task_by_type(&task_type) {
             // Task exists in backing storage
@@ -1573,15 +1588,8 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
             };
             (task_id, task_type)
         };
-        if is_new && is_root {
-            AggregationUpdateQueue::run(
-                AggregationUpdateJob::UpdateAggregationNumber {
-                    task_id,
-                    base_aggregation_number: u32::MAX,
-                    distance: None,
-                },
-                &mut ctx,
-            );
+        if is_new {
+            self.set_initial_aggregation_number(task_id, is_root, is_session_dependent, &mut ctx);
         }
         // Reuse the same ExecuteContext for connect_child
         operation::ConnectChildOperation::run(parent_task, task_id, Some(task_type), ctx);
@@ -1596,6 +1604,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         turbo_tasks: &dyn TurboTasksBackendApi<TurboTasksBackend<B>>,
     ) -> TaskId {
         let is_root = task_type.native_fn.is_root;
+        let is_session_dependent = task_type.native_fn.is_session_dependent;
 
         if let Some(parent_task) = parent_task
             && !parent_task.is_transient()
@@ -1606,16 +1615,17 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
                 /* cell_id */ None,
             );
         }
+        let mut ctx = self.execute_context(turbo_tasks);
         // First check if the task exists in the cache which only uses a read lock.
         // .map(|r| *r) copies the TaskId and drops the DashMap Ref (releasing the read lock)
         // before ConnectChildOperation::run, which may re-enter task_cache with a write lock.
         if let Some(task_id) = self.task_cache.get(&task_type).map(|r| *r) {
             self.track_cache_hit(&task_type);
-            self.connect_child(
+            operation::ConnectChildOperation::run(
                 parent_task,
                 task_id,
                 Some(ArcOrOwned::Owned(task_type)),
-                turbo_tasks,
+                ctx,
             );
             return task_id;
         }
@@ -1625,11 +1635,11 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
                 let task_id = *e.get();
                 drop(e);
                 self.track_cache_hit(&task_type);
-                self.connect_child(
+                operation::ConnectChildOperation::run(
                     parent_task,
                     task_id,
                     Some(ArcOrOwned::Owned(task_type)),
-                    turbo_tasks,
+                    ctx,
                 );
                 task_id
             }
@@ -1639,23 +1649,18 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
                 e.insert(task_type.clone(), task_id);
                 self.track_cache_miss(&task_type);
 
-                if is_root {
-                    let mut ctx = self.execute_context(turbo_tasks);
-                    AggregationUpdateQueue::run(
-                        AggregationUpdateJob::UpdateAggregationNumber {
-                            task_id,
-                            base_aggregation_number: u32::MAX,
-                            distance: None,
-                        },
-                        &mut ctx,
-                    );
-                }
+                self.set_initial_aggregation_number(
+                    task_id,
+                    is_root,
+                    is_session_dependent,
+                    &mut ctx,
+                );
 
-                self.connect_child(
+                operation::ConnectChildOperation::run(
                     parent_task,
                     task_id,
                     Some(ArcOrOwned::Arc(task_type)),
-                    turbo_tasks,
+                    ctx,
                 );
 
                 task_id
@@ -1927,7 +1932,6 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
                     stale: false,
                     once_task,
                     done_event,
-                    session_dependent: false,
                     marked_as_completed: false,
                     new_children: Default::default(),
                 },
@@ -2156,7 +2160,6 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         let &mut InProgressState::InProgress(box InProgressStateInner {
             stale,
             ref mut new_children,
-            session_dependent,
             once_task: is_once_task,
             ..
         }) = in_progress
@@ -2252,7 +2255,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
             // Task was previously marked as immutable
             if !is_immutable
             // Task is not session dependent (session dependent tasks can change between sessions)
-            && !session_dependent
+            && !task.is_session_dependent()
             // Task has no invalidator
             && !task.invalidator()
             // Task has no dependencies on collectibles
@@ -2607,7 +2610,6 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
             done_event,
             once_task: is_once_task,
             stale,
-            session_dependent,
             marked_as_completed: _,
             new_children,
         }) = in_progress
@@ -2646,7 +2648,8 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
             }
         }
 
-        // Compute the new dirty state
+        // Compute and apply the new dirty state, propagating to aggregating ancestors
+        let session_dependent = task.is_session_dependent();
         let new_dirtyness = if session_dependent {
             Some(Dirtyness::SessionDependent)
         } else {
@@ -3041,42 +3044,6 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         );
     }
 
-    fn mark_own_task_as_session_dependent(
-        &self,
-        task_id: TaskId,
-        turbo_tasks: &dyn TurboTasksBackendApi<TurboTasksBackend<B>>,
-    ) {
-        if !self.should_track_dependencies() {
-            // Without dependency tracking we don't need session dependent tasks
-            return;
-        }
-        const SESSION_DEPENDENT_AGGREGATION_NUMBER: u32 = u32::MAX >> 2;
-        let mut ctx = self.execute_context(turbo_tasks);
-        let mut task = ctx.task(task_id, TaskDataCategory::Meta);
-        let aggregation_number = get_aggregation_number(&task);
-        if aggregation_number < SESSION_DEPENDENT_AGGREGATION_NUMBER {
-            drop(task);
-            // We want to use a high aggregation number to avoid large aggregation chains for
-            // session dependent tasks (which change on every run)
-            AggregationUpdateQueue::run(
-                AggregationUpdateJob::UpdateAggregationNumber {
-                    task_id,
-                    base_aggregation_number: SESSION_DEPENDENT_AGGREGATION_NUMBER,
-                    distance: None,
-                },
-                &mut ctx,
-            );
-            task = ctx.task(task_id, TaskDataCategory::Meta);
-        }
-        if let Some(InProgressState::InProgress(box InProgressStateInner {
-            session_dependent,
-            ..
-        })) = task.get_in_progress_mut()
-        {
-            *session_dependent = true;
-        }
-    }
-
     fn mark_own_task_as_finished(
         &self,
         task: TaskId,
@@ -3611,14 +3578,6 @@ impl<B: BackingStorage> Backend for TurboTasksBackend<B> {
         self.0.mark_own_task_as_finished(task_id, turbo_tasks);
     }
 
-    fn mark_own_task_as_session_dependent(
-        &self,
-        task: TaskId,
-        turbo_tasks: &dyn TurboTasksBackendApi<Self>,
-    ) {
-        self.0.mark_own_task_as_session_dependent(task, turbo_tasks);
-    }
-
     fn connect_task(
         &self,
         task: TaskId,

turbopack/crates/turbo-tasks-backend/src/backend/operation/mod.rs

@@ -1050,6 +1050,9 @@ pub trait TaskGuard: Debug + TaskStorageAccessors {
             panic!("Every task must have a task type {self:?}");
         }
     }
+    fn is_session_dependent(&self) -> bool {
+        matches!(self.get_task_type(), TaskTypeRef::Cached(tt) if tt.native_fn.is_session_dependent)
+    }
     fn get_task_desc_fn(&self) -> impl Fn() -> String + Send + Sync + 'static {
         let task_type = self.get_task_type().to_owned();
         let task_id = self.id();

turbopack/crates/turbo-tasks-backend/src/data.rs

@@ -238,7 +238,6 @@ pub struct InProgressStateInner {
     pub stale: bool,
     #[allow(dead_code)]
     pub once_task: bool,
-    pub session_dependent: bool,
     /// Early marking as completed. This is set before the output is available and will ignore full
     /// task completion of the task for strongly consistent reads.
     pub marked_as_completed: bool,

turbopack/crates/turbo-tasks-backend/tests/emptied_cells_session_dependent.rs

@@ -3,7 +3,7 @@
 #![allow(clippy::needless_return)] // tokio macro-generated code doesn't respect this
 
 use anyhow::Result;
-use turbo_tasks::{ResolvedVc, State, Vc, mark_session_dependent};
+use turbo_tasks::{ResolvedVc, State, Vc};
 use turbo_tasks_testing::{Registration, register, run};
 
 static REGISTRATION: Registration = register!();
@@ -76,9 +76,8 @@ async fn inner_compute(input: Vc<ChangingInput>) -> Result<Vc<u32>> {
     Ok(last.unwrap())
 }
 
-#[turbo_tasks::function]
+#[turbo_tasks::function(session_dependent)]
 async fn compute2(input: Vc<u32>) -> Result<Vc<u32>> {
-    mark_session_dependent();
     println!("compute2()");
     let value = *input.await?;
     Ok(Vc::cell(value))

turbopack/crates/turbo-tasks-env/src/command_line.rs

@@ -1,5 +1,5 @@
 use turbo_rcstr::RcStr;
-use turbo_tasks::{FxIndexMap, Vc, mark_session_dependent};
+use turbo_tasks::{FxIndexMap, Vc};
 
 use crate::{GLOBAL_ENV_LOCK, ProcessEnv, TransientEnvMap, sorted_env_vars};
 
@@ -23,9 +23,8 @@ fn env_snapshot() -> FxIndexMap<RcStr, RcStr> {
 
 #[turbo_tasks::value_impl]
 impl ProcessEnv for CommandLineProcessEnv {
-    #[turbo_tasks::function]
+    #[turbo_tasks::function(session_dependent)]
     fn read_all(&self) -> Vc<TransientEnvMap> {
-        mark_session_dependent();
         Vc::cell(env_snapshot())
     }
 }

turbopack/crates/turbo-tasks-fetch/src/client.rs

@@ -3,7 +3,7 @@ use std::{hash::Hash, sync::LazyLock};
 use anyhow::Result;
 use quick_cache::sync::Cache;
 use turbo_rcstr::RcStr;
-use turbo_tasks::{ReadRef, Vc, duration_span, mark_session_dependent};
+use turbo_tasks::{Completion, ReadRef, Vc, duration_span};
 
 use crate::{FetchError, FetchResult, HttpResponse, HttpResponseBody};
 
@@ -121,7 +121,10 @@ impl FetchClientConfig {
             Ok(resp) => Ok(Vc::cell(Ok(resp.resolved_cell()))),
             Err(err) => {
                 // the client failed to construct or the HTTP request failed
-                mark_session_dependent();
+                // Mark session dependent so we get retried in the next sessions
+                // In dev our caller will keep going, but in prod builds this will fail the build
+                // anyway.
+                Completion::session_dependent().as_side_effect().await?;
                 Ok(Vc::cell(Err(
                     FetchError::from_reqwest_error(&err, &url).resolved_cell()
                 )))