Disable local cancellation while within a cycle computation

Author: Veykril

src/function/execute.rs

@@ -109,13 +109,18 @@ where
 
                 (new_value, completed_query)
             }
-            CycleRecoveryStrategy::Fixpoint => self.execute_maybe_iterate(
-                db,
-                opt_old_memo,
-                &mut claim_guard,
-                zalsa_local,
-                memo_ingredient_index,
-            ),
+            CycleRecoveryStrategy::Fixpoint => {
+                let was_disabled = zalsa_local.set_cancellation_disabled(true);
+                let res = self.execute_maybe_iterate(
+                    db,
+                    opt_old_memo,
+                    &mut claim_guard,
+                    zalsa_local,
+                    memo_ingredient_index,
+                );
+                zalsa_local.set_cancellation_disabled(was_disabled);
+                res
+            }
         };
 
         if let Some(old_memo) = opt_old_memo {

src/function/memo.rs

@@ -181,7 +181,11 @@ impl<'db, C: Configuration> Memo<'db, C> {
                     TryClaimHeadsResult::Running(running) => {
                         all_cycles = false;
                         if !running.block_on(zalsa) {
-                            // FIXME: Handle cancellation properly?
+                            // We cannot handle local cancellations in fixpoints
+                            // so we treat it as a general cancellation / panic.
+                            //
+                            // We shouldn't hit this though as we disable local cancellation
+                            // in cycles.
                             crate::Cancelled::PropagatedPanic.throw();
                         }
                     }

src/function/sync.rs

@@ -288,7 +288,7 @@ pub enum SyncOwner {
 
     /// The query's lock ownership has been transferred to another query.
     /// E.g. if `a` transfers its ownership to `b`, then only the thread in the critical path
-    /// to complete b` can claim `a` (in most instances, only the thread owning `b` can claim `a`).
+    /// to complete `b` can claim `a` (in most instances, only the thread owning `b` can claim `a`).
     ///
     /// The thread owning `a` is stored in the `DependencyGraph`.
     ///
@@ -469,7 +469,7 @@ impl<'me> ClaimGuard<'me> {
 impl Drop for ClaimGuard<'_> {
     fn drop(&mut self) {
         if thread::panicking() {
-            if self.zalsa_local.is_cancelled() {
+            if self.zalsa_local.should_trigger_local_cancellation() {
                 self.release_cancelled();
             } else {
                 self.release_panicking();

src/zalsa.rs

@@ -301,7 +301,7 @@ impl Zalsa {
     #[inline]
     pub(crate) fn unwind_if_revision_cancelled(&self, zalsa_local: &ZalsaLocal) {
         self.event(&|| crate::Event::new(crate::EventKind::WillCheckCancellation));
-        if zalsa_local.is_cancelled() {
+        if zalsa_local.should_trigger_local_cancellation() {
             zalsa_local.unwind_cancelled();
         }
         if self.runtime().load_cancellation_flag() {

src/zalsa_local.rs

@@ -3,7 +3,7 @@ use std::fmt;
 use std::fmt::Formatter;
 use std::panic::UnwindSafe;
 use std::ptr::{self, NonNull};
-use std::sync::atomic::Ordering;
+use std::sync::atomic::{AtomicU8, Ordering};
 use std::sync::Arc;
 
 use rustc_hash::FxHashMap;
@@ -47,21 +47,32 @@ pub struct ZalsaLocal {
 
 /// A cancellation token that can be used to cancel a query computation for a specific local `Database`.
 #[derive(Default, Clone, Debug)]
-pub struct CancellationToken(Arc<AtomicBool>);
+pub struct CancellationToken(Arc<AtomicU8>);
 
 impl CancellationToken {
+    const CANCELLED_MASK: u8 = 0b01;
+    const DISABLED_MASK: u8 = 0b10;
+
     /// Inform the database to cancel the current query computation.
     pub fn cancel(&self) {
-        self.0.store(true, Ordering::Relaxed);
+        self.0.fetch_or(Self::CANCELLED_MASK, Ordering::Relaxed);
     }
 
     /// Check if the query computation has been requested to be cancelled.
     pub fn is_cancelled(&self) -> bool {
-        self.0.load(Ordering::Relaxed)
+        self.0.load(Ordering::Relaxed) & Self::CANCELLED_MASK != 0
     }
 
-    pub(crate) fn uncancel(&self) {
-        self.0.store(false, Ordering::Relaxed);
+    fn set_cancellation_disabled(&self, disabled: bool) -> bool {
+        self.0.fetch_or((disabled as u8) << 1, Ordering::Relaxed) & Self::DISABLED_MASK != 0
+    }
+
+    fn should_trigger_local_cancellation(&self) -> bool {
+        self.0.load(Ordering::Relaxed) == Self::CANCELLED_MASK
+    }
+
+    fn reset(&self) {
+        self.0.store(0, Ordering::Relaxed);
     }
 }
 
@@ -433,12 +444,12 @@ impl ZalsaLocal {
 
     #[inline]
     pub(crate) fn uncancel(&self) {
-        self.cancelled.uncancel();
+        self.cancelled.reset();
     }
 
     #[inline]
-    pub fn is_cancelled(&self) -> bool {
-        self.cancelled.0.load(Ordering::Relaxed)
+    pub fn should_trigger_local_cancellation(&self) -> bool {
+        self.cancelled.should_trigger_local_cancellation()
     }
 
     #[cold]
@@ -450,6 +461,10 @@ impl ZalsaLocal {
     pub(crate) fn unwind_cancelled(&self) {
         Cancelled::Local.throw();
     }
+
+    pub(crate) fn set_cancellation_disabled(&self, was_disabled: bool) -> bool {
+        self.cancelled.set_cancellation_disabled(was_disabled)
+    }
 }
 
 // Okay to implement as `ZalsaLocal`` is !Sync