Fix await for queries depending on initial value

Author: MichaReiser

src/function.rs

@@ -5,7 +5,9 @@ use std::ptr::NonNull;
 use std::sync::atomic::Ordering;
 
 use crate::accumulator::accumulated_map::{AccumulatedMap, InputAccumulatedValues};
-use crate::cycle::{CycleHeadKind, CycleHeads, CycleRecoveryAction, CycleRecoveryStrategy};
+use crate::cycle::{
+    empty_cycle_heads, CycleHeadKind, CycleHeads, CycleRecoveryAction, CycleRecoveryStrategy,
+};
 use crate::function::delete::DeletedEntries;
 use crate::function::sync::{ClaimResult, SyncTable};
 use crate::ingredient::Ingredient;
@@ -256,6 +258,12 @@ where
         }
     }
 
+    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`, returning `false` if a cycle occurs.
     fn wait_for(&self, zalsa: &Zalsa, key_index: Id) -> bool {
         match self.sync_table.try_claim(zalsa, key_index) {

src/function/fetch.rs

@@ -202,7 +202,7 @@ where
                 && old_memo.may_be_provisional()
                 && old_memo.verified_at.load() == zalsa.current_revision()
             {
-                old_memo.await_heads(zalsa, database_key_index);
+                old_memo.await_heads(zalsa);
 
                 // It's possible that one of the cycle heads replaced the memo for this ingredient
                 // with fixpoint initial. We ignore that memo because we know it's only a temporary memo
@@ -213,13 +213,9 @@ where
         if let Some(old_memo) = opt_old_memo {
             if old_memo.value.is_some() {
                 let mut cycle_heads = CycleHeads::default();
-                if let VerifyResult::Unchanged(_) = self.deep_verify_memo(
-                    db,
-                    zalsa,
-                    old_memo,
-                    database_key_index,
-                    &mut cycle_heads,
-                ) {
+                if let VerifyResult::Unchanged(_) =
+                    self.deep_verify_memo(db, zalsa, old_memo, database_key_index, &mut cycle_heads)
+                {
                     if cycle_heads.is_empty() {
                         // SAFETY: memo is present in memo_map and we have verified that it is
                         // still valid for the current revision.

src/function/memo.rs

@@ -5,6 +5,7 @@ use std::ptr::NonNull;
 
 use crate::cycle::{empty_cycle_heads, CycleHeadKind, CycleHeads};
 use crate::function::{Configuration, IngredientImpl};
+use crate::hash::FxHashSet;
 use crate::key::DatabaseKeyIndex;
 use crate::revision::AtomicRevision;
 use crate::sync::atomic::Ordering;
@@ -99,6 +100,7 @@ pub struct Memo<V> {
 
 // Memo's are stored a lot, make sure their size is doesn't randomly increase.
 #[cfg(not(feature = "shuttle"))]
+#[cfg(target_pointer_width = "64")]
 const _: [(); std::mem::size_of::<Memo<std::num::NonZeroUsize>>()] =
     [(); std::mem::size_of::<[usize; 13]>()];
 
@@ -141,73 +143,27 @@ impl<V> Memo<V> {
             return false;
         };
 
-        return provisional_retry_cold(zalsa, database_key_index, &self.revisions.cycle_heads);
-
-        #[inline(never)]
-        fn provisional_retry_cold(
-            zalsa: &Zalsa,
-            database_key_index: DatabaseKeyIndex,
-            cycle_heads: &CycleHeads,
-        ) -> bool {
-            let mut retry = false;
-            let mut hit_cycle = false;
-
-            for head in cycle_heads {
-                let head_index = head.database_key_index;
-
-                let ingredient = zalsa.lookup_ingredient(head_index.ingredient_index());
-                let cycle_head_kind = ingredient.cycle_head_kind(zalsa, head_index.key_index());
-                if matches!(
-                    cycle_head_kind,
-                    CycleHeadKind::NotProvisional | CycleHeadKind::FallbackImmediate
-                ) {
-                    // This cycle is already finalized, so we don't need to wait on it;
-                    // keep looping through cycle heads.
-                    retry = true;
-                    tracing::trace!("Dependent cycle head {head_index:?} has been finalized.");
-                } else if ingredient.wait_for(zalsa, head_index.key_index()) {
-                    tracing::trace!("Dependent cycle head {head_index:?} has been released (there's a new memo)");
-                    // There's a new memo available for the cycle head; fetch our own
-                    // updated memo and see if it's still provisional or if the cycle
-                    // has resolved.
-                    tracing::trace!("Dependent cycle head {head_index:?} has been released (there's a new memo)");
-                    retry = true;
-                } else {
-                    // We hit a cycle blocking on the cycle head; this means it's in
-                    // our own active query stack and we are responsible to resolve the
-                    // cycle, so go ahead and return the provisional memo.
-                    tracing::debug!(
-                        "Waiting for {head_index:?} results in a cycle, return {database_key_index:?} once all other cycle heads completed to allow the outer cycle to make progress."
-                    );
-                    hit_cycle = true;
-                }
-            }
-
-            // If `retry` is `true`, all our cycle heads (barring ourself) are complete; re-fetch
-            // and we should get a non-provisional memo. If we get here and `retry` is still
-            // `false`, we have no cycle heads other than ourself, so we are a provisional value of
-            // the cycle head (either initial value, or from a later iteration) and should be
-            // returned to caller to allow fixpoint iteration to proceed. (All cases in the loop
-            // above other than "cycle head is self" are either terminal or set `retry`.)
-            if hit_cycle {
-                false
-            } else if retry {
-                tracing::debug!("Retrying {database_key_index:?}");
-                true
-            } else {
-                false
-            }
+        if self.await_heads(zalsa) {
+            false
+        } else {
+            tracing::debug!(
+                "Retrying provisional memo {database_key_index:?} after awaiting cycle heads."
+            );
+            true
         }
     }
 
-    #[inline(always)]
-    pub(super) fn await_heads(&self, zalsa: &Zalsa, database_key_index: DatabaseKeyIndex) {
-        for head in &self.revisions.cycle_heads {
-            let head_index = head.database_key_index;
+    /// Awaits all cycle heads (recursively) that this memo depends on.
+    ///
+    /// Returns `true` if awaiting the cycle heads resulted in a cycle.
+    pub(super) fn await_heads(&self, zalsa: &Zalsa) -> bool {
+        let mut hit_cycle = false;
 
-            if database_key_index == head_index {
-                continue;
-            }
+        let mut visited = FxHashSet::default();
+        let mut queue: Vec<_> = self.revisions.cycle_heads.iter().collect();
+
+        while let Some(head) = queue.pop() {
+            let head_index = head.database_key_index;
 
             let ingredient = zalsa.lookup_ingredient(head_index.ingredient_index());
             let cycle_head_kind = ingredient.cycle_head_kind(zalsa, head_index.key_index());
@@ -220,13 +176,35 @@ impl<V> Memo<V> {
                 // keep looping through cycle heads.
                 tracing::trace!("Dependent cycle head {head_index:?} has been finalized.");
             } else if ingredient.wait_for(zalsa, head_index.key_index()) {
-                tracing::trace!("Dependent cycle head {head_index:?} has been released");
+                // There's a new memo available for the cycle head; fetch our own
+                // updated memo and see if it's still provisional or if the cycle
+                // has resolved.
+                tracing::trace!(
+                    "Dependent cycle head {head_index:?} has been released (there's a new memo)"
+                );
+                // Recursively wait for all cycle heads that this head depends on.
+                // This is normally not necessary, because cycle heads are transitively added
+                // as query dependencies (they aggregate). The exception to this are queries
+                // that depend on a fixpoint initial value. We don't know all the dependencies of
+                // the query yet, so they can't be carried over. We only know them once the cycle
+                // completes but the cycle heads of the queries don't get updated.
+                // Because of that, recurse here to collect all cycle heads.
+                queue.extend(
+                    ingredient
+                        .cycle_heads(zalsa, head_index.key_index())
+                        .iter()
+                        .filter(|head| visited.insert(head.database_key_index)),
+                );
             } else {
                 // We hit a cycle blocking on the cycle head; this means it's in
                 // our own active query stack and we are responsible to resolve the
-                // cycle
+                // cycle, so go ahead and return the provisional memo.
+                tracing::debug!("Waiting for {head_index:?} results in a cycle");
+                hit_cycle = true;
             }
         }
+
+        hit_cycle
     }
 
     /// Cycle heads that should be propagated to dependent queries.

src/ingredient.rs

@@ -2,7 +2,7 @@ use std::any::{Any, TypeId};
 use std::fmt;
 
 use crate::accumulator::accumulated_map::{AccumulatedMap, InputAccumulatedValues};
-use crate::cycle::{CycleHeadKind, CycleHeads, CycleRecoveryStrategy};
+use crate::cycle::{empty_cycle_heads, CycleHeadKind, CycleHeads, CycleRecoveryStrategy};
 use crate::function::VerifyResult;
 use crate::plumbing::IngredientIndices;
 use crate::sync::Arc;
@@ -67,14 +67,17 @@ pub trait Ingredient: Any + std::fmt::Debug + Send + Sync {
     ) -> VerifyResult;
 
     /// Is the value for `input` in this ingredient a cycle head that is still provisional?
-    ///
-    /// In the case of nested cycles, we are not asking here whether the value is provisional due
-    /// to the outer cycle being unresolved, only whether its own cycle remains provisional.
     fn cycle_head_kind(&self, zalsa: &Zalsa, input: Id) -> CycleHeadKind {
         _ = (zalsa, input);
         CycleHeadKind::NotProvisional
     }
 
+    /// 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`.
     ///
     /// A return value of `true` indicates that a result is now available. A return value of

