Introduce a CancellationToken for cancelling specific computations

Author: Veykril

src/cancelled.rs

@@ -7,15 +7,16 @@ use std::panic::{self, UnwindSafe};
 /// *
 /// *
 /// *
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq)]
 #[non_exhaustive]
 pub enum Cancelled {
+    /// The query was operating but the local database execution has been cancelled.
+    Cancelled,
+
     /// 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::Cancelled => "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();
     }
 
+    /// 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

src/lib.rs

@@ -67,6 +67,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 {

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.take_cancellation() {
+            zalsa_local.unwind_cancelled();
+        }
         if self.runtime().load_cancellation_flag() {
-            zalsa_local.unwind_cancelled(self.current_revision());
+            zalsa_local.unwind_pending_write();
         }
     }
 

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;
@@ -39,13 +41,32 @@ 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)
+    }
 }
 
 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 +422,24 @@ impl ZalsaLocal {
         }
     }
 
+    pub(crate) fn cancellation_token(&self) -> CancellationToken {
+        self.cancelled.clone()
+    }
+
+    #[inline]
+    pub fn take_cancellation(&self) -> bool {
+        self.cancelled.0.swap(false, 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.throw();
+    }
 }
 
 // Okay to implement as `ZalsaLocal`` is !Sync

tests/cancellation_token.rs

@@ -0,0 +1,65 @@
+//! Test that `DeriveWithDb` is correctly derived.
+
+mod common;
+
+use std::{sync::Barrier, thread};
+
+use expect_test::expect;
+use salsa::{Cancelled, Database};
+
+use crate::common::LogDatabase;
+
+#[salsa::input(debug)]
+struct MyInput {
+    field: u32,
+}
+
+#[salsa::tracked]
+fn a(db: &dyn Database, input: MyInput) -> u32 {
+    BARRIER.wait();
+    BARRIER2.wait();
+    b(db, input)
+}
+#[salsa::tracked]
+fn b(db: &dyn Database, input: MyInput) -> u32 {
+    input.field(db)
+}
+
+static BARRIER: Barrier = Barrier::new(2);
+static BARRIER2: Barrier = Barrier::new(2);
+
+#[test]
+fn cancellation_token() {
+    let db = common::EventLoggerDatabase::default();
+    let token = db.cancellation_token();
+    let input = MyInput::new(&db, 22);
+    let res = Cancelled::catch(|| {
+        thread::scope(|s| {
+            s.spawn(|| {
+                BARRIER.wait();
+                token.cancel();
+                BARRIER2.wait();
+            });
+            a(&db, input)
+        })
+    });
+    assert_eq!(res, Err(Cancelled::Cancelled));
+    db.assert_logs(expect![[r#"
+        [
+            "WillCheckCancellation",
+            "WillExecute { database_key: a(Id(0)) }",
+            "WillCheckCancellation",
+        ]"#]]);
+    thread::spawn(|| {
+        BARRIER.wait();
+        BARRIER2.wait();
+    });
+    a(&db, input);
+    db.assert_logs(expect![[r#"
+        [
+            "WillCheckCancellation",
+            "WillExecute { database_key: a(Id(0)) }",
+            "WillCheckCancellation",
+            "WillExecute { database_key: b(Id(0)) }",
+        ]"#]]);
+}

tests/parallel/cancellation_token_recomputes.rs

@@ -0,0 +1,46 @@
+// Shuttle doesn't like panics inside of its runtime.
+#![cfg(not(feature = "shuttle"))]
+
+//! Test for cancellation when another query is blocked on the cancelled thread.
+use salsa::{Cancelled, Database};
+
+use crate::setup::{Knobs, KnobsDatabase};
+
+#[salsa::tracked]
+fn query_a(db: &dyn KnobsDatabase) -> u32 {
+    query_b(db)
+}
+
+#[salsa::tracked]
+fn query_b(db: &dyn KnobsDatabase) -> u32 {
+    db.signal(1);
+    db.wait_for(3);
+    query_c(db)
+}
+
+#[salsa::tracked]
+fn query_c(_db: &dyn KnobsDatabase) -> u32 {
+    1
+}
+
+#[test]
+fn execute() {
+    let db = Knobs::default();
+    let db2 = db.clone();
+    let db_signaler = db.clone();
+    let token = db.cancellation_token();
+
+    let t1 = std::thread::spawn(move || query_a(&db));
+    db_signaler.wait_for(1);
+    db2.signal_on_will_block(2);
+    let t2 = std::thread::spawn(move || query_a(&db2));
+    db_signaler.wait_for(2);
+    token.cancel();
+    db_signaler.signal(3);
+    let (r1, r2) = (t1.join(), t2.join());
+    assert_eq!(
+        *r1.unwrap_err().downcast::<salsa::Cancelled>().unwrap(),
+        Cancelled::Cancelled
+    );
+    assert_eq!(r2.unwrap(), 1);
+}

tests/parallel/main.rs

@@ -3,6 +3,7 @@
 mod setup;
 mod signal;
 
+mod cancellation_token_recomputes;
 mod cycle_a_t1_b_t2;
 mod cycle_a_t1_b_t2_fallback;
 mod cycle_ab_peeping_c;