Introduce a CancellationToken for cancelling specific computations

Veykril · Veykril · commit f97346e40ad1 · 2025-10-26T16:36:33.000+01:00
diff --git a/src/cancelled.rs b/src/cancelled.rs
@@ -1,6 +1,8 @@
 use std::fmt;
 use std::panic::{self, UnwindSafe};
 
+use crate::CancellationToken;
+
 /// A panic payload indicating that execution of a salsa query was cancelled.
 ///
 /// This can occur for a few reasons:
@@ -10,15 +12,25 @@ use std::panic::{self, UnwindSafe};
 #[derive(Debug)]
 #[non_exhaustive]
 pub enum Cancelled {
+    /// The query was operating but the local database execution has been cancelled.
+    Cancelled(UncancelGuard),
+
     /// The query was operating on revision R, but there is a pending write to move to revision R+1.
-    #[non_exhaustive]
     PendingWrite,
 
     /// The query was blocked on another thread, and that thread panicked.
-    #[non_exhaustive]
     PropagatedPanic,
 }
 
+#[derive(Debug)]
+pub struct UncancelGuard(pub(crate) CancellationToken);
+
+impl Drop for UncancelGuard {
+    fn drop(&mut self) {
+        self.0.uncancel();
+    }
+}
+
 impl Cancelled {
     #[cold]
     pub(crate) fn throw(self) -> ! {
@@ -45,6 +57,7 @@ impl Cancelled {
 impl std::fmt::Display for Cancelled {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let why = match self {
+            Cancelled::Cancelled(_) => "canellation request",
             Cancelled::PendingWrite => "pending write",
             Cancelled::PropagatedPanic => "propagated panic",
         };
diff --git a/src/database.rs b/src/database.rs
@@ -3,6 +3,7 @@ use std::ptr::NonNull;
 
 use crate::views::DatabaseDownCaster;
 use crate::zalsa::{IngredientIndex, ZalsaDatabase};
+use crate::zalsa_local::CancellationToken;
 use crate::{Durability, Revision};
 
 #[derive(Copy, Clone)]
@@ -59,14 +60,19 @@ pub trait Database: Send + ZalsaDatabase + AsDynDatabase {
         zalsa_mut.runtime_mut().report_tracked_write(durability);
     }
 
-    /// This method triggers cancellation.
+    /// This method cancels all outstanding computations.
     /// If you invoke it while a snapshot exists, it
     /// will block until that snapshot is dropped -- if that snapshot
     /// is owned by the current thread, this could trigger deadlock.
     fn trigger_cancellation(&mut self) {
         let _ = self.zalsa_mut();
     }
 
+    /// Retrives a [`CancellationToken`] for the current database.
+    fn cancellation_token(&self) -> CancellationToken {
+        self.zalsa_local().cancellation_token()
+    }
+
     /// Reports that the query depends on some state unknown to salsa.
     ///
     /// Queries which report untracked reads will be re-executed in the next
diff --git a/src/function/fetch.rs b/src/function/fetch.rs
@@ -105,7 +105,10 @@ where
     ) -> Option<&'db Memo<'db, C>> {
         let database_key_index = self.database_key_index(id);
         // Try to claim this query: if someone else has claimed it already, go back and start again.
-        let claim_guard = match self.sync_table.try_claim(zalsa, id, Reentrancy::Allow) {
+        let claim_guard = match self
+            .sync_table
+            .try_claim(zalsa, zalsa_local, id, Reentrancy::Allow)
+        {
             ClaimResult::Claimed(guard) => guard,
             ClaimResult::Running(blocked_on) => {
                 blocked_on.block_on(zalsa);
diff --git a/src/function/maybe_changed_after.rs b/src/function/maybe_changed_after.rs
@@ -141,23 +141,24 @@ where
     ) -> Option<VerifyResult> {
         let database_key_index = self.database_key_index(key_index);
 
-        let claim_guard = match self
-            .sync_table
-            .try_claim(zalsa, key_index, Reentrancy::Deny)
-        {
-            ClaimResult::Claimed(guard) => guard,
-            ClaimResult::Running(blocked_on) => {
-                blocked_on.block_on(zalsa);
-                return None;
-            }
-            ClaimResult::Cycle { .. } => {
-                return Some(self.maybe_changed_after_cold_cycle(
-                    zalsa_local,
-                    database_key_index,
-                    cycle_heads,
-                ))
-            }
-        };
+        let claim_guard =
+            match self
+                .sync_table
+                .try_claim(zalsa, zalsa_local, key_index, Reentrancy::Deny)
+            {
+                ClaimResult::Claimed(guard) => guard,
+                ClaimResult::Running(blocked_on) => {
+                    blocked_on.block_on(zalsa);
+                    return None;
+                }
+                ClaimResult::Cycle { .. } => {
+                    return Some(self.maybe_changed_after_cold_cycle(
+                        zalsa_local,
+                        database_key_index,
+                        cycle_heads,
+                    ))
+                }
+            };
         // Load the current memo, if any.
         let Some(old_memo) = self.get_memo_from_table_for(zalsa, key_index, memo_ingredient_index)
         else {
diff --git a/src/function/sync.rs b/src/function/sync.rs
@@ -2,6 +2,7 @@ use rustc_hash::FxHashMap;
 use std::collections::hash_map::OccupiedEntry;
 
 use crate::key::DatabaseKeyIndex;
+use crate::plumbing::ZalsaLocal;
 use crate::runtime::{
     BlockOnTransferredOwner, BlockResult, BlockTransferredResult, Running, WaitResult,
 };
@@ -68,6 +69,7 @@ impl SyncTable {
     pub(crate) fn try_claim<'me>(
         &'me self,
         zalsa: &'me Zalsa,
+        zalsa_local: &'me ZalsaLocal,
         key_index: Id,
         reentrant: Reentrancy,
     ) -> ClaimResult<'me> {
@@ -77,7 +79,12 @@ impl SyncTable {
                 let id = match occupied_entry.get().id {
                     SyncOwner::Thread(id) => id,
                     SyncOwner::Transferred => {
-                        return match self.try_claim_transferred(zalsa, occupied_entry, reentrant) {
+                        return match self.try_claim_transferred(
+                            zalsa,
+                            zalsa_local,
+                            occupied_entry,
+                            reentrant,
+                        ) {
                             Ok(claimed) => claimed,
                             Err(other_thread) => match other_thread.block(write) {
                                 BlockResult::Cycle => ClaimResult::Cycle { inner: false },
@@ -115,6 +122,7 @@ impl SyncTable {
                 ClaimResult::Claimed(ClaimGuard {
                     key_index,
                     zalsa,
+                    zalsa_local,
                     sync_table: self,
                     mode: ReleaseMode::Default,
                 })
@@ -172,6 +180,7 @@ impl SyncTable {
     fn try_claim_transferred<'me>(
         &'me self,
         zalsa: &'me Zalsa,
+        zalsa_local: &'me ZalsaLocal,
         mut entry: OccupiedEntry<Id, SyncState>,
         reentrant: Reentrancy,
     ) -> Result<ClaimResult<'me>, Box<BlockOnTransferredOwner<'me>>> {
@@ -195,6 +204,7 @@ impl SyncTable {
                 Ok(ClaimResult::Claimed(ClaimGuard {
                     key_index,
                     zalsa,
+                    zalsa_local,
                     sync_table: self,
                     mode: ReleaseMode::SelfOnly,
                 }))
@@ -214,6 +224,7 @@ impl SyncTable {
                 Ok(ClaimResult::Claimed(ClaimGuard {
                     key_index,
                     zalsa,
+                    zalsa_local,
                     sync_table: self,
                     mode: ReleaseMode::Default,
                 }))
@@ -295,6 +306,7 @@ pub(crate) struct ClaimGuard<'me> {
     zalsa: &'me Zalsa,
     sync_table: &'me SyncTable,
     mode: ReleaseMode,
+    zalsa_local: &'me ZalsaLocal,
 }
 
 impl<'me> ClaimGuard<'me> {
@@ -319,10 +331,21 @@ impl<'me> ClaimGuard<'me> {
             "Release claim on {:?} due to panic",
             self.database_key_index()
         );
-
         self.release(state, WaitResult::Panicked);
     }
 
+    #[cold]
+    #[inline(never)]
+    fn release_cancelled(&self) {
+        let mut syncs = self.sync_table.syncs.lock();
+        let state = syncs.remove(&self.key_index).expect("key claimed twice?");
+        tracing::debug!(
+            "Release claim on {:?} due to cancellation",
+            self.database_key_index()
+        );
+        self.release(state, WaitResult::Canceled);
+    }
+
     #[inline(always)]
     fn release(&self, state: SyncState, wait_result: WaitResult) {
         let SyncState {
@@ -446,7 +469,11 @@ impl<'me> ClaimGuard<'me> {
 impl Drop for ClaimGuard<'_> {
     fn drop(&mut self) {
         if thread::panicking() {
-            self.release_panicking();
+            if self.zalsa_local.is_cancelled() {
+                self.release_cancelled();
+            } else {
+                self.release_panicking();
+            }
             return;
         }
 
diff --git a/src/interned.rs b/src/interned.rs
@@ -850,6 +850,7 @@ pub struct StructEntry<'db, C>
 where
     C: Configuration,
 {
+    #[allow(dead_code)]
     value: &'db Value<C>,
     key: DatabaseKeyIndex,
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -62,6 +62,7 @@ pub use self::runtime::Runtime;
 pub use self::storage::{Storage, StorageHandle};
 pub use self::update::Update;
 pub use self::zalsa::IngredientIndex;
+pub use self::zalsa_local::CancellationToken;
 pub use crate::attach::{attach, attach_allow_change, with_attached_database};
 
 pub mod prelude {
diff --git a/src/runtime.rs b/src/runtime.rs
@@ -44,6 +44,7 @@ pub struct Runtime {
 pub(super) enum WaitResult {
     Completed,
     Panicked,
+    Canceled,
 }
 
 #[derive(Debug)]
@@ -121,7 +122,7 @@ struct BlockedOnInner<'me> {
 
 impl Running<'_> {
     /// Blocks on the other thread to complete the computation.
-    pub(crate) fn block_on(self, zalsa: &Zalsa) {
+    pub(crate) fn block_on(self, zalsa: &Zalsa) -> bool {
         let BlockedOnInner {
             dg,
             query_mutex_guard,
@@ -151,7 +152,8 @@ impl Running<'_> {
                 // by the other thread and responded to appropriately.
                 Cancelled::PropagatedPanic.throw()
             }
-            WaitResult::Completed => {}
+            WaitResult::Canceled => true,
+            WaitResult::Completed => false,
         }
     }
 }
diff --git a/src/storage.rs b/src/storage.rs
@@ -124,12 +124,15 @@ impl<Db: Database> Storage<Db> {
             .record_unfilled_pages(self.handle.zalsa_impl.table());
         let Self {
             handle,
-            zalsa_local: _,
-        } = &self;
+            zalsa_local,
+        } = &mut self;
         // Avoid rust's annoying destructure prevention rules for `Drop` types
         // SAFETY: We forget `Self` afterwards to discard the original copy, and the destructure
         // above makes sure we won't forget to take into account newly added fields.
         let handle = unsafe { std::ptr::read(handle) };
+        // SAFETY: We forget `Self` afterwards to discard the original copy, and the destructure
+        // above makes sure we won't forget to take into account newly added fields.
+        unsafe { std::ptr::drop_in_place(zalsa_local) };
         std::mem::forget::<Self>(self);
         handle
     }
diff --git a/src/zalsa.rs b/src/zalsa.rs
@@ -301,8 +301,11 @@ 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() {
+            zalsa_local.unwind_cancelled();
+        }
         if self.runtime().load_cancellation_flag() {
-            zalsa_local.unwind_cancelled(self.current_revision());
+            zalsa_local.unwind_pending_write();
         }
     }
 
diff --git a/src/zalsa_local.rs b/src/zalsa_local.rs
@@ -3,6 +3,8 @@ use std::fmt;
 use std::fmt::Formatter;
 use std::panic::UnwindSafe;
 use std::ptr::{self, NonNull};
+use std::sync::atomic::Ordering;
+use std::sync::Arc;
 
 use rustc_hash::FxHashMap;
 use thin_vec::ThinVec;
@@ -13,6 +15,7 @@ use crate::accumulator::{
     Accumulator,
 };
 use crate::active_query::{CompletedQuery, QueryStack};
+use crate::cancelled::UncancelGuard;
 use crate::cycle::{empty_cycle_heads, AtomicIterationCount, CycleHeads, IterationCount};
 use crate::durability::Durability;
 use crate::key::DatabaseKeyIndex;
@@ -39,13 +42,36 @@ pub struct ZalsaLocal {
     /// Stores the most recent page for a given ingredient.
     /// This is thread-local to avoid contention.
     most_recent_pages: UnsafeCell<FxHashMap<IngredientIndex, PageIndex>>,
+
+    cancelled: CancellationToken,
+}
+
+/// 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>);
+
+impl CancellationToken {
+    /// Inform the database to cancel the current query computation.
+    pub fn cancel(&self) {
+        self.0.store(true, Ordering::Relaxed);
+    }
+
+    /// Check if the query computation has been requested to be cancelled.
+    pub fn is_cancelled(&self) -> bool {
+        self.0.load(Ordering::Relaxed)
+    }
+
+    pub(crate) fn uncancel(&self) {
+        self.0.store(false, Ordering::Relaxed);
+    }
 }
 
 impl ZalsaLocal {
     pub(crate) fn new() -> Self {
         ZalsaLocal {
             query_stack: RefCell::new(QueryStack::default()),
             most_recent_pages: UnsafeCell::new(FxHashMap::default()),
+            cancelled: CancellationToken::default(),
         }
     }
 
@@ -401,12 +427,25 @@ impl ZalsaLocal {
         }
     }
 
+    #[inline]
+    pub(crate) fn cancellation_token(&self) -> CancellationToken {
+        self.cancelled.clone()
+    }
+
+    #[inline]
+    pub fn is_cancelled(&self) -> bool {
+        self.cancelled.0.load(Ordering::Relaxed)
+    }
+
     #[cold]
-    pub(crate) fn unwind_cancelled(&self, current_revision: Revision) {
-        // Why is this reporting an untracked read? We do not store the query revisions on unwind do we?
-        self.report_untracked_read(current_revision);
+    pub(crate) fn unwind_pending_write(&self) {
         Cancelled::PendingWrite.throw();
     }
+
+    #[cold]
+    pub(crate) fn unwind_cancelled(&self) {
+        Cancelled::Cancelled(UncancelGuard(self.cancellation_token())).throw();
+    }
 }
 
 // Okay to implement as `ZalsaLocal`` is !Sync
diff --git a/tests/cancellation_token.rs b/tests/cancellation_token.rs
diff --git a/tests/interned-revisions.rs b/tests/interned-revisions.rs
diff --git a/tests/parallel/cancellation_token_recomputes.rs b/tests/parallel/cancellation_token_recomputes.rs
diff --git a/tests/parallel/main.rs b/tests/parallel/main.rs

Original file line number	Diff line number	Diff line change
`@@ -850,6 +850,7 @@ pub struct StructEntry<'db, C>`
`850`	`850`	`where`
`851`	`851`	`C: Configuration,`
`852`	`852`	`{`
	`853`	`+ #[allow(dead_code)]`
`853`	`854`	`value: &'db Value<C>,`
`854`	`855`	`key: DatabaseKeyIndex,`
`855`	`856`	`}`
Original file line number	Diff line number	Diff line change
`@@ -44,6 +44,7 @@ pub struct Runtime {`
`44`	`44`	`pub(super) enum WaitResult {`
`45`	`45`	`Completed,`
`46`	`46`	`Panicked,`
	`47`	`+ Canceled,`
`47`	`48`	`}`
`48`	`49`
`49`	`50`	`#[derive(Debug)]`
`@@ -121,7 +122,7 @@ struct BlockedOnInner<'me> {`
`121`	`122`
`122`	`123`	`impl Running<'_> {`
`123`	`124`	`/// Blocks on the other thread to complete the computation.`
`124`		`- pub(crate) fn block_on(self, zalsa: &Zalsa) {`
	`125`	`+ pub(crate) fn block_on(self, zalsa: &Zalsa) -> bool {`
`125`	`126`	`let BlockedOnInner {`
`126`	`127`	`dg,`
`127`	`128`	`query_mutex_guard,`
`@@ -151,7 +152,8 @@ impl Running<'_> {`
`151`	`152`	`// by the other thread and responded to appropriately.`
`152`	`153`	`Cancelled::PropagatedPanic.throw()`
`153`	`154`	`}`
`154`		`- WaitResult::Completed => {}`
	`155`	`+ WaitResult::Canceled => true,`
	`156`	`+ WaitResult::Completed => false,`
`155`	`157`	`}`
`156`	`158`	`}`
`157`	`159`	`}`
Original file line number	Diff line number	Diff line change
`@@ -301,8 +301,11 @@ impl Zalsa {`
`301`	`301`	`#[inline]`
`302`	`302`	`pub(crate) fn unwind_if_revision_cancelled(&self, zalsa_local: &ZalsaLocal) {`
`303`	`303`	`self.event(&\|\| crate::Event::new(crate::EventKind::WillCheckCancellation));`
	`304`	`+ if zalsa_local.is_cancelled() {`
	`305`	`+ zalsa_local.unwind_cancelled();`
	`306`	`+ }`
`304`	`307`	`if self.runtime().load_cancellation_flag() {`
`305`		`- zalsa_local.unwind_cancelled(self.current_revision());`
	`308`	`+ zalsa_local.unwind_pending_write();`
`306`	`309`	`}`
`307`	`310`	`}`
`308`	`311`