Cancellation cleanup (pointfreeco#1977)

* Clean up effect cancellation.

* wip

* wip

* wip

* wip

* wip

* wip

* wip

* cancel id test

* nb

* benchmark

* wip

* wip

* wip

* wip

* wip

* clean up

* wip

Author: mbrandonw

ComposableArchitecture.xcworkspace/xcshareddata/swiftpm/Package.resolved

@@ -25,7 +25,7 @@ let package = Package(
     .package(url: "https://github.com/pointfreeco/swift-custom-dump", from: "0.9.1"),
     .package(url: "https://github.com/pointfreeco/swift-dependencies", from: "0.2.0"),
     .package(url: "https://github.com/pointfreeco/swift-identified-collections", from: "0.7.0"),
-    .package(url: "https://github.com/pointfreeco/swiftui-navigation", from: "0.7.0"),
+    .package(url: "https://github.com/pointfreeco/swiftui-navigation", from: "0.7.1"),
     .package(url: "https://github.com/pointfreeco/xctest-dynamic-overlay", from: "0.8.4"),
   ],
   targets: [

Package.swift

@@ -45,36 +45,30 @@ extension EffectPublisher {
             _cancellablesLock.lock()
             defer { _cancellablesLock.unlock() }
 
-            let id = _CancelToken(id: id)
             if cancelInFlight {
-              _cancellationCancellables[id]?.forEach { $0.cancel() }
+              _cancellationCancellables.cancel(id: id)
             }
 
             let cancellationSubject = PassthroughSubject<Void, Never>()
 
-            var cancellationCancellable: AnyCancellable!
-            cancellationCancellable = AnyCancellable {
+            var cancellable: AnyCancellable!
+            cancellable = AnyCancellable {
               _cancellablesLock.sync {
                 cancellationSubject.send(())
                 cancellationSubject.send(completion: .finished)
-                _cancellationCancellables[id]?.remove(cancellationCancellable)
-                if _cancellationCancellables[id]?.isEmpty == .some(true) {
-                  _cancellationCancellables[id] = nil
-                }
+                _cancellationCancellables.remove(cancellable, at: id)
               }
             }
 
             return publisher.prefix(untilOutputFrom: cancellationSubject)
               .handleEvents(
                 receiveSubscription: { _ in
-                  _ = _cancellablesLock.sync {
-                    _cancellationCancellables[id, default: []].insert(
-                      cancellationCancellable
-                    )
+                  _cancellablesLock.sync {
+                    _cancellationCancellables.insert(cancellable, at: id)
                   }
                 },
-                receiveCompletion: { _ in cancellationCancellable.cancel() },
-                receiveCancel: cancellationCancellable.cancel
+                receiveCompletion: { _ in cancellable.cancel() },
+                receiveCancel: cancellable.cancel
               )
           }
           .eraseToAnyPublisher()
@@ -113,7 +107,7 @@ extension EffectPublisher {
   public static func cancel(id: AnyHashable) -> Self {
     .fireAndForget {
       _cancellablesLock.sync {
-        _cancellationCancellables[.init(id: id)]?.forEach { $0.cancel() }
+        _cancellationCancellables.cancel(id: id)
       }
     }
   }
@@ -201,22 +195,18 @@ extension EffectPublisher {
     cancelInFlight: Bool = false,
     operation: @Sendable @escaping () async throws -> T
   ) async rethrows -> T {
-    let id = _CancelToken(id: id)
     let (cancellable, task) = _cancellablesLock.sync { () -> (AnyCancellable, Task<T, Error>) in
       if cancelInFlight {
-        _cancellationCancellables[id]?.forEach { $0.cancel() }
+        _cancellationCancellables.cancel(id: id)
       }
       let task = Task { try await operation() }
       let cancellable = AnyCancellable { task.cancel() }
-      _cancellationCancellables[id, default: []].insert(cancellable)
+      _cancellationCancellables.insert(cancellable, at: id)
       return (cancellable, task)
     }
     defer {
       _cancellablesLock.sync {
-        _cancellationCancellables[id]?.remove(cancellable)
-        if _cancellationCancellables[id]?.isEmpty == .some(true) {
-          _cancellationCancellables[id] = nil
-        }
+        _cancellationCancellables.remove(cancellable, at: id)
       }
     }
     do {
@@ -231,22 +221,18 @@ extension EffectPublisher {
     cancelInFlight: Bool = false,
     operation: @Sendable @escaping () async throws -> T
   ) async rethrows -> T {
-    let id = _CancelToken(id: id)
     let (cancellable, task) = _cancellablesLock.sync { () -> (AnyCancellable, Task<T, Error>) in
       if cancelInFlight {
-        _cancellationCancellables[id]?.forEach { $0.cancel() }
+        _cancellationCancellables.cancel(id: id)
       }
       let task = Task { try await operation() }
       let cancellable = AnyCancellable { task.cancel() }
-      _cancellationCancellables[id, default: []].insert(cancellable)
+      _cancellationCancellables.insert(cancellable, at: id)
       return (cancellable, task)
     }
     defer {
       _cancellablesLock.sync {
-        _cancellationCancellables[id]?.remove(cancellable)
-        if _cancellationCancellables[id]?.isEmpty == .some(true) {
-          _cancellationCancellables[id] = nil
-        }
+        _cancellationCancellables.remove(cancellable, at: id)
       }
     }
     do {
@@ -301,7 +287,9 @@ extension Task where Success == Never, Failure == Never {
   ///
   /// - Parameter id: An identifier.
   public static func cancel<ID: Hashable & Sendable>(id: ID) {
-    _cancellablesLock.sync { _cancellationCancellables[.init(id: id)]?.forEach { $0.cancel() } }
+    _cancellablesLock.sync {
+      _cancellationCancellables.cancel(id: id)
+    }
   }
 
   /// Cancel any currently in-flight operation with the given identifier.
@@ -315,7 +303,7 @@ extension Task where Success == Never, Failure == Never {
   }
 }
 
-@_spi(Internals) public struct _CancelToken: Hashable {
+@_spi(Internals) public struct _CancelID: Hashable {
   let id: AnyHashable
   let discriminator: ObjectIdentifier
 
@@ -325,8 +313,8 @@ extension Task where Success == Never, Failure == Never {
   }
 }
 
-@_spi(Internals) public var _cancellationCancellables: [_CancelToken: Set<AnyCancellable>] = [:]
-@_spi(Internals) public let _cancellablesLock = NSRecursiveLock()
+@_spi(Internals) public var _cancellationCancellables = CancellablesCollection()
+private let _cancellablesLock = NSRecursiveLock()
 
 @rethrows
 private protocol _ErrorMechanism {
@@ -346,3 +334,41 @@ extension _ErrorMechanism {
 }
 
 extension Result: _ErrorMechanism {}
+
+@_spi(Internals)
+public class CancellablesCollection {
+  var storage: [_CancelID: Set<AnyCancellable>] = [:]
+
+  func insert(
+    _ cancellable: AnyCancellable,
+    at id: AnyHashable
+  ) {
+    let cancelID = _CancelID(id: id)
+    self.storage[cancelID, default: []].insert(cancellable)
+  }
+
+  func remove(
+    _ cancellable: AnyCancellable,
+    at id: AnyHashable
+  ) {
+    let cancelID = _CancelID(id: id)
+    self.storage[cancelID]?.remove(cancellable)
+    if self.storage[cancelID]?.isEmpty == true {
+      self.storage[cancelID] = nil
+    }
+  }
+
+  func cancel(id: AnyHashable) {
+    let cancelID = _CancelID(id: id)
+    self.storage[cancelID]?.forEach { $0.cancel() }
+    self.storage[cancelID] = nil
+  }
+
+  public func exists(at id: AnyHashable) -> Bool {
+    return self.storage[_CancelID(id: id)] != nil
+  }
+
+  public var count: Int {
+    return self.storage.count
+  }
+}

Sources/ComposableArchitecture/Effects/Cancellation.swift

@@ -0,0 +1,101 @@
+import Benchmark
+import Combine
+@_spi(Internals) import ComposableArchitecture
+import Foundation
+
+let storeSuite = BenchmarkSuite(name: "Store") {
+  var store: StoreOf<Feature>!
+  let levels = 5
+
+  for level in 1...levels {
+    $0.benchmark("Nested send tap: \(level)") {
+      _ = store.send(tap(level: level))
+    } setUp: {
+      store = Store(
+        initialState: state(level: level),
+        reducer: Feature()
+      )
+    } tearDown: {
+      precondition(count(of: store.state.value, level: level) == 1)
+      precondition(_cancellationCancellables.count == 0)
+    }
+  }
+  for level in 1...levels {
+    $0.benchmark("Nested send none: \(level)") {
+      _ = store.send(none(level: level))
+    } setUp: {
+      store = Store(
+        initialState: state(level: level),
+        reducer: Feature()
+      )
+    } tearDown: {
+      precondition(count(of: store.state.value, level: level) == 0)
+      precondition(_cancellationCancellables.count == 0)
+    }
+  }
+}
+
+private struct Feature: ReducerProtocol {
+  struct State {
+    @Box var child: State?
+    var count = 0
+  }
+  enum Action {
+    indirect case child(Action)
+    case tap
+    case none
+  }
+  var body: some ReducerProtocolOf<Self> {
+    Reduce<State, Action> { state, action in
+      switch action {
+      case .child:
+        return .none
+      case .tap:
+        state.count = 1
+        return Empty(completeImmediately: true)
+          .eraseToEffect()
+          .cancellable(id: UUID())
+      case .none:
+        return .none
+      }
+    }
+    .ifLet(\.child, action: /Action.child) {
+      Feature()
+    }
+  }
+}
+
+@propertyWrapper
+private struct Box<Value> {
+  private var value: [Value]
+  var wrappedValue: Value? {
+    get { self.value.first }
+    set { self.value = newValue.map { [$0] } ?? [] }
+  }
+  init(wrappedValue: Value?) {
+    self.value = wrappedValue.map { [$0] } ?? []
+  }
+}
+
+private func state(level: Int) -> Feature.State {
+  Feature.State(
+    child: level == 0
+      ? nil
+      : state(level: level - 1)
+  )
+}
+private func tap(level: Int) -> Feature.Action {
+  level == 0
+    ? .tap
+    : Feature.Action.child(tap(level: level - 1))
+}
+private func none(level: Int) -> Feature.Action {
+  level == 0
+    ? .none
+    : Feature.Action.child(none(level: level - 1))
+}
+private func count(of state: Feature.State?, level: Int) -> Int? {
+  level == 0
+  ? state?.count
+  : count(of: state?.child, level: level - 1)
+}

Sources/swift-composable-architecture-benchmark/StoreSuite.swift

@@ -6,5 +6,6 @@ Benchmark.main([
   dependenciesSuite,
   effectSuite,
   storeScopeSuite,
+  storeSuite,
   viewStoreSuite,
 ])

Sources/swift-composable-architecture-benchmark/main.swift

@@ -4,7 +4,7 @@
   @testable import ComposableArchitecture
 
   @MainActor
-  final class BindingLocalTests: XCTestCase {
+  final class BindingLocalTests: BaseTCATestCase {
     public func testBindingLocalIsActive() {
       XCTAssertFalse(BindingLocal.isActive)
 

Tests/ComposableArchitectureTests/BindingLocalTests.swift

@@ -2,7 +2,7 @@ import ComposableArchitecture
 import XCTest
 
 @MainActor
-final class BindingTests: XCTestCase {
+final class BindingTests: BaseTCATestCase {
   #if swift(>=5.7)
     func testNestedBindingState() {
       struct BindingTest: ReducerProtocol {

Tests/ComposableArchitectureTests/BindingTests.swift

@@ -3,7 +3,7 @@ import ComposableArchitecture
 import XCTest
 
 @MainActor
-final class CompatibilityTests: XCTestCase {
+final class CompatibilityTests: BaseTCATestCase {
   var cancellables: Set<AnyCancellable> = []
 
   // Actions can be re-entrantly sent into the store if an action is sent that holds an object

Tests/ComposableArchitectureTests/CompatibilityTests.swift

@@ -4,7 +4,7 @@ import ComposableArchitecture
 import XCTest
 
 @MainActor
-final class ComposableArchitectureTests: XCTestCase {
+final class ComposableArchitectureTests: BaseTCATestCase {
   var cancellables: Set<AnyCancellable> = []
 
   func testScheduling() async {