Merge pull request groue#1747 from groue/dev/issue-1746

Transaction observers are not impacted by Task Cancellation

Author: groue

GRDB/Core/Database+Statements.swift

@@ -494,11 +494,9 @@ extension Database {
 
         switch ResultCode(rawValue: resultCode) {
         case .SQLITE_INTERRUPT, .SQLITE_ABORT:
-            if suspensionMutex.load().isCancelled {
-                // The only error that a user sees when a Task is cancelled
-                // is CancellationError.
-                throw CancellationError()
-            }
+            // The only error that a user sees when a Task is cancelled
+            // is CancellationError.
+            try suspensionMutex.load().checkCancellation()
         default:
             break
         }

GRDB/Core/Database.swift

@@ -341,10 +341,22 @@ public final class Database: CustomStringConvertible, CustomDebugStringConvertib
 
         /// If true, the database access has been cancelled.
         var isCancelled: Bool
+        
+        /// If true, the database throws an error when it is cancelled.
+        var interruptsWhenCancelled: Bool
+        
+        func checkCancellation() throws {
+            if isCancelled, interruptsWhenCancelled {
+                throw CancellationError()
+            }
+        }
     }
 
     /// Support for `checkForSuspensionViolation(from:)`
-    let suspensionMutex = Mutex(Suspension(isSuspended: false, isCancelled: false))
+    let suspensionMutex = Mutex(Suspension(
+        isSuspended: false,
+        isCancelled: false,
+        interruptsWhenCancelled: true))
 
     // MARK: - Transaction Date
 
@@ -1222,7 +1234,7 @@ public final class Database: CustomStringConvertible, CustomDebugStringConvertib
             }
 
             suspension.isCancelled = true
-            return true
+            return suspension.interruptsWhenCancelled
         }
 
         if needsInterrupt {
@@ -1237,6 +1249,24 @@ public final class Database: CustomStringConvertible, CustomDebugStringConvertib
         }
     }
 
+    /// Within the given closure, Task cancellation does not interrupt
+    /// database accesses.
+    func ignoringCancellation<T>(_ value: () throws -> T) rethrows -> T {
+        let previous = suspensionMutex.withLock {
+            let previous = $0.interruptsWhenCancelled
+            $0.interruptsWhenCancelled = false
+            return previous
+        }
+        
+        defer {
+            suspensionMutex.withLock {
+                $0.interruptsWhenCancelled = previous
+            }
+        }
+        
+        return try value()
+    }
+    
     /// Support for `checkForSuspensionViolation(from:)`
     private func journalMode() throws -> String {
         if let journalMode = journalModeCache {
@@ -1303,7 +1333,7 @@ public final class Database: CustomStringConvertible, CustomDebugStringConvertib
         let interrupt: Interrupt? = try suspensionMutex.withLock { suspension in
             // Check for cancellation first, so that the only error that
             // a user sees when a Task is cancelled is CancellationError.
-            if suspension.isCancelled {
+            if suspension.isCancelled, suspension.interruptsWhenCancelled {
                 return .cancel
             }
 

GRDB/Core/TransactionObserver.swift

@@ -549,8 +549,12 @@ class DatabaseObservationBroker {
         savepointStack.clear()
 
         if !database.isReadOnly {
-            for observation in transactionObservations {
-                observation.databaseDidCommit(database)
+            // Observers must be able to access the database, even if the
+            // task that has performed the commit is cancelled.
+            database.ignoringCancellation {
+                for observation in transactionObservations {
+                    observation.databaseDidCommit(database)
+                }
             }
         }
 
@@ -609,12 +613,18 @@ class DatabaseObservationBroker {
 
     // Called from statementDidExecute or statementDidFail
     private func databaseDidRollback(notifyTransactionObservers: Bool) {
+        #warning("TODO")
         savepointStack.clear()
 
         if notifyTransactionObservers {
             assert(!database.isReadOnly, "Read-only transactions are not notified")
-            for observation in transactionObservations {
-                observation.databaseDidRollback(database)
+            
+            // Observers must be able to access the database, even if the
+            // task that has performed the commit is cancelled.
+            database.ignoringCancellation {
+                for observation in transactionObservations {
+                    observation.databaseDidRollback(database)
+                }
             }
         }
         databaseDidEndTransaction()

Tests/GRDBTests/ValueObservationTests.swift

@@ -1352,4 +1352,58 @@ class ValueObservationTests: GRDBTestCase {
                     XCTAssertEqual(value, true)
                 })
     }
+    
+    // Regression test for <https://github.com/groue/GRDB.swift/discussions/1746>
+    @MainActor func testIssue1746() async throws {
+        let dbQueue = try makeDatabaseQueue()
+        try await dbQueue.write { db in
+            try db.execute(sql: "CREATE TABLE test (id INTEGER PRIMARY KEY)")
+        }
+        
+        // Start a task that waits for writeContinuation before it writes.
+        let (writeStream, writeContinuation) = AsyncStream.makeStream(of: Void.self)
+        let task = Task {
+            for await _ in writeStream { }
+            try? await dbQueue.write { db in
+                XCTAssertFalse(Task.isCancelled) // Required for the test to be meaningful
+                try db.execute(sql: "INSERT INTO test DEFAULT VALUES")
+            }
+        }
+        
+        // A transaction observer that cancels a Task after commit.
+        class CancelObserver: TransactionObserver {
+            let task: Task<Void, Never>
+            init(task: Task<Void, Never>) {
+                self.task = task
+            }
+            func observes(eventsOfKind eventKind: DatabaseEventKind) -> Bool { true }
+            func databaseDidChange(with event: DatabaseEvent) { }
+            func databaseDidRollback(_ db: Database) { }
+            func databaseDidCommit(_ db: Database) {
+                task.cancel()
+            }
+        }
+        
+        // Register CancelObserver first, so no other observer could access
+        // the database before the task is cancelled.
+        dbQueue.add(transactionObserver: CancelObserver(task: task), extent: .databaseLifetime)
+
+        // Start observing.
+        // We expect to see 0, then 1.
+        let values = ValueObservation
+            .tracking(Table("test").fetchCount)
+            .values(in: dbQueue)
+        for try await value in values {
+            if value == 0 {
+                // Perform the write.
+                writeContinuation.finish()
+            } else if value == 1 {
+                // Test passes
+                break
+            } else {
+                XCTFail("Unexpected value \(value)")
+                break
+            }
+        }
+    }
 }