Cancellation cleanup on un-attach

Author: Veykril

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(NonNull::from(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

@@ -1,8 +1,6 @@
 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:
@@ -13,7 +11,7 @@ use crate::CancellationToken;
 #[non_exhaustive]
 pub enum Cancelled {
     /// The query was operating but the local database execution has been cancelled.
-    Cancelled(UncancelGuard),
+    Cancelled,
 
     /// The query was operating on revision R, but there is a pending write to move to revision R+1.
     PendingWrite,
@@ -22,15 +20,6 @@ pub enum Cancelled {
     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) -> ! {
@@ -57,7 +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::Cancelled => "canellation request",
             Cancelled::PendingWrite => "pending write",
             Cancelled::PropagatedPanic => "propagated panic",
         };

src/zalsa_local.rs

@@ -15,7 +15,6 @@ 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;
@@ -432,6 +431,11 @@ impl ZalsaLocal {
         self.cancelled.clone()
     }
 
+    #[inline]
+    pub(crate) fn uncancel(&self) {
+        self.cancelled.uncancel();
+    }
+
     #[inline]
     pub fn is_cancelled(&self) -> bool {
         self.cancelled.0.load(Ordering::Relaxed)
@@ -444,7 +448,7 @@ impl ZalsaLocal {
 
     #[cold]
     pub(crate) fn unwind_cancelled(&self) {
-        Cancelled::Cancelled(UncancelGuard(self.cancellation_token())).throw();
+        Cancelled::Cancelled.throw();
     }
 }
 

tests/cancellation_token.rs

@@ -44,7 +44,7 @@ fn cancellation_token() {
             a(&db, input)
         })
     });
-    assert!(matches!(res, Err(Cancelled::Cancelled(_))), "{res:?}");
+    assert!(matches!(res, Err(Cancelled::Cancelled)), "{res:?}");
     drop(res);
     db.assert_logs(expect![[r#"
         [

tests/parallel/cancellation_token_recomputes.rs

@@ -38,6 +38,6 @@ fn execute() {
     db_signaler.signal(3);
     let (r1, r2) = (t1.join(), t2.join());
     let r1 = *r1.unwrap_err().downcast::<salsa::Cancelled>().unwrap();
-    assert!(matches!(r1, Cancelled::Cancelled(_)), "{r1:?}");
+    assert!(matches!(r1, Cancelled::Cancelled), "{r1:?}");
     assert_eq!(r2.unwrap(), 1);
 }