From 17b3530d3f902e9b98e3b48b47541cab3d3c833e Mon Sep 17 00:00:00 2001 From: pablobaldez Date: Sun, 23 Aug 2026 03:52:33 +0100 Subject: [PATCH 1/2] Fix switchToLatest cancellation hangs Persist iterator cancellation and atomically resume a waiting outer continuation so no replacement child task can outlive a cancelled consumer.\n\nAdd regressions for unfinished latest and non-cooperative outer sequences. --- .../AsyncSwitchToLatestSequence.swift | 32 ++++++-- .../AsyncSwitchToLatestSequenceTests.swift | 81 +++++++++++++++++++ 2 files changed, 107 insertions(+), 6 deletions(-) diff --git a/Sources/Operators/AsyncSwitchToLatestSequence.swift b/Sources/Operators/AsyncSwitchToLatestSequence.swift index 30147b8..595618f 100644 --- a/Sources/Operators/AsyncSwitchToLatestSequence.swift +++ b/Sources/Operators/AsyncSwitchToLatestSequence.swift @@ -85,9 +85,10 @@ where Base.Element: AsyncSequence, Base: Sendable, Base.Element.Element: Sendabl struct State { var childTask: Task? 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) + case immediatelyResume(Task?) 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?, 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,9 @@ 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 { + if state.isCancelled { + return .returnFinish + } else if state.base.isNewAvailableChildIterator { return .pass } else { switch value { @@ -299,10 +306,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?, Never>?, + childTask: Task? + ) = 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) } } } diff --git a/Tests/Operators/AsyncSwitchToLatestSequenceTests.swift b/Tests/Operators/AsyncSwitchToLatestSequenceTests.swift index 8619cde..c22f498 100644 --- a/Tests/Operators/AsyncSwitchToLatestSequenceTests.swift +++ b/Tests/Operators/AsyncSwitchToLatestSequenceTests.swift @@ -57,6 +57,24 @@ private struct LongAsyncSequence: AsyncSequence, AsyncIteratorProtocol } } +private struct NonCooperativeAsyncSequence: 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) in + self.onSuspend() + } + } + } +} + final class AsyncSwitchToLatestSequenceTests: XCTestCase { func testSwitchToLatest_switches_to_latest_asyncSequence_and_cancels_previous_ones() async throws { var asyncSequence1IsCancelled = false @@ -183,4 +201,67 @@ 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>.Continuation! + let outer = AsyncStream> { 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() + first.send(1) + outerContinuation.yield(first) + await fulfillment(of: [receivedFirstValue], timeout: 1) + + let second = AsyncBufferedChannel() + second.send(2) + outerContinuation.yield(second) + await fulfillment(of: [receivedSecondValue], timeout: 1) + + let latest = AsyncBufferedChannel() + 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> { + 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) + } } From ba05c60c9ab9cea2e95215c1cf54abb761c43731 Mon Sep 17 00:00:00 2001 From: pablobaldez Date: Sun, 23 Aug 2026 14:14:29 +0100 Subject: [PATCH 2/2] Release consumed child tasks when switching --- .../AsyncSwitchToLatestSequence.swift | 2 + .../AsyncSwitchToLatestSequenceTests.swift | 65 +++++++++++++++++++ 2 files changed, 67 insertions(+) diff --git a/Sources/Operators/AsyncSwitchToLatestSequence.swift b/Sources/Operators/AsyncSwitchToLatestSequence.swift index 595618f..fa98199 100644 --- a/Sources/Operators/AsyncSwitchToLatestSequence.swift +++ b/Sources/Operators/AsyncSwitchToLatestSequence.swift @@ -265,6 +265,8 @@ where Base.Element: AsyncSequence, Base: Sendable, Base.Element.Element: Sendabl let value = await childTask?.value let decision = state.withCriticalRegion { state -> PostElementDecision in + state.childTask = nil + if state.isCancelled { return .returnFinish } else if state.base.isNewAvailableChildIterator { diff --git a/Tests/Operators/AsyncSwitchToLatestSequenceTests.swift b/Tests/Operators/AsyncSwitchToLatestSequenceTests.swift index c22f498..80118d9 100644 --- a/Tests/Operators/AsyncSwitchToLatestSequenceTests.swift +++ b/Tests/Operators/AsyncSwitchToLatestSequenceTests.swift @@ -75,6 +75,40 @@ private struct NonCooperativeAsyncSequence: AsyncSequence, Se } } +private struct IteratorLifetimeSequence: 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 @@ -264,4 +298,35 @@ final class AsyncSwitchToLatestSequenceTests: XCTestCase { 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>.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() + } }