Merge pull request #75878 from jamieQ/stream-termination

Author: ktoso

stdlib/public/Concurrency/AsyncStreamBuffer.swift

@@ -212,25 +212,21 @@ extension AsyncStream {
       state.onTermination = nil
       state.terminal = true
 
-      if let continuation = state.continuations.first {
-        if state.pending.count > 0 {
-          state.continuations.removeFirst()
-          let toSend = state.pending.removeFirst()
-          unlock()
-          handler?(.finished)
-          continuation.resume(returning: toSend)
-        } else if state.terminal {
-          state.continuations.removeFirst()
-          unlock()
-          handler?(.finished)
-          continuation.resume(returning: nil)
-        } else {
-          unlock()
-          handler?(.finished)
-        }
-      } else {
+      guard !state.continuations.isEmpty else {
         unlock()
         handler?(.finished)
+        return
+      }
+
+      // Hold on to the continuations to resume outside the lock.
+      let continuations = state.continuations
+      state.continuations.removeAll()
+
+      unlock()
+      handler?(.finished)
+
+      for continuation in continuations {
+        continuation.resume(returning: nil)
       }
     }
 

test/Concurrency/Runtime/async_stream.swift

@@ -435,6 +435,40 @@ class NotSendable {}
         expectTrue(expectation.fulfilled)
       }
 
+      // MARK: - Multiple consumers
+
+      tests.test("finish behavior with multiple consumers") {
+        let (stream, continuation) = AsyncStream<Int>.makeStream()
+        let (controlStream, controlContinuation) = AsyncStream<Int>.makeStream()
+        var controlIterator = controlStream.makeAsyncIterator()
+
+        func makeConsumingTaskWithIndex(_ index: Int) -> Task<Void, Never> {
+          Task { @MainActor in
+            controlContinuation.yield(index)
+            for await i in stream {
+              controlContinuation.yield(i)
+            }
+          }
+        }
+
+        // Set up multiple consumers
+        let consumer1 = makeConsumingTaskWithIndex(1)
+        expectEqual(await controlIterator.next(isolation: #isolation), 1)
+
+        let consumer2 = makeConsumingTaskWithIndex(2)
+        expectEqual(await controlIterator.next(isolation: #isolation), 2)
+
+        // Ensure the iterators are suspended
+        await MainActor.run {}
+
+        // Terminate the stream
+        continuation.finish()
+
+        // Ensure the consuming Tasks both complete
+        _ = await consumer1.value
+        _ = await consumer2.value
+      }
+
       await runAllTestsAsync()
     }
   }