Introduce a CancellationToken for cancelling specific computations

src/attach.rs

@@ -70,7 +70,10 @@ impl Attached {
             fn drop(&mut self) {
                 // Reset database to null if we did anything in `DbGuard::new`.
                 if let Some(attached) = self.state {
-                    attached.database.set(None);
+                    if let Some(prev) = attached.database.replace(None) {
+                        // SAFETY: `prev` is a valid pointer to a database.
+                        unsafe { prev.as_ref().zalsa_local().uncancel() };
+                    }
                 }
             }
         }
@@ -85,25 +88,50 @@ impl Attached {
         Db: ?Sized + Database,
     {
         struct DbGuard<'s> {
-            state: &'s Attached,
+            state: Option<&'s Attached>,
             prev: Option<NonNull<dyn Database>>,
         }
 
         impl<'s> DbGuard<'s> {
             #[inline]
             fn new(attached: &'s Attached, db: &dyn Database) -> Self {
-                let prev = attached.database.replace(Some(NonNull::from(db)));
-                Self {
-                    state: attached,
-                    prev,
+                let db = NonNull::from(db);
+                match attached.database.replace(Some(db)) {
+                    Some(prev) => {
+                        if std::ptr::eq(db.as_ptr(), prev.as_ptr()) {
+                            Self {
+                                state: None,
+                                prev: None,
+                            }
+                        } else {
+                            Self {
+                                state: Some(attached),
+                                prev: Some(prev),
+                            }
+                        }
+                    }
+                    None => {
+                        // Otherwise, set the database.
+                        attached.database.set(Some(db));
+                        Self {
+                            state: Some(attached),
+                            prev: None,
+                        }
+                    }
                 }
             }
         }
 
         impl Drop for DbGuard<'_> {
             #[inline]
             fn drop(&mut self) {
-                self.state.database.set(self.prev);
+                // Reset database to null if we did anything in `DbGuard::new`.
+                if let Some(attached) = self.state {
+                    if let Some(prev) = attached.database.replace(self.prev) {
+                        // SAFETY: `prev` is a valid pointer to a database.
+                        unsafe { prev.as_ref().zalsa_local().uncancel() };
+                    }
+                }
             }
         }
 

src/cancelled.rs

@@ -10,12 +10,13 @@ use std::panic::{self, UnwindSafe};
 #[derive(Debug)]
 #[non_exhaustive]
 pub enum Cancelled {
+    /// The query was operating but the local database execution has been cancelled.
+    Local,
+
     /// 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,
 }
 
@@ -45,6 +46,7 @@ impl Cancelled {
 impl std::fmt::Display for Cancelled {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let why = match self {
+            Cancelled::Local => "local canellation request",
             Cancelled::PendingWrite => "pending write",
             Cancelled::PropagatedPanic => "propagated panic",
         };

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();
     }
 
+    /// Retrieves a [`CancellationToken`] for the current database handle.
+    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

src/function/fetch.rs

@@ -105,10 +105,15 @@ 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);
+                if !blocked_on.block_on(zalsa) {
+                    return None;
+                }
 
                 if C::CYCLE_STRATEGY == CycleRecoveryStrategy::FallbackImmediate {
                     let memo = self.get_memo_from_table_for(zalsa, id, memo_ingredient_index);

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 {

src/function/memo.rs

@@ -180,7 +180,10 @@ impl<'db, C: Configuration> Memo<'db, C> {
                     }
                     TryClaimHeadsResult::Running(running) => {
                         all_cycles = false;
-                        running.block_on(zalsa);
+                        if !running.block_on(zalsa) {
+                            // FIXME: Handle cancellation properly?
+                            crate::Cancelled::PropagatedPanic.throw();
+                        }
                     }
                 }
             }
@@ -507,7 +510,7 @@ mod _memory_usage {
     use std::any::TypeId;
     use std::num::NonZeroUsize;
 
-    // Memo's are stored a lot, make sure their size is doesn't randomly increase.
+    // Memo's are stored a lot, make sure their size doesn't randomly increase.
     const _: [(); std::mem::size_of::<super::Memo<DummyConfiguration>>()] =
         [(); std::mem::size_of::<[usize; 6]>()];
 

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,
 };
@@ -21,6 +22,8 @@ pub(crate) struct SyncTable {
 }
 
 pub(crate) enum ClaimResult<'a, Guard = ClaimGuard<'a>> {
+    /// Successfully claimed the query.
+    Claimed(Guard),
     /// Can't claim the query because it is running on an other thread.
     Running(Running<'a>),
     /// Claiming the query results in a cycle.
@@ -30,8 +33,6 @@ pub(crate) enum ClaimResult<'a, Guard = ClaimGuard<'a>> {
         /// [`SyncTable::try_claim`] with [`Reentrant::Allow`].
         inner: bool,
     },
-    /// Successfully claimed the query.
-    Claimed(Guard),
 }
 
 pub(crate) struct SyncState {
@@ -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::Cancelled);
+    }
+
     #[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;
         }
 

src/interned.rs

@@ -850,6 +850,7 @@ pub struct StructEntry<'db, C>
 where
     C: Configuration,
 {
+    #[allow(dead_code)]
     value: &'db Value<C>,
     key: DatabaseKeyIndex,
 }

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 {