Warn if send is called after .run completes (#1900)

* Warn if send called after effect completes

When an Effect.run is converted into a publisher, we already end up discarding `send`s sent after returning. Similarly, `*Task.finish()` will fail to track send()s after other in-flight effects complete. This is effectively UB so let's warn about it.

* Fix tests

* we don't need to print TestAction yet

* Tweak messaging and tests.

* fix tests

* wip

* fix tests part 2

* fix tests part 3

* Wrap testCancellation test in _withMainSerialExecutor

* Update Publisher.swift

---------

Co-authored-by: Brandon Williams <[email protected]>
Co-authored-by: Stephen Celis <[email protected]>
Co-authored-by: Stephen Celis <[email protected]>

Author: 3 people

Sources/ComposableArchitecture/Effects/Publisher.swift

@@ -23,7 +23,32 @@ extension EffectPublisher: Publisher {
       return .create { subscriber in
         let task = Task(priority: priority) { @MainActor in
           defer { subscriber.send(completion: .finished) }
-          let send = Send { subscriber.send($0) }
+          #if DEBUG
+            var isCompleted = false
+            defer { isCompleted = true }
+          #endif
+          let send = Send {
+            #if DEBUG
+              if isCompleted {
+                runtimeWarn(
+                  """
+                  An action was sent from a completed effect:
+
+                    Action:
+                      \(debugCaseOutput($0))
+
+                  Avoid sending actions using the 'send' argument from 'EffectTask.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
+            subscriber.send($0)
+          }
           await operation(send)
         }
         return AnyCancellable {

Sources/ComposableArchitecture/Store.swift

@@ -417,9 +417,35 @@ public final class Store<State, Action> {
         }
       case let .run(priority, operation):
         tasks.wrappedValue.append(
-          Task(priority: priority) {
+          Task(priority: priority) { @MainActor in
+            #if DEBUG
+              var isCompleted = false
+              defer { isCompleted = true }
+            #endif
             await operation(
               EffectTask<Action>.Send {
+                #if DEBUG
+                  if isCompleted {
+                    runtimeWarn(
+                      """
+                      An action was sent from a completed effect:
+
+                        Action:
+                          \(debugCaseOutput($0))
+
+                        Effect returned from:
+                          \(debugCaseOutput(action))
+
+                      Avoid sending actions using the 'send' argument from 'EffectTask.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 = self.send($0, originatingFrom: action) {
                   tasks.wrappedValue.append(task)
                 }

Tests/ComposableArchitectureTests/ComposableArchitectureTests.swift

@@ -112,46 +112,48 @@ final class ComposableArchitectureTests: XCTestCase {
   }
 
   func testCancellation() async {
-    let mainQueue = DispatchQueue.test
-
-    enum Action: Equatable {
-      case cancel
-      case incr
-      case response(Int)
-    }
-
-    let reducer = Reduce<Int, Action> { state, action in
-      enum CancelID {}
-
-      switch action {
-      case .cancel:
-        return .cancel(id: CancelID.self)
-
-      case .incr:
-        state += 1
-        return .task { [state] in
-          try await mainQueue.sleep(for: .seconds(1))
-          return .response(state * state)
+    await _withMainSerialExecutor {
+      let mainQueue = DispatchQueue.test
+      
+      enum Action: Equatable {
+        case cancel
+        case incr
+        case response(Int)
+      }
+      
+      let reducer = Reduce<Int, Action> { state, action in
+        enum CancelID {}
+        
+        switch action {
+        case .cancel:
+          return .cancel(id: CancelID.self)
+          
+        case .incr:
+          state += 1
+          return .task { [state] in
+            try await mainQueue.sleep(for: .seconds(1))
+            return .response(state * state)
+          }
+          .cancellable(id: CancelID.self)
+          
+        case let .response(value):
+          state = value
+          return .none
         }
-        .cancellable(id: CancelID.self)
-
-      case let .response(value):
-        state = value
-        return .none
       }
+      
+      let store = TestStore(
+        initialState: 0,
+        reducer: reducer
+      )
+      
+      await store.send(.incr) { $0 = 1 }
+      await mainQueue.advance(by: .seconds(1))
+      await store.receive(.response(1))
+      
+      await store.send(.incr) { $0 = 2 }
+      await store.send(.cancel)
+      await store.finish()
     }
-
-    let store = TestStore(
-      initialState: 0,
-      reducer: reducer
-    )
-
-    await store.send(.incr) { $0 = 1 }
-    await mainQueue.advance(by: .seconds(1))
-    await store.receive(.response(1))
-
-    await store.send(.incr) { $0 = 2 }
-    await store.send(.cancel)
-    await store.finish()
   }
 }

Tests/ComposableArchitectureTests/EffectRunTests.swift

@@ -117,4 +117,94 @@ final class EffectRunTests: XCTestCase {
     let store = TestStore(initialState: State(), reducer: reducer)
     await store.send(.tapped).finish()
   }
+
+  #if DEBUG
+    func testRunEscapeFailure() async throws {
+      XCTExpectFailure {
+        $0.compactDescription == """
+          An action was sent from a completed effect:
+
+            Action:
+              EffectRunTests.Action.response
+
+            Effect returned from:
+              EffectRunTests.Action.tap
+
+          Avoid sending actions using the 'send' argument from 'EffectTask.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'.
+          """
+      }
+
+      enum Action { case tap, response }
+
+      let queue = DispatchQueue.test
+
+      let store = Store(
+        initialState: 0,
+        reducer: Reduce<Int, Action> { _, action in
+          switch action {
+          case .tap:
+            return .run { send in
+              Task(priority: .userInitiated) {
+                try await queue.sleep(for: .seconds(1))
+                await send(.response)
+              }
+            }
+          case .response:
+            return .none
+          }
+        }
+      )
+
+      let viewStore = ViewStore(store, observe: { $0 })
+      await viewStore.send(.tap).finish()
+      await queue.advance(by: .seconds(1))
+    }
+
+    func testRunEscapeFailurePublisher() async throws {
+      XCTExpectFailure {
+        $0.compactDescription == """
+          An action was sent from a completed effect:
+
+            Action:
+              EffectRunTests.Action.response
+
+          Avoid sending actions using the 'send' argument from 'EffectTask.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'.
+          """
+      }
+
+      enum Action { case tap, response }
+
+      let queue = DispatchQueue.test
+
+      let store = Store(
+        initialState: 0,
+        reducer: Reduce<Int, Action> { _, action in
+          switch action {
+          case .tap:
+            return .run { send in
+              Task(priority: .userInitiated) {
+                try await queue.sleep(for: .seconds(1))
+                await send(.response)
+              }
+            }
+            .eraseToEffect()
+          case .response:
+            return .none
+          }
+        }
+      )
+
+      let viewStore = ViewStore(store, observe: { $0 })
+      await viewStore.send(.tap).finish()
+      await queue.advance(by: .seconds(1))
+    }
+  #endif
 }