refactor: Implement robust flatMapLatest with generation tracking

- Replaces naive implementation with a thread-safe approach using ManagedCriticalState.
- Introduces generation tracking to prevent race conditions where cancelled inner sequences could yield stale values.
- Adds test_interleaving_race_condition to verify correctness under concurrent load.
- Ensures Swift 6 Sendable compliance.

Author: peterfriese

Sources/AsyncAlgorithms/AsyncFlatMapLatestSequence.swift

@@ -26,43 +26,81 @@ extension AsyncSequence where Self: Sendable {
   ) -> AsyncThrowingStream<T.Element, Error>
   where T.Element: Sendable {
 
-    AsyncThrowingStream { continuation in
-      let outerIterationTask = Task {
-        var innerIterationTask: Task<Void, Never>? = nil
-        
+    // Explicitly specify the type of the stream
+    return AsyncThrowingStream<T.Element, Error> { continuation in
+      let state = ManagedCriticalState(FlatMapLatestState())
+      
+      let outerTask = Task {
         do {
           for try await element in self {
-            innerIterationTask?.cancel()
-            
             let innerSequence = transform(element)
 
-            innerIterationTask = Task {
+            // Increment generation and get the new value
+            let currentGeneration = state.withCriticalRegion { state -> Int in
+              state.innerTask?.cancel()
+              state.generation += 1
+              return state.generation
+            }
+            
+            let innerTask = Task {
               do {
                 for try await innerElement in innerSequence {
-                  try Task.checkCancellation()
-                  continuation.yield(innerElement)
+                  // Check if we are still the latest generation
+                  let shouldYield = state.withCriticalRegion { state in
+                    state.generation == currentGeneration
+                  }
+                  
+                  if shouldYield {
+                    continuation.yield(innerElement)
+                  } else {
+                    // If we are not the latest, we should stop
+                    return
+                  }
                 }
               } catch is CancellationError {
-                // Inner task was cancelled, this is normal
+                // Normal cancellation
               } catch {
-                // Inner sequence threw an error
-                continuation.finish(throwing: error)
+                // If an error occurs, we only propagate it if we are the latest generation
+                let shouldPropagate = state.withCriticalRegion { state in
+                  state.generation == currentGeneration
+                }
+                if shouldPropagate {
+                  continuation.finish(throwing: error)
+                }
+              }
+            }
+            
+            state.withCriticalRegion { state in
+              // Only update the inner task if the generation hasn't changed again
+              if state.generation == currentGeneration {
+                state.innerTask = innerTask
               }
             }
           }
+          
+          // Outer sequence finished
+          // Wait for the last inner task to finish
+          let lastInnerTask = state.withCriticalRegion { $0.innerTask }
+          _ = await lastInnerTask?.result
+          continuation.finish()
+          
         } catch {
-          // Outer sequence threw an error
           continuation.finish(throwing: error)
         }
-        
-        // Outer sequence finished
-        await innerIterationTask?.value
-        continuation.finish()
       }
 
       continuation.onTermination = { @Sendable _ in
-        outerIterationTask.cancel()
+        outerTask.cancel()
+        state.withCriticalRegion { state in
+          state.innerTask?.cancel()
+        }
       }
     }
   }
 }
+
+@available(AsyncAlgorithms 1.0, *)
+private struct FlatMapLatestState: Sendable {
+  var generation: Int = 0
+  var innerTask: Task<Void, Never>? = nil
+}

Tests/AsyncAlgorithmsTests/TestFlatMapLatest.swift

@@ -35,4 +35,60 @@ final class TestFlatMapLatest: XCTestCase {
     XCTAssertTrue(expected.isEmpty)
   }
 
+  func test_interleaving_race_condition() async throws {
+    // This test simulates a scenario where the inner sequence is slow.
+    // In a naive implementation (without generation tracking), the inner task for '1'
+    // might wake up and yield AFTER '2' has already started, causing interleaving.
+    
+    let source = [1, 2, 3].async
+    let transformed = source.flatMapLatest { intValue -> AsyncStream<Int> in
+      AsyncStream { continuation in
+        Task {
+          // Yield the value immediately
+          continuation.yield(intValue)
+          
+          // Sleep for a bit to allow the outer sequence to move on
+          try? await Task.sleep(nanoseconds: 10_000_000) // 10ms
+          
+          // Yield a second value - this should be ignored if a new outer value has arrived
+          continuation.yield(intValue * 10)
+          continuation.finish()
+        }
+      }
+    }
+
+    // We expect:
+    // 1 arrives -> starts inner(1) -> yields 1 -> sleeps
+    // 2 arrives -> cancels inner(1) -> starts inner(2) -> yields 2 -> sleeps
+    // 3 arrives -> cancels inner(2) -> starts inner(3) -> yields 3 -> sleeps
+    // inner(3) finishes sleep -> yields 30
+    //
+    // Ideally, we should NOT see 10 or 20.
+    // However, without strict synchronization, we might see them.
+    // The strict expectation for flatMapLatest is that once a new value arrives,
+    // the old one produces NO MORE values.
+    
+    // Note: This test is probabilistic in the naive implementation. 
+    // It might pass or fail depending on scheduling. 
+    // But with a correct implementation, it should ALWAYS pass.
+    
+    var expected = [3, 30] // We only want the latest
+    
+    // We'll collect all results to see what happened
+    var results: [Int] = []
+    
+    for try await element in transformed {
+      results.append(element)
+    }
+    
+    // In the naive implementation, we might get [1, 2, 3, 10, 20, 30] or similar.
+    // We want strictly [3, 30] (or [1, 2, 3, 30] depending on how fast the outer sequence is consumed vs produced)
+    // Actually, if the outer sequence is consumed fast, we might see intermediate "first" values (1, 2).
+    // But we should NEVER see "second" values (10, 20) from cancelled sequences.
+    
+    // Let's relax the check to: "Must not contain 10 or 20"
+    XCTAssertFalse(results.contains(10), "Should not contain 10 (from cancelled sequence 1)")
+    XCTAssertFalse(results.contains(20), "Should not contain 20 (from cancelled sequence 2)")
+    XCTAssertTrue(results.contains(30), "Should contain 30 (from final sequence)")
+  }
 }