Avoid simultaneous access crash on double-cancel (#137)

* Avoid simultaneous access crash on double-cancel

* fixes

* fix

Author: stephencelis

Sources/ComposableArchitecture/Effects/Cancellation.swift

@@ -29,40 +29,34 @@ 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)
-          }
+          subject.send(completion: .finished)
         }
+        .store(in: &cancellationCancellables[id, default: []])
       }
 
-      func cleanup() {
-        isCleaningUp = true
+      func cleanUp() {
         cancellablesLock.sync {
-          cancellationCancellables[id]?[uuid] = nil
-          if cancellationCancellables[id]?.isEmpty == true {
-            cancellationCancellables[id] = nil
-          }
+          guard !isCancelling.contains(id) else { return }
+          isCancelling.insert(id)
+          defer { isCancelling.remove(id) }
+          cancellationCancellables[id] = nil
         }
       }
 
       return subject.handleEvents(
-        receiveCompletion: { _ in cleanup() },
-        receiveCancel: cleanup
+        receiveCompletion: { _ in cleanUp() },
+        receiveCancel: cleanUp
       )
     }
     .eraseToEffect()
@@ -76,12 +70,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

@@ -229,11 +229,32 @@ final class EffectCancellationTests: XCTestCase {
 
     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] = [:]
+  cancellablesLock.sync {
+    for (id, _) in cancellationCancellables {
+      cancellationCancellables[id] = []
+    }
+    cancellationCancellables = [:]
   }
-  cancellationCancellables = [:]
 }