Drop excessive cells after task reexecution (vercel/turborepo#8170)

### Description

When cells become unused after recomputation of a task, drop them.

### Testing Instructions

<!--
  Give a quick description of steps to test your changes.
-->

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,
  );