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 {}