turbo-tasks: task-storage memory wins

* Migrate Arc<CachedTaskType> to triomphe::Arc via CachedTaskTypeArc newtype.
  Saves one usize per allocation (no weak count) and avoids the weak-count
  CAS in drop_slow.
* Niche-encode cell_dependencies / cell_dependents elements as a CellDependency
  enum that reuses the niche on ValueTypeId (NonZero<u16>) for the variant
  tag. Element size drops from 32 B to 24 B; LazyField from 56 B to 48 B.
* Replace TaskStorage::lazy: Vec<LazyField> with a custom 16 B TinyVec
  (u8 len + u8 cap, the schema's max field count is well under 255).
  Drops size_of::<TaskStorage>() from 136 B to 128 B. The included
  micro-benchmarks show TinyVec push is 11-32% faster than Vec across
  realistic sizes and iter is neutral.
* Add CellDependency::into_parts() and use it in iter_cell_dependents /
  iter_cell_dependencies hot loops to avoid checking the enum discriminant
  twice via back-to-back cell_ref() + key() calls.

Cumulative MaxRSS reduction is ~2.9% on a representative Next.js build
(measured on a real project, 15 runs each side, t=-13.7 vs canary).

Author: lukesandberg

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

@@ -34,8 +34,8 @@ use turbo_tasks::{
     TaskId, TaskPersistence, TaskPriority, TraitTypeId, TurboTasksBackendApi, TurboTasksPanic,
     ValueTypeId,
     backend::{
-        Backend, CachedTaskType, CellContent, CellHash, TaskExecutionSpec, TransientTaskType,
-        TurboTaskContextError, TurboTaskLocalContextError, TurboTasksError,
+        Backend, CachedTaskType, CachedTaskTypeArc, CellContent, CellHash, TaskExecutionSpec,
+        TransientTaskType, TurboTaskContextError, TurboTaskLocalContextError, TurboTasksError,
         TurboTasksExecutionError, TurboTasksExecutionErrorMessage, TypedCellContent,
         VerificationMode,
     },
@@ -70,8 +70,8 @@ use crate::{
     },
     backing_storage::{BackingStorage, SnapshotItem, compute_task_type_hash},
     data::{
-        ActivenessState, CellRef, CollectibleRef, CollectiblesRef, Dirtyness, InProgressCellState,
-        InProgressState, InProgressStateInner, OutputValue, TransientTask,
+        ActivenessState, CellDependency, CellRef, CollectibleRef, CollectiblesRef, Dirtyness,
+        InProgressCellState, InProgressState, InProgressStateInner, OutputValue, TransientTask,
     },
     error::TaskError,
     utils::{
@@ -785,7 +785,8 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
                 && (!task.immutable() || cfg!(feature = "verify_immutable"))
             {
                 let reader = reader.unwrap();
-                let _ = task.add_cell_dependents((cell, key, reader));
+                let _ = task
+                    .add_cell_dependents(CellDependency::new(CellRef { task: reader, cell }, key));
                 drop(task);
 
                 // Note: We use `task_pair` earlier to lock the task and its reader at the same
@@ -797,8 +798,9 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
                     task: task_id,
                     cell,
                 };
-                if !reader_task.remove_outdated_cell_dependencies(&(target, key)) {
-                    let _ = reader_task.add_cell_dependencies((target, key));
+                let dep = CellDependency::new(target, key);
+                if !reader_task.remove_outdated_cell_dependencies(&dep) {
+                    let _ = reader_task.add_cell_dependencies(dep);
                 }
                 drop(reader_task);
             }
@@ -1526,7 +1528,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
                     // Only now do we force the allocation.
                     // NOTE: if our caller had to perform resolution, then this will have already
                     // been boxed and take_box just takes it.
-                    let task_type = Arc::new(CachedTaskType {
+                    let task_type = CachedTaskTypeArc::new(CachedTaskType {
                         native_fn,
                         this,
                         arg: arg.take_box(),
@@ -1757,7 +1759,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         }
     }
 
-    fn debug_get_cached_task_type(&self, task_id: TaskId) -> Option<Arc<CachedTaskType>> {
+    fn debug_get_cached_task_type(&self, task_id: TaskId) -> Option<CachedTaskTypeArc> {
         let task = self.storage.access_mut(task_id);
         task.get_persistent_task_type().cloned()
     }
@@ -2197,7 +2199,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
             Some(
                 // Collect all dependencies on tasks to check if all dependencies are immutable
                 task.iter_output_dependencies()
-                    .chain(task.iter_cell_dependencies().map(|(target, _key)| target.task))
+                    .chain(task.iter_cell_dependencies().map(|dep| dep.cell_ref().task))
                     .collect::<FxHashSet<_>>(),
             )
         } else {
@@ -2236,7 +2238,7 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
             // breaking dependency tracking.
             old_edges.extend(
                 task.iter_outdated_cell_dependencies()
-                    .map(|(target, key)| OutdatedEdge::CellDependency(target, key)),
+                    .map(OutdatedEdge::CellDependency),
             );
             old_edges.extend(
                 task.iter_outdated_output_dependencies()

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

@@ -17,7 +17,7 @@ use crate::{
         },
         storage_schema::TaskStorageAccessors,
     },
-    data::{CellRef, CollectibleRef, CollectiblesRef},
+    data::{CellDependency, CellRef, CollectibleRef, CollectiblesRef},
 };
 
 #[derive(Encode, Decode, Clone)]
@@ -48,7 +48,7 @@ impl Default for CleanupOldEdgesOperation {
 pub enum OutdatedEdge {
     Child(TaskId),
     Collectible(CollectibleRef, i32),
-    CellDependency(CellRef, Option<u64>),
+    CellDependency(CellDependency),
     OutputDependency(TaskId),
     CollectiblesDependency(CollectiblesRef),
 }
@@ -166,27 +166,32 @@ impl CleanupOldEdgesOperation {
                                     AggregatedDataUpdate::new().collectibles_update(collectibles),
                                 ));
                             }
-                            OutdatedEdge::CellDependency(
-                                CellRef {
-                                    task: cell_task_id,
-                                    cell,
-                                },
-                                key,
-                            ) => {
+                            OutdatedEdge::CellDependency(dep) => {
+                                let (
+                                    CellRef {
+                                        task: cell_task_id,
+                                        cell,
+                                    },
+                                    key,
+                                ) = dep.into_parts();
                                 {
                                     let mut task = ctx.task(cell_task_id, TaskDataCategory::Data);
-                                    task.remove_cell_dependents(&(cell, key, task_id));
-                                }
-                                {
-                                    let mut task = ctx.task(task_id, TaskDataCategory::Data);
-                                    task.remove_cell_dependencies(&(
+                                    // The `cell_dependents` set on the dependee task is keyed
+                                    // by `(cell, task_id, key)`. Re-encode that as a
+                                    // `CellDependency` whose `CellRef.task` is the dependent
+                                    // task — same bits, opposite direction.
+                                    task.remove_cell_dependents(&CellDependency::new(
                                         CellRef {
-                                            task: cell_task_id,
+                                            task: task_id,
                                             cell,
                                         },
                                         key,
                                     ));
                                 }
+                                {
+                                    let mut task = ctx.task(task_id, TaskDataCategory::Data);
+                                    task.remove_cell_dependencies(&dep);
+                                }
                             }
                             OutdatedEdge::OutputDependency(output_task_id) => {
                                 #[cfg(feature = "trace_task_output_dependencies")]

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

@@ -19,7 +19,7 @@ use tracing::info_span;
 use tracing::trace_span;
 use turbo_tasks::{
     CellId, DynTaskInputs, FxIndexMap, RawVc, SharedReference, TaskExecutionReason, TaskId,
-    TaskPriority, TurboTasksBackendApi, TurboTasksCallApi, backend::CachedTaskType,
+    TaskPriority, TurboTasksBackendApi, TurboTasksCallApi, backend::CachedTaskTypeArc,
     macro_helpers::NativeFunction,
 };
 
@@ -103,15 +103,16 @@ pub trait ExecuteContext<'e>: Sized {
     /// Uses hash-based lookup which may return multiple candidates due to hash collisions,
     /// then verifies each candidate by comparing the stored `persistent_task_type`.
     /// Returns `Some((task_id, task_type))` if a matching task is found, where `task_type` is
-    /// the existing `Arc<CachedTaskType>` from storage (avoiding a duplicate allocation).
+    /// the existing `CachedTaskTypeArc` from storage (avoiding a duplicate
+    /// allocation).
     ///
     /// Accepts exploded components so the caller does not need to box the argument before calling.
     fn task_by_type(
         &mut self,
         native_fn: &'static NativeFunction,
         this: Option<RawVc>,
         arg: &dyn DynTaskInputs,
-    ) -> Option<(TaskId, Arc<CachedTaskType>)>;
+    ) -> Option<(TaskId, CachedTaskTypeArc)>;
     fn debug_get_task_description(&self, task_id: TaskId) -> String;
 }
 
@@ -606,7 +607,7 @@ struct TaskRestoreEntry {
     /// Another thread claimed the meta restore; we must wait in Phase 3.
     wait_meta: bool,
     /// Task type discovered during Phase 1c data restore (used to update task cache in Phase 2).
-    task_type: Option<Arc<CachedTaskType>>,
+    task_type: Option<CachedTaskTypeArc>,
     /// This thread performed the restore for at least one category (set in Phase 1c).
     self_restored: bool,
 }
@@ -985,7 +986,7 @@ impl<'e, B: BackingStorage> ExecuteContext<'e> for ExecuteContextImpl<'e, B> {
         native_fn: &'static NativeFunction,
         this: Option<RawVc>,
         arg: &dyn DynTaskInputs,
-    ) -> Option<(TaskId, Arc<CachedTaskType>)> {
+    ) -> Option<(TaskId, CachedTaskTypeArc)> {
         if !self.backend.should_restore() {
             return None;
         }
@@ -1032,14 +1033,14 @@ impl<'e, B: BackingStorage> ChildExecuteContext<'e> for ChildExecuteContextImpl<
 }
 
 pub enum TaskTypeRef<'l> {
-    Cached(&'l Arc<CachedTaskType>),
+    Cached(&'l CachedTaskTypeArc),
     Transient(&'l Arc<TransientTask>),
 }
 
 impl TaskTypeRef<'_> {
     pub fn to_owned(&self) -> TaskType {
         match self {
-            TaskTypeRef::Cached(ty) => TaskType::Cached(Arc::clone(ty)),
+            TaskTypeRef::Cached(ty) => TaskType::Cached((*ty).clone()),
             TaskTypeRef::Transient(ty) => TaskType::Transient(Arc::clone(ty)),
         }
     }
@@ -1056,7 +1057,7 @@ impl Display for TaskTypeRef<'_> {
 
 #[derive(Debug)]
 pub enum TaskType {
-    Cached(Arc<CachedTaskType>),
+    Cached(CachedTaskTypeArc),
     Transient(Arc<TransientTask>),
 }
 
@@ -1398,7 +1399,7 @@ impl TaskGuard for TaskGuardImpl<'_> {
             .map(|target| (target, TaskDataCategory::Meta))
             .chain(
                 self.iter_cell_dependencies()
-                    .map(|(target, _key)| (target.task, TaskDataCategory::All)),
+                    .map(|dep| (dep.cell_ref().task, TaskDataCategory::All)),
             )
             .chain(
                 self.iter_collectibles_dependencies()

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

@@ -20,7 +20,7 @@ use crate::{
         },
         storage_schema::TaskStorageAccessors,
     },
-    data::CellRef,
+    data::{CellDependency, CellRef},
 };
 
 #[derive(Encode, Decode, Clone, Default)]
@@ -137,17 +137,24 @@ impl UpdateCellOperation {
             let mut dependent_tasks: FxIndexMap<TaskId, SmallVec<[Option<u64>; 2]>> =
                 FxIndexMap::default();
             if !skip_invalidation {
-                let tasks_with_keys =
-                    task.iter_cell_dependents()
-                        .filter_map(|(dependent_cell, key, task)| {
-                            (dependent_cell == cell
-                                && key.is_none_or(|key_hash| {
-                                    updated_key_hashes_set
-                                        .as_ref()
-                                        .is_none_or(|set| set.contains(&key_hash))
-                                }))
-                            .then_some((task, key))
-                        });
+                let tasks_with_keys = task.iter_cell_dependents().filter_map(|dep| {
+                    // `cell_dependents` reuses `CellDependency`, but with the dependent task
+                    // stored in `CellRef.task` and the storing task's cell in `CellRef.cell`.
+                    let (
+                        CellRef {
+                            task: dependent_task,
+                            cell: dependent_cell,
+                        },
+                        key,
+                    ) = dep.into_parts();
+                    (dependent_cell == cell
+                        && key.is_none_or(|key_hash| {
+                            updated_key_hashes_set
+                                .as_ref()
+                                .is_none_or(|set| set.contains(&key_hash))
+                        }))
+                    .then_some((dependent_task, key))
+                });
                 for (task, key) in tasks_with_keys {
                     dependent_tasks.entry(task).or_default().push(key);
                 }
@@ -276,14 +283,15 @@ impl Operation for UpdateCellOperation {
                         let mut make_stale = false;
                         let dependent = ctx.task(dependent_task_id, TaskDataCategory::All);
                         for key in keys.iter().copied() {
-                            if dependent.outdated_cell_dependencies_contains(&(cell_ref, key)) {
+                            let dep = CellDependency::new(cell_ref, key);
+                            if dependent.outdated_cell_dependencies_contains(&dep) {
                                 // cell dependency is outdated, so it hasn't read the cell yet
                                 // and doesn't need to be invalidated.
                                 // We do not need to make the task stale in this case.
                                 // But importantly we still need to make the task dirty as it should
                                 // no longer be considered as
                                 // "recomputation".
-                            } else if !dependent.cell_dependencies_contains(&(cell_ref, key)) {
+                            } else if !dependent.cell_dependencies_contains(&dep) {
                                 // cell dependency has been removed, so the task doesn't depend on
                                 // the cell anymore and doesn't need
                                 // to be invalidated

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

@@ -12,7 +12,7 @@ use std::{
 use thread_local::ThreadLocal;
 use tracing::span::Id;
 use turbo_bincode::TurboBincodeBuffer;
-use turbo_tasks::{FxDashMap, TaskId, backend::CachedTaskType, event::Event, parallel};
+use turbo_tasks::{FxDashMap, TaskId, backend::CachedTaskTypeArc, event::Event, parallel};
 
 use crate::{
     backend::storage_schema::{
@@ -183,7 +183,7 @@ pub struct Storage {
     /// This is backed by the TaskCache table in the database.
     ///
     /// LockOrdering: See the comments on [map].
-    pub task_cache: FxDashMap<Arc<CachedTaskType>, TaskId>,
+    pub task_cache: FxDashMap<CachedTaskTypeArc, TaskId>,
 }
 
 impl Storage {
@@ -253,7 +253,7 @@ impl Storage {
     /// Mark a newly allocated task as restored (skip DB queries) and new (include in persistence
     /// snapshots). Optionally sets the `persistent_task_type` eagerly so it's available for
     /// persistence snapshots without needing to propagate it through `connect_child`.
-    pub fn initialize_new_task(&self, task_id: TaskId, task_type: Option<Arc<CachedTaskType>>) {
+    pub fn initialize_new_task(&self, task_id: TaskId, task_type: Option<CachedTaskTypeArc>) {
         let mut task = self.access_mut(task_id);
         task.flags.set_restored(TaskDataCategory::All);
         task.flags.set_new_task(true);
@@ -516,7 +516,7 @@ impl Storage {
             // was contended. We defer them until after the map shard lock is released to
             // avoid a lock cycle with get_or_create_persistent_task, which takes task_cache
             // before map. Allocated lazily on first conflict.
-            let mut deferred_task_cache_removals: Vec<Arc<CachedTaskType>> = Vec::new();
+            let mut deferred_task_cache_removals: Vec<CachedTaskTypeArc> = Vec::new();
             // SAFETY: We hold the write lock for the duration of iteration.
             for bucket in unsafe { shard.iter() } {
                 // SAFETY: The write lock guard outlives the bucket reference.