handle recomputation of cells on read

Author: sokra

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

@@ -30,11 +30,10 @@ use turbo_tasks::{
 };
 
 use crate::{
-    cell::RecomputingCell,
     edges_set::{TaskEdge, TaskEdgesSet},
     gc::{GcQueue, PERCENTAGE_IDLE_TARGET_MEMORY, PERCENTAGE_TARGET_MEMORY},
     output::Output,
-    task::{Task, DEPENDENCIES_TO_TRACK},
+    task::{ReadCellError, Task, DEPENDENCIES_TO_TRACK},
     task_statistics::TaskStatisticsApi,
 };
 
@@ -412,23 +411,17 @@ impl Backend for MemoryBackend {
         } else {
             Task::add_dependency_to_current(TaskEdge::Cell(task_id, index));
             self.with_task(task_id, |task| {
-                match task.access_cell_for_read(index, self.gc_queue.as_ref(), |cell| {
-                    cell.map(|cell| {
-                        cell.read_content(
-                            reader,
-                            move || format!("{task_id} {index}"),
-                            move || format!("reading {} {} from {}", task_id, index, reader),
-                        )
-                    })
-                }) {
-                    None => Err(anyhow!("Cell doesn't exist (anymore)")),
-                    Some(Ok(content)) => Ok(Ok(content)),
-                    Some(Err(RecomputingCell { listener, schedule })) => {
-                        if schedule {
-                            task.recompute(self, turbo_tasks);
-                        }
-                        Ok(Err(listener))
-                    }
+                match task.read_cell(
+                    index,
+                    self.gc_queue.as_ref(),
+                    move || format!("reading {} {} from {}", task_id, index, reader),
+                    Some(reader),
+                    self,
+                    turbo_tasks,
+                ) {
+                    Ok(content) => Ok(Ok(content)),
+                    Err(ReadCellError::Recomputing(listener)) => Ok(Err(listener)),
+                    Err(ReadCellError::CellRemoved) => Err(anyhow!("Cell {index:?} doesn't exist")),
                 }
             })
         }
@@ -452,22 +445,17 @@ impl Backend for MemoryBackend {
         turbo_tasks: &dyn TurboTasksBackendApi<MemoryBackend>,
     ) -> Result<Result<CellContent, EventListener>> {
         self.with_task(task_id, |task| {
-            match task.access_cell_for_read(index, self.gc_queue.as_ref(), |cell| {
-                cell.map(|cell| {
-                    cell.read_content_untracked(
-                        move || format!("{task_id} {index}"),
-                        move || format!("reading {} {} untracked", task_id, index),
-                    )
-                })
-            }) {
-                None => Err(anyhow!("Cell doesn't exist (anymore)")),
-                Some(Ok(content)) => Ok(Ok(content)),
-                Some(Err(RecomputingCell { listener, schedule })) => {
-                    if schedule {
-                        task.recompute(self, turbo_tasks);
-                    }
-                    Ok(Err(listener))
-                }
+            match task.read_cell(
+                index,
+                self.gc_queue.as_ref(),
+                move || format!("reading {} {} untracked", task_id, index),
+                None,
+                self,
+                turbo_tasks,
+            ) {
+                Ok(content) => Ok(Ok(content)),
+                Err(ReadCellError::Recomputing(listener)) => Ok(Err(listener)),
+                Err(ReadCellError::CellRemoved) => Err(anyhow!("Cell {index:?} doesn't exist")),
             }
         })
     }

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

@@ -21,7 +21,7 @@ use tokio::task_local;
 use tracing::Span;
 use turbo_prehash::PreHashed;
 use turbo_tasks::{
-    backend::{PersistentTaskType, TaskExecutionSpec},
+    backend::{CellContent, PersistentTaskType, TaskExecutionSpec},
     event::{Event, EventListener},
     get_invalidator, registry, CellId, Invalidator, RawVc, TaskId, TaskIdSet, TraitTypeId,
     TurboTasksBackendApi, ValueTypeId,
@@ -32,7 +32,7 @@ use crate::{
         aggregation_data, handle_new_edge, prepare_aggregation_data, query_root_info,
         AggregationDataGuard, PreparedOperation,
     },
-    cell::Cell,
+    cell::{Cell, RecomputingCell},
     edges_set::{TaskEdge, TaskEdgesList, TaskEdgesSet},
     gc::{GcQueue, GcTaskState},
     output::{Output, OutputContent},
@@ -448,6 +448,11 @@ pub enum GcResult {
     Unloaded,
 }
 
+pub enum ReadCellError {
+    CellRemoved,
+    Recomputing(EventListener),
+}
+
 impl Task {
     pub(crate) fn new_persistent(
         id: TaskId,
@@ -1286,26 +1291,82 @@ impl Task {
         }
     }
 
-    /// Access to a cell.
-    pub(crate) fn access_cell_for_read<T>(
+    /// Read a cell.
+    pub(crate) fn read_cell(
         &self,
         index: CellId,
         gc_queue: Option<&GcQueue>,
-        func: impl FnOnce(Option<&mut Cell>) -> T,
-    ) -> T {
+        note: impl Fn() -> String + Sync + Send + 'static,
+        reader: Option<TaskId>,
+        backend: &MemoryBackend,
+        turbo_tasks: &dyn TurboTasksBackendApi<MemoryBackend>,
+    ) -> Result<CellContent, ReadCellError> {
+        let task_id = self.id;
         let mut state = self.full_state_mut();
         if let Some(gc_queue) = gc_queue {
             let generation = gc_queue.generation();
             if state.gc.on_read(generation) {
                 let _ = gc_queue.task_accessed(self.id);
             }
         }
-        let list = state.cells.entry(index.type_id).or_default();
-        let i = index.index as usize;
-        if list.len() <= i {
-            func(None)
-        } else {
-            func(Some(&mut list[i]))
+        match state.state_type {
+            Done { .. } | InProgress(..) | InProgressDirty { .. } => {
+                let is_done = matches!(state.state_type, Done { .. });
+                let list = state.cells.entry(index.type_id).or_default();
+                let i = index.index as usize;
+                if list.len() <= i {
+                    if is_done {
+                        return Err(ReadCellError::CellRemoved);
+                    } else {
+                        list.resize_with(i + 1, Default::default);
+                    }
+                }
+                let cell = &mut list[i];
+                let description = move || format!("{task_id} {index}");
+                let read_result = if let Some(reader) = reader {
+                    cell.read_content(reader, description, note)
+                } else {
+                    cell.read_content_untracked(description, note)
+                };
+                drop(state);
+                match read_result {
+                    Ok(content) => Ok(content),
+                    Err(RecomputingCell { listener, schedule }) => {
+                        if schedule {
+                            self.recompute(backend, turbo_tasks);
+                        }
+                        Err(ReadCellError::Recomputing(listener))
+                    }
+                }
+            }
+            Dirty {
+                ref mut outdated_edges,
+            } => {
+                let mut aggregation_context = TaskAggregationContext::new(turbo_tasks, backend);
+                let description = self.get_event_description();
+                let event = Event::new(move || format!("TaskState({})::event", description()));
+                let listener = event.listen_with_note(note);
+                state.state_type = Scheduled {
+                    event,
+                    outdated_edges: take(outdated_edges),
+                    clean: false,
+                };
+                let change_job = state.aggregation_node.apply_change(
+                    &aggregation_context,
+                    TaskChange {
+                        dirty_tasks_update: vec![(self.id, -1)],
+                        ..Default::default()
+                    },
+                );
+                drop(state);
+                turbo_tasks.schedule(self.id);
+                change_job.apply(&aggregation_context);
+                aggregation_context.apply_queued_updates();
+                Err(ReadCellError::Recomputing(listener))
+            }
+            Scheduled { ref event, .. } => {
+                Err(ReadCellError::Recomputing(event.listen_with_note(note)))
+            }
         }
     }