Split Dirty into Dirtyness and CleanInSession

Author: sokra

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

@@ -61,7 +61,7 @@ use crate::{
     backing_storage::BackingStorage,
     data::{
         ActivenessState, AggregationNumber, CachedDataItem, CachedDataItemKey, CachedDataItemType,
-        CachedDataItemValueRef, CellRef, CollectibleRef, CollectiblesRef, DirtyState,
+        CachedDataItemValueRef, CellRef, CollectibleRef, CollectiblesRef, Dirtyness,
         InProgressCellState, InProgressState, InProgressStateInner, OutputValue, RootType,
     },
     utils::{
@@ -567,8 +567,9 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
                 }
             }
 
-            let is_dirty =
-                get!(task, Dirty).map_or(false, |dirty_state| dirty_state.get(self.session_id));
+            let is_dirty = task
+                .dirty_state()
+                .map_or(false, |dirty_state| dirty_state.get(self.session_id));
 
             // Check the dirty count of the root node
             let dirty_tasks = get!(task, AggregatedDirtyContainerCount)
@@ -624,7 +625,8 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
                             visited: &mut FxHashSet<TaskId>,
                         ) -> String {
                             let task = ctx.task(task_id, TaskDataCategory::Data);
-                            let is_dirty = get!(task, Dirty)
+                            let is_dirty = task
+                                .dirty_state()
                                 .map_or(false, |dirty_state| dirty_state.get(ctx.session_id()));
                             let in_progress =
                                 get!(task, InProgress).map_or("not in progress", |p| match p {
@@ -2369,57 +2371,60 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
         ));
 
         // Update the dirty state
-        let old_dirty_state = get!(task, Dirty).copied();
+        let old_dirtyness = task.dirtyness_and_session();
 
-        let new_dirty_state = if session_dependent {
-            Some(DirtyState {
-                clean_in_session: Some(self.session_id),
-            })
+        let new_dirtyness = if session_dependent {
+            Some((Dirtyness::SessionDependent, Some(self.session_id)))
         } else {
             None
         };
 
-        let dirty_changed = old_dirty_state != new_dirty_state;
+        let dirty_changed = old_dirtyness != new_dirtyness;
         let data_update = if dirty_changed {
-            if let Some(new_dirty_state) = new_dirty_state {
-                task.insert(CachedDataItem::Dirty {
-                    value: new_dirty_state,
-                });
-            } else {
+            if let Some((value, _)) = new_dirtyness {
+                task.insert(CachedDataItem::Dirty { value });
+            } else if old_dirtyness.is_some() {
                 task.remove(&CachedDataItemKey::Dirty {});
             }
+            if let Some(session_id) = new_dirtyness.and_then(|t| t.1) {
+                task.insert(CachedDataItem::CleanInSession { value: session_id });
+            } else if old_dirtyness.is_some_and(|t| t.1.is_some()) {
+                task.remove(&CachedDataItemKey::CleanInSession {});
+            }
 
-            if old_dirty_state.is_some() || new_dirty_state.is_some() {
-                let mut dirty_containers = get!(task, AggregatedDirtyContainerCount)
-                    .cloned()
-                    .unwrap_or_default();
-                if let Some(old_dirty_state) = old_dirty_state {
-                    dirty_containers.update_with_dirty_state(&old_dirty_state);
+            let mut dirty_containers = get!(task, AggregatedDirtyContainerCount)
+                .cloned()
+                .unwrap_or_default();
+            if let Some((old_dirtyness, old_clean_in_session)) = old_dirtyness {
+                dirty_containers
+                    .update_with_dirtyness_and_session(old_dirtyness, old_clean_in_session);
+            }
+            let aggregated_update = match (old_dirtyness, new_dirtyness) {
+                (None, None) => unreachable!(),
+                (Some(old), None) => {
+                    dirty_containers.undo_update_with_dirtyness_and_session(old.0, old.1)
                 }
-                let aggregated_update = match (old_dirty_state, new_dirty_state) {
-                    (None, None) => unreachable!(),
-                    (Some(old), None) => dirty_containers.undo_update_with_dirty_state(&old),
-                    (None, Some(new)) => dirty_containers.update_with_dirty_state(&new),
-                    (Some(old), Some(new)) => dirty_containers.replace_dirty_state(&old, &new),
-                };
-                if !aggregated_update.is_zero() {
-                    if aggregated_update.get(self.session_id) < 0
-                        && let Some(activeness_state) = get_mut!(task, Activeness)
-                    {
-                        activeness_state.all_clean_event.notify(usize::MAX);
-                        activeness_state.unset_active_until_clean();
-                        if activeness_state.is_empty() {
-                            task.remove(&CachedDataItemKey::Activeness {});
-                        }
+                (None, Some(new)) => {
+                    dirty_containers.update_with_dirtyness_and_session(new.0, new.1)
+                }
+                (Some(old), Some(new)) => {
+                    dirty_containers.replace_dirtyness_and_session(old.0, old.1, new.0, new.1)
+                }
+            };
+            if !aggregated_update.is_zero() {
+                if aggregated_update.get(self.session_id) < 0
+                    && let Some(activeness_state) = get_mut!(task, Activeness)
+                {
+                    activeness_state.all_clean_event.notify(usize::MAX);
+                    activeness_state.unset_active_until_clean();
+                    if activeness_state.is_empty() {
+                        task.remove(&CachedDataItemKey::Activeness {});
                     }
-                    AggregationUpdateJob::data_update(
-                        &mut task,
-                        AggregatedDataUpdate::new()
-                            .dirty_container_update(task_id, aggregated_update),
-                    )
-                } else {
-                    None
                 }
+                AggregationUpdateJob::data_update(
+                    &mut task,
+                    AggregatedDataUpdate::new().dirty_container_update(task_id, aggregated_update),
+                )
             } else {
                 None
             }
@@ -2887,7 +2892,9 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
 
         let mut ctx = self.execute_context(turbo_tasks);
         let mut task = ctx.task(task_id, TaskDataCategory::All);
-        let is_dirty = get!(task, Dirty).map_or(false, |dirty| dirty.get(self.session_id));
+        let is_dirty = task
+            .dirty_state()
+            .map_or(false, |dirty| dirty.get(self.session_id));
         let has_dirty_containers = get!(task, AggregatedDirtyContainerCount)
             .map_or(false, |dirty_containers| {
                 dirty_containers.get(self.session_id) > 0
@@ -2980,7 +2987,9 @@ impl<B: BackingStorage> TurboTasksBackendInner<B> {
                     }
                 }
 
-                let is_dirty = get!(task, Dirty).is_some_and(|dirty| dirty.get(self.session_id));
+                let is_dirty = task
+                    .dirty_state()
+                    .is_some_and(|dirty| dirty.get(self.session_id));
                 let has_dirty_container = get!(task, AggregatedDirtyContainerCount)
                     .is_some_and(|count| count.get(self.session_id) > 0);
                 let should_be_in_upper = is_dirty || has_dirty_container;

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

@@ -249,8 +249,8 @@ impl AggregatedDataUpdate {
                 collectibles_update.push((collectible, 1));
             }
         }
-        if let Some(dirty) = get!(task, Dirty) {
-            dirty_container_count.update_with_dirty_state(dirty);
+        if let Some((dirtyness, clean_in_session)) = task.dirtyness_and_session() {
+            dirty_container_count.update_with_dirtyness_and_session(dirtyness, clean_in_session);
         }
 
         let mut result = Self::new().collectibles_update(collectibles_update);
@@ -323,7 +323,7 @@ impl AggregatedDataUpdate {
             );
 
             if !aggregated_update.is_zero() {
-                let dirty_state = get!(task, Dirty).copied();
+                let dirty_state = task.dirty_state();
                 let task_id = task.id();
                 update!(task, AggregatedDirtyContainerCount, |old: Option<
                     DirtyContainerCount,
@@ -1209,7 +1209,7 @@ impl AggregationUpdateQueue {
     ) {
         let session_id = ctx.session_id();
         // Task need to be scheduled if it's dirty or doesn't have output
-        let dirty = get!(task, Dirty).map_or(false, |d| d.get(session_id));
+        let dirty = task.dirty_state().map_or(false, |d| d.get(session_id));
         let should_schedule = if dirty {
             Some(TaskExecutionReason::ActivateDirty)
         } else if !task.has_key(&CachedDataItemKey::Output {}) {

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

@@ -11,11 +11,11 @@ use crate::{
                 AggregatedDataUpdate, AggregationUpdateJob, AggregationUpdateQueue,
             },
         },
-        storage::{get, get_mut},
+        storage::{get, get_mut, remove},
     },
     data::{
-        CachedDataItem, CachedDataItemKey, CachedDataItemValue, DirtyState, InProgressState,
-        InProgressStateInner,
+        CachedDataItem, CachedDataItemKey, CachedDataItemValue, DirtyState, Dirtyness,
+        InProgressState, InProgressStateInner,
     },
 };
 
@@ -232,15 +232,11 @@ pub fn make_task_dirty_internal(
         *stale = true;
     }
     let old = task.insert(CachedDataItem::Dirty {
-        value: DirtyState {
-            clean_in_session: None,
-        },
+        value: Dirtyness::Dirty,
     });
     let mut dirty_container = match old {
         Some(CachedDataItemValue::Dirty {
-            value: DirtyState {
-                clean_in_session: None,
-            },
+            value: Dirtyness::Dirty,
         }) => {
             #[cfg(feature = "trace_task_dirty")]
             let _span = tracing::trace_span!(
@@ -254,16 +250,30 @@ pub fn make_task_dirty_internal(
             return;
         }
         Some(CachedDataItemValue::Dirty {
-            value: DirtyState {
-                clean_in_session: Some(session_id),
-            },
+            value: Dirtyness::SessionDependent,
         }) => {
-            // Got dirty in that one session only
-            let mut dirty_container = get!(task, AggregatedDirtyContainerCount)
-                .cloned()
-                .unwrap_or_default();
-            dirty_container.update_session_dependent(session_id, 1);
-            dirty_container
+            let old = remove!(task, CleanInSession);
+            match old {
+                None => {
+                    #[cfg(feature = "trace_task_dirty")]
+                    let _span = tracing::trace_span!(
+                        "session-dependent task already dirty",
+                        name = ctx.get_task_description(task_id),
+                        cause = %TaskDirtyCauseInContext::new(&cause, ctx)
+                    )
+                    .entered();
+                    // already dirty
+                    return;
+                }
+                Some(session_id) => {
+                    // Got dirty in that one session only
+                    let mut dirty_container = get!(task, AggregatedDirtyContainerCount)
+                        .cloned()
+                        .unwrap_or_default();
+                    dirty_container.update_session_dependent(session_id, 1);
+                    dirty_container
+                }
+            }
         }
         None => {
             // Get dirty for all sessions

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

@@ -20,12 +20,12 @@ use crate::{
     backend::{
         OperationGuard, TaskDataCategory, TransientTask, TurboTasksBackend, TurboTasksBackendInner,
         TurboTasksBackendJob,
-        storage::{SpecificTaskDataCategory, StorageWriteGuard, iter_many},
+        storage::{SpecificTaskDataCategory, StorageWriteGuard, get, iter_many},
     },
     backing_storage::BackingStorage,
     data::{
         CachedDataItem, CachedDataItemKey, CachedDataItemType, CachedDataItemValue,
-        CachedDataItemValueRef, CachedDataItemValueRefMut,
+        CachedDataItemValueRef, CachedDataItemValueRefMut, DirtyState, Dirtyness,
     },
 };
 
@@ -415,6 +415,30 @@ pub trait TaskGuard: Debug {
     fn invalidate_serialization(&mut self);
     fn prefetch(&mut self) -> Option<FxIndexMap<TaskId, bool>>;
     fn is_immutable(&self) -> bool;
+    fn dirtyness_and_session(&self) -> Option<(Dirtyness, Option<SessionId>)> {
+        match get!(self, Dirty)? {
+            Dirtyness::Dirty => Some((Dirtyness::Dirty, None)),
+            Dirtyness::SessionDependent => Some((
+                Dirtyness::SessionDependent,
+                get!(self, CleanInSession).copied(),
+            )),
+        }
+    }
+    fn dirty_state(&self) -> Option<DirtyState> {
+        match get!(self, Dirty)? {
+            Dirtyness::Dirty => Some(DirtyState {
+                clean_in_session: None,
+            }),
+            Dirtyness::SessionDependent => match get!(self, CleanInSession) {
+                None => Some(DirtyState {
+                    clean_in_session: None,
+                }),
+                Some(session) => Some(DirtyState {
+                    clean_in_session: Some(*session),
+                }),
+            },
+        }
+    }
 }
 
 pub struct TaskGuardImpl<'a, B: BackingStorage> {