Avoid simultaneous access crash on double-cancel

Sources/ComposableArchitecture/Effects/Cancellation.swift

@@ -29,40 +29,25 @@ extension Effect {
   public func cancellable(id: AnyHashable, cancelInFlight: Bool = false) -> Effect {
     return Deferred { () -> Publishers.HandleEvents<PassthroughSubject<Output, Failure>> in
       let subject = PassthroughSubject<Output, Failure>()
-      let uuid = UUID()
-
-      var isCleaningUp = false
 
       cancellablesLock.sync {
         if cancelInFlight {
-          cancellationCancellables[id]?.forEach { _, cancellable in cancellable.cancel() }
+          cancellationCancellables[id]?.forEach { cancellable in cancellable.cancel() }
           cancellationCancellables[id] = nil
         }
 
         let cancellable = self.subscribe(subject)
 
-        cancellationCancellables[id] = cancellationCancellables[id] ?? [:]
-        cancellationCancellables[id]?[uuid] = AnyCancellable {
+        AnyCancellable {
           cancellable.cancel()
-          if !isCleaningUp {
-            subject.send(completion: .finished)
-          }
-        }
-      }
-
-      func cleanup() {
-        isCleaningUp = true
-        cancellablesLock.sync {
-          cancellationCancellables[id]?[uuid] = nil
-          if cancellationCancellables[id]?.isEmpty == true {
-            cancellationCancellables[id] = nil
-          }
+          subject.send(completion: .finished)
         }
+        .store(in: &cancellationCancellables[id, default: []])
       }
 
       return subject.handleEvents(
-        receiveCompletion: { _ in cleanup() },
-        receiveCancel: cleanup
+        receiveCompletion: { _ in cancellablesLock.sync { cancellationCancellables[id] = nil } },
+        receiveCancel: { cancellablesLock.sync { cancellationCancellables[id] = nil } }
       )
     }
     .eraseToEffect()
@@ -76,12 +61,13 @@ extension Effect {
   public static func cancel(id: AnyHashable) -> Effect {
     .fireAndForget {
       cancellablesLock.sync {
-        cancellationCancellables[id]?.forEach { _, cancellable in cancellable.cancel() }
+        cancellationCancellables[id]?.forEach { cancellable in cancellable.cancel() }
         cancellationCancellables[id] = nil
       }
     }
   }
 }
 
-var cancellationCancellables: [AnyHashable: [UUID: AnyCancellable]] = [:]
+var cancellationCancellables: [AnyHashable: Set<AnyCancellable>] = [:]
 let cancellablesLock = NSRecursiveLock()
+var isCancelling: Set<AnyHashable> = []

Sources/ComposableArchitecture/Internal/Locking.swift

@@ -11,9 +11,9 @@ extension UnsafeMutablePointer where Pointee == os_unfair_lock_s {
 
 extension NSRecursiveLock {
   @inlinable
-  func sync(work: () -> Void) {
+  func sync<R>(work: () -> R) -> R {
     self.lock()
     defer { self.unlock() }
-    work()
+    return work()
   }
 }

Tests/ComposableArchitectureTests/EffectCancellationTests.swift

@@ -223,17 +223,40 @@ final class EffectCancellationTests: XCTestCase {
 
     let expectation = self.expectation(description: "wait")
     effect
-      .sink(receiveCompletion: { _ in expectation.fulfill() }, receiveValue: { _ in })
+      .sink(receiveCompletion: { _ in
+        print("completed")
+        expectation.fulfill()
+
+      }, receiveValue: { _ in })
       .store(in: &self.cancellables)
     self.wait(for: [expectation], timeout: 999)
 
     XCTAssertTrue(cancellationCancellables.isEmpty)
   }
+
+  func testNestedCancels() {
+    var effect = Empty<Void, Never>(completeImmediately: false)
+      .eraseToEffect()
+      .cancellable(id: 1)
+
+    for _ in 1 ... .random(in: 1...1_000) {
+      effect = effect.cancellable(id: 1)
+    }
+
+    effect
+      .sink(receiveValue: { _ in })
+      .store(in: &cancellables)
+
+    cancellables.removeAll()
+
+    XCTAssertEqual([:], cancellationCancellables)
+    XCTAssertEqual([], isCancelling)
+  }
 }
 
 func resetCancellables() {
   for (id, _) in cancellationCancellables {
-    cancellationCancellables[id] = [:]
+    cancellationCancellables[id] = []
   }
   cancellationCancellables = [:]
 }