Fix switchToLatest cancellation hangs

Sources/Operators/AsyncSwitchToLatestSequence.swift

@@ -85,9 +85,10 @@ where Base.Element: AsyncSequence, Base: Sendable, Base.Element.Element: Sendabl
     struct State {
       var childTask: Task<ChildValue?, Never>?
       var base: BaseState
+      var isCancelled: Bool
 
       static var initial: State {
-        State(childTask: nil, base: .notStarted)
+        State(childTask: nil, base: .notStarted, isCancelled: false)
       }
     }
 
@@ -97,7 +98,7 @@ where Base.Element: AsyncSequence, Base: Sendable, Base.Element.Element: Sendabl
     }
 
     enum NextDecision {
-      case immediatelyResume(Task<ChildValue?, Never>)
+      case immediatelyResume(Task<ChildValue?, Never>?)
       case suspend
     }
 
@@ -140,6 +141,8 @@ where Base.Element: AsyncSequence, Base: Sendable, Base.Element.Element: Sendabl
           for try await child in base {
             let childIterator = child.makeAsyncIterator()
             let decision = state.withCriticalRegion { state -> BaseDecision in
+              guard !state.isCancelled else { return .cancelPreviousChildTask(nil) }
+
               switch state.base {
                 case .waitingForChildIterator(let continuation):
                   state.base = .processingChildIterator(.success(childIterator))
@@ -225,6 +228,8 @@ where Base.Element: AsyncSequence, Base: Sendable, Base.Element.Element: Sendabl
         while true {
           let childTask = await withUnsafeContinuation { [state] (continuation: UnsafeContinuation<Task<ChildValue?, Never>?, Never>) in
             let decision = state.withCriticalRegion { state -> NextDecision in
+              guard !state.isCancelled else { return .immediatelyResume(nil) }
+
               switch state.base {
                 case .newChildIteratorAvailable(let childIterator):
                   state.base = .processingChildIterator(childIterator)
@@ -260,7 +265,11 @@ where Base.Element: AsyncSequence, Base: Sendable, Base.Element.Element: Sendabl
           let value = await childTask?.value
 
           let decision = state.withCriticalRegion { state -> PostElementDecision in
-            if state.base.isNewAvailableChildIterator {
+            state.childTask = nil
+
+            if state.isCancelled {
+              return .returnFinish
+            } else if state.base.isNewAvailableChildIterator {
               return .pass
             } else {
               switch value {
@@ -299,10 +308,23 @@ where Base.Element: AsyncSequence, Base: Sendable, Base.Element.Element: Sendabl
           }
         }
       } onCancel: { [baseTask, state] in
-        baseTask?.cancel()
-        state.withCriticalRegion {
-          $0.childTask?.cancel()
+        let cancellation: (
+          continuation: UnsafeContinuation<Task<ChildValue?, Never>?, Never>?,
+          childTask: Task<ChildValue?, Never>?
+        ) = state.withCriticalRegion { state in
+          state.isCancelled = true
+
+          if case .waitingForChildIterator(let continuation) = state.base {
+            state.base = .finished(nil)
+            return (continuation, state.childTask)
+          } else {
+            return (nil, state.childTask)
+          }
         }
+
+        baseTask?.cancel()
+        cancellation.childTask?.cancel()
+        cancellation.continuation?.resume(returning: nil)
       }
     }
   }

Tests/Operators/AsyncSwitchToLatestSequenceTests.swift

@@ -57,6 +57,58 @@ private struct LongAsyncSequence<Element>: AsyncSequence, AsyncIteratorProtocol
   }
 }
 
+private struct NonCooperativeAsyncSequence<Element: Sendable>: AsyncSequence, Sendable {
+  let onSuspend: @Sendable () -> Void
+
+  func makeAsyncIterator() -> Iterator {
+    Iterator(onSuspend: self.onSuspend)
+  }
+
+  struct Iterator: AsyncIteratorProtocol, Sendable {
+    let onSuspend: @Sendable () -> Void
+
+    mutating func next() async -> Element? {
+      await withUnsafeContinuation { (_: UnsafeContinuation<Element?, Never>) in
+        self.onSuspend()
+      }
+    }
+  }
+}
+
+private struct IteratorLifetimeSequence<Element: Sendable>: AsyncSequence, Sendable {
+  let element: Element
+  let onIteratorCreated: @Sendable () -> Void
+  let onIteratorReleased: @Sendable () -> Void
+
+  func makeAsyncIterator() -> Iterator {
+    self.onIteratorCreated()
+    return Iterator(element: self.element, onReleased: self.onIteratorReleased)
+  }
+
+  final class Iterator: AsyncIteratorProtocol, Sendable {
+    let element: Element
+    let onReleased: @Sendable () -> Void
+    let hasEmitted = ManagedCriticalState(false)
+
+    init(element: Element, onReleased: @escaping @Sendable () -> Void) {
+      self.element = element
+      self.onReleased = onReleased
+    }
+
+    deinit {
+      self.onReleased()
+    }
+
+    func next() async -> Element? {
+      self.hasEmitted.withCriticalRegion { hasEmitted in
+        guard !hasEmitted else { return nil }
+        hasEmitted = true
+        return self.element
+      }
+    }
+  }
+}
+
 final class AsyncSwitchToLatestSequenceTests: XCTestCase {
   func testSwitchToLatest_switches_to_latest_asyncSequence_and_cancels_previous_ones() async throws {
     var asyncSequence1IsCancelled = false
@@ -183,4 +235,98 @@ final class AsyncSwitchToLatestSequenceTests: XCTestCase {
 
     wait(for: [taskHasFinishedExpectation], timeout: 5) // task has been cancelled and has finished
   }
+
+  func testSwitchToLatest_finishes_when_awaiting_an_unfinished_latest_sequence_and_task_is_cancelled() async {
+    let receivedFirstValue = expectation(description: "The first sequence emitted")
+    let receivedSecondValue = expectation(description: "The second sequence emitted")
+    let receivedLatestValue = expectation(description: "The latest sequence emitted")
+    let collectionFinished = expectation(description: "The collection task finished")
+
+    var outerContinuation: AsyncStream<AsyncBufferedChannel<Int>>.Continuation!
+    let outer = AsyncStream<AsyncBufferedChannel<Int>> { continuation in
+      outerContinuation = continuation
+    }
+
+    let collectionTask = Task {
+      for await element in outer.switchToLatest() {
+        switch element {
+          case 1: receivedFirstValue.fulfill()
+          case 2: receivedSecondValue.fulfill()
+          case 4: receivedLatestValue.fulfill()
+          default: XCTFail("Received unexpected element: \(element)")
+        }
+      }
+      collectionFinished.fulfill()
+    }
+
+    let first = AsyncBufferedChannel<Int>()
+    first.send(1)
+    outerContinuation.yield(first)
+    await fulfillment(of: [receivedFirstValue], timeout: 1)
+
+    let second = AsyncBufferedChannel<Int>()
+    second.send(2)
+    outerContinuation.yield(second)
+    await fulfillment(of: [receivedSecondValue], timeout: 1)
+
+    let latest = AsyncBufferedChannel<Int>()
+    latest.send(4)
+    outerContinuation.yield(latest)
+    await fulfillment(of: [receivedLatestValue], timeout: 1)
+
+    collectionTask.cancel()
+
+    await fulfillment(of: [collectionFinished], timeout: 1)
+  }
+
+  func testSwitchToLatest_finishes_when_awaiting_a_non_cooperative_outer_sequence_and_task_is_cancelled() async {
+    let outerSequenceIsSuspended = expectation(description: "The outer sequence is suspended")
+    let collectionFinished = expectation(description: "The collection task finished")
+    let outer = NonCooperativeAsyncSequence<AsyncBufferedChannel<Int>> {
+      outerSequenceIsSuspended.fulfill()
+    }
+
+    let collectionTask = Task {
+      for await _ in outer.switchToLatest() {}
+      collectionFinished.fulfill()
+    }
+
+    await fulfillment(of: [outerSequenceIsSuspended], timeout: 1)
+    await Task.yield()
+
+    collectionTask.cancel()
+
+    await fulfillment(of: [collectionFinished], timeout: 1)
+  }
+
+  func testSwitchToLatest_releases_previous_iterator_when_new_sequence_arrives_between_downstream_calls() async throws {
+    let firstIteratorCreated = expectation(description: "The first iterator was created")
+    let firstIteratorReleased = expectation(description: "The first iterator was released")
+    let secondIteratorCreated = expectation(description: "The second iterator was created")
+    let (outer, continuation) = AsyncStream<IteratorLifetimeSequence<Int>>.makeStream()
+    var iterator = outer.switchToLatest().makeAsyncIterator()
+
+    continuation.yield(
+      IteratorLifetimeSequence(
+        element: 1,
+        onIteratorCreated: { firstIteratorCreated.fulfill() },
+        onIteratorReleased: { firstIteratorReleased.fulfill() }
+      )
+    )
+
+    let firstValue = await iterator.next()
+    XCTAssertEqual(firstValue, 1)
+    await fulfillment(of: [firstIteratorCreated], timeout: 1)
+
+    continuation.yield(
+      IteratorLifetimeSequence(
+        element: 2,
+        onIteratorCreated: { secondIteratorCreated.fulfill() },
+        onIteratorReleased: {}
+      )
+    )
+
+    await fulfillment(of: [secondIteratorCreated, firstIteratorReleased], timeout: 1)
+    continuation.finish()
+  }
 }