Merge pull request swiftlang#38582 from DougGregor/async-stream-catch-up

Author: DougGregor

stdlib/public/Concurrency/AsyncStream.swift

@@ -49,10 +49,50 @@ import Swift
 @available(SwiftStdlib 5.5, *)
 public struct AsyncStream<Element> {
   public struct Continuation: Sendable {
+    /// Indication of the type of termination informed to
+    /// `onTermination`.
     public enum Termination {
+      
+      /// The stream was finished via the `finish` method
       case finished
+      
+      /// The stream was cancelled
       case cancelled
     }
+    
+    /// A result of yielding values.
+    public enum YieldResult {
+      
+      /// When a value is successfully enqueued, either buffered
+      /// or immediately consumed to resume a pending call to next
+      /// and a count of remaining slots available in the buffer at
+      /// the point in time of yielding. Note: transacting upon the
+      /// remaining count is only valid when then calls to yield are
+      /// mutually exclusive.
+      case enqueued(remaining: Int)
+      
+      /// Yielding resulted in not buffering an element because the
+      /// buffer was full. The element is the dropped value.
+      case dropped(Element)
+      
+      /// Indication that the continuation was yielded when the
+      /// stream was already in a terminal state: either by cancel or
+      /// by finishing.
+      case terminated
+    }
+    
+    /// A strategy that handles exhaustion of a buffer’s capacity.
+    public enum BufferingPolicy {
+      case unbounded
+      
+      /// When the buffer is full, discard the newly received element.
+      /// This enforces keeping the specified amount of oldest values.
+      case bufferingOldest(Int)
+      
+      /// When the buffer is full, discard the oldest element in the buffer.
+      /// This enforces keeping the specified amount of newest values.
+      case bufferingNewest(Int)
+    }
 
     let storage: _Storage
 
@@ -64,7 +104,8 @@ public struct AsyncStream<Element> {
     ///
     /// This can be called more than once and returns to the caller immediately
     /// without blocking for any awaiting consumption from the iteration.
-    public func yield(_ value: __owned Element) {
+    @discardableResult
+    public func yield(_ value: __owned Element) -> YieldResult {
       storage.yield(value)
     }
 
@@ -105,7 +146,8 @@ public struct AsyncStream<Element> {
   /// - Parameter elementType: The type the AsyncStream will produce.
   /// - Parameter maxBufferedElements: The maximum number of elements to
   ///   hold in the buffer past any checks for continuations being resumed.
-  /// - Parameter build: The work associated with yielding values to the AsyncStream.
+  /// - Parameter build: The work associated with yielding values to the 
+  ///   AsyncStream.
   ///
   /// The maximum number of pending elements limited by dropping the oldest
   /// value when a new value comes in if the buffer would exceed the limit
@@ -117,13 +159,34 @@ public struct AsyncStream<Element> {
   /// concurrent contexts could result in out of order delivery.
   public init(
     _ elementType: Element.Type = Element.self,
-    maxBufferedElements limit: Int = .max,
+    bufferingPolicy limit: Continuation.BufferingPolicy = .unbounded,
     _ build: (Continuation) -> Void
   ) {
     let storage: _Storage = .create(limit: limit)
-    produce = storage.next
+    self.init(unfolding: storage.next)
     build(Continuation(storage: storage))
   }
+
+  
+  public init(
+    unfolding produce: @escaping () async -> Element?, 
+    onCancel: (@Sendable () -> Void)? = nil
+  ) {
+    let storage: _AsyncStreamCriticalStorage<Optional<() async -> Element?>>
+      = .create(produce)
+    self.produce = {
+      return await Task.withCancellationHandler {
+        storage.value = nil
+        onCancel?()
+      } operation: {
+        guard let result = await storage.value?() else {
+          storage.value = nil
+          return nil
+        }
+        return result
+      }
+    }
+  }
 }
 
 @available(SwiftStdlib 5.5, *)
@@ -139,9 +202,9 @@ extension AsyncStream: AsyncSequence {
 
     /// The next value from the AsyncStream.
     ///
-    /// When next returns nil this signifies the end of the AsyncStream. Any such
-    /// case that next is invoked concurrently and contends with another call to
-    /// next is a programmer error and will fatalError.
+    /// When next returns nil this signifies the end of the AsyncStream. Any 
+    /// such case that next is invoked concurrently and contends with another 
+    /// call to next is a programmer error and will fatalError.
     ///
     /// If the task this iterator is running in is canceled while next is
     /// awaiting a value, this will terminate the AsyncStream and next may return nil
@@ -167,12 +230,13 @@ extension AsyncStream.Continuation {
   ///
   /// This can be called more than once and returns to the caller immediately
   /// without blocking for any awaiting consumption from the iteration.
+  @discardableResult
   public func yield(
     with result: Result<Element, Never>
-  ) {
+  ) -> YieldResult {
     switch result {
       case .success(let val):
-        storage.yield(val)
+        return storage.yield(val)
     }
   }
 
@@ -182,7 +246,9 @@ extension AsyncStream.Continuation {
   ///
   /// This can be called more than once and returns to the caller immediately
   /// without blocking for any awaiting consumption from the iteration.
-  public func yield() where Element == Void {
-    storage.yield(())
+  /// without blocking for any awaiting consuption from the iteration.
+  @discardableResult
+  public func yield() -> YieldResult where Element == Void {
+    return storage.yield(())
   }
 }