From 08c4e58aa835d60305e293bad91b6873ef6927b2 Mon Sep 17 00:00:00 2001 From: Tobias Koppers Date: Fri, 10 Jan 2025 17:57:56 +0100 Subject: [PATCH] Improve memory usage of the new backend (#74740) ### What? Improves memory usage by changing the way data is stored in the new backend. Instead of keeping a map `enum CachedDataItemKey` -> `enum CachedDataItemValue`, it uses a Vec `enum CachedDataItemStorage` which keeps a single storage for every `enum CachedDataItemType`. The storage is strongly types and e. g. stores a map `TaskId` -> `()`, which makes it more memory efficient while still not wasting memory for unused types. It also calls `shrink_to_fit` after task completion and `shrink_amortized` after removing items, to reduce the unnecessary capacity. --- .../src/backend/indexed.rs | 4 - .../turbo-tasks-backend/src/backend/mod.rs | 224 +++---- .../backend/operation/aggregation_update.rs | 26 +- .../src/backend/operation/mod.rs | 99 ++- .../src/backend/operation/update_cell.rs | 4 +- .../src/backend/operation/update_output.rs | 27 +- .../src/backend/storage.rs | 323 ++++------ .../crates/turbo-tasks-backend/src/data.rs | 108 +--- .../src/kv_backing_storage.rs | 3 +- .../src/derive/key_value_pair_macro.rs | 587 +++++++++++++++++- .../crates/turbo-tasks/src/key_value_pair.rs | 10 + 11 files changed, 854 insertions(+), 561 deletions(-) delete mode 100644 turbopack/crates/turbo-tasks-backend/src/backend/indexed.rs diff --git a/turbopack/crates/turbo-tasks-backend/src/backend/indexed.rs b/turbopack/crates/turbo-tasks-backend/src/backend/indexed.rs deleted file mode 100644 index b5598224c2eef..0000000000000 --- a/turbopack/crates/turbo-tasks-backend/src/backend/indexed.rs +++ /dev/null @@ -1,4 +0,0 @@ -pub trait Indexed { - type Index: Clone + PartialEq + Eq + std::hash::Hash; - fn index(&self) -> Self::Index; -} diff --git a/turbopack/crates/turbo-tasks-backend/src/backend/mod.rs b/turbopack/crates/turbo-tasks-backend/src/backend/mod.rs index c8bec9feaa506..56b9c5bed0988 100644 --- a/turbopack/crates/turbo-tasks-backend/src/backend/mod.rs +++ b/turbopack/crates/turbo-tasks-backend/src/backend/mod.rs @@ -1,4 +1,3 @@ -pub mod indexed; mod operation; mod persisted_storage_log; mod storage; @@ -49,9 +48,9 @@ use crate::{ }, backing_storage::BackingStorage, data::{ - ActiveType, AggregationNumber, CachedDataItem, CachedDataItemIndex, CachedDataItemKey, - CachedDataItemValue, CachedDataUpdate, CellRef, CollectibleRef, CollectiblesRef, - DirtyState, InProgressCellState, InProgressState, OutputValue, RootState, + ActiveType, AggregationNumber, CachedDataItem, CachedDataItemKey, CachedDataItemType, + CachedDataItemValue, CachedDataItemValueRef, CachedDataUpdate, CellRef, CollectibleRef, + CollectiblesRef, DirtyState, InProgressCellState, InProgressState, OutputValue, RootState, }, utils::{bi_map::BiMap, chunked_vec::ChunkedVec, ptr_eq_arc::PtrEqArc, sharded::Sharded}, }; @@ -151,7 +150,7 @@ struct TurboTasksBackendInner { persisted_storage_data_log: Option, persisted_storage_meta_log: Option, - storage: Storage, + storage: Storage, /// Number of executing operations + Highest bit is set when snapshot is /// requested. When that bit is set, operations should pause until the @@ -459,7 +458,7 @@ impl TurboTasksBackendInner { AggregatedDirtyContainer { task } count if count.get(self.session_id) > 0 => { - *task + task } ); if is_dirty { @@ -1001,13 +1000,8 @@ impl TurboTasksBackendInner { Current(TaskId), Outdated(TaskId), } - let children = task - .iter(CachedDataItemIndex::Children) - .filter_map(|(key, _)| match *key { - CachedDataItemKey::Child { task } => Some(Child::Current(task)), - CachedDataItemKey::OutdatedChild { task } => Some(Child::Outdated(task)), - _ => None, - }) + let children = iter_many!(task, Child { task } => Child::Current(task)) + .chain(iter_many!(task, OutdatedChild { task } => Child::Outdated(task))) .collect::>(); for child in children { match child { @@ -1030,18 +1024,8 @@ impl TurboTasksBackendInner { Current(CollectibleRef, i32), Outdated(CollectibleRef), } - let collectibles = task - .iter(CachedDataItemIndex::Collectibles) - .filter_map(|(key, value)| match (key, value) { - ( - &CachedDataItemKey::Collectible { collectible }, - &CachedDataItemValue::Collectible { value }, - ) => Some(Collectible::Current(collectible, value)), - (&CachedDataItemKey::OutdatedCollectible { collectible }, _) => { - Some(Collectible::Outdated(collectible)) - } - _ => None, - }) + let collectibles = iter_many!(task, Collectible { collectible } value => Collectible::Current(collectible, *value)) + .chain(iter_many!(task, OutdatedCollectible { collectible } => Collectible::Outdated(collectible))) .collect::>(); for collectible in collectibles { match collectible { @@ -1068,23 +1052,10 @@ impl TurboTasksBackendInner { OutdatedCell(CellRef), OutdatedOutput(TaskId), } - let dependencies = task - .iter(CachedDataItemIndex::Dependencies) - .filter_map(|(key, _)| match *key { - CachedDataItemKey::CellDependency { target } => { - Some(Dep::CurrentCell(target)) - } - CachedDataItemKey::OutputDependency { target } => { - Some(Dep::CurrentOutput(target)) - } - CachedDataItemKey::OutdatedCellDependency { target } => { - Some(Dep::OutdatedCell(target)) - } - CachedDataItemKey::OutdatedOutputDependency { target } => { - Some(Dep::OutdatedOutput(target)) - } - _ => None, - }) + let dependencies = iter_many!(task, CellDependency { target } => Dep::CurrentCell(target)) + .chain(iter_many!(task, OutputDependency { target } => Dep::CurrentOutput(target))) + .chain(iter_many!(task, OutdatedCellDependency { target } => Dep::OutdatedCell(target))) + .chain(iter_many!(task, OutdatedOutputDependency { target } => Dep::OutdatedOutput(target))) .collect::>(); for dep in dependencies { match dep { @@ -1237,7 +1208,7 @@ impl TurboTasksBackendInner { // handle cell counters: update max index and remove cells that are no longer used let mut old_counters: HashMap<_, _> = - get_many!(task, CellTypeMaxIndex { cell_type } max_index => (*cell_type, *max_index)); + get_many!(task, CellTypeMaxIndex { cell_type } max_index => (cell_type, *max_index)); for (&cell_type, &max_index) in cell_counters.iter() { if let Some(old_max_index) = old_counters.remove(&cell_type) { if old_max_index != max_index { @@ -1257,128 +1228,69 @@ impl TurboTasksBackendInner { task.remove(&CachedDataItemKey::CellTypeMaxIndex { cell_type }); } - let mut removed_data = Vec::new(); + let mut removed_data: Vec = Vec::new(); let mut old_edges = Vec::new(); // Remove no longer existing cells and notify in progress cells // find all outdated data items (removed cells, outdated edges) - if task.is_indexed() { - removed_data.extend(task.extract_if( - CachedDataItemIndex::InProgressCell, - |key, value| { - match (key, value) { - ( - &CachedDataItemKey::InProgressCell { cell }, - CachedDataItemValue::InProgressCell { value }, - ) if cell_counters - .get(&cell.type_id) - .is_none_or(|start_index| cell.index >= *start_index) => - { - value.event.notify(usize::MAX); - true - } - _ => false, - } - }, - )); - removed_data.extend(task.extract_if(CachedDataItemIndex::CellData, |key, _| { - matches!(key, &CachedDataItemKey::CellData { cell } if cell_counters - .get(&cell.type_id).is_none_or(|start_index| cell.index >= *start_index)) - })); - if self.should_track_children() { - old_edges.extend(task.iter(CachedDataItemIndex::Children).filter_map( - |(key, _)| match *key { - CachedDataItemKey::OutdatedChild { task } => { - Some(OutdatedEdge::Child(task)) - } - _ => None, - }, - )); - old_edges.extend(task.iter(CachedDataItemIndex::Collectibles).filter_map( - |(key, value)| match (key, value) { - ( - CachedDataItemKey::OutdatedCollectible { collectible }, - CachedDataItemValue::OutdatedCollectible { value }, - ) => Some(OutdatedEdge::Collectible(*collectible, *value)), - _ => None, - }, - )); - } - if self.should_track_dependencies() { - old_edges.extend(task.iter(CachedDataItemIndex::Dependencies).filter_map( - |(key, _): (&CachedDataItemKey, &CachedDataItemValue)| match *key { - CachedDataItemKey::OutdatedCellDependency { target } => { - Some(OutdatedEdge::CellDependency(target)) - } - CachedDataItemKey::OutdatedOutputDependency { target } => { - Some(OutdatedEdge::OutputDependency(target)) - } - _ => None, - }, - )); - old_edges.extend(task.iter(CachedDataItemIndex::CellDependent).filter_map( - |(key, _)| { - match *key { - CachedDataItemKey::CellDependent { cell, task } - if cell_counters - .get(&cell.type_id) - .is_none_or(|start_index| cell.index >= *start_index) => - { - Some(OutdatedEdge::RemovedCellDependent(task, cell.type_id)) - } - _ => None, - } - }, - )); - } - } else { - removed_data.extend(task.extract_if_all(|key, value| { + removed_data.extend( + task.extract_if(CachedDataItemType::InProgressCell, |key, value| { match (key, value) { ( - &CachedDataItemKey::InProgressCell { cell }, - CachedDataItemValue::InProgressCell { value }, + CachedDataItemKey::InProgressCell { cell }, + CachedDataItemValueRef::InProgressCell { value }, ) if cell_counters .get(&cell.type_id) .is_none_or(|start_index| cell.index >= *start_index) => { value.event.notify(usize::MAX); - return true; - } - (&CachedDataItemKey::CellData { cell }, _) - if cell_counters - .get(&cell.type_id) - .is_none_or(|start_index| cell.index >= *start_index) => - { - return true; - } - (&CachedDataItemKey::OutdatedChild { task }, _) => { - old_edges.push(OutdatedEdge::Child(task)); - } - ( - &CachedDataItemKey::OutdatedCollectible { collectible }, - &CachedDataItemValue::OutdatedCollectible { value }, - ) => old_edges.push(OutdatedEdge::Collectible(collectible, value)), - (&CachedDataItemKey::OutdatedCellDependency { target }, _) => { - old_edges.push(OutdatedEdge::CellDependency(target)); - } - (&CachedDataItemKey::OutdatedOutputDependency { target }, _) => { - old_edges.push(OutdatedEdge::OutputDependency(target)); + true } - (&CachedDataItemKey::OutdatedCollectiblesDependency { target }, _) => { - old_edges.push(OutdatedEdge::CollectiblesDependency(target)); - } - (&CachedDataItemKey::CellDependent { cell, task }, _) - if cell_counters - .get(&cell.type_id) - .is_none_or(|start_index| cell.index >= *start_index) => - { - old_edges.push(OutdatedEdge::RemovedCellDependent(task, cell.type_id)); - } - _ => {} + _ => false, } - false - })); - }; + }), + ); + removed_data.extend(task.extract_if(CachedDataItemType::CellData, |key, _| { + matches!(key, CachedDataItemKey::CellData { cell } if cell_counters + .get(&cell.type_id).is_none_or(|start_index| cell.index >= *start_index)) + })); + if self.should_track_children() { + old_edges.extend(task.iter(CachedDataItemType::OutdatedChild).filter_map( + |(key, _)| match key { + CachedDataItemKey::OutdatedChild { task } => Some(OutdatedEdge::Child(task)), + _ => None, + }, + )); + old_edges.extend( + task.iter(CachedDataItemType::OutdatedCollectible) + .filter_map(|(key, value)| match (key, value) { + ( + CachedDataItemKey::OutdatedCollectible { collectible }, + CachedDataItemValueRef::OutdatedCollectible { value }, + ) => Some(OutdatedEdge::Collectible(collectible, *value)), + _ => None, + }), + ); + } + if self.should_track_dependencies() { + old_edges.extend(iter_many!(task, OutdatedCellDependency { target } => OutdatedEdge::CellDependency(target))); + old_edges.extend(iter_many!(task, OutdatedOutputDependency { target } => OutdatedEdge::OutputDependency(target))); + old_edges.extend(task.iter(CachedDataItemType::CellDependent).filter_map( + |(key, _)| { + match key { + CachedDataItemKey::CellDependent { cell, task } + if cell_counters + .get(&cell.type_id) + .is_none_or(|start_index| cell.index >= *start_index) => + { + Some(OutdatedEdge::RemovedCellDependent(task, cell.type_id)) + } + _ => None, + } + }, + )); + } + drop(task); // Remove outdated edges first, before removing in_progress+dirty flag. @@ -1479,6 +1391,14 @@ impl TurboTasksBackendInner { drop(removed_data); + let mut task = ctx.task(task_id, TaskDataCategory::All); + task.shrink_to_fit(CachedDataItemType::CellData); + task.shrink_to_fit(CachedDataItemType::CellTypeMaxIndex); + task.shrink_to_fit(CachedDataItemType::CellDependency); + task.shrink_to_fit(CachedDataItemType::OutputDependency); + task.shrink_to_fit(CachedDataItemType::CollectiblesDependency); + drop(task); + false } diff --git a/turbopack/crates/turbo-tasks-backend/src/backend/operation/aggregation_update.rs b/turbopack/crates/turbo-tasks-backend/src/backend/operation/aggregation_update.rs index ac615c63f70ba..15a9ec258ac68 100644 --- a/turbopack/crates/turbo-tasks-backend/src/backend/operation/aggregation_update.rs +++ b/turbopack/crates/turbo-tasks-backend/src/backend/operation/aggregation_update.rs @@ -51,9 +51,9 @@ fn get_followers_with_aggregation_number( aggregation_number: u32, ) -> Vec { if is_aggregating_node(aggregation_number) { - get_many!(task, Follower { task } count if *count > 0 => *task) + get_many!(task, Follower { task } count if *count > 0 => task) } else { - get_many!(task, Child { task } => *task) + get_many!(task, Child { task } => task) } } @@ -66,12 +66,12 @@ fn get_followers(task: &impl TaskGuard) -> Vec { /// Returns a list of tasks that are considered as "upper" tasks of the task. The upper tasks are /// aggregating over the task. pub fn get_uppers(task: &impl TaskGuard) -> Vec { - get_many!(task, Upper { task } count if *count > 0 => *task) + get_many!(task, Upper { task } count if *count > 0 => task) } /// Returns an iterator of tasks that are considered as "upper" tasks of the task. See `get_uppers` fn iter_uppers<'a>(task: &'a (impl TaskGuard + 'a)) -> impl Iterator + 'a { - iter_many!(task, Upper { task } count if *count > 0 => *task) + iter_many!(task, Upper { task } count if *count > 0 => task) } /// Returns the aggregation number of the task. @@ -84,7 +84,7 @@ pub fn get_aggregation_number(task: &impl TaskGuard) -> u32 { /// A job in the job queue for updating something in the aggregated graph. #[derive(Serialize, Deserialize, Clone, Debug)] pub enum AggregationUpdateJob { - /// Update the aggregation number of a task. This might result in balancing needed to update + /// Update the aggregation number of a task. This might result in balancing needed to update /// "upper" and "follower" edges. UpdateAggregationNumber { task_id: TaskId, @@ -175,7 +175,7 @@ impl AggregatedDataUpdate { let aggregation = get_aggregation_number(task); let mut dirty_container_count = Default::default(); let mut collectibles_update: Vec<_> = - get_many!(task, Collectible { collectible } count => (*collectible, *count)); + get_many!(task, Collectible { collectible } count => (collectible, *count)); if is_aggregating_node(aggregation) { dirty_container_count = get!(task, AggregatedDirtyContainerCount) .cloned() @@ -185,7 +185,7 @@ impl AggregatedDataUpdate { AggregatedCollectible { collectible } count if *count > 0 => { - *collectible + collectible } ); for collectible in collectibles { @@ -321,8 +321,8 @@ impl AggregatedDataUpdate { CollectiblesDependent { collectible_type, task, - } if *collectible_type == ty => { - *task + } if collectible_type == ty => { + task } ); if !dependent.is_empty() { @@ -1027,7 +1027,7 @@ impl AggregationUpdateQueue { // if it has an `AggregateRoot` we can skip visiting the nested nodes since // this would already be scheduled by the `AggregateRoot` if !task.has_key(&CachedDataItemKey::AggregateRoot {}) { - let dirty_containers: Vec<_> = get_many!(task, AggregatedDirtyContainer { task } count if count.get(session_id) > 0 => *task); + let dirty_containers: Vec<_> = get_many!(task, AggregatedDirtyContainer { task } count if count.get(session_id) > 0 => task); if !dirty_containers.is_empty() || dirty { task.insert(CachedDataItem::AggregateRoot { value: RootState::new(ActiveType::CachedActiveUntilClean, task_id), @@ -1613,7 +1613,7 @@ impl AggregationUpdateQueue { if !is_aggregating_node(old) && is_aggregating_node(aggregation_number) { // When converted from leaf to aggregating node, all children become // followers - let children: Vec<_> = get_many!(task, Child { task } => *task); + let children: Vec<_> = get_many!(task, Child { task } => task); for child_id in children { task.add_new(CachedDataItem::Follower { task: child_id, @@ -1625,7 +1625,7 @@ impl AggregationUpdateQueue { if is_aggregating_node(aggregation_number) { // followers might become inner nodes when the aggregation number is // increased - let followers = iter_many!(task, Follower { task } count if *count > 0 => *task); + let followers = iter_many!(task, Follower { task } count if *count > 0 => task); for follower_id in followers { self.push(AggregationUpdateJob::BalanceEdge { upper_id: task_id, @@ -1637,7 +1637,7 @@ impl AggregationUpdateQueue { self.push(AggregationUpdateJob::BalanceEdge { upper_id, task_id }); } } else { - let children = iter_many!(task, Child { task } => *task); + let children = iter_many!(task, Child { task } => task); for child_id in children { self.push(AggregationUpdateJob::UpdateAggregationNumber { task_id: child_id, diff --git a/turbopack/crates/turbo-tasks-backend/src/backend/operation/mod.rs b/turbopack/crates/turbo-tasks-backend/src/backend/operation/mod.rs index 7b0c72b3e03e0..99d84b563ff1d 100644 --- a/turbopack/crates/turbo-tasks-backend/src/backend/operation/mod.rs +++ b/turbopack/crates/turbo-tasks-backend/src/backend/operation/mod.rs @@ -21,7 +21,10 @@ use crate::{ TurboTasksBackend, TurboTasksBackendInner, }, backing_storage::BackingStorage, - data::{CachedDataItem, CachedDataItemIndex, CachedDataItemKey, CachedDataItemValue}, + data::{ + CachedDataItem, CachedDataItemKey, CachedDataItemType, CachedDataItemValue, + CachedDataItemValueRef, CachedDataItemValueRefMut, + }, }; pub trait Operation: @@ -361,35 +364,31 @@ pub trait TaskGuard: Debug { fn insert(&mut self, item: CachedDataItem) -> Option; fn update( &mut self, - key: &CachedDataItemKey, + key: CachedDataItemKey, update: impl FnOnce(Option) -> Option, ); fn remove(&mut self, key: &CachedDataItemKey) -> Option; - fn get(&self, key: &CachedDataItemKey) -> Option<&CachedDataItemValue>; - fn get_mut(&mut self, key: &CachedDataItemKey) -> Option<&mut CachedDataItemValue>; + fn get(&self, key: &CachedDataItemKey) -> Option>; + fn get_mut(&mut self, key: &CachedDataItemKey) -> Option>; fn has_key(&self, key: &CachedDataItemKey) -> bool; - fn is_indexed(&self) -> bool; fn iter( &self, - index: CachedDataItemIndex, - ) -> impl Iterator; - fn iter_all(&self) -> impl Iterator; + ty: CachedDataItemType, + ) -> impl Iterator)>; + fn shrink_to_fit(&mut self, ty: CachedDataItemType); fn extract_if<'l, F>( &'l mut self, - index: CachedDataItemIndex, + ty: CachedDataItemType, f: F, ) -> impl Iterator where - F: for<'a, 'b> FnMut(&'a CachedDataItemKey, &'b CachedDataItemValue) -> bool + 'l; - fn extract_if_all<'l, F>(&'l mut self, f: F) -> impl Iterator - where - F: for<'a, 'b> FnMut(&'a CachedDataItemKey, &'b CachedDataItemValue) -> bool + 'l; + F: for<'a> FnMut(CachedDataItemKey, CachedDataItemValueRef<'a>) -> bool + 'l; fn invalidate_serialization(&mut self); } struct TaskGuardImpl<'a, B: BackingStorage> { task_id: TaskId, - task: StorageWriteGuard<'a, TaskId, CachedDataItem>, + task: StorageWriteGuard<'a>, backend: &'a TurboTasksBackendInner, #[cfg(debug_assertions)] category: TaskDataCategory, @@ -474,10 +473,9 @@ impl TaskGuard for TaskGuardImpl<'_, B> { self.task .insert(CachedDataItem::from_key_and_value(key, value)) } else if value.is_persistent() { - let old = self.task.insert(CachedDataItem::from_key_and_value( - key.clone(), - value.clone(), - )); + let old = self + .task + .insert(CachedDataItem::from_key_and_value(key, value.clone())); self.task.persistance_state_mut().add_persisting_item(); self.backend .persisted_storage_log(key.category()) @@ -491,7 +489,7 @@ impl TaskGuard for TaskGuardImpl<'_, B> { ); old } else { - let item = CachedDataItem::from_key_and_value(key.clone(), value); + let item = CachedDataItem::from_key_and_value(key, value); if let Some(old) = self.task.insert(item) { if old.is_persistent() { self.task.persistance_state_mut().add_persisting_item(); @@ -509,7 +507,7 @@ impl TaskGuard for TaskGuardImpl<'_, B> { fn update( &mut self, - key: &CachedDataItemKey, + key: CachedDataItemKey, update: impl FnOnce(Option) -> Option, ) { self.check_access(key.category()); @@ -538,7 +536,7 @@ impl TaskGuard for TaskGuardImpl<'_, B> { add_persisting_item = true; backend.persisted_storage_log(key.category()).unwrap().push( *task_id, - key.clone(), + key, Some(old_value), None, ); @@ -547,7 +545,7 @@ impl TaskGuard for TaskGuardImpl<'_, B> { add_persisting_item = true; backend.persisted_storage_log(key.category()).unwrap().push( *task_id, - key.clone(), + key, old_value, new.clone(), ); @@ -570,14 +568,13 @@ impl TaskGuard for TaskGuardImpl<'_, B> { && key.is_persistent() && value.is_persistent() { - let key = key.clone(); self.task.persistance_state_mut().add_persisting_item(); self.backend .persisted_storage_log(key.category()) .unwrap() .push( self.task_id, - key, + *key, value.is_persistent().then(|| value.clone()), None, ); @@ -588,12 +585,12 @@ impl TaskGuard for TaskGuardImpl<'_, B> { } } - fn get(&self, key: &CachedDataItemKey) -> Option<&CachedDataItemValue> { + fn get(&self, key: &CachedDataItemKey) -> Option> { self.check_access(key.category()); self.task.get(key) } - fn get_mut(&mut self, key: &CachedDataItemKey) -> Option<&mut CachedDataItemValue> { + fn get_mut(&mut self, key: &CachedDataItemKey) -> Option> { self.check_access(key.category()); self.task.get_mut(key) } @@ -603,52 +600,29 @@ impl TaskGuard for TaskGuardImpl<'_, B> { self.task.has_key(key) } - fn is_indexed(&self) -> bool { - self.task.is_indexed() - } - fn iter( &self, - index: CachedDataItemIndex, - ) -> impl Iterator { - self.task.iter(Some(index)) + ty: CachedDataItemType, + ) -> impl Iterator)> { + self.task.iter(ty) } - fn iter_all(&self) -> impl Iterator { - self.task.iter_all() + fn shrink_to_fit(&mut self, ty: CachedDataItemType) { + self.task.shrink_to_fit(ty) } fn extract_if<'l, F>( &'l mut self, - index: CachedDataItemIndex, + ty: CachedDataItemType, f: F, ) -> impl Iterator where - F: for<'a, 'b> FnMut(&'a CachedDataItemKey, &'b CachedDataItemValue) -> bool + 'l, - { - if !self.backend.should_persist() || self.task_id.is_transient() { - return Either::Left(self.task.extract_if(Some(index), f)); - } - Either::Right(self.task.extract_if(Some(index), f).inspect(|item| { - if item.is_persistent() { - let key = item.key(); - let value = item.value(); - self.backend - .persisted_storage_log(key.category()) - .unwrap() - .push(self.task_id, key, Some(value), None); - } - })) - } - - fn extract_if_all<'l, F>(&'l mut self, f: F) -> impl Iterator - where - F: for<'a, 'b> FnMut(&'a CachedDataItemKey, &'b CachedDataItemValue) -> bool + 'l, + F: for<'a> FnMut(CachedDataItemKey, CachedDataItemValueRef<'a>) -> bool + 'l, { if !self.backend.should_persist() || self.task_id.is_transient() { - return Either::Left(self.task.extract_if_all(f)); + return Either::Left(self.task.extract_if(ty, f)); } - Either::Right(self.task.extract_if_all(f).inspect(|item| { + Either::Right(self.task.extract_if(ty, f).inspect(|item| { if item.is_persistent() { let key = item.key(); let value = item.value(); @@ -666,12 +640,15 @@ impl TaskGuard for TaskGuardImpl<'_, B> { } let mut count = 0; let cell_data = self - .iter(CachedDataItemIndex::CellData) + .iter(CachedDataItemType::CellData) .filter_map(|(key, value)| match (key, value) { - (CachedDataItemKey::CellData { cell }, CachedDataItemValue::CellData { value }) => { + ( + CachedDataItemKey::CellData { cell }, + CachedDataItemValueRef::CellData { value }, + ) => { count += 1; Some(CachedDataItem::CellData { - cell: *cell, + cell, value: value.clone(), }) } diff --git a/turbopack/crates/turbo-tasks-backend/src/backend/operation/update_cell.rs b/turbopack/crates/turbo-tasks-backend/src/backend/operation/update_cell.rs index ffd2daa2ab459..3767c8aa8cb53 100644 --- a/turbopack/crates/turbo-tasks-backend/src/backend/operation/update_cell.rs +++ b/turbopack/crates/turbo-tasks-backend/src/backend/operation/update_cell.rs @@ -41,8 +41,8 @@ impl UpdateCellOperation { let dependent = get_many!( task, CellDependent { cell: dependent_cell, task } - if *dependent_cell == cell - => *task + if dependent_cell == cell + => task ); drop(task); diff --git a/turbopack/crates/turbo-tasks-backend/src/backend/operation/update_output.rs b/turbopack/crates/turbo-tasks-backend/src/backend/operation/update_output.rs index ca108a9782a21..79637c5862a0b 100644 --- a/turbopack/crates/turbo-tasks-backend/src/backend/operation/update_output.rs +++ b/turbopack/crates/turbo-tasks-backend/src/backend/operation/update_output.rs @@ -13,10 +13,7 @@ use crate::{ storage::{get, get_many}, TaskDataCategory, }, - data::{ - CachedDataItem, CachedDataItemKey, CachedDataItemValue, CellRef, InProgressState, - OutputValue, - }, + data::{CachedDataItem, CachedDataItemKey, CellRef, InProgressState, OutputValue}, }; #[derive(Serialize, Deserialize, Clone, Default)] @@ -50,13 +47,10 @@ impl UpdateOutputOperation { return; } let old_error = task.remove(&CachedDataItemKey::Error {}); - let current_output = task.get(&CachedDataItemKey::Output {}); + let current_output = get!(task, Output); let output_value = match output { Ok(Ok(RawVc::TaskOutput(output_task_id))) => { - if let Some(CachedDataItemValue::Output { - value: OutputValue::Output(current_task_id), - }) = current_output - { + if let Some(OutputValue::Output(current_task_id)) = current_output { if *current_task_id == output_task_id { return; } @@ -64,13 +58,10 @@ impl UpdateOutputOperation { OutputValue::Output(output_task_id) } Ok(Ok(RawVc::TaskCell(output_task_id, cell))) => { - if let Some(CachedDataItemValue::Output { - value: - OutputValue::Cell(CellRef { - task: current_task_id, - cell: current_cell, - }), - }) = current_output + if let Some(OutputValue::Cell(CellRef { + task: current_task_id, + cell: current_cell, + })) = current_output { if *current_task_id == output_task_id && *current_cell == cell { return; @@ -113,11 +104,11 @@ impl UpdateOutputOperation { let dependent_tasks = ctx .should_track_dependencies() - .then(|| get_many!(task, OutputDependent { task } => *task)) + .then(|| get_many!(task, OutputDependent { task } => task)) .unwrap_or_default(); let children = ctx .should_track_children() - .then(|| get_many!(task, Child { task } => *task)) + .then(|| get_many!(task, Child { task } => task)) .unwrap_or_default(); let mut queue = AggregationUpdateQueue::new(); diff --git a/turbopack/crates/turbo-tasks-backend/src/backend/storage.rs b/turbopack/crates/turbo-tasks-backend/src/backend/storage.rs index e7a633f5ea12b..328cca2d53c2b 100644 --- a/turbopack/crates/turbo-tasks-backend/src/backend/storage.rs +++ b/turbopack/crates/turbo-tasks-backend/src/backend/storage.rs @@ -1,19 +1,19 @@ use std::{ hash::{BuildHasherDefault, Hash}, - mem::take, ops::{Deref, DerefMut}, - panic, thread::available_parallelism, }; -use auto_hash_map::{map::Entry, AutoMap}; use dashmap::DashMap; use either::Either; use rustc_hash::FxHasher; -use turbo_tasks::KeyValuePair; +use turbo_tasks::{KeyValuePair, TaskId}; use crate::{ - backend::indexed::Indexed, + data::{ + CachedDataItem, CachedDataItemKey, CachedDataItemStorage, CachedDataItemType, + CachedDataItemValue, CachedDataItemValueRef, CachedDataItemValueRefMut, + }, utils::dash_map_multi::{get_multiple_mut, RefMut}, }; @@ -119,246 +119,155 @@ impl PersistanceState { } } -const INDEX_THRESHOLD: usize = 128; - -type IndexedMap = AutoMap< - <::Key as Indexed>::Index, - AutoMap<::Key, ::Value>, ->; +type InnerStorageMap = Vec; -pub enum InnerStorage -where - T::Key: Indexed, -{ - Plain { - map: AutoMap, - persistance_state: PersistanceState, - }, - Indexed { - map: IndexedMap, - persistance_state: PersistanceState, - }, +pub struct InnerStorage { + map: InnerStorageMap, + persistance_state: PersistanceState, } -impl InnerStorage -where - T::Key: Indexed, -{ +impl InnerStorage { fn new() -> Self { - Self::Plain { - map: AutoMap::new(), + Self { + map: Default::default(), persistance_state: PersistanceState::default(), } } pub fn persistance_state(&self) -> &PersistanceState { - match self { - InnerStorage::Plain { - persistance_state, .. - } => persistance_state, - InnerStorage::Indexed { - persistance_state, .. - } => persistance_state, - } + &self.persistance_state } pub fn persistance_state_mut(&mut self) -> &mut PersistanceState { - match self { - InnerStorage::Plain { - persistance_state, .. - } => persistance_state, - InnerStorage::Indexed { - persistance_state, .. - } => persistance_state, - } + &mut self.persistance_state } - fn check_threshold(&mut self) { - let InnerStorage::Plain { - map: plain_map, - persistance_state, - } = self - else { - return; - }; - if plain_map.len() >= INDEX_THRESHOLD { - let mut map: IndexedMap = AutoMap::new(); - for (key, value) in take(plain_map).into_iter() { - let index = key.index(); - map.entry(index).or_default().insert(key, value); - } - *self = InnerStorage::Indexed { - map, - persistance_state: take(persistance_state), - }; + fn get_or_create_map_mut(&mut self, ty: CachedDataItemType) -> &mut CachedDataItemStorage { + let i = self.map.iter().position(|m| m.ty() == ty); + if let Some(i) = i { + &mut self.map[i] + } else { + self.map.push(CachedDataItemStorage::new(ty)); + self.map.last_mut().unwrap() } } - fn get_or_create_map_mut(&mut self, key: &T::Key) -> &mut AutoMap { - self.check_threshold(); - match self { - InnerStorage::Plain { map, .. } => map, - InnerStorage::Indexed { map, .. } => map.entry(key.index()).or_default(), - } + fn get_map_mut(&mut self, ty: CachedDataItemType) -> Option<&mut CachedDataItemStorage> { + self.map.iter_mut().find(|m| m.ty() == ty) } - fn get_map_mut(&mut self, key: &T::Key) -> Option<&mut AutoMap> { - self.check_threshold(); - match self { - InnerStorage::Plain { map, .. } => Some(map), - InnerStorage::Indexed { map, .. } => map.get_mut(&key.index()), - } + fn get_map_index(&mut self, ty: CachedDataItemType) -> Option { + self.map.iter_mut().position(|m| m.ty() == ty) } - fn get_map(&self, key: &T::Key) -> Option<&AutoMap> { - match self { - InnerStorage::Plain { map, .. } => Some(map), - InnerStorage::Indexed { map, .. } => map.get(&key.index()), - } + fn get_map(&self, ty: CachedDataItemType) -> Option<&CachedDataItemStorage> { + self.map.iter().find(|m| m.ty() == ty) } - fn index_map(&self, index: ::Index) -> Option<&AutoMap> { - match self { - InnerStorage::Plain { map, .. } => Some(map), - InnerStorage::Indexed { map, .. } => map.get(&index), - } + pub fn add(&mut self, item: CachedDataItem) -> bool { + let ty = item.ty(); + self.get_or_create_map_mut(ty).add(item) } - fn index_map_mut( - &mut self, - index: ::Index, - ) -> Option<&mut AutoMap> { - match self { - InnerStorage::Plain { map, .. } => Some(map), - InnerStorage::Indexed { map, .. } => map.get_mut(&index), - } + pub fn insert(&mut self, item: CachedDataItem) -> Option { + let ty = item.ty(); + self.get_or_create_map_mut(ty).insert(item) } - pub fn add(&mut self, item: T) -> bool { - let (key, value) = item.into_key_and_value(); - match self.get_or_create_map_mut(&key).entry(key) { - Entry::Occupied(_) => false, - Entry::Vacant(e) => { - e.insert(value); - true + pub fn remove(&mut self, key: &CachedDataItemKey) -> Option { + self.get_map_index(key.ty()).and_then(|i| { + let storage = &mut self.map[i]; + let result = storage.remove(key); + if result.is_some() && storage.is_empty() { + self.map.swap_remove(i); } - } - } - - pub fn insert(&mut self, item: T) -> Option { - let (key, value) = item.into_key_and_value(); - self.get_or_create_map_mut(&key).insert(key, value) - } - - pub fn remove(&mut self, key: &T::Key) -> Option { - self.get_map_mut(key).and_then(|m| m.remove(key)) + result + }) } - pub fn get(&self, key: &T::Key) -> Option<&T::Value> { - self.get_map(key).and_then(|m| m.get(key)) + pub fn get(&self, key: &CachedDataItemKey) -> Option { + self.get_map(key.ty()).and_then(|m| m.get(key)) } - pub fn get_mut(&mut self, key: &T::Key) -> Option<&mut T::Value> { - self.get_map_mut(key).and_then(|m| m.get_mut(key)) + pub fn get_mut(&mut self, key: &CachedDataItemKey) -> Option { + self.get_map_mut(key.ty()).and_then(|m| m.get_mut(key)) } - pub fn has_key(&self, key: &T::Key) -> bool { - self.get_map(key) + pub fn has_key(&self, key: &CachedDataItemKey) -> bool { + self.get_map(key.ty()) .map(|m| m.contains_key(key)) .unwrap_or_default() } - pub fn is_indexed(&self) -> bool { - matches!(self, InnerStorage::Indexed { .. }) - } - pub fn iter( &self, - index: ::Index, - ) -> impl Iterator { - self.index_map(index) - .map(|m| m.iter()) - .into_iter() - .flatten() + ty: CachedDataItemType, + ) -> impl Iterator)> { + self.get_map(ty).map(|m| m.iter()).into_iter().flatten() } - pub fn iter_all(&self) -> impl Iterator { - match self { - InnerStorage::Plain { map, .. } => Either::Left(map.iter()), - InnerStorage::Indexed { map, .. } => { - Either::Right(map.iter().flat_map(|(_, m)| m.iter())) - } - } + pub fn iter_all( + &self, + ) -> impl Iterator)> { + self.map.iter().flat_map(|m| m.iter()) } pub fn extract_if<'l, F>( &'l mut self, - index: ::Index, + ty: CachedDataItemType, mut f: F, - ) -> impl Iterator + use<'l, T, F> + ) -> impl Iterator + use<'l, F> where - F: for<'a, 'b> FnMut(&'a T::Key, &'b T::Value) -> bool + 'l, + F: for<'a> FnMut(CachedDataItemKey, CachedDataItemValueRef<'a>) -> bool + 'l, { - self.index_map_mut(index) - .map(move |m| m.extract_if(move |k, v| f(k, v))) + // TODO this could be more efficient when the storage would support extract_if directly. + // This requires some macro magic to make it work... + // But we could potentially avoid the two temporary Vecs. + let Some(i) = self.get_map_index(ty) else { + return Either::Left(std::iter::empty()); + }; + let storage = &mut self.map[i]; + let items_to_extract = storage + .iter() + .filter(|(k, v)| f(*k, *v)) + .map(|(key, _)| key) + .collect::>(); + let items = items_to_extract .into_iter() - .flatten() - .map(|(key, value)| T::from_key_and_value(key, value)) - } - - pub fn extract_if_all<'l, F>(&'l mut self, mut f: F) -> impl Iterator + use<'l, T, F> - where - F: for<'a, 'b> FnMut(&'a T::Key, &'b T::Value) -> bool + 'l, - { - match self { - InnerStorage::Plain { map, .. } => map - .extract_if(move |k, v| f(k, v)) - .map(|(key, value)| T::from_key_and_value(key, value)), - InnerStorage::Indexed { .. } => { - panic!("Do not use extract_if_all with indexed storage") - } + .map(move |key| { + let value = storage.remove(&key).unwrap(); + CachedDataItem::from_key_and_value(key, value) + }) + .collect::>(); + if self.map[i].is_empty() { + self.map.swap_remove(i); } + Either::Right(items.into_iter()) } -} -impl InnerStorage -where - T::Key: Indexed, - T::Value: Default, - T::Key: Clone, -{ pub fn update( &mut self, - key: &T::Key, - update: impl FnOnce(Option) -> Option, + key: CachedDataItemKey, + update: impl FnOnce(Option) -> Option, ) { - let map = self.get_or_create_map_mut(key); - if let Some(value) = map.get_mut(key) { - let v = take(value); - if let Some(v) = update(Some(v)) { - *value = v; - } else { - map.remove(key); - } - } else if let Some(v) = update(None) { - map.insert(key.clone(), v); + let map = self.get_or_create_map_mut(key.ty()); + if let Some(v) = update(map.remove(&key)) { + map.insert(CachedDataItem::from_key_and_value(key, v)); + } + } + + pub fn shrink_to_fit(&mut self, ty: CachedDataItemType) { + if let Some(map) = self.get_map_mut(ty) { + map.shrink_to_fit(); } } } -pub struct Storage -where - T::Key: Indexed, -{ - map: DashMap, BuildHasherDefault>, +pub struct Storage { + map: DashMap>, } -impl Storage -where - T: KeyValuePair, - T::Key: Indexed, - K: Eq + std::hash::Hash + Clone, -{ +impl Storage { pub fn new() -> Self { let shard_amount = (available_parallelism().map_or(4, |v| v.get()) * 64).next_power_of_two(); @@ -371,7 +280,7 @@ where } } - pub fn access_mut(&self, key: K) -> StorageWriteGuard<'_, K, T> { + pub fn access_mut(&self, key: TaskId) -> StorageWriteGuard<'_> { let inner = match self.map.entry(key) { dashmap::mapref::entry::Entry::Occupied(e) => e.into_ref(), dashmap::mapref::entry::Entry::Vacant(e) => e.insert(InnerStorage::new()), @@ -383,10 +292,10 @@ where pub fn access_pair_mut( &self, - key1: K, - key2: K, - ) -> (StorageWriteGuard<'_, K, T>, StorageWriteGuard<'_, K, T>) { - let (a, b) = get_multiple_mut(&self.map, key1, key2, || InnerStorage::new()); + key1: TaskId, + key2: TaskId, + ) -> (StorageWriteGuard<'_>, StorageWriteGuard<'_>) { + let (a, b) = get_multiple_mut(&self.map, key1, key2, InnerStorage::new); ( StorageWriteGuard { inner: a }, StorageWriteGuard { inner: b }, @@ -394,33 +303,19 @@ where } } -pub struct StorageWriteGuard<'a, K, T> -where - T: KeyValuePair, - T::Key: Indexed, -{ - inner: RefMut<'a, K, InnerStorage>, +pub struct StorageWriteGuard<'a> { + inner: RefMut<'a, TaskId, InnerStorage>, } -impl Deref for StorageWriteGuard<'_, K, T> -where - T: KeyValuePair, - T::Key: Indexed, - K: Eq + Hash, -{ - type Target = InnerStorage; +impl Deref for StorageWriteGuard<'_> { + type Target = InnerStorage; fn deref(&self) -> &Self::Target { &self.inner } } -impl DerefMut for StorageWriteGuard<'_, K, T> -where - T: KeyValuePair, - T::Key: Indexed, - K: Eq + Hash, -{ +impl DerefMut for StorageWriteGuard<'_> { fn deref_mut(&mut self) -> &mut Self::Target { &mut self.inner } @@ -430,9 +325,9 @@ macro_rules! get { ($task:ident, $key:ident $input:tt) => {{ #[allow(unused_imports)] use $crate::backend::operation::TaskGuard; - if let Some($crate::data::CachedDataItemValue::$key { + if let Some($crate::data::CachedDataItemValueRef::$key { value, - }) = $task.get(&$crate::data::CachedDataItemKey::$key $input).as_ref() { + }) = $task.get(&$crate::data::CachedDataItemKey::$key $input) { Some(value) } else { None @@ -447,7 +342,7 @@ macro_rules! get_mut { ($task:ident, $key:ident $input:tt) => {{ #[allow(unused_imports)] use $crate::backend::operation::TaskGuard; - if let Some($crate::data::CachedDataItemValue::$key { + if let Some($crate::data::CachedDataItemValueRefMut::$key { value, }) = $task.get_mut(&$crate::data::CachedDataItemKey::$key $input).as_mut() { let () = $crate::data::allow_mut_access::$key; @@ -471,7 +366,7 @@ macro_rules! iter_many { #[allow(unused_imports)] use $crate::backend::operation::TaskGuard; $task - .iter($crate::data::indicies::$key) + .iter($crate::data::CachedDataItemType::$key) .filter_map(|(key, _)| match key { $crate::data::CachedDataItemKey::$key $key_pattern $(if $cond)? => Some( $iter_item @@ -483,11 +378,11 @@ macro_rules! iter_many { #[allow(unused_imports)] use $crate::backend::operation::TaskGuard; $task - .iter($crate::data::indicies::$key) + .iter($crate::data::CachedDataItemType::$key) .filter_map(|(key, value)| match (key, value) { ( $crate::data::CachedDataItemKey::$key $input, - $crate::data::CachedDataItemValue::$key { value: $value_pattern } + $crate::data::CachedDataItemValueRef::$key { value: $value_pattern } ) $(if $cond)? => Some($iter_item), _ => None, }) @@ -510,7 +405,7 @@ macro_rules! update { use $crate::backend::operation::TaskGuard; #[allow(unused_mut)] let mut update = $update; - $task.update(&$crate::data::CachedDataItemKey::$key $input, |old| { + $task.update($crate::data::CachedDataItemKey::$key $input, |old| { update(old.and_then(|old| { if let $crate::data::CachedDataItemValue::$key { value } = old { Some(value) diff --git a/turbopack/crates/turbo-tasks-backend/src/data.rs b/turbopack/crates/turbo-tasks-backend/src/data.rs index 0eb709750e63b..7cecd0caca081 100644 --- a/turbopack/crates/turbo-tasks-backend/src/data.rs +++ b/turbopack/crates/turbo-tasks-backend/src/data.rs @@ -8,7 +8,7 @@ use turbo_tasks::{ CellId, KeyValuePair, SessionId, TaskId, TraitTypeId, TypedSharedReference, ValueTypeId, }; -use crate::backend::{indexed::Indexed, TaskDataCategory}; +use crate::backend::TaskDataCategory; // this traits are needed for the transient variants of `CachedDataItem` // transient variants are never cloned or compared @@ -623,100 +623,6 @@ pub mod allow_mut_access { pub const AggregateRoot: () = (); } -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub enum CachedDataItemIndex { - Children, - Collectibles, - Follower, - Upper, - AggregatedDirtyContainer, - AggregatedCollectible, - CellData, - CellTypeMaxIndex, - CellDependent, - OutputDependent, - CollectiblesDependent, - Dependencies, - InProgressCell, -} - -#[allow(non_upper_case_globals, dead_code)] -pub mod indicies { - use super::CachedDataItemIndex; - - pub const Child: CachedDataItemIndex = CachedDataItemIndex::Children; - pub const OutdatedChild: CachedDataItemIndex = CachedDataItemIndex::Children; - pub const Collectible: CachedDataItemIndex = CachedDataItemIndex::Collectibles; - pub const OutdatedCollectible: CachedDataItemIndex = CachedDataItemIndex::Collectibles; - pub const Follower: CachedDataItemIndex = CachedDataItemIndex::Follower; - pub const Upper: CachedDataItemIndex = CachedDataItemIndex::Upper; - pub const AggregatedDirtyContainer: CachedDataItemIndex = - CachedDataItemIndex::AggregatedDirtyContainer; - pub const AggregatedCollectible: CachedDataItemIndex = - CachedDataItemIndex::AggregatedCollectible; - pub const CellData: CachedDataItemIndex = CachedDataItemIndex::CellData; - pub const CellTypeMaxIndex: CachedDataItemIndex = CachedDataItemIndex::CellTypeMaxIndex; - pub const CellDependent: CachedDataItemIndex = CachedDataItemIndex::CellDependent; - pub const OutputDependent: CachedDataItemIndex = CachedDataItemIndex::OutputDependent; - pub const CollectiblesDependent: CachedDataItemIndex = - CachedDataItemIndex::CollectiblesDependent; - pub const OutputDependency: CachedDataItemIndex = CachedDataItemIndex::Dependencies; - pub const CellDependency: CachedDataItemIndex = CachedDataItemIndex::Dependencies; - pub const CollectibleDependency: CachedDataItemIndex = CachedDataItemIndex::Dependencies; - pub const OutdatedOutputDependency: CachedDataItemIndex = CachedDataItemIndex::Dependencies; - pub const OutdatedCellDependency: CachedDataItemIndex = CachedDataItemIndex::Dependencies; - pub const OutdatedCollectiblesDependency: CachedDataItemIndex = - CachedDataItemIndex::Dependencies; - pub const OutdatedCollectibleDependency: CachedDataItemIndex = - CachedDataItemIndex::Dependencies; - pub const InProgressCell: CachedDataItemIndex = CachedDataItemIndex::InProgressCell; -} - -impl Indexed for CachedDataItemKey { - type Index = Option; - - fn index(&self) -> Option { - match self { - CachedDataItemKey::Child { .. } => Some(CachedDataItemIndex::Children), - CachedDataItemKey::OutdatedChild { .. } => Some(CachedDataItemIndex::Children), - CachedDataItemKey::Collectible { .. } => Some(CachedDataItemIndex::Collectibles), - CachedDataItemKey::OutdatedCollectible { .. } => { - Some(CachedDataItemIndex::Collectibles) - } - CachedDataItemKey::Follower { .. } => Some(CachedDataItemIndex::Follower), - CachedDataItemKey::Upper { .. } => Some(CachedDataItemIndex::Upper), - CachedDataItemKey::AggregatedDirtyContainer { .. } => { - Some(CachedDataItemIndex::AggregatedDirtyContainer) - } - CachedDataItemKey::AggregatedCollectible { .. } => { - Some(CachedDataItemIndex::AggregatedCollectible) - } - CachedDataItemKey::CellData { .. } => Some(CachedDataItemIndex::CellData), - CachedDataItemKey::CellTypeMaxIndex { .. } => { - Some(CachedDataItemIndex::CellTypeMaxIndex) - } - CachedDataItemKey::CellDependent { .. } => Some(CachedDataItemIndex::CellDependent), - CachedDataItemKey::OutputDependent { .. } => Some(CachedDataItemIndex::OutputDependent), - CachedDataItemKey::OutputDependency { .. } => Some(CachedDataItemIndex::Dependencies), - CachedDataItemKey::CellDependency { .. } => Some(CachedDataItemIndex::Dependencies), - CachedDataItemKey::CollectiblesDependency { .. } => { - Some(CachedDataItemIndex::Dependencies) - } - CachedDataItemKey::OutdatedOutputDependency { .. } => { - Some(CachedDataItemIndex::Dependencies) - } - CachedDataItemKey::OutdatedCellDependency { .. } => { - Some(CachedDataItemIndex::Dependencies) - } - CachedDataItemKey::OutdatedCollectiblesDependency { .. } => { - Some(CachedDataItemIndex::Dependencies) - } - CachedDataItemKey::InProgressCell { .. } => Some(CachedDataItemIndex::InProgressCell), - _ => None, - } - } -} - impl CachedDataItemValue { pub fn is_persistent(&self) -> bool { match self { @@ -742,3 +648,15 @@ pub enum CachedDataUpdate { /// An item was replaced. This is step 2 and tells about the new value. Replace2 { value: CachedDataItemValue }, } + +#[cfg(test)] +mod tests { + #[test] + fn test_sizes() { + assert_eq!(std::mem::size_of::(), 40); + assert_eq!(std::mem::size_of::(), 20); + assert_eq!(std::mem::size_of::(), 32); + assert_eq!(std::mem::size_of::(), 48); + assert_eq!(std::mem::size_of::(), 48); + } +} diff --git a/turbopack/crates/turbo-tasks-backend/src/kv_backing_storage.rs b/turbopack/crates/turbo-tasks-backend/src/kv_backing_storage.rs index 50d7aedd83a4f..283c353d90ad1 100644 --- a/turbopack/crates/turbo-tasks-backend/src/kv_backing_storage.rs +++ b/turbopack/crates/turbo-tasks-backend/src/kv_backing_storage.rs @@ -812,7 +812,8 @@ struct SerializeLikeCachedDataItem<'l>(&'l CachedDataItemKey, &'l CachedDataItem impl Serialize for SerializeLikeCachedDataItem<'_> { fn serialize(&self, serializer: S) -> Result { - let item = CachedDataItem::from_key_and_value(self.0.clone(), self.1.clone()); + // TODO add CachedDataItemRef to avoid cloning + let item = CachedDataItem::from_key_and_value(*self.0, self.1.clone()); item.serialize(serializer) } } diff --git a/turbopack/crates/turbo-tasks-macros/src/derive/key_value_pair_macro.rs b/turbopack/crates/turbo-tasks-macros/src/derive/key_value_pair_macro.rs index 3753a920a79ff..05692a23053f7 100644 --- a/turbopack/crates/turbo-tasks-macros/src/derive/key_value_pair_macro.rs +++ b/turbopack/crates/turbo-tasks-macros/src/derive/key_value_pair_macro.rs @@ -7,8 +7,13 @@ pub fn derive_key_value_pair(input: TokenStream) -> TokenStream { let ident = &input.ident; let vis = &input.vis; + let type_name = Ident::new(&format!("{}Type", input.ident), input.ident.span()); let key_name = Ident::new(&format!("{}Key", input.ident), input.ident.span()); let value_name = Ident::new(&format!("{}Value", input.ident), input.ident.span()); + let value_ref_name = Ident::new(&format!("{}ValueRef", input.ident), input.ident.span()); + let value_ref_mut_name = Ident::new(&format!("{}ValueRefMut", input.ident), input.ident.span()); + let iter_name = Ident::new(&format!("{}Iter", input.ident), input.ident.span()); + let storage_name = Ident::new(&format!("{}Storage", input.ident), input.ident.span()); let variant_names = input .variants @@ -58,10 +63,312 @@ pub fn derive_key_value_pair(input: TokenStream) -> TokenStream { let value_pat = patterns(&value_fields); let value_clone_fields = clone_fields(&value_fields); + let value_ref_decl = ref_field_declarations(&value_fields); + let value_ref_mut_decl = mut_ref_field_declarations(&value_fields); + let value_ref_fields = ref_fields(&value_fields); + + let storage = key_fields + .iter() + .zip(value_fields.iter()).zip(variant_names.iter()) + .map(|((key_fields, value_fields), variant_name)| { + let value_types = value_fields + .iter() + .map(|field| { + let ty = &field.ty; + quote! { + #ty + } + }) + .collect::>(); + let key_types = key_fields + .iter() + .map(|field| { + let ty = &field.ty; + quote! { + #ty + } + }) + .collect::>(); + let value_fields = value_fields + .iter() + .map(|field| { + let ident = field.ident.as_ref().unwrap(); + quote! { + #ident + } + }) + .collect::>(); + let key_fields = key_fields + .iter() + .map(|field| { + let ident = field.ident.as_ref().unwrap(); + quote! { + #ident + } + }) + .collect::>(); + match (key_types.len(), value_types.len()) { + (0, 1) => { + StorageDecl { + decl: quote! { + storage: Option<#(#value_types)*>, + }, + add: quote! { + (#storage_name::#variant_name { storage }, #ident::#variant_name { #(#value_fields)* }) => { + if storage.is_none() { + *storage = Some(#(#value_fields)*); + true + } else { + false + } + } + }, + insert: quote! { + (#storage_name::#variant_name { storage }, #ident::#variant_name { #(#value_fields)* }) => { + std::mem::replace(storage, Some(#(#value_fields)*)) + .map(|#(#value_fields)*| #value_name::#variant_name { #(#value_fields)* }) + } + }, + remove: quote! { + (#storage_name::#variant_name { storage }, #key_name::#variant_name {}) => { + storage.take() + .map(|#(#value_fields)*| #value_name::#variant_name { #(#value_fields)* }) + } + }, + contains_key: quote! { + (#storage_name::#variant_name { storage }, #key_name::#variant_name {}) => { + storage.is_some() + } + }, + get: quote! { + (#storage_name::#variant_name { storage }, #key_name::#variant_name {}) => { + storage.as_ref() + .map(|#(#value_fields)*| #value_ref_name::#variant_name { #(#value_fields)* }) + } + }, + get_mut: quote! { + (#storage_name::#variant_name { storage }, #key_name::#variant_name {}) => { + storage.as_mut() + .map(|#(#value_fields)*| #value_ref_mut_name::#variant_name { #(#value_fields)* }) + } + }, + shrink_to_fit: quote! { + #storage_name::#variant_name { .. } => { + // nothing to do + } + }, + is_empty: quote! { + #storage_name::#variant_name { storage } => { + storage.is_none() + } + }, + iter: quote! { + #storage_name::#variant_name { storage } => { + #iter_name::#variant_name(storage.iter()) + } + }, + iterator: quote! { + std::option::Iter<'l, #(#value_types)*> + }, + iterator_next: quote! { + iter.next().map(|#(#value_fields)*| (#key_name::#variant_name {}, #value_ref_name::#variant_name { #(#value_fields)* })) + }, + } + } + (1, 1) => { + StorageDecl { + decl: quote! { + storage: auto_hash_map::AutoMap<#(#key_types)*, #(#value_types)*, std::hash::BuildHasherDefault, 1>, + }, + add: quote! { + (#storage_name::#variant_name { storage }, #ident::#variant_name { #(#key_fields)*, #(#value_fields)* }) => { + match storage.entry(#(#key_fields)*) { + auto_hash_map::map::Entry::Occupied(_) => false, + auto_hash_map::map::Entry::Vacant(e) => { + e.insert(#(#value_fields)*); + true + } + } + } + }, + insert: quote! { + (#storage_name::#variant_name { storage }, #ident::#variant_name { #(#key_fields)*, #(#value_fields)* }) => { + storage.insert(#(#key_fields)*, #(#value_fields)*) + .map(|#(#value_fields)*| #value_name::#variant_name { #(#value_fields)* }) + } + }, + remove: quote! { + (#storage_name::#variant_name { storage }, #key_name::#variant_name { #(#key_fields)* }) => { + let result = storage.remove(#(#key_fields)*) + .map(|#(#value_fields)*| #value_name::#variant_name { #(#value_fields)* }); + if result.is_some() { + storage.shrink_amortized(); + } + result + } + }, + contains_key: quote! { + (#storage_name::#variant_name { storage }, #key_name::#variant_name { #(#key_fields)* }) => { + storage.contains_key(#(#key_fields)*) + } + }, + get: quote! { + (#storage_name::#variant_name { storage }, #key_name::#variant_name { #(#key_fields)* }) => { + storage.get(#(#key_fields)*) + .map(|#(#value_fields)*| #value_ref_name::#variant_name { #(#value_fields)* }) + } + }, + get_mut: quote! { + (#storage_name::#variant_name { storage }, #key_name::#variant_name { #(#key_fields)* }) => { + storage.get_mut(#(#key_fields)*) + .map(|#(#value_fields)*| #value_ref_mut_name::#variant_name { #(#value_fields)* }) + } + }, + shrink_to_fit: quote! { + #storage_name::#variant_name { storage } => { + storage.shrink_to_fit() + } + }, + is_empty: quote! { + #storage_name::#variant_name { storage } => { + storage.is_empty() + } + }, + iter: quote! { + #storage_name::#variant_name { storage } => { + #iter_name::#variant_name(storage.iter()) + } + }, + iterator: quote! { + auto_hash_map::map::Iter<'l, #(#key_types)*, #(#value_types)*> + }, + iterator_next: quote! { + iter.next() + .map(|(#(#key_fields)*, #(#value_fields)*)| (#key_name::#variant_name { #(#key_fields: *#key_fields)* }, #value_ref_name::#variant_name { #(#value_fields)* })) + }, + } + } + (_, 1) => { + StorageDecl { + decl: quote! { + storage: auto_hash_map::AutoMap<(#(#key_types),*), #(#value_types)*, std::hash::BuildHasherDefault, 1>, + }, + add: quote! { + (#storage_name::#variant_name { storage }, #ident::#variant_name { #(#key_fields),*, #(#value_fields)* }) => { + match storage.entry((#(#key_fields),*)) { + auto_hash_map::map::Entry::Occupied(_) => false, + auto_hash_map::map::Entry::Vacant(e) => { + e.insert(#(#value_fields)*); + true + } + } + } + }, + insert: quote! { + (#storage_name::#variant_name { storage }, #ident::#variant_name { #(#key_fields),*, #(#value_fields)* }) => { + storage.insert((#(#key_fields),*), #(#value_fields)*) + .map(|#(#value_fields)*| #value_name::#variant_name { #(#value_fields)* }) + } + }, + remove: quote! { + (#storage_name::#variant_name { storage }, &#key_name::#variant_name { #(#key_fields),* }) => { + let result = storage.remove(&(#(#key_fields),*)) + .map(|#(#value_fields)*| #value_name::#variant_name { #(#value_fields)* }); + if result.is_some() { + storage.shrink_amortized(); + } + result + } + }, + contains_key: quote! { + (#storage_name::#variant_name { storage }, &#key_name::#variant_name { #(#key_fields),* }) => { + storage.contains_key(&(#(#key_fields),*)) + } + }, + get: quote! { + (#storage_name::#variant_name { storage }, &#key_name::#variant_name { #(#key_fields),* }) => { + storage.get(&(#(#key_fields),*)) + .map(|#(#value_fields)*| #value_ref_name::#variant_name { #(#value_fields)* }) + } + }, + get_mut: quote! { + (#storage_name::#variant_name { storage }, &#key_name::#variant_name { #(#key_fields),* }) => { + storage.get_mut(&(#(#key_fields),*)) + .map(|#(#value_fields)*| #value_ref_mut_name::#variant_name { #(#value_fields)* }) + } + }, + shrink_to_fit: quote! { + #storage_name::#variant_name { storage } => { + storage.shrink_to_fit() + } + }, + is_empty: quote! { + #storage_name::#variant_name { storage } => { + storage.is_empty() + } + }, + iter: quote! { + #storage_name::#variant_name { storage } => { + #iter_name::#variant_name(storage.iter()) + } + }, + iterator: quote! { + auto_hash_map::map::Iter<'l, (#(#key_types),*), #(#value_types)*> + }, + iterator_next: quote! { + iter.next() + .map(|((#(#key_fields),*), #(#value_fields)*)| (#key_name::#variant_name { #(#key_fields: *#key_fields),* }, #value_ref_name::#variant_name { #(#value_fields)* })) + }, + } + } + _ => unreachable!() + } + }) + .collect::>(); + + let storage_decl = storage.iter().map(|decl| &decl.decl).collect::>(); + let storage_add = storage.iter().map(|decl| &decl.add).collect::>(); + let storage_insert = storage.iter().map(|decl| &decl.insert).collect::>(); + let storage_remove = storage.iter().map(|decl| &decl.remove).collect::>(); + let storage_contains_key = storage + .iter() + .map(|decl| &decl.contains_key) + .collect::>(); + let storage_get = storage.iter().map(|decl| &decl.get).collect::>(); + let storage_get_mut = storage.iter().map(|decl| &decl.get_mut).collect::>(); + let storage_shrink_to_fit = storage + .iter() + .map(|decl| &decl.shrink_to_fit) + .collect::>(); + let storage_is_empty = storage + .iter() + .map(|decl| &decl.is_empty) + .collect::>(); + let storage_iter = storage.iter().map(|decl| &decl.iter).collect::>(); + let storage_iterator = storage + .iter() + .map(|decl| &decl.iterator) + .collect::>(); + let storage_iterator_next = storage + .iter() + .map(|decl| &decl.iterator_next) + .collect::>(); + quote! { impl turbo_tasks::KeyValuePair for #ident { + type Type = #type_name; type Key = #key_name; type Value = #value_name; + type ValueRef<'l> = #value_ref_name<'l> where Self: 'l; + type ValueRefMut<'l> = #value_ref_mut_name<'l> where Self: 'l; + + fn ty(&self) -> #type_name { + match self { + #( + #ident::#variant_names { .. } => #type_name::#variant_names, + )* + } + } fn key(&self) -> #key_name { match self { @@ -79,6 +386,22 @@ pub fn derive_key_value_pair(input: TokenStream) -> TokenStream { } } + fn value_ref(&self) -> #value_ref_name<'_> { + match self { + #( + #ident::#variant_names { #value_pat .. } => #value_ref_name::#variant_names { #value_ref_fields }, + )* + } + } + + fn value_mut(&mut self) -> #value_ref_mut_name<'_> { + match self { + #( + #ident::#variant_names { #value_pat .. } => #value_ref_mut_name::#variant_names { #value_ref_fields }, + )* + } + } + fn from_key_and_value(key: #key_name, value: #value_name) -> Self { match (key, value) { #( @@ -97,7 +420,14 @@ pub fn derive_key_value_pair(input: TokenStream) -> TokenStream { } } - #[derive(Debug, Clone, PartialEq, Eq, Hash)] + #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] + #vis enum #type_name { + #( + #variant_names, + )* + } + + #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] #vis enum #key_name { #( #variant_names { @@ -116,6 +446,178 @@ pub fn derive_key_value_pair(input: TokenStream) -> TokenStream { #[default] Reserved, } + + #[derive(Debug, Copy, Clone, PartialEq, Eq)] + #vis enum #value_ref_name<'l> { + #( + #variant_names { + #value_ref_decl + }, + )* + } + + #[derive(Debug, PartialEq, Eq)] + #vis enum #value_ref_mut_name<'l> { + #( + #variant_names { + #value_ref_mut_decl + }, + )* + } + + impl #key_name { + pub fn ty(&self) -> #type_name { + match self { + #( + #key_name::#variant_names { .. } => #type_name::#variant_names, + )* + } + } + } + + impl #value_name { + pub fn as_ref(&self) -> #value_ref_name<'_> { + match self { + #( + #value_name::#variant_names { #value_pat .. } => #value_ref_name::#variant_names { #value_ref_fields }, + )* + #value_name::Reserved => unreachable!(), + } + } + + pub fn as_mut(&mut self) -> #value_ref_mut_name<'_> { + match self { + #( + #value_name::#variant_names { #value_pat .. } => #value_ref_mut_name::#variant_names { #value_ref_fields }, + )* + #value_name::Reserved => unreachable!(), + } + } + } + + #vis enum #storage_name { + #( + #variant_names { + #storage_decl + }, + )* + } + + impl #storage_name { + pub fn new(ty: #type_name) -> Self { + match ty { + #( + #type_name::#variant_names => #storage_name::#variant_names { storage: Default::default() }, + )* + } + } + + pub fn ty(&self) -> #type_name { + match self { + #( + #storage_name::#variant_names { .. } => #type_name::#variant_names, + )* + } + } + + pub fn add(&mut self, item: #ident) -> bool { + match (self, item) { + #( + #storage_add + )* + _ => unreachable!(), + } + } + + pub fn insert(&mut self, item: #ident) -> Option<#value_name> { + match (self, item) { + #( + #storage_insert + )* + _ => unreachable!(), + } + } + + pub fn remove(&mut self, key: &#key_name) -> Option<#value_name> { + match (self, key) { + #( + #storage_remove + )* + _ => unreachable!(), + } + } + + pub fn contains_key(&self, key: &#key_name) -> bool { + match (self, key) { + #( + #storage_contains_key + )* + _ => unreachable!(), + } + } + + pub fn get(&self, key: &#key_name) -> Option<#value_ref_name> { + match (self, key) { + #( + #storage_get + )* + _ => unreachable!(), + } + } + + pub fn get_mut(&mut self, key: &#key_name) -> Option<#value_ref_mut_name> { + match (self, key) { + #( + #storage_get_mut + )* + _ => unreachable!(), + } + } + + pub fn shrink_to_fit(&mut self) { + match self { + #( + #storage_shrink_to_fit + )* + } + } + + pub fn is_empty(&self) -> bool { + match self { + #( + #storage_is_empty + )* + } + } + + pub fn iter(&self) -> #iter_name { + match self { + #( + #storage_iter + )* + } + } + } + + #vis enum #iter_name<'l> { + #( + #variant_names(#storage_iterator), + )* + } + + impl<'l> Iterator for #iter_name<'l> { + type Item = (#key_name, #value_ref_name<'l>); + + fn next(&mut self) -> Option { + match self { + #( + #iter_name::#variant_names(iter) => { + #storage_iterator_next + } + )* + } + } + } } .into() } @@ -162,6 +664,27 @@ fn clone_fields(fields: &[Vec<&syn::Field>]) -> Vec { variant_pat } +fn ref_fields(fields: &[Vec<&syn::Field>]) -> Vec { + let variant_pat = fields + .iter() + .map(|fields| { + let pat = fields + .iter() + .map(|field| { + let ident = field.ident.as_ref().unwrap(); + quote! { + #ident + } + }) + .collect::>(); + quote! { + #(#pat,)* + } + }) + .collect::>(); + variant_pat +} + fn field_declarations(fields: &[Vec<&syn::Field>]) -> Vec { fields .iter() @@ -184,3 +707,65 @@ fn field_declarations(fields: &[Vec<&syn::Field>]) -> Vec>() } + +fn ref_field_declarations(fields: &[Vec<&syn::Field>]) -> Vec { + fields + .iter() + .map(|fields| { + let fields = fields + .iter() + .map(|field| { + let ty = &field.ty; + let ident = field.ident.as_ref().unwrap(); + let attrs = &field.attrs; + quote! { + #(#attrs)* + #ident: &'l #ty + } + }) + .collect::>(); + quote! { + #(#fields),* + } + }) + .collect::>() +} + +fn mut_ref_field_declarations(fields: &[Vec<&syn::Field>]) -> Vec { + fields + .iter() + .map(|fields| { + let fields = fields + .iter() + .map(|field| { + let ty = &field.ty; + let ident = field.ident.as_ref().unwrap(); + let attrs = &field.attrs; + quote! { + #(#attrs)* + #ident: &'l mut #ty + } + }) + .collect::>(); + quote! { + #(#fields),* + } + }) + .collect::>() +} + +struct StorageDecl { + decl: proc_macro2::TokenStream, + add: proc_macro2::TokenStream, + insert: proc_macro2::TokenStream, + remove: proc_macro2::TokenStream, + contains_key: proc_macro2::TokenStream, + get: proc_macro2::TokenStream, + get_mut: proc_macro2::TokenStream, + shrink_to_fit: proc_macro2::TokenStream, + is_empty: proc_macro2::TokenStream, + iter: proc_macro2::TokenStream, + + iterator: proc_macro2::TokenStream, + iterator_next: proc_macro2::TokenStream, +} diff --git a/turbopack/crates/turbo-tasks/src/key_value_pair.rs b/turbopack/crates/turbo-tasks/src/key_value_pair.rs index dfd0e7bdb92e0..80fe85afbd102 100644 --- a/turbopack/crates/turbo-tasks/src/key_value_pair.rs +++ b/turbopack/crates/turbo-tasks/src/key_value_pair.rs @@ -1,10 +1,20 @@ use std::fmt::Debug; pub trait KeyValuePair { + type Type: Debug + Copy + Clone + PartialEq + Eq + std::hash::Hash; type Key: Debug + Clone + PartialEq + Eq + std::hash::Hash; type Value: Debug + Clone + Default + PartialEq + Eq; + type ValueRef<'l>: Debug + Copy + Clone + PartialEq + Eq + 'l + where + Self: 'l; + type ValueRefMut<'l>: Debug + PartialEq + Eq + 'l + where + Self: 'l; + fn ty(&self) -> Self::Type; fn key(&self) -> Self::Key; fn value(&self) -> Self::Value; + fn value_ref(&self) -> Self::ValueRef<'_>; + fn value_mut(&mut self) -> Self::ValueRefMut<'_>; fn from_key_and_value(key: Self::Key, value: Self::Value) -> Self; fn into_key_and_value(self) -> (Self::Key, Self::Value); }