Drop excessive cells after task reexecution

crates/turbo-tasks-memory/src/cell.rs

@@ -201,11 +201,44 @@ impl Cell {
         }
     }
 
+    pub fn empty(
+        &mut self,
+        clean: bool,
+        turbo_tasks: &dyn TurboTasksBackendApi<MemoryBackend>,
+    ) -> Option<CellContent> {
+        let content = match replace(&mut self.state, CellState::Empty) {
+            CellState::TrackedValueless | CellState::Empty => None,
+            CellState::Computing { event } => {
+                event.notify(usize::MAX);
+                if clean {
+                    // We can assume that the task is deterministic and produces the same content
+                    // again. No need to notify dependent tasks.
+                    return None;
+                }
+                None
+            }
+            CellState::Value { content } => Some(content),
+        };
+        // Assigning to a cell will invalidate all dependent tasks as the content might
+        // have changed.
+        if !self.dependent_tasks.is_empty() {
+            turbo_tasks.schedule_notify_tasks_set(&self.dependent_tasks);
+            self.dependent_tasks.clear();
+        }
+        content
+    }
+
     /// Reduces memory needs to the minimum.
     pub fn shrink_to_fit(&mut self) {
         self.dependent_tasks.shrink_to_fit();
     }
 
+    /// Returns true if the cell is current not used and could be dropped from
+    /// the array.
+    pub fn is_unused(&self) -> bool {
+        self.dependent_tasks.is_empty() && matches!(self.state, CellState::Empty)
+    }
+
     /// Takes the content out of the cell. Make sure to drop the content outside
     /// of the task state lock.
     #[must_use]

crates/turbo-tasks-memory/src/memory_backend.rs

@@ -13,6 +13,7 @@ use std::{
 };
 
 use anyhow::{anyhow, bail, Result};
+use auto_hash_map::AutoMap;
 use dashmap::{mapref::entry::Entry, DashMap};
 use rustc_hash::FxHasher;
 use tokio::task::futures::TaskLocalFuture;
@@ -25,7 +26,7 @@ use turbo_tasks::{
     },
     event::EventListener,
     util::{IdFactoryWithReuse, NoMoveVec},
-    CellId, RawVc, TaskId, TaskIdSet, TraitTypeId, TurboTasksBackendApi, Unused,
+    CellId, RawVc, TaskId, TaskIdSet, TraitTypeId, TurboTasksBackendApi, Unused, ValueTypeId,
 };
 
 use crate::{
@@ -324,6 +325,7 @@ impl Backend for MemoryBackend {
         task_id: TaskId,
         duration: Duration,
         memory_usage: usize,
+        cell_counters: AutoMap<ValueTypeId, u32, BuildHasherDefault<FxHasher>, 8>,
         stateful: bool,
         turbo_tasks: &dyn TurboTasksBackendApi<MemoryBackend>,
     ) -> bool {
@@ -339,6 +341,7 @@ impl Backend for MemoryBackend {
                     duration,
                     memory_usage,
                     generation,
+                    cell_counters,
                     stateful,
                     self,
                     turbo_tasks,

crates/turbo-tasks-memory/src/task.rs

@@ -944,6 +944,7 @@ impl Task {
         duration: Duration,
         memory_usage: usize,
         generation: NonZeroU32,
+        cell_counters: AutoMap<ValueTypeId, u32, BuildHasherDefault<FxHasher>, 8>,
         stateful: bool,
         backend: &MemoryBackend,
         turbo_tasks: &dyn TurboTasksBackendApi<MemoryBackend>,
@@ -953,6 +954,7 @@ impl Task {
         {
             let mut change_job = None;
             let mut remove_job = None;
+            let mut drained_cells = SmallVec::<[Cell; 8]>::new();
             let dependencies = DEPENDENCIES_TO_TRACK.with(|deps| deps.take());
             {
                 let mut state = self.full_state_mut();
@@ -961,21 +963,43 @@ impl Task {
                     .gc
                     .execution_completed(duration, memory_usage, generation);
 
+                let TaskState {
+                    ref mut cells,
+                    ref mut state_type,
+                    ..
+                } = *state;
+
                 let InProgress(box InProgressState {
                     ref mut done_event,
                     count_as_finished,
                     ref mut outdated_edges,
                     ref mut outdated_collectibles,
                     ref mut new_children,
-                    clean: _,
+                    clean,
                     stale,
-                }) = state.state_type
+                }) = *state_type
                 else {
                     panic!(
                         "Task execution completed in unexpected state {}",
                         Task::state_string(&state)
                     )
                 };
+                for (value_type, cells) in cells.iter_mut() {
+                    let counter =
+                        cell_counters.get(value_type).copied().unwrap_or_default() as usize;
+                    let mut is_unused = true;
+                    while counter < cells.len() {
+                        let last = cells.last_mut().unwrap();
+                        last.empty(clean, turbo_tasks);
+                        if is_unused {
+                            if last.is_unused() {
+                                drained_cells.push(cells.pop().unwrap());
+                            } else {
+                                is_unused = false;
+                            }
+                        }
+                    }
+                }
                 let done_event = done_event.take();
                 let outdated_collectibles = outdated_collectibles.take_collectibles();
                 let mut outdated_edges = take(outdated_edges);
@@ -1064,6 +1088,9 @@ impl Task {
                     self.clear_dependencies(outdated_edges, backend, turbo_tasks);
                 }
             }
+            for cell in drained_cells {
+                cell.gc_drop(turbo_tasks);
+            }
             change_job.apply(&aggregation_context);
             remove_job.apply(&aggregation_context);
         }

crates/turbo-tasks/src/backend.rs

@@ -429,6 +429,7 @@ pub trait Backend: Sync + Send {
         task: TaskId,
         duration: Duration,
         memory_usage: usize,
+        cell_counters: AutoMap<ValueTypeId, u32, BuildHasherDefault<FxHasher>, 8>,
         stateful: bool,
         turbo_tasks: &dyn TurboTasksBackendApi<Self>,
     ) -> bool;

crates/turbo-tasks/src/manager.rs

@@ -518,10 +518,13 @@ impl<B: Backend + 'static> TurboTasks<B> {
                                 });
                                 this.backend.task_execution_result(task_id, result, &*this);
                                 let stateful = this.finish_current_task_state();
+                                let cell_counters =
+                                    CELL_COUNTERS.with(|cc| take(&mut *cc.borrow_mut()));
                                 let schedule_again = this.backend.task_execution_completed(
                                     task_id,
                                     duration,
                                     memory_usage,
+                                    cell_counters,
                                     stateful,
                                     &*this,
                                 );