Turbopack: avoid race condition when updating cells

Author: sokra

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

@@ -2,7 +2,10 @@ use std::mem::take;
 
 use serde::{Deserialize, Serialize};
 use smallvec::SmallVec;
-use turbo_tasks::{CellId, TaskId, backend::CellContent};
+use turbo_tasks::{
+    CellId, TaskId, TypedSharedReference,
+    backend::{CellContent, TypedCellContent},
+};
 
 #[cfg(feature = "trace_task_dirty")]
 use crate::backend::operation::invalidate::TaskDirtyCause;
@@ -24,6 +27,12 @@ pub enum UpdateCellOperation {
     InvalidateWhenCellDependency {
         cell_ref: CellRef,
         dependent_tasks: SmallVec<[TaskId; 4]>,
+        content: Option<TypedSharedReference>,
+        queue: AggregationUpdateQueue,
+    },
+    FinalCellChange {
+        cell_ref: CellRef,
+        content: Option<TypedSharedReference>,
         queue: AggregationUpdateQueue,
     },
     AggregationUpdate {
@@ -36,29 +45,22 @@ pub enum UpdateCellOperation {
 impl UpdateCellOperation {
     pub fn run(task_id: TaskId, cell: CellId, content: CellContent, mut ctx: impl ExecuteContext) {
         let mut task = ctx.task(task_id, TaskDataCategory::All);
-        let old_content = if let CellContent(Some(new_content)) = content {
-            task.insert(CachedDataItem::CellData {
-                cell,
-                value: new_content.into_typed(cell.type_id),
-            })
-        } else {
-            task.remove(&CachedDataItemKey::CellData { cell })
-        };
-
-        if let Some(in_progress) = remove!(task, InProgressCell { cell }) {
-            in_progress.event.notify(usize::MAX);
-        }
 
         // We need to detect recomputation, because here the content has not actually changed (even
         // if it's not equal to the old content, as not all values implement Eq). We have to
         // assume that tasks are deterministic and pure.
+        let should_invalidate = ctx.should_track_dependencies()
+            && (task.has_key(&CachedDataItemKey::Dirty {}) ||
+            // This is a hack for the streaming hack. Stateful tasks are never recomputed, so this forces invalidation for them in case of this hack.
+            task.has_key(&CachedDataItemKey::Stateful {}));
+
+        if should_invalidate {
+            // Slow path: We need to invalidate tasks depending on this cell.
+            // To avoid a race condition, we need to remove the old content first,
+            // then invalidate dependent tasks and only then update the cell content.
+
+            let old_content = task.remove(&CachedDataItemKey::CellData { cell });
 
-        if ctx.should_track_dependencies()
-            && (task.has_key(&CachedDataItemKey::Dirty {})
-                ||
-                // This is a hack for the streaming hack. Stateful tasks are never recomputed, so this forces invalidation for them in case of this hack.
-                task.has_key(&CachedDataItemKey::Stateful {}))
-        {
             let dependent_tasks = get_many!(
                 task,
                 CellDependent { cell: dependent_cell, task }
@@ -69,18 +71,44 @@ impl UpdateCellOperation {
             drop(task);
             drop(old_content);
 
+            let content = if let CellContent(Some(new_content)) = content {
+                Some(new_content.into_typed(cell.type_id))
+            } else {
+                None
+            };
+
             UpdateCellOperation::InvalidateWhenCellDependency {
                 cell_ref: CellRef {
                     task: task_id,
                     cell,
                 },
                 dependent_tasks,
+                content,
                 queue: AggregationUpdateQueue::new(),
             }
             .execute(&mut ctx);
         } else {
+            // Fast path: We don't need to invalidate anything.
+            // So we can just update the cell content.
+
+            let old_content = if let CellContent(Some(new_content)) = content {
+                let new_content = new_content.into_typed(cell.type_id);
+                task.insert(CachedDataItem::CellData {
+                    cell,
+                    value: new_content,
+                })
+            } else {
+                task.remove(&CachedDataItemKey::CellData { cell })
+            };
+
+            let in_progress_cell = remove!(task, InProgressCell { cell });
+
             drop(task);
             drop(old_content);
+
+            if let Some(in_progress) = in_progress_cell {
+                in_progress.event.notify(usize::MAX);
+            }
         }
     }
 }
@@ -93,6 +121,7 @@ impl Operation for UpdateCellOperation {
                 UpdateCellOperation::InvalidateWhenCellDependency {
                     cell_ref,
                     ref mut dependent_tasks,
+                    ref mut content,
                     ref mut queue,
                 } => {
                     if let Some(dependent_task_id) = dependent_tasks.pop() {
@@ -129,9 +158,37 @@ impl Operation for UpdateCellOperation {
                         );
                     }
                     if dependent_tasks.is_empty() {
-                        self = UpdateCellOperation::AggregationUpdate { queue: take(queue) };
+                        self = UpdateCellOperation::FinalCellChange {
+                            cell_ref,
+                            content: take(content),
+                            queue: take(queue),
+                        };
                     }
                 }
+                UpdateCellOperation::FinalCellChange {
+                    cell_ref: CellRef { task, cell },
+                    content,
+                    ref mut queue,
+                } => {
+                    let mut task = ctx.task(task, TaskDataCategory::Data);
+
+                    if let Some(content) = content {
+                        task.add_new(CachedDataItem::CellData {
+                            cell,
+                            value: content,
+                        })
+                    }
+
+                    let in_progress_cell = remove!(task, InProgressCell { cell });
+
+                    drop(task);
+
+                    if let Some(in_progress) = in_progress_cell {
+                        in_progress.event.notify(usize::MAX);
+                    }
+
+                    self = UpdateCellOperation::AggregationUpdate { queue: take(queue) };
+                }
                 UpdateCellOperation::AggregationUpdate { ref mut queue } => {
                     if queue.process(ctx) {
                         self = UpdateCellOperation::Done