[Turbopack] fix incorrect task access for Persistent Caching

turbopack/crates/turbo-tasks-backend/Cargo.toml

@@ -15,6 +15,7 @@ workspace = true
 [features]
 default = []
 verify_serialization = []
+verify_task_access = []
 trace_aggregation_update = []
 trace_find_and_schedule = []
 lmdb = ["dep:lmdb-rkv"]

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

@@ -1592,7 +1592,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         let mut ctx = self.execute_context(turbo_tasks);
         let mut collectibles = AutoMap::default();
         {
-            let mut task = ctx.task(task_id, TaskDataCategory::Data);
+            let mut task = ctx.task(task_id, TaskDataCategory::All);
             // Ensure it's an root node
             loop {
                 let aggregation_number = get_aggregation_number(&task);

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

@@ -126,7 +126,7 @@ impl ConnectChildOperation {
             drop(parent_task);
 
             {
-                let mut task = ctx.task(child_task_id, TaskDataCategory::Data);
+                let mut task = ctx.task(child_task_id, TaskDataCategory::Meta);
                 if !task.has_key(&CachedDataItemKey::Output {}) {
                     let description = ctx.get_task_desc_fn(child_task_id);
                     let should_schedule = task.add(CachedDataItem::new_scheduled(description));

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

@@ -204,6 +204,8 @@
             task,
             task_id,
             backend: self.backend,
+            #[cfg(feature = "verify_task_access")]
+            category,
         }
     }
 
@@ -258,11 +260,15 @@
                 task: task1,
                 task_id: task_id1,
                 backend: self.backend,
+                #[cfg(feature = "verify_task_access")]
+                category,
             },
             TaskGuardImpl {
                 task: task2,
                 task_id: task_id2,
                 backend: self.backend,
+                #[cfg(feature = "verify_task_access")]
+                category,
             },
         )
     }
@@ -385,6 +391,34 @@
     task_id: TaskId,
     task: StorageWriteGuard<'a, TaskId, CachedDataItem>,
     backend: &'a TurboTasksBackendInner<B>,
+    #[cfg(feature = "verify_task_access")]
+    category: TaskDataCategory,
+}
+
+impl<B: BackingStorage> TaskGuardImpl<'_, B> {
+    #[inline]
+    fn check_access(&self, category: TaskDataCategory) {
+        #[cfg(feature = "verify_task_access")]
+        {
+            match category {
+                TaskDataCategory::All => {
+                    // This category is used for non-persisted data
+                }
+                TaskDataCategory::Data => {
+                    assert!(
+                        self.category == TaskDataCategory::Data
+                            || self.category == TaskDataCategory::All
+                    );
+                }
+                TaskDataCategory::Meta => {
+                    assert!(
+                        self.category == TaskDataCategory::Meta
+                            || self.category == TaskDataCategory::All
+                    );
+                }
+            }
+        }
+    }
 }
 
 impl<B: BackingStorage> Debug for TaskGuardImpl<'_, B> {
