Fix Store leak when async effect is in flight (#2643)

* Fix Store leak when async effect is in flight A [Slack conversation](https://pointfreecommunity.slack.com/archives/C04KQQ7NXHV/p1702049135565039) pointed to a potential leak of the `Store` when it was executing an effect. This PR allows cancellation to properly propagate during deinitialization, and ensures that async effects don't strongly capture the store. * wip * feedback * wip
pointfreeco · Dec 12, 2023 · b86012c · b86012c
1 parent c6b0a6c
commit b86012c
Show file tree

Hide file tree

Showing 6 changed files with 385 additions and 297 deletions.
diff --git a/Sources/ComposableArchitecture/Store.swift b/Sources/ComposableArchitecture/Store.swift
@@ -485,6 +485,10 @@ public final class Store<State, Action> {
     for id in self.children.keys {
       self.invalidateChild(id: id)
     }
+    self.effectCancellables.values.forEach { cancellable in
+      cancellable.cancel()
+    }
+    self.effectCancellables.removeAll()
   }
 
   fileprivate func invalidateChild(id: ScopeID<State, Action>) {
@@ -526,14 +530,14 @@ public final class Store<State, Action> {
       defer { index += 1 }
       let action = self.bufferedActions[index]
       let effect = self.reducer.reduce(into: &currentState, action: action)
+      let uuid = UUID()
 
       switch effect.operation {
       case .none:
         break
       case let .publisher(publisher):
         var didComplete = false
         let boxedTask = Box<Task<Void, Never>?>(wrappedValue: nil)
-        let uuid = UUID()
         let effectCancellable = withEscapedDependencies { continuation in
           publisher
             .handleEvents(
@@ -571,45 +575,48 @@ public final class Store<State, Action> {
         }
       case let .run(priority, operation):
         withEscapedDependencies { continuation in
-          tasks.wrappedValue.append(
-            Task(priority: priority) { @MainActor in
-              #if DEBUG
-                let isCompleted = LockIsolated(false)
-                defer { isCompleted.setValue(true) }
-              #endif
-              await operation(
-                Send { effectAction in
-                  #if DEBUG
-                    if isCompleted.value {
-                      runtimeWarn(
-                        """
-                        An action was sent from a completed effect:
-
-                          Action:
-                            \(debugCaseOutput(effectAction))
-
-                          Effect returned from:
-                            \(debugCaseOutput(action))
-
-                        Avoid sending actions using the 'send' argument from 'Effect.run' after \
-                        the effect has completed. This can happen if you escape the 'send' \
-                        argument in an unstructured context.
-
-                        To fix this, make sure that your 'run' closure does not return until \
-                        you're done calling 'send'.
-                        """
-                      )
-                    }
-                  #endif
-                  if let task = continuation.yield({
-                    self.send(effectAction, originatingFrom: action)
-                  }) {
-                    tasks.wrappedValue.append(task)
+          let task = Task(priority: priority) { @MainActor [weak self] in
+            #if DEBUG
+              let isCompleted = LockIsolated(false)
+              defer { isCompleted.setValue(true) }
+            #endif
+            await operation(
+              Send { effectAction in
+                #if DEBUG
+                  if isCompleted.value {
+                    runtimeWarn(
+                      """
+                      An action was sent from a completed effect:
+
+                        Action:
+                          \(debugCaseOutput(effectAction))
+
+                        Effect returned from:
+                          \(debugCaseOutput(action))
+
+                      Avoid sending actions using the 'send' argument from 'Effect.run' after \
+                      the effect has completed. This can happen if you escape the 'send' \
+                      argument in an unstructured context.
+
+                      To fix this, make sure that your 'run' closure does not return until \
+                      you're done calling 'send'.
+                      """
+                    )
                   }
+                #endif
+                if let task = continuation.yield({
+                  self?.send(effectAction, originatingFrom: action)
+                }) {
+                  tasks.wrappedValue.append(task)
                 }
-              )
-            }
-          )
+              }
+            )
+            self?.effectCancellables[uuid] = nil
+          }
+          tasks.wrappedValue.append(task)
+          self.effectCancellables[uuid] = AnyCancellable {
+            task.cancel()
+          }
         }
       }
     }