Skip to content

Commit

Permalink
Improve memory usage of the new backend (#74740)
Browse files Browse the repository at this point in the history 
### 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.
  • Loading branch information
@sokra
sokra authored Jan 10, 2025
1 parent dd79127 commit 08c4e58
Show file tree
Hide file tree
Showing 11 changed files with 854 additions and 561 deletions.
4 changes: 0 additions & 4 deletions turbopack/crates/turbo-tasks-backend/src/backend/indexed.rs
View file

This file was deleted.

224 changes: 72 additions & 152 deletions turbopack/crates/turbo-tasks-backend/src/backend/mod.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
pub mod indexed;
mod operation;
mod persisted_storage_log;
mod storage;
Expand Down Expand Up @@ -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},
};
Expand Down Expand Up @@ -151,7 +150,7 @@ struct TurboTasksBackendInner<B: BackingStorage> {

persisted_storage_data_log: Option<PersistedStorageLog>,
persisted_storage_meta_log: Option<PersistedStorageLog>,
storage: Storage<TaskId, CachedDataItem>,
storage: Storage,

/// Number of executing operations + Highest bit is set when snapshot is
/// requested. When that bit is set, operations should pause until the
Expand Down Expand Up @@ -459,7 +458,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
AggregatedDirtyContainer {
task
} count if count.get(self.session_id) > 0 => {
*task
task
}
);
if is_dirty {
Expand Down Expand Up @@ -1001,13 +1000,8 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
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::<Vec<_>>();
for child in children {
match child {
Expand All @@ -1030,18 +1024,8 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
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::<Vec<_>>();
for collectible in collectibles {
match collectible {
Expand All @@ -1068,23 +1052,10 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
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::<Vec<_>>();
for dep in dependencies {
match dep {
Expand Down Expand Up @@ -1237,7 +1208,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {

// 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 {
Expand All @@ -1257,128 +1228,69 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
task.remove(&CachedDataItemKey::CellTypeMaxIndex { cell_type });
}

let mut removed_data = Vec::new();
let mut removed_data: Vec<CachedDataItem> = 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.
Expand Down Expand Up @@ -1479,6 +1391,14 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {

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
}

Expand Down
Loading

0 comments on commit 08c4e58

Please sign in to comment.