Fix cache invalidation when cycle head becomes non-head (#1014)

* Fix cache invalidation when cycle head becomes non-head

* Discard changes to src/function/fetch.rs

* Inline comment

Author: MichaReiser

src/cycle.rs

@@ -346,6 +346,7 @@ impl CycleHeads {
 
             if *removed {
                 *removed = false;
+                existing.iteration_count.store_mut(iteration_count);
 
                 true
             } else {
@@ -468,15 +469,25 @@ pub(crate) fn empty_cycle_heads() -> &'static CycleHeads {
     EMPTY_CYCLE_HEADS.get_or_init(|| CycleHeads(ThinVec::new()))
 }
 
-#[derive(Debug, PartialEq, Eq)]
-pub enum ProvisionalStatus {
+#[derive(Debug)]
+pub enum ProvisionalStatus<'db> {
     Provisional {
         iteration: IterationCount,
         verified_at: Revision,
+        cycle_heads: &'db CycleHeads,
     },
     Final {
         iteration: IterationCount,
         verified_at: Revision,
     },
     FallbackImmediate,
 }
+
+impl<'db> ProvisionalStatus<'db> {
+    pub(crate) fn cycle_heads(&self) -> &'db CycleHeads {
+        match self {
+            ProvisionalStatus::Provisional { cycle_heads, .. } => cycle_heads,
+            _ => empty_cycle_heads(),
+        }
+    }
+}

src/function.rs

@@ -7,10 +7,7 @@ use std::ptr::NonNull;
 use std::sync::atomic::Ordering;
 use std::sync::OnceLock;
 
-use crate::cycle::{
-    empty_cycle_heads, CycleHeads, CycleRecoveryAction, CycleRecoveryStrategy, IterationCount,
-    ProvisionalStatus,
-};
+use crate::cycle::{CycleRecoveryAction, CycleRecoveryStrategy, IterationCount, ProvisionalStatus};
 use crate::database::RawDatabase;
 use crate::function::delete::DeletedEntries;
 use crate::hash::{FxHashSet, FxIndexSet};
@@ -357,7 +354,11 @@ where
     ///
     /// Otherwise, the value is still provisional. For both final and provisional, it also
     /// returns the iteration in which this memo was created (always 0 except for cycle heads).
-    fn provisional_status(&self, zalsa: &Zalsa, input: Id) -> Option<ProvisionalStatus> {
+    fn provisional_status<'db>(
+        &self,
+        zalsa: &'db Zalsa,
+        input: Id,
+    ) -> Option<ProvisionalStatus<'db>> {
         let memo =
             self.get_memo_from_table_for(zalsa, input, self.memo_ingredient_index(zalsa, input))?;
 
@@ -377,6 +378,7 @@ where
             ProvisionalStatus::Provisional {
                 iteration,
                 verified_at: memo.verified_at.load(),
+                cycle_heads: memo.cycle_heads(),
             }
         })
     }
@@ -416,12 +418,6 @@ where
         self.sync_table.mark_as_transfer_target(key_index)
     }
 
-    fn cycle_heads<'db>(&self, zalsa: &'db Zalsa, input: Id) -> &'db CycleHeads {
-        self.get_memo_from_table_for(zalsa, input, self.memo_ingredient_index(zalsa, input))
-            .map(|memo| memo.cycle_heads())
-            .unwrap_or(empty_cycle_heads())
-    }
-
     /// Attempts to claim `key_index` without blocking.
     ///
     /// * [`WaitForResult::Running`] if the `key_index` is running on another thread. It's up to the caller to block on the other thread

src/function/execute.rs

@@ -248,9 +248,11 @@ where
                 let ingredient =
                     zalsa.lookup_ingredient(head.database_key_index.ingredient_index());
 
