store task data and aggregation separately

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

@@ -22,6 +22,7 @@ use dashmap::DashMap;
 use parking_lot::{Condvar, Mutex};
 use rustc_hash::FxHasher;
 use smallvec::smallvec;
+pub use storage::TaskDataCategory;
 use turbo_tasks::{
     backend::{
         Backend, BackendJobId, CachedTaskType, CellContent, TaskExecutionSpec, TransientTaskRoot,
@@ -101,7 +102,8 @@ pub struct TurboTasksBackend {
     task_cache: BiMap<Arc<CachedTaskType>, TaskId>,
     transient_tasks: DashMap<TaskId, Arc<TransientTask>>,
 
-    persisted_storage_log: Mutex<ChunkedVec<CachedDataUpdate>>,
+    persisted_storage_data_log: Mutex<ChunkedVec<CachedDataUpdate>>,
+    persisted_storage_meta_log: Mutex<ChunkedVec<CachedDataUpdate>>,
     storage: Storage<TaskId, CachedDataItem>,
 
     /// Number of executing operations + Highest bit is set when snapshot is
@@ -140,7 +142,8 @@ impl TurboTasksBackend {
             persisted_task_cache_log: Mutex::new(ChunkedVec::new()),
             task_cache: BiMap::new(),
             transient_tasks: DashMap::new(),
-            persisted_storage_log: Mutex::new(ChunkedVec::new()),
+            persisted_storage_data_log: Mutex::new(ChunkedVec::new()),
+            persisted_storage_meta_log: Mutex::new(ChunkedVec::new()),
             storage: Storage::new(),
             in_progress_operations: AtomicUsize::new(0),
             snapshot_request: Mutex::new(SnapshotRequest::new()),
@@ -205,6 +208,17 @@ impl TurboTasksBackend {
         }
         OperationGuard { backend: self }
     }
+
+    fn persisted_storage_log(
+        &self,
+        category: TaskDataCategory,
+    ) -> &Mutex<ChunkedVec<CachedDataUpdate>> {
+        match category {
+            TaskDataCategory::Data => &self.persisted_storage_data_log,
+            TaskDataCategory::Meta => &self.persisted_storage_meta_log,
+            TaskDataCategory::All => unreachable!(),
+        }
+    }
 }
 
 pub(crate) struct OperationGuard<'a> {
@@ -246,7 +260,7 @@ impl TurboTasksBackend {
         turbo_tasks: &dyn TurboTasksBackendApi<Self>,
     ) -> Result<Result<RawVc, EventListener>> {
         let ctx = self.execute_context(turbo_tasks);
-        let mut task = ctx.task(task_id);
+        let mut task = ctx.task(task_id, TaskDataCategory::All);
 
         if let Some(in_progress) = get!(task, InProgress) {
             match in_progress {
@@ -281,7 +295,7 @@ impl TurboTasksBackend {
                     },
                     &ctx,
                 );
-                task = ctx.task(task_id);
+                task = ctx.task(task_id, TaskDataCategory::All);
             }
 
             // Check the dirty count of the root node
@@ -328,7 +342,7 @@ impl TurboTasksBackend {
                     });
                     drop(task);
 
-                    let mut reader_task = ctx.task(reader);
+                    let mut reader_task = ctx.task(reader, TaskDataCategory::Data);
                     if reader_task
                         .remove(&CachedDataItemKey::OutdatedOutputDependency { target: task_id })
                         .is_none()
@@ -370,7 +384,7 @@ impl TurboTasksBackend {
         turbo_tasks: &dyn TurboTasksBackendApi<Self>,
     ) -> Result<Result<TypedCellContent, EventListener>> {
         let ctx = self.execute_context(turbo_tasks);
-        let mut task = ctx.task(task_id);
+        let mut task = ctx.task(task_id, TaskDataCategory::Data);
         if let Some(content) = get!(task, CellData { cell }) {
             let content = content.clone();
             if let Some(reader) = reader {
@@ -381,7 +395,7 @@ impl TurboTasksBackend {
                 });
                 drop(task);
 
-                let mut reader_task = ctx.task(reader);
+                let mut reader_task = ctx.task(reader, TaskDataCategory::Data);
                 let target = CellRef {
                     task: task_id,
                     cell,
@@ -494,7 +508,8 @@ impl TurboTasksBackend {
             .map(|op| op.arc().clone())
             .collect::<Vec<_>>();
         drop(snapshot_request);
-        let persisted_storage_log = take(&mut *self.persisted_storage_log.lock());
+        let persisted_storage_meta_log = take(&mut *self.persisted_storage_meta_log.lock());
+        let persisted_storage_data_log = take(&mut *self.persisted_storage_data_log.lock());
         let persisted_task_cache_log = take(&mut *self.persisted_task_cache_log.lock());
         let mut snapshot_request = self.snapshot_request.lock();
         snapshot_request.snapshot_requested = false;
@@ -505,18 +520,25 @@ impl TurboTasksBackend {
         drop(snapshot_request);
 
         let mut counts: HashMap<TaskId, u32> = HashMap::new();
-        for CachedDataUpdate { task, .. } in persisted_storage_log.iter() {
+        for CachedDataUpdate { task, .. } in persisted_storage_data_log
+            .iter()
+            .chain(persisted_storage_meta_log.iter())
+        {
             *counts.entry(*task).or_default() += 1;
         }
 
         let mut new_items = false;
 
-        if !persisted_task_cache_log.is_empty() || !persisted_storage_log.is_empty() {
+        if !persisted_task_cache_log.is_empty()
+            || !persisted_storage_meta_log.is_empty()
+            || !persisted_storage_data_log.is_empty()
+        {
             new_items = true;
             if let Err(err) = self.backing_storage.save_snapshot(
                 suspended_operations,
                 persisted_task_cache_log,
-                persisted_storage_log,
+                persisted_storage_meta_log,
+                persisted_storage_data_log,
             ) {
                 println!("Persising failed: {:#?}", err);
                 return None;
@@ -655,7 +677,7 @@ impl Backend for TurboTasksBackend {
             return;
         }
         let ctx = self.execute_context(turbo_tasks);
-        let mut task = ctx.task(task_id);
+        let mut task = ctx.task(task_id, TaskDataCategory::Data);
         task.invalidate_serialization();
     }
 
@@ -696,7 +718,7 @@ impl Backend for TurboTasksBackend {
         };
         {
             let ctx = self.execute_context(turbo_tasks);
-            let mut task = ctx.task(task_id);
+            let mut task = ctx.task(task_id, TaskDataCategory::Data);
             let in_progress = remove!(task, InProgress)?;
             let InProgressState::Scheduled { done_event } = in_progress else {
                 task.add_new(CachedDataItem::InProgress { value: in_progress });
@@ -887,7 +909,7 @@ impl Backend for TurboTasksBackend {
         turbo_tasks: &dyn TurboTasksBackendApi<Self>,
     ) -> bool {
         let ctx = self.execute_context(turbo_tasks);
-        let mut task = ctx.task(task_id);
+        let mut task = ctx.task(task_id, TaskDataCategory::All);
         let Some(CachedDataItemValue::InProgress { value: in_progress }) =
             task.remove(&CachedDataItemKey::InProgress {})
         else {
@@ -1126,7 +1148,7 @@ impl Backend for TurboTasksBackend {
         turbo_tasks: &dyn TurboTasksBackendApi<Self>,
     ) -> Result<TypedCellContent> {
         let ctx = self.execute_context(turbo_tasks);
-        let task = ctx.task(task_id);
+        let task = ctx.task(task_id, TaskDataCategory::Data);
         if let Some(content) = get!(task, CellData { cell }) {
             Ok(CellContent(Some(content.1.clone())).into_typed(cell.type_id))
         } else {

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

@@ -7,6 +7,7 @@ use crate::{
     backend::{
         operation::{ExecuteContext, Operation, TaskGuard},
         storage::{get, get_many, iter_many, remove, update, update_count},
+        TaskDataCategory,
     },
     data::{ActiveType, AggregationNumber, CachedDataItem, CachedDataItemKey, RootState},
 };
@@ -385,7 +386,7 @@ impl AggregationUpdateQueue {
     }
 
     fn balance_edge(&mut self, ctx: &ExecuteContext, upper_id: TaskId, task_id: TaskId) {
-        let (mut upper, mut task) = ctx.task_pair(upper_id, task_id);
+        let (mut upper, mut task) = ctx.task_pair(upper_id, task_id, TaskDataCategory::Meta);
         let upper_aggregation_number = get_aggregation_number(&upper);
         let task_aggregation_number = get_aggregation_number(&task);
 
@@ -490,7 +491,7 @@ impl AggregationUpdateQueue {
             self.push(AggregationUpdateJob::FindAndScheduleDirty { task_ids });
         }
         if let Some(task_id) = popped {
-            let mut task = ctx.task(task_id);
+            let mut task = ctx.task(task_id, TaskDataCategory::Meta);
             #[allow(clippy::collapsible_if, reason = "readablility")]
             if task.has_key(&CachedDataItemKey::Dirty {}) {
                 let description = ctx.backend.get_task_desc_fn(task_id);
@@ -522,7 +523,7 @@ impl AggregationUpdateQueue {
         update: AggregatedDataUpdate,
     ) {
         for upper_id in upper_ids {
-            let mut upper = ctx.task(upper_id);
+            let mut upper = ctx.task(upper_id, TaskDataCategory::Meta);
             let diff = update.apply(&mut upper, self);
             if !diff.is_empty() {
                 let upper_ids = get_uppers(&upper);
@@ -542,7 +543,7 @@ impl AggregationUpdateQueue {
         lost_follower_id: TaskId,
         mut upper_ids: Vec<TaskId>,
     ) {
-        let mut follower = ctx.task(lost_follower_id);
+        let mut follower = ctx.task(lost_follower_id, TaskDataCategory::Meta);
         let mut follower_in_upper_ids = Vec::new();
         upper_ids.retain(|&upper_id| {
             let mut keep_upper = false;
@@ -571,7 +572,7 @@ impl AggregationUpdateQueue {
             if !data.is_empty() {
                 for upper_id in upper_ids.iter() {
                     // remove data from upper
-                    let mut upper = ctx.task(*upper_id);
+                    let mut upper = ctx.task(*upper_id, TaskDataCategory::Meta);
                     let diff = data.apply(&mut upper, self);
                     if !diff.is_empty() {
                         let upper_ids = get_uppers(&upper);
@@ -593,7 +594,7 @@ impl AggregationUpdateQueue {
         }
 
         for upper_id in follower_in_upper_ids {
-            let mut upper = ctx.task(upper_id);
+            let mut upper = ctx.task(upper_id, TaskDataCategory::Meta);
             if update_count!(
                 upper,
                 Follower {
@@ -617,12 +618,12 @@ impl AggregationUpdateQueue {
         mut upper_ids: Vec<TaskId>,
     ) {
         let follower_aggregation_number = {
-            let follower = ctx.task(new_follower_id);
+            let follower = ctx.task(new_follower_id, TaskDataCategory::Meta);
             get_aggregation_number(&follower)
         };
         let mut upper_ids_as_follower = Vec::new();
         upper_ids.retain(|&upper_id| {
-            let upper = ctx.task(upper_id);
+            let upper = ctx.task(upper_id, TaskDataCategory::Meta);
             // decide if it should be an inner or follower
             let upper_aggregation_number = get_aggregation_number(&upper);
 
@@ -638,7 +639,7 @@ impl AggregationUpdateQueue {
             }
         });
         if !upper_ids.is_empty() {
-            let mut follower = ctx.task(new_follower_id);
+            let mut follower = ctx.task(new_follower_id, TaskDataCategory::Meta);
             upper_ids.retain(|&upper_id| {
                 if update_count!(follower, Upper { task: upper_id }, 1) {
                     // It's a new upper
@@ -656,7 +657,7 @@ impl AggregationUpdateQueue {
                 if !data.is_empty() {
                     for upper_id in upper_ids.iter() {
                         // add data to upper
-                        let mut upper = ctx.task(*upper_id);
+                        let mut upper = ctx.task(*upper_id, TaskDataCategory::Meta);
                         let diff = data.apply(&mut upper, self);
                         if !diff.is_empty() {
                             let upper_ids = get_uppers(&upper);
@@ -678,7 +679,7 @@ impl AggregationUpdateQueue {
             }
         }
         upper_ids_as_follower.retain(|&upper_id| {
-            let mut upper = ctx.task(upper_id);
+            let mut upper = ctx.task(upper_id, TaskDataCategory::Meta);
             update_count!(
                 upper,
                 Follower {
@@ -704,14 +705,14 @@ impl AggregationUpdateQueue {
         let mut followers_with_aggregation_number = new_follower_ids
             .into_iter()
             .map(|new_follower_id| {
-                let follower = ctx.task(new_follower_id);
+                let follower = ctx.task(new_follower_id, TaskDataCategory::Meta);
                 (new_follower_id, get_aggregation_number(&follower))
             })
             .collect::<Vec<_>>();
 
         let mut followers_of_upper = Vec::new();
         {
-            let upper = ctx.task(upper_id);
+            let upper = ctx.task(upper_id, TaskDataCategory::Meta);
             // decide if it should be an inner or follower
             let upper_aggregation_number = get_aggregation_number(&upper);
 
@@ -734,7 +735,7 @@ impl AggregationUpdateQueue {
         let mut upper_data_updates = Vec::new();
         let mut upper_new_followers = Vec::new();
         for (follower_id, _) in followers_with_aggregation_number {
-            let mut follower = ctx.task(follower_id);
+            let mut follower = ctx.task(follower_id, TaskDataCategory::Meta);
             if update_count!(follower, Upper { task: upper_id }, 1) {
                 // It's a new upper
                 let data = AggregatedDataUpdate::from_task(&mut follower);
@@ -755,7 +756,7 @@ impl AggregationUpdateQueue {
         }
         if !upper_data_updates.is_empty() {
             // add data to upper
-            let mut upper = ctx.task(upper_id);
+            let mut upper = ctx.task(upper_id, TaskDataCategory::Meta);
             let diffs = upper_data_updates
                 .into_iter()
                 .filter_map(|data| {
@@ -782,7 +783,7 @@ impl AggregationUpdateQueue {
             }
         }
         if !followers_of_upper.is_empty() {
-            let mut upper = ctx.task(upper_id);
+            let mut upper = ctx.task(upper_id, TaskDataCategory::Meta);
             followers_of_upper
                 .retain(|follower_id| update_count!(upper, Follower { task: *follower_id }, 1));
             if !followers_of_upper.is_empty() {
@@ -801,7 +802,7 @@ impl AggregationUpdateQueue {
         base_effective_distance: Option<std::num::NonZero<u32>>,
         base_aggregation_number: u32,
     ) {
-        let mut task = ctx.task(task_id);
+        let mut task = ctx.task(task_id, TaskDataCategory::Meta);
         let current = get!(task, AggregationNumber).copied().unwrap_or_default();
         // The wanted new distance is either the provided one or the old distance
         let distance = base_effective_distance.map_or(current.distance, |d| d.get());

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

@@ -4,13 +4,16 @@ use serde::{Deserialize, Serialize};
 use turbo_tasks::TaskId;
 
 use crate::{
-    backend::operation::{
-        aggregation_update::{
-            get_aggregation_number, get_uppers, is_aggregating_node, AggregationUpdateJob,
-            AggregationUpdateQueue,
+    backend::{
+        operation::{
+            aggregation_update::{
+                get_aggregation_number, get_uppers, is_aggregating_node, AggregationUpdateJob,
+                AggregationUpdateQueue,
+            },
+            invalidate::make_task_dirty,
+            ExecuteContext, Operation,
         },
-        invalidate::make_task_dirty,
-        ExecuteContext, Operation,
+        TaskDataCategory,
     },
     data::{CachedDataItemKey, CellRef},
 };
@@ -69,7 +72,7 @@ impl Operation for CleanupOldEdgesOperation {
                     if let Some(edge) = outdated.pop() {
                         match edge {
                             OutdatedEdge::Child(child_id) => {
-                                let mut task = ctx.task(task_id);
+                                let mut task = ctx.task(task_id, TaskDataCategory::All);
                                 task.remove(&CachedDataItemKey::Child { task: child_id });
                                 if is_aggregating_node(get_aggregation_number(&task)) {
                                     queue.push(AggregationUpdateJob::InnerLostFollower {
@@ -89,14 +92,14 @@ impl Operation for CleanupOldEdgesOperation {
                                 cell,
                             }) => {
                                 {
-                                    let mut task = ctx.task(cell_task_id);
+                                    let mut task = ctx.task(cell_task_id, TaskDataCategory::Data);
                                     task.remove(&CachedDataItemKey::CellDependent {
                                         cell,
                                         task: task_id,
                                     });
                                 }
                                 {
-                                    let mut task = ctx.task(task_id);
+                                    let mut task = ctx.task(task_id, TaskDataCategory::Data);
                                     task.remove(&CachedDataItemKey::CellDependency {
                                         target: CellRef {
                                             task: cell_task_id,
@@ -107,13 +110,13 @@ impl Operation for CleanupOldEdgesOperation {
                             }
                             OutdatedEdge::OutputDependency(output_task_id) => {
                                 {
-                                    let mut task = ctx.task(output_task_id);
+                                    let mut task = ctx.task(output_task_id, TaskDataCategory::Data);
                                     task.remove(&CachedDataItemKey::OutputDependent {
                                         task: task_id,
                                     });
                                 }
                                 {
-                                    let mut task = ctx.task(task_id);
+                                    let mut task = ctx.task(task_id, TaskDataCategory::Data);
                                     task.remove(&CachedDataItemKey::OutputDependency {
                                         target: output_task_id,
                                     });