@@ -408,6 +442,7 @@
 
     #[must_use]
     fn add(&mut self, item: CachedDataItem) -> bool {
+        self.check_access(item.category());
         if !self.backend.should_persist() || self.task_id.is_transient() || !item.is_persistent() {
             self.task.add(item)
         } else if self.task.add(item.clone()) {
@@ -424,11 +459,13 @@
     }
 
     fn add_new(&mut self, item: CachedDataItem) {
+        self.check_access(item.category());
         let added = self.add(item);
         assert!(added, "Item already exists");
     }
 
     fn insert(&mut self, item: CachedDataItem) -> Option<CachedDataItemValue> {
+        self.check_access(item.category());
         let (key, value) = item.into_key_and_value();
         if !self.backend.should_persist() || self.task_id.is_transient() || !key.is_persistent() {
             self.task
@@ -472,6 +509,7 @@
         key: &CachedDataItemKey,
         update: impl FnOnce(Option<CachedDataItemValue>) -> Option<CachedDataItemValue>,
     ) {
+        self.check_access(key.category());
         if !self.backend.should_persist() || self.task_id.is_transient() || !key.is_persistent() {
             self.task.update(key, update);
             return;
@@ -480,6 +518,8 @@
             task,
             task_id,
             backend,
+            #[cfg(feature = "verify_task_access")]
+                category: _,
         } = self;
         let mut add_persisting_item = false;
         task.update(key, |old| {
@@ -519,6 +559,7 @@
     }
 
     fn remove(&mut self, key: &CachedDataItemKey) -> Option<CachedDataItemValue> {
+        self.check_access(key.category());
         let old_value = self.task.remove(key);
         if let Some(value) = old_value {
             if self.backend.should_persist()
@@ -545,14 +586,17 @@
     }
 
     fn get(&self, key: &CachedDataItemKey) -> Option<&CachedDataItemValue> {
+        self.check_access(key.category());
         self.task.get(key)
     }
 
     fn get_mut(&mut self, key: &CachedDataItemKey) -> Option<&mut CachedDataItemValue> {
+        self.check_access(key.category());
         self.task.get_mut(key)
     }
 
     fn has_key(&self, key: &CachedDataItemKey) -> bool {
+        self.check_access(key.category());
         self.task.has_key(key)
     }
 

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

@@ -43,7 +43,7 @@ impl UpdateOutputOperation {
         output: Result<Result<RawVc>, Option<Cow<'static, str>>>,
         mut ctx: impl ExecuteContext,
     ) {
-        let mut task = ctx.task(task_id, TaskDataCategory::Data);
+        let mut task = ctx.task(task_id, TaskDataCategory::Meta);
         if let Some(InProgressState::InProgress { stale: true, .. }) = get!(task, InProgress) {
             // Skip updating the output when the task is stale
             return;
@@ -173,7 +173,7 @@ impl Operation for UpdateOutputOperation {
                     ref mut queue,
                 } => {
                     if let Some(child_id) = children.pop() {
-                        let mut child_task = ctx.task(child_id, TaskDataCategory::Data);
+                        let mut child_task = ctx.task(child_id, TaskDataCategory::Meta);
                         if !child_task.has_key(&CachedDataItemKey::Output {}) {
                             make_task_dirty_internal(
                                 &mut child_task,

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

@@ -496,6 +496,42 @@ impl CachedDataItem {
             listener,
         )
     }
+
+    pub fn category(&self) -> TaskDataCategory {
+        match self {
+            Self::Collectible { .. }
+            | Self::Child { .. }
+            | Self::ChildrenCount { .. }
+            | Self::CellData { .. }
+            | Self::CellTypeMaxIndex { .. }
+            | Self::OutputDependency { .. }
+            | Self::CellDependency { .. }
+            | Self::CollectiblesDependency { .. }
+            | Self::OutputDependent { .. }
+            | Self::CellDependent { .. }
+            | Self::CollectiblesDependent { .. } => TaskDataCategory::Data,
+
+            Self::Output { .. }
+            | Self::AggregationNumber { .. }
+            | Self::Dirty { .. }
+            | Self::Follower { .. }
+            | Self::Upper { .. }
+            | Self::PersistentUpperCount { .. }
+            | Self::AggregatedDirtyContainer { .. }
+            | Self::AggregatedCollectible { .. }
+            | Self::AggregatedDirtyContainerCount { .. } => TaskDataCategory::Meta,
+
+            Self::OutdatedCollectible { .. }
+            | Self::OutdatedOutputDependency { .. }
+            | Self::OutdatedCellDependency { .. }
+            | Self::OutdatedCollectiblesDependency { .. }
+            | Self::OutdatedChild { .. }
+            | Self::InProgressCell { .. }
+            | Self::InProgress { .. }
+            | Self::Error { .. }
+            | Self::AggregateRoot { .. } => TaskDataCategory::All,
+        }
+    }
 }
 
 impl CachedDataItemKey {
@@ -543,36 +579,37 @@ impl CachedDataItemKey {
 
     pub fn category(&self) -> TaskDataCategory {
         match self {
-            CachedDataItemKey::Collectible { .. }
-            | CachedDataItemKey::Child { .. }
-            | CachedDataItemKey::ChildrenCount { .. }
-            | CachedDataItemKey::CellData { .. }
-            | CachedDataItemKey::CellTypeMaxIndex { .. }
-            | CachedDataItemKey::OutputDependency { .. }
-            | CachedDataItemKey::CellDependency { .. }
-            | CachedDataItemKey::CollectiblesDependency { .. }
-            | CachedDataItemKey::OutputDependent { .. }
-            | CachedDataItemKey::CellDependent { .. }
-            | CachedDataItemKey::CollectiblesDependent { .. }
-            | CachedDataItemKey::InProgress { .. }
-            | CachedDataItemKey::InProgressCell { .. }
-            | CachedDataItemKey::OutdatedCollectible { .. }
-            | CachedDataItemKey::OutdatedOutputDependency { .. }
-            | CachedDataItemKey::OutdatedCellDependency { .. }
-            | CachedDataItemKey::OutdatedCollectiblesDependency { .. }
-            | CachedDataItemKey::OutdatedChild { .. }
-            | CachedDataItemKey::Error { .. } => TaskDataCategory::Data,
-
-            CachedDataItemKey::Output { .. }
-            | CachedDataItemKey::AggregationNumber { .. }
-            | CachedDataItemKey::Dirty { .. }
-            | CachedDataItemKey::Follower { .. }
-            | CachedDataItemKey::Upper { .. }
-            | CachedDataItemKey::PersistentUpperCount { .. }
-            | CachedDataItemKey::AggregatedDirtyContainer { .. }
-            | CachedDataItemKey::AggregatedCollectible { .. }
-            | CachedDataItemKey::AggregatedDirtyContainerCount { .. }
-            | CachedDataItemKey::AggregateRoot { .. } => TaskDataCategory::Meta,
+            Self::Collectible { .. }
+            | Self::Child { .. }
+            | Self::ChildrenCount { .. }
+            | Self::CellData { .. }
+            | Self::CellTypeMaxIndex { .. }
+            | Self::OutputDependency { .. }
+            | Self::CellDependency { .. }
+            | Self::CollectiblesDependency { .. }
+            | Self::OutputDependent { .. }
+            | Self::CellDependent { .. }
+            | Self::CollectiblesDependent { .. } => TaskDataCategory::Data,
+
+            Self::Output { .. }
+            | Self::AggregationNumber { .. }
+            | Self::Dirty { .. }
+            | Self::Follower { .. }
+            | Self::Upper { .. }
+            | Self::PersistentUpperCount { .. }
+            | Self::AggregatedDirtyContainer { .. }
+            | Self::AggregatedCollectible { .. }
+            | Self::AggregatedDirtyContainerCount { .. } => TaskDataCategory::Meta,
+
+            Self::OutdatedCollectible { .. }
+            | Self::OutdatedOutputDependency { .. }
+            | Self::OutdatedCellDependency { .. }
+            | Self::OutdatedCollectiblesDependency { .. }
+            | Self::OutdatedChild { .. }
+            | Self::InProgressCell { .. }
+            | Self::InProgress { .. }
+            | Self::Error { .. }
+            | Self::AggregateRoot { .. } => TaskDataCategory::All,
         }
     }
 }