-                for nested_head in
-                    ingredient.cycle_heads(zalsa, head.database_key_index.key_index())
-                {
+                let provisional_status = ingredient
+                    .provisional_status(zalsa, head.database_key_index.key_index())
+                    .expect("cycle head memo must have been created during the execution");
+
+                for nested_head in provisional_status.cycle_heads() {
                     let nested_as_tuple = (
                         nested_head.database_key_index,
                         nested_head.iteration_count.load(),
@@ -442,6 +444,8 @@ where
 
             // Update the iteration count of this cycle head, but only after restoring
             // the cycle heads array (or this becomes a no-op).
+            // We don't call the same method on `cycle_heads` because that one doens't update
+            // the `memo.iteration_count`
             completed_query.revisions.set_cycle_heads(cycle_heads);
             completed_query
                 .revisions

src/function/maybe_changed_after.rs

@@ -484,8 +484,9 @@ where
 
         // Always return `false` for cycle initial values "unless" they are running in the same thread.
         if cycle_heads
-            .iter()
-            .all(|head| head.database_key_index == memo_database_key_index)
+            .iter_not_eq(memo_database_key_index)
+            .next()
+            .is_none()
         {
             // SAFETY: We do not access the query stack reentrantly.
             let on_stack = unsafe {
@@ -508,6 +509,8 @@ where
                     head_iteration_count,
                     memo_iteration_count: current_iteration_count,
                     verified_at: head_verified_at,
+                    cycle_heads,
+                    database_key_index: head_database_key,
                 } => {
                     if head_verified_at != memo_verified_at {
                         return false;
@@ -516,6 +519,27 @@ where
                     if head_iteration_count != current_iteration_count {
                         return false;
                     }
+
+                    // Check if the memo is still a cycle head and hasn't changed
+                    // to a normal cycle participant. This is to force re-execution in
+                    // a scenario like this:
+                    //
+                    // * There's a nested cycle with the outermost query A
+                    // * B participates in the cycle and is a cycle head in the first few iterations
+                    // * B becomes a non-cycle head in a later iteration
+                    // * There's a query `C` that has `B` as its cycle head
+                    //
+                    // The crucial point is that `B` switches from being a cycle head to being a regular cycle participant.
+                    // The issue with that is that `A` doesn't update `B`'s `iteration_count `when the iteration completes
+                    // because it only does that for cycle heads (and collecting all queries participating in a query would be sort of expensive?).
+                    //
+                    // When we now pull `C` in a later iteration, `validate_same_iteration` iterates over all its cycle heads (`B`),
+                    // and check if the iteration count still matches. Which is the case because `A` didn't update `B`'s iteration count.
+                    //
+                    // That's why we also check if `B` is still a cycle head in the current iteration.
+                    if !cycle_heads.contains(&head_database_key) {
+                        return false;
+                    }
                 }
                 _ => {
                     return false;

src/function/memo.rs

@@ -409,9 +409,11 @@ mod persistence {
 pub(super) enum TryClaimHeadsResult<'me> {
     /// Claiming the cycle head results in a cycle.
     Cycle {
+        database_key_index: DatabaseKeyIndex,
         head_iteration_count: IterationCount,
         memo_iteration_count: IterationCount,
         verified_at: Revision,
+        cycle_heads: &'me CycleHeads,
     },
 
     /// The cycle head is not finalized, but it can be claimed.
@@ -458,23 +460,28 @@ impl<'me> Iterator for TryClaimCycleHeadsIter<'me> {
                 let provisional_status = ingredient
                     .provisional_status(self.zalsa, head_key_index)
                     .expect("cycle head memo to exist");
-                let (current_iteration_count, verified_at) = match provisional_status {
+                let (current_iteration_count, verified_at, cycle_heads) = match provisional_status {
                     ProvisionalStatus::Provisional {
                         iteration,
                         verified_at,
-                    }
-                    | ProvisionalStatus::Final {
+                        cycle_heads,
+                    } => (iteration, verified_at, cycle_heads),
+                    ProvisionalStatus::Final {
                         iteration,
                         verified_at,
-                    } => (iteration, verified_at),
-                    ProvisionalStatus::FallbackImmediate => {
-                        (IterationCount::initial(), self.zalsa.current_revision())
-                    }
+                    } => (iteration, verified_at, empty_cycle_heads()),
+                    ProvisionalStatus::FallbackImmediate => (
+                        IterationCount::initial(),
+                        self.zalsa.current_revision(),
+                        empty_cycle_heads(),
+                    ),
                 };
 
                 Some(TryClaimHeadsResult::Cycle {
+                    database_key_index: head_database_key,
                     memo_iteration_count: current_iteration_count,
                     head_iteration_count: head.iteration_count.load(),
+                    cycle_heads,
                     verified_at,
                 })
             }

src/ingredient.rs

@@ -1,7 +1,7 @@
 use std::any::{Any, TypeId};
 use std::fmt;
 
-use crate::cycle::{empty_cycle_heads, CycleHeads, IterationCount, ProvisionalStatus};
+use crate::cycle::{IterationCount, ProvisionalStatus};
 use crate::database::RawDatabase;
 use crate::function::{VerifyCycleHeads, VerifyResult};
 use crate::hash::{FxHashSet, FxIndexSet};
@@ -74,16 +74,14 @@ pub trait Ingredient: Any + std::fmt::Debug + Send + Sync {
     /// Is it a provisional value or has it been finalized and in which iteration.
     ///
     /// Returns `None` if `input` doesn't exist.
-    fn provisional_status(&self, _zalsa: &Zalsa, _input: Id) -> Option<ProvisionalStatus> {
+    fn provisional_status<'db>(
+        &self,
+        _zalsa: &'db Zalsa,
+        _input: Id,
+    ) -> Option<ProvisionalStatus<'db>> {
         unreachable!("provisional_status should only be called on cycle heads and only functions can be cycle heads");
     }
 
-    /// Returns the cycle heads for this ingredient.
-    fn cycle_heads<'db>(&self, zalsa: &'db Zalsa, input: Id) -> &'db CycleHeads {
-        _ = (zalsa, input);
-        empty_cycle_heads()
-    }
-
     /// Invoked when the current thread needs to wait for a result for the given `key_index`.
     /// This call doesn't block the current thread. Instead, it's up to the caller to block
     /// in case `key_index` is [running](`WaitForResult::Running`) on another thread.