diff --git a/Sources/AsyncAlgorithms/FlatMapLatest/FlatMapLatestStateMachine.swift b/Sources/AsyncAlgorithms/FlatMapLatest/FlatMapLatestStateMachine.swift index 71ee36c4..8891f4d2 100644 --- a/Sources/AsyncAlgorithms/FlatMapLatest/FlatMapLatestStateMachine.swift +++ b/Sources/AsyncAlgorithms/FlatMapLatest/FlatMapLatestStateMachine.swift @@ -125,12 +125,39 @@ where Base.Element: Sendable, Inner.Element: Sendable { let innerTask, let innerCont, let downstreamCont, - let buffer, + var buffer, let generation, let outerFinished ): precondition(downstreamCont == nil, "Already have downstream continuation") - precondition(buffer.isEmpty, "Buffer should be empty if suspending") + + // Producers may buffer elements or finish between `next()` releasing + // the lock and this call; serve the demand from the current state. + if let result = buffer.popFirst() { + state = .running( + outerTask: outerTask, + outerContinuation: outerCont, + innerTask: innerTask, + innerContinuation: innerCont, + downstreamContinuation: nil, + buffer: buffer, + generation: generation, + outerFinished: outerFinished + ) + switch result { + case .success(let element): + continuation.resume(returning: element) + case .failure(let error): + continuation.resume(throwing: error) + } + return .none + } + + if outerFinished && innerTask == nil { + state = .finished + continuation.resume(returning: nil) + return .none + } state = .running( outerTask: outerTask, diff --git a/Sources/AsyncAlgorithms/FlatMapLatest/FlatMapLatestStorage.swift b/Sources/AsyncAlgorithms/FlatMapLatest/FlatMapLatestStorage.swift index dfe5aa57..b621a717 100644 --- a/Sources/AsyncAlgorithms/FlatMapLatest/FlatMapLatestStorage.swift +++ b/Sources/AsyncAlgorithms/FlatMapLatest/FlatMapLatestStorage.swift @@ -124,6 +124,9 @@ where Base.Element: Sendable, Inner.Element: Sendable { } private func startInnerTask(_ inner: Inner, generation: Int) { + // Hold the lock so the task is registered before its body can take the + // lock and transition the state machine (matching `startOuterTask`). + lock.lock() let task = Task { var iterator = inner.makeAsyncIterator() @@ -165,6 +168,7 @@ where Base.Element: Sendable, Inner.Element: Sendable { } } stateMachine.innerTaskStarted(task, generation: generation) + lock.unlock() } private func handleAction(_ action: FlatMapLatestStateMachine.Action) { diff --git a/Tests/AsyncAlgorithmsTests/TestFlatMapLatest.swift b/Tests/AsyncAlgorithmsTests/TestFlatMapLatest.swift index 2cfab1cc..6cfe61c7 100644 --- a/Tests/AsyncAlgorithmsTests/TestFlatMapLatest.swift +++ b/Tests/AsyncAlgorithmsTests/TestFlatMapLatest.swift @@ -226,6 +226,39 @@ final class TestFlatMapLatest: XCTestCase { // Determine success by running without crashing for try await _ in combined {} } + + func test_rapid_switching_single_consumer_stress() async throws { + // Regression test: rapid inner sequence switching used to trap in + // `next(for:)` or hang; the timeout guards against the hang variant. + let finished = expectation(description: "all iterations finished") + let consumer = Task { + for _ in 0..<1000 { + let toggles = AsyncStream { continuation in + for index in 0..<20 { + continuation.yield(index.isMultiple(of: 2)) + } + continuation.finish() + } + + let sequence = toggles.flatMapLatest { enabled in + AsyncStream { continuation in + if enabled { + for value in 0..<200 { + continuation.yield(value) + } + } + continuation.finish() + } + } + + for await _ in sequence {} + } + finished.fulfill() + } + defer { consumer.cancel() } + + await fulfillment(of: [finished], timeout: 120) + } } private struct FlatMapLatestFailure: Error, Equatable {}