Fix forEach(cancelPrevious:) leaving its in-flight body running after cancel

CHANGELOG.md

@@ -12,6 +12,12 @@ All notable changes are documented here. The format follows [Keep a Changelog](h
   - Once inside the context lock, the memoize's update closure now checks whether the model has been destructed, and if so serves its last produced value instead of running the producer. The result would have been discarded anyway, because teardown clears the cache.
   - `MemoizeDuringTeardownTests` reproduces the downstream shape. Without the fix, 200 iterations ran the producer on a torn-down model 293 times and recorded 186 unattributed issues; with it, zero and zero.
 
+- **Cancelling a `forEach(cancelPrevious: true)` subscription could leave its in-flight body running.** A child task is spawned first and keyed into its parent's cancellation context afterwards (`inheritCancellationContext()`), so the parent's context could be cancelled in between. Its cancel had then already run and never reached the child. For `forEach(cancelPrevious:)` that meant `subscription.cancel()`, landing right after a body started, cancelled the outer loop but left the body running to completion. The same hole applied to any `task { … }.inheritCancellationContext()` inside a task whose context was just cancelled. `InheritCancellationContextTests.testForEachCancelPreviousInheritsContext` hit it intermittently in CI and 7 times in 10 when run alone locally.
+  - One-shot contexts — the anonymous `cancellationContext { }`, which every `node.task` is wrapped in — now use a `ContextToken` key that stays cancelled once cancelled.
+  - Closing the token and taking the snapshot of its registrations happen in one critical section under the store's lock. Keying a cancellable into a closed context (or registering one inside it) takes the same lock and cancels it at once. So a registration racing the cancel either lands in the snapshot or sees the context closed, never neither.
+  - User keys (`cancellationContext(for: key)`, `cancel(for:)`, `cancelInFlight`) stay reusable: they are not one-shot.
+  - Validation: with the race window artificially widened, the test failed every run before the fix and passes every run after it. Alone it went from 3/10 passing to 20/20. `CancelledContextKeyingTests` pins the behaviour deterministically; its public-API case fails without the fix.
+
 ---
 
 ## [1.1.0] — Signals (`onSignal` / `signal` / `onTeardown`) + `TestPredicate` `==` no longer leaks into app code

Sources/SwiftModel/Internal/Cancellations.swift

@@ -20,9 +20,17 @@ final class Cancellations: @unchecked Sendable {
             cancelAll(for: key)
         }
 
-        lock {
-            guard registered[c.id] != nil else { return }
-            keyed[.init(key: key), default: []].append(c.id)
+        let cancelNow: Bool = lock {
+            guard registered[c.id] != nil else { return false }
+            let key = CancellableKey(key: key)
+            // Keying into a context that has already been cancelled: its cancel has run
+            // and will never reach this one — cancel it now instead (see `ContextToken`).
+            if key.contextToken?.isCancelled == true { return true }
+            keyed[key, default: []].append(c.id)
+            return false
+        }
+        if cancelNow {
+            cancel(c)
         }
     }
 
@@ -42,14 +50,27 @@ final class Cancellations: @unchecked Sendable {
     }
 
     func register(_ c: InternalCancellable) {
+        let contexts = AnyCancellable.contexts
+        var inCancelledContext = false
         let shouldImmediatelyCancel: Bool = lock {
             if _sealed { return true }
+            // Registered inside a context that has already been cancelled (e.g. a task
+            // started by a task whose context was just cancelled): cancel at once, or
+            // nothing ever would (see `ContextToken`).
+            if contexts.contains(where: { $0.contextToken?.isCancelled == true }) {
+                inCancelledContext = true
+                return false
+            }
             registered[c.id] = c
-            for key in AnyCancellable.contexts {
+            for key in contexts {
                 keyed[key, default: []].append(c.id)
             }
             return false
         }
+        if inCancelledContext {
+            c.onCancel()
+            return
+        }
         if shouldImmediatelyCancel {
             // Not while holding `lock` — `reportIssue` must never run inside a
             // context/cancellations critical section (see the AB-BA discussion
@@ -93,8 +114,13 @@ final class Cancellations: @unchecked Sendable {
     }
 
     func cancelAll(for key: some Hashable&Sendable) {
+        let key = CancellableKey(key: key)
         let cancellables = lock {
-            (keyed.removeValue(forKey: .init(key: key)) ?? []).compactMap { id in
+            // Close a one-shot context in the SAME critical section that takes the
+            // snapshot, so a concurrent registration under it either lands in this
+            // snapshot or sees it closed — never neither.
+            key.contextToken?.markCancelled()
+            return (keyed.removeValue(forKey: key) ?? []).compactMap { id in
                 registered.removeValue(forKey: id)
             }
         }
@@ -165,9 +191,34 @@ enum ContextCancellationKey {
     case onActivate
 }
 
+/// The key of a one-shot cancellation context: the anonymous `cancellationContext { }`,
+/// and the context `node.task` wraps every task in. Unlike a user key, such a context is
+/// never reused, so once cancelled it stays cancelled — and a cancellable registered
+/// under it AFTER its cancel has run must be cancelled immediately rather than left
+/// running with nothing left to cancel it.
+///
+/// That window is real: a child task is spawned first and keyed into its parent's
+/// context afterwards (`inheritCancellationContext()`), so the parent can be cancelled
+/// in between — `forEach(cancelPrevious:)` then left its in-flight body running after
+/// the subscription was cancelled. `isCancelled` is set and read under the
+/// `Cancellations` lock that also guards `keyed`, which makes "register" and "cancel"
+/// linearizable per store.
+final class ContextToken: Hashable, @unchecked Sendable {
+    private let lock = NSLock()
+    private var _isCancelled = false
+
+    var isCancelled: Bool { lock { _isCancelled } }
+    func markCancelled() { lock { _isCancelled = true } }
+
+    static func == (lhs: ContextToken, rhs: ContextToken) -> Bool { lhs === rhs }
+    func hash(into hasher: inout Hasher) { hasher.combine(ObjectIdentifier(self)) }
+}
+
 struct CancellableKey: Hashable, @unchecked Sendable {
     var key: AnyHashable
 
+    var contextToken: ContextToken? { key.base as? ContextToken }
+
     init<Key: Hashable&Sendable>(key: Key) {
         if let key = key as? CancellableKey {
             self.key = key.key

Sources/SwiftModel/ModelNode+Cancellation.swift

@@ -96,7 +96,7 @@ public extension ModelNode {
     func cancellationContext(perform: () throws -> Void) rethrows -> Cancellable {
         guard let cancellations = enforcedContext()?.cancellations else { return EmptyCancellable() }
 
-        let key = UUID()
+        let key = ContextToken()   // one-shot: stays cancelled once cancelled
         try AnyCancellable.$contexts.withValue(AnyCancellable.contexts + [CancellableKey(key: key)]) {
             try perform()
         }
@@ -112,7 +112,7 @@ public extension ModelNode {
     func cancellationContext(perform: () async throws -> Void) async rethrows -> Cancellable {
         guard let cancellations = enforcedContext()?.cancellations else { return EmptyCancellable() }
 
-        let key = UUID()
+        let key = ContextToken()   // one-shot: stays cancelled once cancelled
         let cancellable = AnyCancellable(cancellations: cancellations) { [weak cancellations] in
             cancellations?.cancelAll(for: key)
         }

Tests/SwiftModelTests/InheritCancellationContextTests.swift

@@ -1,6 +1,7 @@
 import Testing
 import AsyncAlgorithms
 import Foundation
+import ConcurrencyExtras
 @testable import SwiftModel
 
 struct InheritCancellationContextTests {
@@ -149,6 +150,82 @@ struct InheritCancellationContextTests {
     }
 }
 
+// MARK: - Keying into an already-cancelled context
+
+/// A child is spawned first and keyed into its parent's context afterwards, so the
+/// parent's context can be cancelled in between. Its cancel has then already run; the
+/// late child must be cancelled at once, not left running with nothing to cancel it.
+/// (`forEach(cancelPrevious:)` hit exactly this: cancelling the subscription while a body
+/// had just started left the body running — `testForEachCancelPreviousInheritsContext`
+/// timed out when `subscription.cancel()` landed in that window.)
+struct CancelledContextKeyingTests {
+    @Test func inheritingAnAlreadyCancelledContextCancelsImmediately() async throws {
+        let model = InheritModel().withAnchor().testNode
+        let started = AsyncChannel<Void>()
+        let childCancelled = LockIsolated(false)
+        let parentDone = LockIsolated(false)
+
+        let parent = model.task {
+            await started.send(())
+            while !Task.isCancelled { await Task.yield() }   // our context is now cancelled
+            model.onCancel { childCancelled.setValue(true) }.inheritCancellationContext()
+            parentDone.setValue(true)
+        }
+
+        var it = started.makeAsyncIterator()
+        await it.next()
+        parent.cancel()
+        try await waitUntil(parentDone.value)
+        #expect(childCancelled.value)
+    }
+
+    @Test func registeringUnderACancelledOneShotContextCancelsImmediately() {
+        let cancellations = Cancellations()
+        let context = ContextToken()
+        let inside = Recording(id: cancellations.nextId)
+        AnyCancellable.$contexts.withValue([CancellableKey(key: context)]) {
+            cancellations.register(inside)
+        }
+
+        cancellations.cancelAll(for: context)
+        #expect(inside.cancelled.value)
+
+        // Keyed in after the cancel (`inheritCancellationContext` / `cancel(for:)`).
+        let keyedLate = Recording(id: cancellations.nextId)
+        cancellations.register(keyedLate)
+        cancellations.cancel(keyedLate, for: context, cancelInFlight: false)
+        #expect(keyedLate.cancelled.value)
+
+        // Registered inside the cancelled context.
+        let registeredLate = Recording(id: cancellations.nextId)
+        AnyCancellable.$contexts.withValue([CancellableKey(key: context)]) {
+            cancellations.register(registeredLate)
+        }
+        #expect(registeredLate.cancelled.value)
+    }
+
+    @Test func userKeysStayReusableAfterCancelAll() {
+        let cancellations = Cancellations()
+        let first = Recording(id: cancellations.nextId)
+        cancellations.register(first)
+        cancellations.cancel(first, for: "reload", cancelInFlight: false)
+        cancellations.cancelAll(for: "reload")
+        #expect(first.cancelled.value)
+
+        let second = Recording(id: cancellations.nextId)
+        cancellations.register(second)
+        cancellations.cancel(second, for: "reload", cancelInFlight: false)
+        #expect(!second.cancelled.value)   // a user key is not one-shot
+    }
+}
+
+private final class Recording: InternalCancellable, Sendable {
+    let id: Int
+    let cancelled = LockIsolated(false)
+    init(id: Int) { self.id = id }
+    func onCancel() { cancelled.setValue(true) }
+}
+
 // MARK: - Supporting types
 
 enum InheritKey { case outer }