Correct lifetime based cancellation of AsyncStream and AsyncThrowingStream (#40087)

Author: phausler

stdlib/public/Concurrency/AsyncStream.swift

@@ -139,7 +139,22 @@ public struct AsyncStream<Element> {
     }
   }
 
-  let produce: () async -> Element?
+  final class _Context {
+    let storage: _Storage?
+    let produce: () async -> Element?
+
+    init(storage: _Storage? = nil, produce: @escaping () async -> Element?) {
+      self.storage = storage
+      self.produce = produce
+    }
+
+    deinit {
+      storage?.cancel()
+    }
+  }
+
+  let context: _Context
+  
 
   /// Construct a AsyncStream buffering given an Element type.
   ///
@@ -163,7 +178,7 @@ public struct AsyncStream<Element> {
     _ build: (Continuation) -> Void
   ) {
     let storage: _Storage = .create(limit: limit)
-    self.init(unfolding: storage.next)
+    context = _Context(storage: storage, produce: storage.next)
     build(Continuation(storage: storage))
   }
 
@@ -174,7 +189,7 @@ public struct AsyncStream<Element> {
   ) {
     let storage: _AsyncStreamCriticalStorage<Optional<() async -> Element?>>
       = .create(produce)
-    self.produce = {
+    context = _Context {
       return await withTaskCancellationHandler {
         guard let result = await storage.value?() else {
           storage.value = nil
@@ -198,7 +213,7 @@ extension AsyncStream: AsyncSequence {
   /// concurrently and contends with another call to next is a programmer error
   /// and will fatalError.
   public struct Iterator: AsyncIteratorProtocol {
-    let produce: () async -> Element?
+    let context: _Context
 
     /// The next value from the AsyncStream.
     ///
@@ -210,13 +225,13 @@ extension AsyncStream: AsyncSequence {
     /// awaiting a value, this will terminate the AsyncStream and next may return nil
     /// immediately (or will return nil on subsequent calls)
     public mutating func next() async -> Element? {
-      await produce()
+      await context.produce()
     }
   }
 
   /// Construct an iterator.
   public func makeAsyncIterator() -> Iterator {
-    return Iterator(produce: produce)
+    return Iterator(context: context)
   }
 }
 

stdlib/public/Concurrency/AsyncThrowingStream.swift

@@ -108,7 +108,21 @@ public struct AsyncThrowingStream<Element, Failure: Error> {
     }
   }
 
-  let produce: () async throws -> Element?
+  final class _Context {
+    let storage: _Storage?
+    let produce: () async throws -> Element?
+
+    init(storage: _Storage? = nil, produce: @escaping () async throws -> Element?) {
+      self.storage = storage
+      self.produce = produce
+    }
+
+    deinit {
+      storage?.cancel()
+    }
+  }
+
+  let context: _Context
 
   /// Construct a AsyncThrowingStream buffering given an Element type.
   ///
@@ -132,14 +146,26 @@ public struct AsyncThrowingStream<Element, Failure: Error> {
     _ build: (Continuation) -> Void
   ) where Failure == Error {
     let storage: _Storage = .create(limit: limit)
-    self.init(unfolding: storage.next)
+    context = _Context(storage: storage, produce: storage.next)
     build(Continuation(storage: storage))
   }
 
   public init(
     unfolding produce: @escaping () async throws -> Element?
   ) where Failure == Error {
-    self.produce = produce
+    let storage: _AsyncStreamCriticalStorage<Optional<() async throws -> Element?>>
+      = .create(produce)
+    context = _Context {
+      return try await withTaskCancellationHandler {
+        guard let result = try await storage.value?() else {
+          storage.value = nil
+          return nil
+        }
+        return result
+      } onCancel: {
+        storage.value = nil
+      }
+    }
   }
 }
 
@@ -152,16 +178,16 @@ extension AsyncThrowingStream: AsyncSequence {
   /// concurrently and contends with another call to next is a programmer error
   /// and will fatalError.
   public struct Iterator: AsyncIteratorProtocol {
-    let produce: () async throws -> Element?
+    let context: _Context
 
     public mutating func next() async throws -> Element? {
-      return try await produce()
+      return try await context.produce()
     }
   }
 
   /// Construct an iterator.
   public func makeAsyncIterator() -> Iterator {
-    return Iterator(produce: produce)
+    return Iterator(context: context)
   }
 }