Short circuit if entire query runs on single thread

Author: MichaReiser

src/function/accumulated.rs

@@ -95,8 +95,9 @@ where
         db: &'db C::DbView,
         key: Id,
     ) -> (Option<&'db AccumulatedMap>, InputAccumulatedValues) {
+        let (zalsa, zalsa_local) = db.zalsas();
         // NEXT STEP: stash and refactor `fetch` to return an `&Memo` so we can make this work
-        let memo = self.refresh_memo(db, db.zalsa(), key);
+        let memo = self.refresh_memo(db, zalsa, zalsa_local, key);
         (
             memo.revisions.accumulated.as_deref(),
             memo.revisions.accumulated_inputs.load(),

src/function/fetch.rs

@@ -3,7 +3,7 @@ use crate::function::memo::Memo;
 use crate::function::sync::ClaimResult;
 use crate::function::{Configuration, IngredientImpl, VerifyResult};
 use crate::zalsa::{MemoIngredientIndex, Zalsa, ZalsaDatabase};
-use crate::zalsa_local::QueryRevisions;
+use crate::zalsa_local::{QueryRevisions, ZalsaLocal};
 use crate::Id;
 
 impl<C> IngredientImpl<C>
@@ -18,7 +18,7 @@ where
         #[cfg(debug_assertions)]
         let _span = tracing::debug_span!("fetch", query = ?database_key_index).entered();
 
-        let memo = self.refresh_memo(db, zalsa, id);
+        let memo = self.refresh_memo(db, zalsa, zalsa_local, id);
         // SAFETY: We just refreshed the memo so it is guaranteed to contain a value now.
         let memo_value = unsafe { memo.value.as_ref().unwrap_unchecked() };
 
@@ -41,13 +41,16 @@ where
         &'db self,
         db: &'db C::DbView,
         zalsa: &'db Zalsa,
+        zalsa_local: &'db ZalsaLocal,
         id: Id,
     ) -> &'db Memo<C::Output<'db>> {
         let memo_ingredient_index = self.memo_ingredient_index(zalsa, id);
         loop {
             if let Some(memo) = self
                 .fetch_hot(zalsa, id, memo_ingredient_index)
-                .or_else(|| self.fetch_cold_with_retry(zalsa, db, id, memo_ingredient_index))
+                .or_else(|| {
+                    self.fetch_cold_with_retry(zalsa, zalsa_local, db, id, memo_ingredient_index)
+                })
             {
                 return memo;
             }
@@ -84,11 +87,12 @@ where
     fn fetch_cold_with_retry<'db>(
         &'db self,
         zalsa: &'db Zalsa,
+        zalsa_local: &'db ZalsaLocal,
         db: &'db C::DbView,
         id: Id,
         memo_ingredient_index: MemoIngredientIndex,
     ) -> Option<&'db Memo<C::Output<'db>>> {
-        let memo = self.fetch_cold(zalsa, db, id, memo_ingredient_index)?;
+        let memo = self.fetch_cold(zalsa, zalsa_local, db, id, memo_ingredient_index)?;
 
         // If we get back a provisional cycle memo, and it's provisional on any cycle heads
         // that are claimed by a different thread, we can't propagate the provisional memo
@@ -98,7 +102,7 @@ where
         // That is only correct for fixpoint cycles, though: `FallbackImmediate` cycles
         // never have provisional entries.
         if C::CYCLE_STRATEGY == CycleRecoveryStrategy::FallbackImmediate
-            || !memo.provisional_retry(zalsa, self.database_key_index(id))
+            || !memo.provisional_retry(zalsa, zalsa_local, self.database_key_index(id))
         {
             Some(memo)
         } else {
@@ -109,6 +113,7 @@ where
     fn fetch_cold<'db>(
         &'db self,
         zalsa: &'db Zalsa,
+        zalsa_local: &'db ZalsaLocal,
         db: &'db C::DbView,
         id: Id,
         memo_ingredient_index: MemoIngredientIndex,
@@ -161,7 +166,7 @@ where
                         tracing::debug!(
                             "hit a `FallbackImmediate` cycle at {database_key_index:#?}"
                         );
-                        let active_query = db.zalsa_local().push_query(database_key_index, 0);
+                        let active_query = zalsa_local.push_query(database_key_index, 0);
                         let fallback_value = C::cycle_initial(db, C::id_to_input(db, id));
                         let mut revisions = active_query.pop();
                         revisions.cycle_heads = CycleHeads::initial(database_key_index);
@@ -202,7 +207,7 @@ where
                 && old_memo.may_be_provisional()
                 && old_memo.verified_at.load() == zalsa.current_revision()
             {
-                old_memo.await_heads(zalsa);
+                old_memo.await_heads(zalsa, zalsa_local);
 
                 // 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
@@ -227,7 +232,7 @@ where
 
         let memo = self.execute(
             db,
-            db.zalsa_local().push_query(database_key_index, 0),
+            zalsa_local.push_query(database_key_index, 0),
             opt_old_memo,
         );
 

src/function/maybe_changed_after.rs

@@ -4,10 +4,9 @@ use crate::function::memo::Memo;
 use crate::function::sync::ClaimResult;
 use crate::function::{Configuration, IngredientImpl};
 use crate::key::DatabaseKeyIndex;
-use crate::plumbing::ZalsaLocal;
 use crate::sync::atomic::Ordering;
 use crate::zalsa::{MemoIngredientIndex, Zalsa, ZalsaDatabase};
-use crate::zalsa_local::{QueryEdge, QueryOrigin};
+use crate::zalsa_local::{QueryEdge, QueryOrigin, ZalsaLocal};
 use crate::{AsDynDatabase as _, Id, Revision};
 
 /// Result of memo validation.
@@ -303,20 +302,7 @@ where
             memo = memo.tracing_debug()
         );
 
-        let cycle_heads = &memo.revisions.cycle_heads;
-        if cycle_heads.is_empty() {
-            return true;
-        }
-
-        zalsa_local.with_query_stack(|stack| {
-            cycle_heads.iter().all(|cycle_head| {
-                stack
-                    .iter()
-                    .rev()
-                    .find(|query| query.database_key_index == cycle_head.database_key_index)
-                    .is_some_and(|query| query.iteration_count() == cycle_head.iteration_count)
-            })
-        })
+        memo.validate_same_iteration(zalsa_local)
     }
 
     /// VerifyResult::Unchanged if the memo's value and `changed_at` time is up-to-date in the

src/function/memo.rs

@@ -11,7 +11,7 @@ use crate::revision::AtomicRevision;
 use crate::sync::atomic::Ordering;
 use crate::table::memo::MemoTableWithTypesMut;
 use crate::zalsa::{MemoIngredientIndex, Zalsa};
-use crate::zalsa_local::{QueryOrigin, QueryRevisions};
+use crate::zalsa_local::{QueryOrigin, QueryRevisions, ZalsaLocal};
 use crate::{Event, EventKind, Id, Revision};
 
 impl<C: Configuration> IngredientImpl<C> {
@@ -133,6 +133,7 @@ impl<V> Memo<V> {
     pub(super) fn provisional_retry(
         &self,
         zalsa: &Zalsa,
+        zalsa_local: &ZalsaLocal,
         database_key_index: DatabaseKeyIndex,
     ) -> bool {
         if self.revisions.cycle_heads.is_empty() {
@@ -143,7 +144,7 @@ impl<V> Memo<V> {
             return false;
         };
 
-        if self.await_heads(zalsa) {
+        if self.await_heads(zalsa, zalsa_local) {
             // If we get here, 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.
@@ -162,7 +163,14 @@ impl<V> Memo<V> {
     ///
     /// Returns `true` if awaiting the cycle heads resulted in a cycle.
     #[inline(never)]
-    pub(super) fn await_heads(&self, zalsa: &Zalsa) -> bool {
+    pub(super) fn await_heads(&self, zalsa: &Zalsa, zalsa_local: &ZalsaLocal) -> bool {
+        // The most common case is that the entire cycle is running in the same thread.
+        // If that's the case, short circuit and return `true` immediately.
+        if self.validate_same_iteration(zalsa_local) {
+            return true;
+        }
+
+        // Otherwise, await all cycle heads, recursively.
         let mut queue: Vec<_> = self
             .revisions
             .cycle_heads
@@ -220,6 +228,26 @@ impl<V> Memo<V> {
         hit_cycle
     }
 
+    /// If this is a provisional memo, validate that it was cached in the same iteration of the
+    /// same cycle(s) that we are still executing. If so, it is valid for reuse. This avoids
+    /// runaway re-execution of the same queries within a fixpoint iteration.
+    pub(super) fn validate_same_iteration(&self, zalsa_local: &ZalsaLocal) -> bool {
+        let cycle_heads = &self.revisions.cycle_heads;
+        if cycle_heads.is_empty() {
+            return true;
+        }
+
+        zalsa_local.with_query_stack(|stack| {
+            cycle_heads.iter().all(|cycle_head| {
+                stack
+                    .iter()
+                    .rev()
+                    .find(|query| query.database_key_index == cycle_head.database_key_index)
+                    .is_some_and(|query| query.iteration_count() == cycle_head.iteration_count)
+            })
+        })
+    }
+
     /// Cycle heads that should be propagated to dependent queries.
     #[inline(always)]
     pub(super) fn cycle_heads(&self) -> &CycleHeads {