fix: fix data race with task cancellation

Author: soumyamahunt

.github/workflows/main.yml

@@ -5,8 +5,8 @@ on:
     branches: [ main ]
   pull_request:
     branches: [ main ]
-  schedule:
-    - cron: '0 0 * * *'
+  # schedule:
+  #   - cron: '0 0 * * *'
   workflow_dispatch:
     inputs:
       release:

AsyncObjects.podspec

@@ -41,5 +41,6 @@ Pod::Spec.new do |s|
 
   s.test_spec do |ts|
     ts.source_files = "Tests/#{s.name}Tests/**/*.swift"
+    ts.scheme = { :parallelizable => true }
   end
 end

AsyncObjects.xcodeproj/project.pbxproj

@@ -22,7 +22,10 @@ import OrderedCollections
 public actor AsyncCountdownEvent: AsyncObject {
     /// The suspended tasks continuation type.
     @usableFromInline
-    typealias Continuation = GlobalContinuation<Void, Error>
+    typealias Continuation = SafeContinuation<GlobalContinuation<Void, Error>>
+    /// The platform dependent lock used to synchronize continuations tracking.
+    @usableFromInline
+    let locker: Locker = .init()
     /// The continuations stored with an associated key for all the suspended task that are waiting to be resumed.
     @usableFromInline
     private(set) var continuations: OrderedDictionary<UUID, Continuation> = [:]
@@ -47,6 +50,12 @@ public actor AsyncCountdownEvent: AsyncObject {
 
     // MARK: Internal
 
+    /// Checks whether to wait for countdown to signal.
+    ///
+    /// - Returns: Whether to wait to be resumed later.
+    @inlinable
+    func _wait() -> Bool { !isSet || !continuations.isEmpty }
+
     /// Resume provided continuation with additional changes based on the associated flags.
     ///
     /// - Parameter continuation: The queued continuation to resume.
@@ -66,10 +75,8 @@ public actor AsyncCountdownEvent: AsyncObject {
         _ continuation: Continuation,
         withKey key: UUID
     ) {
-        guard !isSet, continuations.isEmpty else {
-            _resumeContinuation(continuation)
-            return
-        }
+        guard !continuation.resumed else { return }
+        guard _wait() else { _resumeContinuation(continuation); return }
         continuations[key] = continuation
     }
 
@@ -79,8 +86,7 @@ public actor AsyncCountdownEvent: AsyncObject {
     /// - Parameter key: The key in the map.
     @inlinable
     func _removeContinuation(withKey key: UUID) {
-        let continuation = continuations.removeValue(forKey: key)
-        continuation?.cancel()
+        continuations.removeValue(forKey: key)
     }
 
     /// Decrements countdown count by the provided number.
@@ -113,15 +119,13 @@ public actor AsyncCountdownEvent: AsyncObject {
     @inlinable
     nonisolated func _withPromisedContinuation() async throws {
         let key = UUID()
-        try await withTaskCancellationHandler { [weak self] in
+        try await Continuation.withCancellation(synchronizedWith: locker) {
             Task { [weak self] in
                 await self?._removeContinuation(withKey: key)
             }
-        } operation: { () -> Continuation.Success in
-            try await Continuation.with { continuation in
-                Task { [weak self] in
-                    await self?._addContinuation(continuation, withKey: key)
-                }
+        } operation: { continuation in
+            Task { [weak self] in
+                await self?._addContinuation(continuation, withKey: key)
             }
         }
     }
@@ -205,7 +209,7 @@ public actor AsyncCountdownEvent: AsyncObject {
     /// Use this to wait for high priority tasks completion to start low priority ones.
     @Sendable
     public func wait() async {
-        guard !isSet else { currentCount += 1; return }
+        guard _wait() else { currentCount += 1; return }
         try? await _withPromisedContinuation()
     }
 }

Sources/AsyncObjects/AsyncCountdownEvent.swift

@@ -13,7 +13,10 @@ import Foundation
 public actor AsyncEvent: AsyncObject {
     /// The suspended tasks continuation type.
     @usableFromInline
-    typealias Continuation = GlobalContinuation<Void, Error>
+    typealias Continuation = SafeContinuation<GlobalContinuation<Void, Error>>
+    /// The platform dependent lock used to synchronize continuations tracking.
+    @usableFromInline
+    let locker: Locker = .init()
     /// The continuations stored with an associated key for all the suspended task that are waiting for event signal.
     @usableFromInline
     private(set) var continuations: [UUID: Continuation] = [:]
@@ -33,6 +36,7 @@ public actor AsyncEvent: AsyncObject {
         _ continuation: Continuation,
         withKey key: UUID
     ) {
+        guard !continuation.resumed else { return }
         guard !signalled else { continuation.resume(); return }
         continuations[key] = continuation
     }
@@ -43,8 +47,7 @@ public actor AsyncEvent: AsyncObject {
     /// - Parameter key: The key in the map.
     @inlinable
     func _removeContinuation(withKey key: UUID) {
-        let continuation = continuations.removeValue(forKey: key)
-        continuation?.cancel()
+        continuations.removeValue(forKey: key)
     }
 
     /// Suspends the current task, then calls the given closure with a throwing continuation for the current task.
@@ -58,15 +61,13 @@ public actor AsyncEvent: AsyncObject {
     @inlinable
     nonisolated func _withPromisedContinuation() async throws {
         let key = UUID()
-        try await withTaskCancellationHandler { [weak self] in
+        try await Continuation.withCancellation(synchronizedWith: locker) {
             Task { [weak self] in
                 await self?._removeContinuation(withKey: key)
             }
-        } operation: { () -> Continuation.Success in
-            try await Continuation.with { continuation in
-                Task { [weak self] in
-                    await self?._addContinuation(continuation, withKey: key)
-                }
+        } operation: { continuation in
+            Task { [weak self] in
+                await self?._addContinuation(continuation, withKey: key)
             }
         }
     }

Sources/AsyncObjects/AsyncEvent.swift

@@ -16,14 +16,17 @@ import OrderedCollections
 public actor AsyncSemaphore: AsyncObject {
     /// The suspended tasks continuation type.
     @usableFromInline
-    typealias Continuation = GlobalContinuation<Void, Error>
+    typealias Continuation = SafeContinuation<GlobalContinuation<Void, Error>>
+    /// The platform dependent lock used to synchronize continuations tracking.
+    @usableFromInline
+    let locker: Locker = .init()
     /// The continuations stored with an associated key for all the suspended task that are waiting for access to resource.
     @usableFromInline
     private(set) var continuations: OrderedDictionary<UUID, Continuation> = [:]
     /// Pool size for concurrent resource access.
     /// Has value provided during initialization incremented by one.
     @usableFromInline
-    private(set) var limit: UInt
+    let limit: UInt
     /// Current count of semaphore.
     /// Can have maximum value up to `limit`.
     @usableFromInline
@@ -41,6 +44,8 @@ public actor AsyncSemaphore: AsyncObject {
         _ continuation: Continuation,
         withKey key: UUID
     ) {
+        count -= 1
+        guard !continuation.resumed else { return }
         guard count <= 0 else { continuation.resume(); return }
         continuations[key] = continuation
     }
@@ -51,8 +56,7 @@ public actor AsyncSemaphore: AsyncObject {
     /// - Parameter key: The key in the map.
     @inlinable
     func _removeContinuation(withKey key: UUID) {
-        let continuation = continuations.removeValue(forKey: key)
-        continuation?.cancel()
+        continuations.removeValue(forKey: key)
         _incrementCount()
     }
 
@@ -74,15 +78,13 @@ public actor AsyncSemaphore: AsyncObject {
     @inlinable
     nonisolated func _withPromisedContinuation() async throws {
         let key = UUID()
-        try await withTaskCancellationHandler { [weak self] in
+        try await Continuation.withCancellation(synchronizedWith: locker) {
             Task { [weak self] in
                 await self?._removeContinuation(withKey: key)
             }
-        } operation: { () -> Continuation.Success in
-            try await Continuation.with { continuation in
-                Task { [weak self] in
-                    await self?._addContinuation(continuation, withKey: key)
-                }
+        } operation: { continuation in
+            Task { [weak self] in
+                await self?._addContinuation(continuation, withKey: key)
             }
         }
     }
@@ -123,8 +125,7 @@ public actor AsyncSemaphore: AsyncObject {
     /// current task is suspended until a signal occurs.
     @Sendable
     public func wait() async {
-        count -= 1
-        guard count <= 0 else { return }
+        guard count <= 1 else { count -= 1; return }
         try? await _withPromisedContinuation()
     }
 }