refactor(turbo-tasks): Use an execution id instead of the parent task id to prevent local Vc escapes (#78487)

Re-using the task id just for the `Vc` scope escape check isn't ideal:

- A task can execute multiple times. Just using task ids doesn't guarantee that the `Vc` hasn't leaked across multiple executions of the same task.
- This is just a fallback check to help with debugging, so allocating 32 bits to this id seems wasteful. We can do a good job with 16 bits. Nothing else really needs the task id, so we can get rid of it later in this PR stack.

A version of this idea was part of the original local tasks idea: https://www.notion.so/vercel/Resolved-Vcs-Vc-Lifetimes-Local-Vcs-and-Vc-Refcounts-49d666d3f9594017b5b312b87ddc5bff?pvs=4#da95dbdc390144a9bb1b3d2fb73867b7

Even with this extra 4 bits, the overall enum size is still 16 bits due to alignment and the size of the other enum variants, so this shouldn't regress memory consumption.

The next PR (#78561) removes `TaskId` entirely...

Author: bgw

turbopack/crates/turbo-tasks-testing/src/lib.rs

@@ -20,8 +20,8 @@ use turbo_tasks::{
     registry,
     test_helpers::with_turbo_tasks_for_testing,
     util::{SharedError, StaticOrArc},
-    CellId, InvalidationReason, LocalTaskId, MagicAny, RawVc, ReadCellOptions, ReadConsistency,
-    TaskId, TaskPersistence, TraitTypeId, TurboTasksApi, TurboTasksCallApi,
+    CellId, ExecutionId, InvalidationReason, LocalTaskId, MagicAny, RawVc, ReadCellOptions,
+    ReadConsistency, TaskId, TaskPersistence, TraitTypeId, TurboTasksApi, TurboTasksCallApi,
 };
 
 pub use crate::run::{run, run_with_tt, run_without_cache_check, Registration};
@@ -57,9 +57,11 @@ impl VcStorage {
             i
         };
         let task_id = TaskId::try_from(u32::try_from(i + 1).unwrap()).unwrap();
+        let execution_id = ExecutionId::try_from(u16::try_from(i + 1).unwrap()).unwrap();
         handle.spawn(with_turbo_tasks_for_testing(
             this.clone(),
             task_id,
+            execution_id,
             async move {
                 let result = AssertUnwindSafe(future).catch_unwind().await;
 
@@ -230,7 +232,7 @@ impl TurboTasksApi for VcStorage {
 
     fn try_read_local_output(
         &self,
-        _parent_task_id: TaskId,
+        _execution_id: ExecutionId,
         _local_task_id: LocalTaskId,
     ) -> Result<Result<RawVc, EventListener>> {
         unimplemented!()
@@ -322,6 +324,7 @@ impl VcStorage {
                 ..Default::default()
             }),
             TaskId::MAX,
+            ExecutionId::MIN,
             f,
         )
     }

turbopack/crates/turbo-tasks/src/id.rs

@@ -125,6 +125,13 @@ define_id!(
     serde(transparent),
     doc = "Represents the nth `local` function call inside a task.",
 );
+define_id!(
+    ExecutionId: u16,
+    derive(Debug, Serialize, Deserialize),
+    serde(transparent),
+    doc = "An identifier for a specific task execution. Used to assert that local `Vc`s don't \
+        leak. This value may overflow and re-use old values.",
+);
 
 impl Debug for TaskId {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

turbopack/crates/turbo-tasks/src/id_factory.rs

@@ -96,10 +96,36 @@ where
             ),
         }
     }
+
+    /// Returns an id, potentially allowing an overflow. This may cause ids to be silently re-used.
+    /// Used for [`crate::id::ExecutionId`].
+    ///
+    /// If id re-use is desired only for "freed" ids, use [`IdFactoryWithReuse`] instead.
+    pub fn wrapping_get(&self) -> T {
+        let count = self.counter.fetch_add(1, Ordering::Relaxed);
+
+        let new_id_u64 = (count % self.max_count) + self.id_offset;
+        // Safety:
+        // - `id_offset` is a non-zero value.
+        // - `id_offset + max_count < u64::MAX`.
+        let new_id = unsafe { NonZeroU64::new_unchecked(new_id_u64) };
+
+        match new_id.try_into() {
+            Ok(id) => id,
+            Err(_) => panic!(
+                "Failed to convert NonZeroU64 value of {} into {}",
+                new_id,
+                type_name::<T>()
+            ),
+        }
+    }
 }
 
 /// An [`IdFactory`], but extended with a free list to allow for id reuse for ids such as
 /// [`BackendJobId`][crate::backend::BackendJobId].
+///
+/// If silent untracked re-use of ids is okay, consider using the cheaper
+/// [`IdFactory::wrapping_get`] method.
 pub struct IdFactoryWithReuse<T> {
     factory: IdFactory<T>,
     free_ids: ConcurrentQueue<T>,

turbopack/crates/turbo-tasks/src/lib.rs

@@ -91,7 +91,8 @@ pub use completion::{Completion, Completions};
 pub use display::ValueToString;
 pub use effect::{apply_effects, effect, get_effects, Effects};
 pub use id::{
-    FunctionId, LocalTaskId, SessionId, TaskId, TraitTypeId, ValueTypeId, TRANSIENT_TASK_BIT,
+    ExecutionId, FunctionId, LocalTaskId, SessionId, TaskId, TraitTypeId, ValueTypeId,
+    TRANSIENT_TASK_BIT,
 };
 pub use invalidation::{
     get_invalidator, DynamicEqHash, InvalidationReason, InvalidationReasonKind,

turbopack/crates/turbo-tasks/src/manager.rs

@@ -31,7 +31,7 @@ use crate::{
     },
     capture_future::{self, CaptureFuture},
     event::{Event, EventListener},
-    id::{BackendJobId, FunctionId, LocalTaskId, TraitTypeId, TRANSIENT_TASK_BIT},
+    id::{BackendJobId, ExecutionId, FunctionId, LocalTaskId, TraitTypeId, TRANSIENT_TASK_BIT},
     id_factory::IdFactoryWithReuse,
     magic_any::MagicAny,
     raw_vc::{CellId, RawVc},
@@ -41,10 +41,11 @@ use crate::{
     task_statistics::TaskStatisticsApi,
     trace::TraceRawVcs,
     trait_helpers::get_trait_method,
-    util::StaticOrArc,
+    util::{IdFactory, StaticOrArc},
     vc::ReadVcFuture,
-    Completion, InvalidationReason, InvalidationReasonSet, ReadCellOptions, ResolvedVc,
-    SharedReference, TaskId, TaskIdSet, ValueTypeId, Vc, VcRead, VcValueTrait, VcValueType,
+    Completion, InvalidationReason, InvalidationReasonSet, OutputContent, ReadCellOptions,
+    ResolvedVc, SharedReference, TaskId, TaskIdSet, ValueTypeId, Vc, VcRead, VcValueTrait,
+    VcValueType,
 };
 
 pub trait TurboTasksCallApi: Sync + Send {
@@ -139,7 +140,7 @@ pub trait TurboTasksApi: TurboTasksCallApi + Sync + Send {
     /// `OperationVc`s, which should never be local tasks.
     fn try_read_local_output(
         &self,
-        parent_task_id: TaskId,
+        execution_id: ExecutionId,
         local_task_id: LocalTaskId,
     ) -> Result<Result<RawVc, EventListener>>;
 
@@ -347,6 +348,7 @@ pub struct TurboTasks<B: Backend + 'static> {
     backend: B,
     task_id_factory: IdFactoryWithReuse<TaskId>,
     transient_task_id_factory: IdFactoryWithReuse<TaskId>,
+    execution_id_factory: IdFactory<ExecutionId>,
     stopped: AtomicBool,
     currently_scheduled_tasks: AtomicUsize,
     currently_scheduled_foreground_jobs: AtomicUsize,
@@ -371,6 +373,7 @@ pub struct TurboTasks<B: Backend + 'static> {
 /// - The backend is aware of.
 struct CurrentTaskState {
     task_id: TaskId,
+    execution_id: ExecutionId,
 
     /// Affected tasks, that are tracked during task execution. These tasks will
     /// be invalidated when the execution finishes or before reading a cell
@@ -397,9 +400,14 @@ struct CurrentTaskState {
 }
 
 impl CurrentTaskState {
-    fn new(task_id: TaskId, backend_state: Box<dyn Any + Send + Sync>) -> Self {
+    fn new(
+        task_id: TaskId,
+        execution_id: ExecutionId,
+        backend_state: Box<dyn Any + Send + Sync>,
+    ) -> Self {
         Self {
             task_id,
+            execution_id,
             tasks_to_notify: Vec::new(),
             stateful: false,
             cell_counters: Some(AutoMap::default()),
@@ -409,10 +417,11 @@ impl CurrentTaskState {
         }
     }
 
-    fn assert_task_id(&self, expected_task_id: TaskId) {
-        if self.task_id != expected_task_id {
-            unimplemented!(
-                "Local tasks can currently only be scheduled/awaited within their parent task"
+    fn assert_execution_id(&self, expected_execution_id: ExecutionId) {
+        if self.execution_id != expected_execution_id {
+            panic!(
+                "Local tasks can only be scheduled/awaited within the same execution of the \
+                 parent task that created them"
             );
         }
     }
@@ -458,11 +467,13 @@ impl<B: Backend + 'static> TurboTasks<B> {
         );
         let transient_task_id_factory =
             IdFactoryWithReuse::new(TaskId::try_from(TRANSIENT_TASK_BIT).unwrap(), TaskId::MAX);
+        let execution_id_factory = IdFactory::new(ExecutionId::MIN, ExecutionId::MAX);
         let this = Arc::new_cyclic(|this| Self {
             this: this.clone(),
             backend,
             task_id_factory,
             transient_task_id_factory,
+            execution_id_factory,
             stopped: AtomicBool::new(false),
             currently_scheduled_tasks: AtomicUsize::new(0),
             currently_scheduled_background_jobs: AtomicUsize::new(0),
@@ -644,8 +655,11 @@ impl<B: Backend + 'static> TurboTasks<B> {
             let mut schedule_again = true;
             while schedule_again {
                 let backend_state = this.backend.new_task_state(task_id);
+                // it's okay for execution ids to overflow and wrap, they're just used for an assert
+                let execution_id = this.execution_id_factory.wrapping_get();
                 let current_task_state = Arc::new(RwLock::new(CurrentTaskState::new(
                     task_id,
+                    execution_id,
                     Box::new(backend_state),
                 )));
                 let single_execution_future = async {
@@ -722,22 +736,31 @@ impl<B: Backend + 'static> TurboTasks<B> {
         &self,
         ty: LocalTaskType,
         // if this is a `LocalTaskType::Resolve*`, we may spawn another task with this persistence,
-        // if this is a `LocalTaskType::Native`, persistence is unused (there's no caching).
+        // if this is a `LocalTaskType::Native`, persistence is unused.
+        //
+        // TODO: In the rare case that we're crossing a transient->persistent boundary, we should
+        // force `LocalTaskType::Native` to be spawned as real tasks, so that any cells they create
+        // have the correct persistence. This is not an issue for resolution stub task, as they
+        // don't end up owning any cells.
         persistence: TaskPersistence,
     ) -> RawVc {
-        use crate::OutputContent;
-
         let ty = Arc::new(ty);
-        let (global_task_state, local_task_id, parent_task_id) = CURRENT_TASK_STATE.with(|gts| {
-            let mut gts_write = gts.write().unwrap();
-            let local_task_id = gts_write.create_local_task(LocalTask::Scheduled {
-                done_event: Event::new({
-                    let ty = Arc::clone(&ty);
-                    move || format!("LocalTask({})::done_event", ty)
-                }),
+        let (global_task_state, parent_task_id, execution_id, local_task_id) = CURRENT_TASK_STATE
+            .with(|gts| {
+                let mut gts_write = gts.write().unwrap();
+                let local_task_id = gts_write.create_local_task(LocalTask::Scheduled {
+                    done_event: Event::new({
+                        let ty = Arc::clone(&ty);
+                        move || format!("LocalTask({})::done_event", ty)
+                    }),
+                });
+                (
+                    Arc::clone(gts),
+                    gts_write.task_id,
+                    gts_write.execution_id,
+                    local_task_id,
+                )
             });
-            (Arc::clone(gts), local_task_id, gts_write.task_id)
-        });
 
         #[cfg(feature = "tokio_tracing")]
         let description = format!(
@@ -799,7 +822,7 @@ impl<B: Backend + 'static> TurboTasks<B> {
         #[cfg(not(feature = "tokio_tracing"))]
         tokio::task::spawn(future);
 
-        RawVc::LocalOutput(parent_task_id, persistence, local_task_id)
+        RawVc::LocalOutput(parent_task_id, persistence, execution_id, local_task_id)
     }
 
     fn begin_primary_job(&self) {
@@ -1277,17 +1300,17 @@ impl<B: Backend + 'static> TurboTasksApi for TurboTasks<B> {
 
     fn try_read_local_output(
         &self,
-        parent_task_id: TaskId,
+        execution_id: ExecutionId,
         local_task_id: LocalTaskId,
     ) -> Result<Result<RawVc, EventListener>> {
         CURRENT_TASK_STATE.with(|gts| {
             let gts_read = gts.read().unwrap();
 
-            // Local Vcs are local to their parent task, and do not exist outside of it. This is
-            // weakly enforced at compile time using the `NonLocalValue` marker trait. This
-            // assertion exists to handle any potential escapes that the compile-time checks cannot
-            // capture.
-            gts_read.assert_task_id(parent_task_id);
+            // Local Vcs are local to their parent task's current execution, and do not exist
+            // outside of it. This is weakly enforced at compile time using the `NonLocalValue`
+            // marker trait. This assertion exists to handle any potential escapes that the
+            // compile-time checks cannot capture.
+            gts_read.assert_execution_id(execution_id);
 
             match gts_read.get_local_task(local_task_id) {
                 LocalTask::Scheduled { done_event } => Ok(Err(done_event.listen())),
@@ -1632,13 +1655,15 @@ pub fn turbo_tasks_future_scope<T>(
 pub fn with_turbo_tasks_for_testing<T>(
     tt: Arc<dyn TurboTasksApi>,
     current_task: TaskId,
+    execution_id: ExecutionId,
     f: impl Future<Output = T>,
 ) -> impl Future<Output = T> {
     TURBO_TASKS.scope(
         tt,
         CURRENT_TASK_STATE.scope(
             Arc::new(RwLock::new(CurrentTaskState::new(
                 current_task,
+                execution_id,
                 Box::new(()),
             ))),
             f,
@@ -1969,11 +1994,11 @@ pub fn find_cell_by_type(ty: ValueTypeId) -> CurrentCellRef {
 
 pub(crate) async fn read_local_output(
     this: &dyn TurboTasksApi,
-    parent_task_id: TaskId,
+    execution_id: ExecutionId,
     local_task_id: LocalTaskId,
 ) -> Result<RawVc> {
     loop {
-        match this.try_read_local_output(parent_task_id, local_task_id)? {
+        match this.try_read_local_output(execution_id, local_task_id)? {
             Ok(raw_vc) => return Ok(raw_vc),
             Err(event_listener) => event_listener.await,
         }

turbopack/crates/turbo-tasks/src/raw_vc.rs

@@ -8,7 +8,7 @@ use thiserror::Error;
 use crate::{
     backend::{CellContent, TypedCellContent},
     event::EventListener,
-    id::LocalTaskId,
+    id::{ExecutionId, LocalTaskId},
     manager::{read_local_output, read_task_cell, read_task_output, with_turbo_tasks},
     registry::{self, get_value_type},
     turbo_tasks, CollectiblesSource, ReadCellOptions, ReadConsistency, ResolvedVc, TaskId,
@@ -57,7 +57,7 @@ impl Display for CellId {
 pub enum RawVc {
     TaskOutput(TaskId),
     TaskCell(TaskId, CellId),
-    LocalOutput(TaskId, TaskPersistence, LocalTaskId),
+    LocalOutput(TaskId, TaskPersistence, ExecutionId, LocalTaskId),
 }
 
 impl RawVc {
@@ -148,8 +148,8 @@ impl RawVc {
                         return Err(ResolveTypeError::NoContent);
                     }
                 }
-                RawVc::LocalOutput(task_id, _persistence, local_task_id) => {
-                    current = read_local_output(&*tt, task_id, local_task_id)
+                RawVc::LocalOutput(_task_id, _persistence, execution_id, local_task_id) => {
+                    current = read_local_output(&*tt, execution_id, local_task_id)
                         .await
                         .map_err(|source| ResolveTypeError::TaskError { source })?;
                 }
@@ -185,9 +185,9 @@ impl RawVc {
                     consistency = ReadConsistency::Eventual;
                 }
                 RawVc::TaskCell(_, _) => return Ok(current),
-                RawVc::LocalOutput(task_id, _persistence, local_task_id) => {
+                RawVc::LocalOutput(_task_id, _persistence, execution_id, local_task_id) => {
                     debug_assert_eq!(consistency, ReadConsistency::Eventual);
-                    current = read_local_output(&*tt, task_id, local_task_id).await?;
+                    current = read_local_output(&*tt, execution_id, local_task_id).await?;
                 }
             }
         }
@@ -200,8 +200,8 @@ impl RawVc {
         let mut current = self;
         loop {
             match current {
-                RawVc::LocalOutput(task_id, _persistence, local_task_id) => {
-                    current = read_local_output(&*tt, task_id, local_task_id).await?;
+                RawVc::LocalOutput(_task_id, _persistence, execution_id, local_task_id) => {
+                    current = read_local_output(&*tt, execution_id, local_task_id).await?;
                 }
                 non_local => return Ok(non_local),
             }
@@ -215,14 +215,14 @@ impl RawVc {
 
     pub fn get_task_id(&self) -> TaskId {
         match self {
-            RawVc::TaskOutput(t) | RawVc::TaskCell(t, _) | RawVc::LocalOutput(t, ..) => *t,
+            RawVc::TaskOutput(t) | RawVc::TaskCell(t, ..) | RawVc::LocalOutput(t, ..) => *t,
         }
     }
 
     pub fn try_get_type_id(&self) -> Option<ValueTypeId> {
         match self {
             RawVc::TaskCell(_, CellId { type_id, .. }) => Some(*type_id),
-            RawVc::TaskOutput(_) | RawVc::LocalOutput(..) => None,
+            RawVc::TaskOutput(..) | RawVc::LocalOutput(..) => None,
         }
     }
 
@@ -360,9 +360,9 @@ impl Future for ReadRawVcFuture {
                             Err(err) => return Poll::Ready(Err(err)),
                         }
                     }
-                    RawVc::LocalOutput(task_id, _persistence, local_output_id) => {
+                    RawVc::LocalOutput(_task_id, _persistence, execution_id, local_output_id) => {
                         debug_assert_eq!(this.consistency, ReadConsistency::Eventual);
-                        let read_result = tt.try_read_local_output(task_id, local_output_id);
+                        let read_result = tt.try_read_local_output(execution_id, local_output_id);
                         match read_result {
                             Ok(Ok(vc)) => {
                                 this.current = vc;