tests/parallel/cycle_a_t1_b_t2_fallback.rs

@@ -12,7 +12,7 @@
 //!  +--------------------+
 //! ```
 use crate::sync::thread;
-use crate::{Knobs, KnobsDatabase};
+use crate::KnobsDatabase;
 
 const FALLBACK_A: u32 = 0b01;
 const FALLBACK_B: u32 = 0b10;
@@ -53,6 +53,9 @@ fn cycle_result_b(_db: &dyn KnobsDatabase) -> u32 {
 #[test_log::test]
 #[cfg(not(feature = "shuttle"))] // This test is currently failing.
 fn the_test() {
+    use crate::sync::thread;
+    use crate::Knobs;
+
     crate::sync::check(|| {
         let db_t1 = Knobs::default();
         let db_t2 = db_t1.clone();

tests/parallel/cycle_nested_deep.rs

@@ -3,8 +3,9 @@
 //! The trick is that different threads call into the same cycle from different entry queries.
 //!
 //! * Thread 1: `a` -> b -> c (which calls back into d, e, b, a)
-//! * Thread 2: `d` -> `c`
-//! * Thread 3: `e` -> `c`
+//! * Thread 2: `b`
+//! * Thread 3: `d` -> `c`
+//! * Thread 4: `e` -> `c`
 use crate::sync::thread;
 use crate::{Knobs, KnobsDatabase};
 
@@ -65,6 +66,7 @@ fn the_test() {
         let db_t1 = Knobs::default();
         let db_t2 = db_t1.clone();
         let db_t3 = db_t1.clone();
+        let db_t4 = db_t1.clone();
 
         let t1 = thread::spawn(move || {
             let _span = tracing::debug_span!("t1", thread_id = ?thread::current().id()).entered();
@@ -73,11 +75,16 @@ fn the_test() {
             result
         });
         let t2 = thread::spawn(move || {
+            let _span = tracing::debug_span!("t4", thread_id = ?thread::current().id()).entered();
+            db_t4.wait_for(1);
+            query_b(&db_t4)
+        });
+        let t3 = thread::spawn(move || {
             let _span = tracing::debug_span!("t2", thread_id = ?thread::current().id()).entered();
             db_t2.wait_for(1);
             query_d(&db_t2)
         });
-        let t3 = thread::spawn(move || {
+        let t4 = thread::spawn(move || {
             let _span = tracing::debug_span!("t3", thread_id = ?thread::current().id()).entered();
             db_t3.wait_for(1);
             query_e(&db_t3)
@@ -86,7 +93,8 @@ fn the_test() {
         let r_t1 = t1.join().unwrap();
         let r_t2 = t2.join().unwrap();
         let r_t3 = t3.join().unwrap();
+        let r_t4 = t4.join().unwrap();
 
-        assert_eq!((r_t1, r_t2, r_t3), (MAX, MAX, MAX));
+        assert_eq!((r_t1, r_t2, r_t3, r_t4), (MAX, MAX, MAX, MAX));
     });
 }