From 81dc1e4bd7e90c8e9c5f0c2e978d947c3a410ad2 Mon Sep 17 00:00:00 2001 From: Guilherme Souza Date: Fri, 17 Oct 2025 10:21:51 -0300 Subject: [PATCH] feat: pull OpenAPIURLSession changes --- .../BufferedStream/BufferedStream.swift | 3274 +++++++++-------- .../BufferedStream/Lock.swift | 193 +- ...rectionalStreamingURLSessionDelegate.swift | 73 +- .../HTTPBodyOutputStreamBridge.swift | 39 +- .../URLSessionTransport.swift | 25 +- 5 files changed, 1847 insertions(+), 1757 deletions(-) diff --git a/Sources/HTTPClient/HTTPClientFoundation/BufferedStream/BufferedStream.swift b/Sources/HTTPClient/HTTPClientFoundation/BufferedStream/BufferedStream.swift index 6b30d21..5c2ae53 100644 --- a/Sources/HTTPClient/HTTPClientFoundation/BufferedStream/BufferedStream.swift +++ b/Sources/HTTPClient/HTTPClientFoundation/BufferedStream/BufferedStream.swift @@ -11,7 +11,7 @@ // SPDX-License-Identifier: Apache-2.0 // //===----------------------------------------------------------------------===// -// swift-format-ignore-file + //===----------------------------------------------------------------------===// // // This source file is part of the Swift.org open source project @@ -131,102 +131,102 @@ import DequeModule @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *) @usableFromInline internal struct BufferedStream { + @usableFromInline + final class _Backing: Sendable { @usableFromInline - final class _Backing: Sendable { - @usableFromInline - let storage: _BackPressuredStorage - - @usableFromInline - init(storage: _BackPressuredStorage) { - self.storage = storage - } + let storage: _BackPressuredStorage - deinit { - self.storage.sequenceDeinitialized() - } + @usableFromInline + init(storage: _BackPressuredStorage) { + self.storage = storage } - @usableFromInline - enum _Implementation: Sendable { - /// This is the implementation with backpressure based on the Source - case backpressured(_Backing) + deinit { + self.storage.sequenceDeinitialized() } + } - @usableFromInline - let implementation: _Implementation + @usableFromInline + enum _Implementation: Sendable { + /// This is the implementation with backpressure based on the Source + case backpressured(_Backing) + } + + @usableFromInline + let implementation: _Implementation } @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *) extension BufferedStream: AsyncSequence { - /// The asynchronous iterator for iterating an asynchronous stream. - /// - /// This type is not `Sendable`. Don't use it from multiple - /// concurrent contexts. It is a programmer error to invoke `next()` from a - /// concurrent context that contends with another such call, which - /// results in a call to `fatalError()`. + /// The asynchronous iterator for iterating an asynchronous stream. + /// + /// This type is not `Sendable`. Don't use it from multiple + /// concurrent contexts. It is a programmer error to invoke `next()` from a + /// concurrent context that contends with another such call, which + /// results in a call to `fatalError()`. + @usableFromInline + internal struct Iterator: AsyncIteratorProtocol { @usableFromInline - internal struct Iterator: AsyncIteratorProtocol { - @usableFromInline - final class _Backing { - @usableFromInline - let storage: _BackPressuredStorage - - @usableFromInline - init(storage: _BackPressuredStorage) { - self.storage = storage - self.storage.iteratorInitialized() - } - - deinit { - self.storage.iteratorDeinitialized() - } - } - - @usableFromInline - enum _Implementation { - /// This is the implementation with backpressure based on the Source - case backpressured(_Backing) - } + final class _Backing { + @usableFromInline + let storage: _BackPressuredStorage + + @usableFromInline + init(storage: _BackPressuredStorage) { + self.storage = storage + self.storage.iteratorInitialized() + } + + deinit { + self.storage.iteratorDeinitialized() + } + } - @usableFromInline - var implementation: _Implementation + @usableFromInline + enum _Implementation { + /// This is the implementation with backpressure based on the Source + case backpressured(_Backing) + } - @usableFromInline - init(implementation: _Implementation) { - self.implementation = implementation - } + @usableFromInline + var implementation: _Implementation - /// The next value from the asynchronous stream. - /// - /// When `next()` returns `nil`, this signifies the end of the - /// `BufferedStream`. - /// - /// It is a programmer error to invoke `next()` from a concurrent context - /// that contends with another such call, which results in a call to - /// `fatalError()`. - /// - /// If you cancel the task this iterator is running in while `next()` is - /// awaiting a value, the `BufferedStream` terminates. In this case, - /// `next()` may return `nil` immediately, or else return `nil` on - /// subsequent calls. - @inlinable - internal mutating func next() async throws -> Element? { - switch self.implementation { - case .backpressured(let backing): - return try await backing.storage.next() - } - } + @usableFromInline + init(implementation: _Implementation) { + self.implementation = implementation } - /// Creates the asynchronous iterator that produces elements of this - /// asynchronous sequence. + /// The next value from the asynchronous stream. + /// + /// When `next()` returns `nil`, this signifies the end of the + /// `BufferedStream`. + /// + /// It is a programmer error to invoke `next()` from a concurrent context + /// that contends with another such call, which results in a call to + /// `fatalError()`. + /// + /// If you cancel the task this iterator is running in while `next()` is + /// awaiting a value, the `BufferedStream` terminates. In this case, + /// `next()` may return `nil` immediately, or else return `nil` on + /// subsequent calls. @inlinable - internal func makeAsyncIterator() -> Iterator { - switch self.implementation { - case .backpressured(let backing): - return Iterator(implementation: .backpressured(.init(storage: backing.storage))) - } + internal mutating func next() async throws -> Element? { + switch self.implementation { + case .backpressured(let backing): + return try await backing.storage.next() + } + } + } + + /// Creates the asynchronous iterator that produces elements of this + /// asynchronous sequence. + @inlinable + internal func makeAsyncIterator() -> Iterator { + switch self.implementation { + case .backpressured(let backing): + return Iterator(implementation: .backpressured(.init(storage: backing.storage))) } + } } @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *) @@ -234,1740 +234,1750 @@ extension BufferedStream: Sendable where Element: Sendable {} @usableFromInline internal struct _ManagedCriticalState: @unchecked Sendable { - @usableFromInline - let lock: LockedValueBox - - @usableFromInline - internal init(_ initial: State) { - self.lock = .init(initial) - } - - @inlinable - internal func withCriticalRegion( - _ critical: (inout State) throws -> R - ) rethrows -> R { - try self.lock.withLockedValue(critical) - } + @usableFromInline + let lock: LockedValueBox + + @usableFromInline + internal init(_ initial: State) { + self.lock = .init(initial) + } + + @inlinable + internal func withCriticalRegion( + _ critical: (inout State) throws -> R + ) rethrows -> R { + try self.lock.withLockedValue(critical) + } } @usableFromInline internal struct AlreadyFinishedError: Error { - @usableFromInline - init() {} + @usableFromInline + init() {} } @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *) extension BufferedStream { - /// A mechanism to interface between producer code and an asynchronous stream. - /// - /// Use this source to provide elements to the stream by calling one of the `write` methods, then terminate the stream normally - /// by calling the `finish()` method. You can also use the source's `finish(throwing:)` method to terminate the stream by - /// throwing an error. + /// A mechanism to interface between producer code and an asynchronous stream. + /// + /// Use this source to provide elements to the stream by calling one of the `write` methods, then terminate the stream normally + /// by calling the `finish()` method. You can also use the source's `finish(throwing:)` method to terminate the stream by + /// throwing an error. + @usableFromInline + internal struct Source: Sendable { + /// A strategy that handles the backpressure of the asynchronous stream. @usableFromInline - internal struct Source: Sendable { - /// A strategy that handles the backpressure of the asynchronous stream. - @usableFromInline - internal struct BackPressureStrategy: Sendable { - /// When the high watermark is reached producers will be suspended. All producers will be resumed again once - /// the low watermark is reached. The current watermark is the number of elements in the buffer. - @inlinable - internal static func watermark(low: Int, high: Int) -> BackPressureStrategy { - BackPressureStrategy( - internalBackPressureStrategy: .watermark(.init(low: low, high: high)) - ) - } - - /// When the high watermark is reached producers will be suspended. All producers will be resumed again once - /// the low watermark is reached. The current watermark is computed using the given closure. - static func customWatermark( - low: Int, - high: Int, - waterLevelForElement: @escaping @Sendable (Element) -> Int - ) -> BackPressureStrategy where Element: RandomAccessCollection { - BackPressureStrategy( - internalBackPressureStrategy: .watermark(.init(low: low, high: high, waterLevelForElement: waterLevelForElement)) - ) - } + internal struct BackPressureStrategy: Sendable { + /// When the high watermark is reached producers will be suspended. All producers will be resumed again once + /// the low watermark is reached. The current watermark is the number of elements in the buffer. + @inlinable + internal static func watermark(low: Int, high: Int) -> BackPressureStrategy { + BackPressureStrategy( + internalBackPressureStrategy: .watermark(.init(low: low, high: high)) + ) + } + + /// When the high watermark is reached producers will be suspended. All producers will be resumed again once + /// the low watermark is reached. The current watermark is computed using the given closure. + static func customWatermark( + low: Int, + high: Int, + waterLevelForElement: @escaping @Sendable (Element) -> Int + ) -> BackPressureStrategy where Element: RandomAccessCollection { + BackPressureStrategy( + internalBackPressureStrategy: .watermark( + .init(low: low, high: high, waterLevelForElement: waterLevelForElement) + ) + ) + } - @usableFromInline - init(internalBackPressureStrategy: _InternalBackPressureStrategy) { - self._internalBackPressureStrategy = internalBackPressureStrategy - } + @usableFromInline + init(internalBackPressureStrategy: _InternalBackPressureStrategy) { + self._internalBackPressureStrategy = internalBackPressureStrategy + } - @usableFromInline - let _internalBackPressureStrategy: _InternalBackPressureStrategy - } + @usableFromInline + let _internalBackPressureStrategy: _InternalBackPressureStrategy + } - /// A type that indicates the result of writing elements to the source. - @frozen + /// A type that indicates the result of writing elements to the source. + @frozen + @usableFromInline + internal enum WriteResult: Sendable { + /// A token that is returned when the asynchronous stream's backpressure strategy indicated that production should + /// be suspended. Use this token to enqueue a callback by calling the ``enqueueCallback(_:)`` method. + @usableFromInline + internal struct CallbackToken: Sendable { @usableFromInline - internal enum WriteResult: Sendable { - /// A token that is returned when the asynchronous stream's backpressure strategy indicated that production should - /// be suspended. Use this token to enqueue a callback by calling the ``enqueueCallback(_:)`` method. - @usableFromInline - internal struct CallbackToken: Sendable { - @usableFromInline - let id: UInt - @usableFromInline - init(id: UInt) { - self.id = id - } - } - - /// Indicates that more elements should be produced and written to the source. - case produceMore - - /// Indicates that a callback should be enqueued. - /// - /// The associated token should be passed to the ``enqueueCallback(_:)`` method. - case enqueueCallback(CallbackToken) - } - - /// Backing class for the source used to hook a deinit. + let id: UInt @usableFromInline - final class _Backing: Sendable { - @usableFromInline - let storage: _BackPressuredStorage + init(id: UInt) { + self.id = id + } + } - @usableFromInline - init(storage: _BackPressuredStorage) { - self.storage = storage - } + /// Indicates that more elements should be produced and written to the source. + case produceMore - deinit { - self.storage.sourceDeinitialized() - } - } + /// Indicates that a callback should be enqueued. + /// + /// The associated token should be passed to the ``enqueueCallback(_:)`` method. + case enqueueCallback(CallbackToken) + } - /// A callback to invoke when the stream finished. - /// - /// The stream finishes and calls this closure in the following cases: - /// - No iterator was created and the sequence was deinited - /// - An iterator was created and deinited - /// - After ``finish(throwing:)`` was called and all elements have been consumed - /// - The consuming task got cancelled - @inlinable - internal var onTermination: (@Sendable () -> Void)? { - set { - self._backing.storage.onTermination = newValue - } - get { - self._backing.storage.onTermination - } - } + /// Backing class for the source used to hook a deinit. + @usableFromInline + final class _Backing: Sendable { + @usableFromInline + let storage: _BackPressuredStorage - @usableFromInline - var _backing: _Backing + @usableFromInline + init(storage: _BackPressuredStorage) { + self.storage = storage + } - @usableFromInline - internal init(storage: _BackPressuredStorage) { - self._backing = .init(storage: storage) - } + deinit { + self.storage.sourceDeinitialized() + } + } - /// Writes new elements to the asynchronous stream. - /// - /// If there is a task consuming the stream and awaiting the next element then the task will get resumed with the - /// first element of the provided sequence. If the asynchronous stream already terminated then this method will throw an error - /// indicating the failure. - /// - /// - Parameter sequence: The elements to write to the asynchronous stream. - /// - Returns: The result that indicates if more elements should be produced at this time. - @inlinable - internal func write(contentsOf sequence: S) throws -> WriteResult - where Element == S.Element, S: Sequence { - try self._backing.storage.write(contentsOf: sequence) - } + /// A callback to invoke when the stream finished. + /// + /// The stream finishes and calls this closure in the following cases: + /// - No iterator was created and the sequence was deinited + /// - An iterator was created and deinited + /// - After ``finish(throwing:)`` was called and all elements have been consumed + /// - The consuming task got cancelled + @inlinable + internal var onTermination: (@Sendable () -> Void)? { + set { + self._backing.storage.onTermination = newValue + } + get { + self._backing.storage.onTermination + } + } - /// Write the element to the asynchronous stream. - /// - /// If there is a task consuming the stream and awaiting the next element then the task will get resumed with the - /// provided element. If the asynchronous stream already terminated then this method will throw an error - /// indicating the failure. - /// - /// - Parameter element: The element to write to the asynchronous stream. - /// - Returns: The result that indicates if more elements should be produced at this time. - @inlinable - internal func write(_ element: Element) throws -> WriteResult { - try self._backing.storage.write(contentsOf: CollectionOfOne(element)) - } + @usableFromInline + var _backing: _Backing - /// Enqueues a callback that will be invoked once more elements should be produced. - /// - /// Call this method after ``write(contentsOf:)`` or ``write(:)`` returned ``WriteResult/enqueueCallback(_:)``. - /// - /// - Important: Enqueueing the same token multiple times is not allowed. - /// - /// - Parameters: - /// - callbackToken: The callback token. - /// - onProduceMore: The callback which gets invoked once more elements should be produced. - @inlinable - internal func enqueueCallback( - callbackToken: WriteResult.CallbackToken, - onProduceMore: @escaping @Sendable (Result) -> Void - ) { - self._backing.storage.enqueueProducer( - callbackToken: callbackToken, - onProduceMore: onProduceMore - ) - } + @usableFromInline + internal init(storage: _BackPressuredStorage) { + self._backing = .init(storage: storage) + } - /// Cancel an enqueued callback. - /// - /// Call this method to cancel a callback enqueued by the ``enqueueCallback(callbackToken:onProduceMore:)`` method. - /// - /// - Note: This methods supports being called before ``enqueueCallback(callbackToken:onProduceMore:)`` is called and - /// will mark the passed `callbackToken` as cancelled. - /// - /// - Parameter callbackToken: The callback token. - @inlinable - internal func cancelCallback(callbackToken: WriteResult.CallbackToken) { - self._backing.storage.cancelProducer(callbackToken: callbackToken) - } + /// Writes new elements to the asynchronous stream. + /// + /// If there is a task consuming the stream and awaiting the next element then the task will get resumed with the + /// first element of the provided sequence. If the asynchronous stream already terminated then this method will throw an error + /// indicating the failure. + /// + /// - Parameter sequence: The elements to write to the asynchronous stream. + /// - Returns: The result that indicates if more elements should be produced at this time. + @inlinable + internal func write(contentsOf sequence: S) throws -> WriteResult + where Element == S.Element, S: Sequence { + try self._backing.storage.write(contentsOf: sequence) + } - /// Write new elements to the asynchronous stream and provide a callback which will be invoked once more elements should be produced. - /// - /// If there is a task consuming the stream and awaiting the next element then the task will get resumed with the - /// first element of the provided sequence. If the asynchronous stream already terminated then `onProduceMore` will be invoked with - /// a `Result.failure`. - /// - /// - Parameters: - /// - sequence: The elements to write to the asynchronous stream. - /// - onProduceMore: The callback which gets invoked once more elements should be produced. This callback might be - /// invoked during the call to ``write(contentsOf:onProduceMore:)``. - @inlinable - internal func write( - contentsOf sequence: S, - onProduceMore: @escaping @Sendable (Result) -> Void - ) where Element == S.Element, S: Sequence { - do { - let writeResult = try self.write(contentsOf: sequence) - - switch writeResult { - case .produceMore: - onProduceMore(Result.success(())) - - case .enqueueCallback(let callbackToken): - self.enqueueCallback(callbackToken: callbackToken, onProduceMore: onProduceMore) - } - } catch { - onProduceMore(.failure(error)) - } - } + /// Write the element to the asynchronous stream. + /// + /// If there is a task consuming the stream and awaiting the next element then the task will get resumed with the + /// provided element. If the asynchronous stream already terminated then this method will throw an error + /// indicating the failure. + /// + /// - Parameter element: The element to write to the asynchronous stream. + /// - Returns: The result that indicates if more elements should be produced at this time. + @inlinable + internal func write(_ element: Element) throws -> WriteResult { + try self._backing.storage.write(contentsOf: CollectionOfOne(element)) + } - /// Writes the element to the asynchronous stream. - /// - /// If there is a task consuming the stream and awaiting the next element then the task will get resumed with the - /// provided element. If the asynchronous stream already terminated then `onProduceMore` will be invoked with - /// a `Result.failure`. - /// - /// - Parameters: - /// - sequence: The element to write to the asynchronous stream. - /// - onProduceMore: The callback which gets invoked once more elements should be produced. This callback might be - /// invoked during the call to ``write(_:onProduceMore:)``. - @inlinable - internal func write( - _ element: Element, - onProduceMore: @escaping @Sendable (Result) -> Void - ) { - self.write(contentsOf: CollectionOfOne(element), onProduceMore: onProduceMore) - } + /// Enqueues a callback that will be invoked once more elements should be produced. + /// + /// Call this method after ``write(contentsOf:)`` or ``write(:)`` returned ``WriteResult/enqueueCallback(_:)``. + /// + /// - Important: Enqueueing the same token multiple times is not allowed. + /// + /// - Parameters: + /// - callbackToken: The callback token. + /// - onProduceMore: The callback which gets invoked once more elements should be produced. + @inlinable + internal func enqueueCallback( + callbackToken: WriteResult.CallbackToken, + onProduceMore: @escaping @Sendable (Result) -> Void + ) { + self._backing.storage.enqueueProducer( + callbackToken: callbackToken, + onProduceMore: onProduceMore + ) + } - /// Write new elements to the asynchronous stream. - /// - /// If there is a task consuming the stream and awaiting the next element then the task will get resumed with the - /// first element of the provided sequence. If the asynchronous stream already terminated then this method will throw an error - /// indicating the failure. - /// - /// This method returns once more elements should be produced. - /// - /// - Parameters: - /// - sequence: The elements to write to the asynchronous stream. - @inlinable - internal func write(contentsOf sequence: S) async throws - where Element == S.Element, S: Sequence { - let writeResult = try { try self.write(contentsOf: sequence) }() - - switch writeResult { - case .produceMore: - return - - case .enqueueCallback(let callbackToken): - try await withTaskCancellationHandler { - try await withCheckedThrowingContinuation { continuation in - self.enqueueCallback( - callbackToken: callbackToken, - onProduceMore: { result in - switch result { - case .success(): - continuation.resume(returning: ()) - case .failure(let error): - continuation.resume(throwing: error) - } - } - ) - } - } onCancel: { - self.cancelCallback(callbackToken: callbackToken) - } - } - } + /// Cancel an enqueued callback. + /// + /// Call this method to cancel a callback enqueued by the ``enqueueCallback(callbackToken:onProduceMore:)`` method. + /// + /// - Note: This methods supports being called before ``enqueueCallback(callbackToken:onProduceMore:)`` is called and + /// will mark the passed `callbackToken` as cancelled. + /// + /// - Parameter callbackToken: The callback token. + @inlinable + internal func cancelCallback(callbackToken: WriteResult.CallbackToken) { + self._backing.storage.cancelProducer(callbackToken: callbackToken) + } - /// Write new element to the asynchronous stream. - /// - /// If there is a task consuming the stream and awaiting the next element then the task will get resumed with the - /// provided element. If the asynchronous stream already terminated then this method will throw an error - /// indicating the failure. - /// - /// This method returns once more elements should be produced. - /// - /// - Parameters: - /// - sequence: The element to write to the asynchronous stream. - @inlinable - internal func write(_ element: Element) async throws { - try await self.write(contentsOf: CollectionOfOne(element)) - } + /// Write new elements to the asynchronous stream and provide a callback which will be invoked once more elements should be produced. + /// + /// If there is a task consuming the stream and awaiting the next element then the task will get resumed with the + /// first element of the provided sequence. If the asynchronous stream already terminated then `onProduceMore` will be invoked with + /// a `Result.failure`. + /// + /// - Parameters: + /// - sequence: The elements to write to the asynchronous stream. + /// - onProduceMore: The callback which gets invoked once more elements should be produced. This callback might be + /// invoked during the call to ``write(contentsOf:onProduceMore:)``. + @inlinable + internal func write( + contentsOf sequence: S, + onProduceMore: @escaping @Sendable (Result) -> Void + ) where Element == S.Element, S: Sequence { + do { + let writeResult = try self.write(contentsOf: sequence) + + switch writeResult { + case .produceMore: + onProduceMore(Result.success(())) + + case .enqueueCallback(let callbackToken): + self.enqueueCallback(callbackToken: callbackToken, onProduceMore: onProduceMore) + } + } catch { + onProduceMore(.failure(error)) + } + } - /// Write the elements of the asynchronous sequence to the asynchronous stream. - /// - /// This method returns once the provided asynchronous sequence or the the asynchronous stream finished. - /// - /// - Important: This method does not finish the source if consuming the upstream sequence terminated. - /// - /// - Parameters: - /// - sequence: The elements to write to the asynchronous stream. - @inlinable - internal func write(contentsOf sequence: S) async throws - where Element == S.Element, S: AsyncSequence { - for try await element in sequence { - try await self.write(contentsOf: CollectionOfOne(element)) - } - } + /// Writes the element to the asynchronous stream. + /// + /// If there is a task consuming the stream and awaiting the next element then the task will get resumed with the + /// provided element. If the asynchronous stream already terminated then `onProduceMore` will be invoked with + /// a `Result.failure`. + /// + /// - Parameters: + /// - sequence: The element to write to the asynchronous stream. + /// - onProduceMore: The callback which gets invoked once more elements should be produced. This callback might be + /// invoked during the call to ``write(_:onProduceMore:)``. + @inlinable + internal func write( + _ element: Element, + onProduceMore: @escaping @Sendable (Result) -> Void + ) { + self.write(contentsOf: CollectionOfOne(element), onProduceMore: onProduceMore) + } - /// Indicates that the production terminated. - /// - /// After all buffered elements are consumed the next iteration point will return `nil` or throw an error. - /// - /// Calling this function more than once has no effect. After calling finish, the stream enters a terminal state and doesn't accept - /// new elements. - /// - /// - Parameters: - /// - error: The error to throw, or `nil`, to finish normally. - @inlinable - internal func finish(throwing error: (any Error)?) { - self._backing.storage.finish(error) + /// Write new elements to the asynchronous stream. + /// + /// If there is a task consuming the stream and awaiting the next element then the task will get resumed with the + /// first element of the provided sequence. If the asynchronous stream already terminated then this method will throw an error + /// indicating the failure. + /// + /// This method returns once more elements should be produced. + /// + /// - Parameters: + /// - sequence: The elements to write to the asynchronous stream. + @inlinable + internal func write(contentsOf sequence: S) async throws + where Element == S.Element, S: Sequence { + let writeResult = try { try self.write(contentsOf: sequence) }() + + switch writeResult { + case .produceMore: + return + + case .enqueueCallback(let callbackToken): + try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { continuation in + self.enqueueCallback( + callbackToken: callbackToken, + onProduceMore: { result in + switch result { + case .success(): + continuation.resume(returning: ()) + case .failure(let error): + continuation.resume(throwing: error) + } + } + ) + } + } onCancel: { + self.cancelCallback(callbackToken: callbackToken) } + } } - /// Initializes a new ``BufferedStream`` and an ``BufferedStream/Source``. + /// Write new element to the asynchronous stream. + /// + /// If there is a task consuming the stream and awaiting the next element then the task will get resumed with the + /// provided element. If the asynchronous stream already terminated then this method will throw an error + /// indicating the failure. + /// + /// This method returns once more elements should be produced. /// /// - Parameters: - /// - elementType: The element type of the stream. - /// - failureType: The failure type of the stream. - /// - backPressureStrategy: The backpressure strategy that the stream should use. - /// - Returns: A tuple containing the stream and its source. The source should be passed to the - /// producer while the stream should be passed to the consumer. + /// - sequence: The element to write to the asynchronous stream. @inlinable - internal static func makeStream( - of elementType: Element.Type = Element.self, - throwing failureType: any Error.Type = (any Error).self, - backPressureStrategy: Source.BackPressureStrategy - ) -> (`Self`, Source) where any Error == any Error { - let storage = _BackPressuredStorage( - backPressureStrategy: backPressureStrategy._internalBackPressureStrategy - ) - let source = Source(storage: storage) + internal func write(_ element: Element) async throws { + try await self.write(contentsOf: CollectionOfOne(element)) + } - return (.init(storage: storage), source) + /// Write the elements of the asynchronous sequence to the asynchronous stream. + /// + /// This method returns once the provided asynchronous sequence or the the asynchronous stream finished. + /// + /// - Important: This method does not finish the source if consuming the upstream sequence terminated. + /// + /// - Parameters: + /// - sequence: The elements to write to the asynchronous stream. + @inlinable + internal func write(contentsOf sequence: S) async throws + where Element == S.Element, S: AsyncSequence { + for try await element in sequence { + try await self.write(contentsOf: CollectionOfOne(element)) + } } - @usableFromInline - init(storage: _BackPressuredStorage) { - self.implementation = .backpressured(.init(storage: storage)) + /// Indicates that the production terminated. + /// + /// After all buffered elements are consumed the next iteration point will return `nil` or throw an error. + /// + /// Calling this function more than once has no effect. After calling finish, the stream enters a terminal state and doesn't accept + /// new elements. + /// + /// - Parameters: + /// - error: The error to throw, or `nil`, to finish normally. + @inlinable + internal func finish(throwing error: (any Error)?) { + self._backing.storage.finish(error) } + } + + /// Initializes a new ``BufferedStream`` and an ``BufferedStream/Source``. + /// + /// - Parameters: + /// - elementType: The element type of the stream. + /// - failureType: The failure type of the stream. + /// - backPressureStrategy: The backpressure strategy that the stream should use. + /// - Returns: A tuple containing the stream and its source. The source should be passed to the + /// producer while the stream should be passed to the consumer. + @inlinable + internal static func makeStream( + of elementType: Element.Type = Element.self, + throwing failureType: any Error.Type = (any Error).self, + backPressureStrategy: Source.BackPressureStrategy + ) -> (`Self`, Source) where any Error == any Error { + let storage = _BackPressuredStorage( + backPressureStrategy: backPressureStrategy._internalBackPressureStrategy + ) + let source = Source(storage: storage) + + return (.init(storage: storage), source) + } + + @usableFromInline + init(storage: _BackPressuredStorage) { + self.implementation = .backpressured(.init(storage: storage)) + } } @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *) extension BufferedStream { + @usableFromInline + struct _WatermarkBackPressureStrategy: Sendable { + /// The low watermark where demand should start. @usableFromInline - struct _WatermarkBackPressureStrategy: Sendable { - /// The low watermark where demand should start. - @usableFromInline - let _low: Int - /// The high watermark where demand should be stopped. - @usableFromInline - let _high: Int - /// The current watermark. - @usableFromInline - private(set) var _current: Int - /// Function to compute the contribution to the water level for a given element. - @usableFromInline - let _waterLevelForElement: (@Sendable (Element) -> Int)? - - /// Initializes a new ``_WatermarkBackPressureStrategy``. - /// - /// - Parameters: - /// - low: The low watermark where demand should start. - /// - high: The high watermark where demand should be stopped. - /// - waterLevelForElement: Function to compute the contribution to the water level for a given element. - @usableFromInline - init(low: Int, high: Int, waterLevelForElement: (@Sendable (Element) -> Int)? = nil) { - precondition(low <= high) - self._low = low - self._high = high - self._current = 0 - self._waterLevelForElement = waterLevelForElement - } + let _low: Int + /// The high watermark where demand should be stopped. + @usableFromInline + let _high: Int + /// The current watermark. + @usableFromInline + private(set) var _current: Int + /// Function to compute the contribution to the water level for a given element. + @usableFromInline + let _waterLevelForElement: (@Sendable (Element) -> Int)? - @usableFromInline - mutating func didYield(elements: Deque.SubSequence) -> Bool { - if let waterLevelForElement = self._waterLevelForElement { - self._current += elements.reduce(0) { $0 + waterLevelForElement($1) } - } else { - self._current += elements.count - } - precondition(self._current >= 0, "Watermark below zero") - // We are demanding more until we reach the high watermark - return self._current < self._high - } + /// Initializes a new ``_WatermarkBackPressureStrategy``. + /// + /// - Parameters: + /// - low: The low watermark where demand should start. + /// - high: The high watermark where demand should be stopped. + /// - waterLevelForElement: Function to compute the contribution to the water level for a given element. + @usableFromInline + init(low: Int, high: Int, waterLevelForElement: (@Sendable (Element) -> Int)? = nil) { + precondition(low <= high) + self._low = low + self._high = high + self._current = 0 + self._waterLevelForElement = waterLevelForElement + } - @usableFromInline - mutating func didConsume(elements: Deque.SubSequence) -> Bool { - if let waterLevelForElement = self._waterLevelForElement { - self._current -= elements.reduce(0) { $0 + waterLevelForElement($1) } - } else { - self._current -= elements.count - } - precondition(self._current >= 0, "Watermark below zero") - // We start demanding again once we are below the low watermark - return self._current < self._low - } + @usableFromInline + mutating func didYield(elements: Deque.SubSequence) -> Bool { + if let waterLevelForElement = self._waterLevelForElement { + self._current += elements.reduce(0) { $0 + waterLevelForElement($1) } + } else { + self._current += elements.count + } + precondition(self._current >= 0, "Watermark below zero") + // We are demanding more until we reach the high watermark + return self._current < self._high + } - @usableFromInline - mutating func didConsume(element: Element) -> Bool { - if let waterLevelForElement = self._waterLevelForElement { - self._current -= waterLevelForElement(element) - } else { - self._current -= 1 - } - precondition(self._current >= 0, "Watermark below zero") - // We start demanding again once we are below the low watermark - return self._current < self._low - } + @usableFromInline + mutating func didConsume(elements: Deque.SubSequence) -> Bool { + if let waterLevelForElement = self._waterLevelForElement { + self._current -= elements.reduce(0) { $0 + waterLevelForElement($1) } + } else { + self._current -= elements.count + } + precondition(self._current >= 0, "Watermark below zero") + // We start demanding again once we are below the low watermark + return self._current < self._low } @usableFromInline - enum _InternalBackPressureStrategy: Sendable { - case watermark(_WatermarkBackPressureStrategy) - - @inlinable - mutating func didYield(elements: Deque.SubSequence) -> Bool { - switch self { - case .watermark(var strategy): - let result = strategy.didYield(elements: elements) - self = .watermark(strategy) - return result - } - } + mutating func didConsume(element: Element) -> Bool { + if let waterLevelForElement = self._waterLevelForElement { + self._current -= waterLevelForElement(element) + } else { + self._current -= 1 + } + precondition(self._current >= 0, "Watermark below zero") + // We start demanding again once we are below the low watermark + return self._current < self._low + } + } - @usableFromInline - mutating func didConsume(elements: Deque.SubSequence) -> Bool { - switch self { - case .watermark(var strategy): - let result = strategy.didConsume(elements: elements) - self = .watermark(strategy) - return result - } - } + @usableFromInline + enum _InternalBackPressureStrategy: Sendable { + case watermark(_WatermarkBackPressureStrategy) - @usableFromInline - mutating func didConsume(element: Element) -> Bool { - switch self { - case .watermark(var strategy): - let result = strategy.didConsume(element: element) - self = .watermark(strategy) - return result - } - } + @inlinable + mutating func didYield(elements: Deque.SubSequence) -> Bool { + switch self { + case .watermark(var strategy): + let result = strategy.didYield(elements: elements) + self = .watermark(strategy) + return result + } } + + @usableFromInline + mutating func didConsume(elements: Deque.SubSequence) -> Bool { + switch self { + case .watermark(var strategy): + let result = strategy.didConsume(elements: elements) + self = .watermark(strategy) + return result + } + } + + @usableFromInline + mutating func didConsume(element: Element) -> Bool { + switch self { + case .watermark(var strategy): + let result = strategy.didConsume(element: element) + self = .watermark(strategy) + return result + } + } + } } @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *) extension BufferedStream { - // We are unchecked Sendable since we are protecting our state with a lock. + // We are unchecked Sendable since we are protecting our state with a lock. + @usableFromInline + final class _BackPressuredStorage: Sendable { + /// The state machine @usableFromInline - final class _BackPressuredStorage: Sendable { - /// The state machine - @usableFromInline - let _stateMachine: _ManagedCriticalState<_StateMachine> + let _stateMachine: _ManagedCriticalState<_StateMachine> - @usableFromInline - var onTermination: (@Sendable () -> Void)? { - set { - self._stateMachine.withCriticalRegion { - $0._onTermination = newValue - } - } - get { - self._stateMachine.withCriticalRegion { - $0._onTermination - } - } + @usableFromInline + var onTermination: (@Sendable () -> Void)? { + set { + self._stateMachine.withCriticalRegion { + $0._onTermination = newValue } - - @usableFromInline - init( - backPressureStrategy: _InternalBackPressureStrategy - ) { - self._stateMachine = .init(.init(backPressureStrategy: backPressureStrategy)) + } + get { + self._stateMachine.withCriticalRegion { + $0._onTermination } + } + } - @inlinable - func sequenceDeinitialized() { - let action = self._stateMachine.withCriticalRegion { - $0.sequenceDeinitialized() - } + @usableFromInline + init( + backPressureStrategy: _InternalBackPressureStrategy + ) { + self._stateMachine = .init(.init(backPressureStrategy: backPressureStrategy)) + } - switch action { - case .callOnTermination(let onTermination): - onTermination?() + @inlinable + func sequenceDeinitialized() { + let action = self._stateMachine.withCriticalRegion { + $0.sequenceDeinitialized() + } - case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): - for producerContinuation in producerContinuations { - producerContinuation(.failure(AlreadyFinishedError())) - } - onTermination?() + switch action { + case .callOnTermination(let onTermination): + onTermination?() - case .none: - break - } + case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): + for producerContinuation in producerContinuations { + producerContinuation(.failure(AlreadyFinishedError())) } + onTermination?() - @inlinable - func iteratorInitialized() { - self._stateMachine.withCriticalRegion { - $0.iteratorInitialized() - } - } + case .none: + break + } + } - @inlinable - func iteratorDeinitialized() { - let action = self._stateMachine.withCriticalRegion { - $0.iteratorDeinitialized() - } + @inlinable + func iteratorInitialized() { + self._stateMachine.withCriticalRegion { + $0.iteratorInitialized() + } + } - switch action { - case .callOnTermination(let onTermination): - onTermination?() + @inlinable + func iteratorDeinitialized() { + let action = self._stateMachine.withCriticalRegion { + $0.iteratorDeinitialized() + } - case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): - for producerContinuation in producerContinuations { - producerContinuation(.failure(AlreadyFinishedError())) - } - onTermination?() + switch action { + case .callOnTermination(let onTermination): + onTermination?() - case .none: - break - } + case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): + for producerContinuation in producerContinuations { + producerContinuation(.failure(AlreadyFinishedError())) } + onTermination?() - @inlinable - func sourceDeinitialized() { - let action = self._stateMachine.withCriticalRegion { - $0.sourceDeinitialized() - } + case .none: + break + } + } - switch action { - case .callOnTermination(let onTermination): - onTermination?() - - case .failProducersAndCallOnTermination( - let consumer, - let producerContinuations, - let onTermination - ): - consumer?.resume(returning: nil) - for producerContinuation in producerContinuations { - producerContinuation(.failure(AlreadyFinishedError())) - } - onTermination?() + @inlinable + func sourceDeinitialized() { + let action = self._stateMachine.withCriticalRegion { + $0.sourceDeinitialized() + } + + switch action { + case .callOnTermination(let onTermination): + onTermination?() + + case .failProducersAndCallOnTermination( + let consumer, + let producerContinuations, + let onTermination + ): + consumer?.resume(returning: nil) + for producerContinuation in producerContinuations { + producerContinuation(.failure(AlreadyFinishedError())) + } + onTermination?() + + case .failProducers(let producerContinuations): + for producerContinuation in producerContinuations { + producerContinuation(.failure(AlreadyFinishedError())) + } + + case .none: + break + } + } - case .failProducers(let producerContinuations): - for producerContinuation in producerContinuations { - producerContinuation(.failure(AlreadyFinishedError())) - } + @inlinable + func write( + contentsOf sequence: some Sequence + ) throws -> Source.WriteResult { + let action = self._stateMachine.withCriticalRegion { + return $0.write(sequence) + } + + switch action { + case .returnProduceMore: + return .produceMore + + case .returnEnqueue(let callbackToken): + return .enqueueCallback(callbackToken) + + case .resumeConsumerAndReturnProduceMore(let continuation, let element): + continuation.resume(returning: element) + return .produceMore + + case .resumeConsumerAndReturnEnqueue(let continuation, let element, let callbackToken): + continuation.resume(returning: element) + return .enqueueCallback(callbackToken) + + case .throwFinishedError: + throw AlreadyFinishedError() + } + } - case .none: - break - } - } + @inlinable + func enqueueProducer( + callbackToken: Source.WriteResult.CallbackToken, + onProduceMore: @escaping @Sendable (Result) -> Void + ) { + let action = self._stateMachine.withCriticalRegion { + $0.enqueueProducer(callbackToken: callbackToken, onProduceMore: onProduceMore) + } + + switch action { + case .resumeProducer(let onProduceMore): + onProduceMore(Result.success(())) + + case .resumeProducerWithError(let onProduceMore, let error): + onProduceMore(Result.failure(error)) + + case .none: + break + } + } - @inlinable - func write( - contentsOf sequence: some Sequence - ) throws -> Source.WriteResult { - let action = self._stateMachine.withCriticalRegion { - return $0.write(sequence) - } + @inlinable + func cancelProducer(callbackToken: Source.WriteResult.CallbackToken) { + let action = self._stateMachine.withCriticalRegion { + $0.cancelProducer(callbackToken: callbackToken) + } + + switch action { + case .resumeProducerWithCancellationError(let onProduceMore): + onProduceMore(Result.failure(CancellationError())) + + case .none: + break + } + } - switch action { - case .returnProduceMore: - return .produceMore + @inlinable + func finish(_ failure: (any Error)?) { + let action = self._stateMachine.withCriticalRegion { + $0.finish(failure) + } - case .returnEnqueue(let callbackToken): - return .enqueueCallback(callbackToken) + switch action { + case .callOnTermination(let onTermination): + onTermination?() - case .resumeConsumerAndReturnProduceMore(let continuation, let element): - continuation.resume(returning: element) - return .produceMore + case .resumeConsumerAndCallOnTermination( + let consumerContinuation, + let failure, + let onTermination + ): + switch failure { + case .some(let error): + consumerContinuation.resume(throwing: error) + case .none: + consumerContinuation.resume(returning: nil) + } - case .resumeConsumerAndReturnEnqueue(let continuation, let element, let callbackToken): - continuation.resume(returning: element) - return .enqueueCallback(callbackToken) + onTermination?() - case .throwFinishedError: - throw AlreadyFinishedError() - } + case .resumeProducers(let producerContinuations): + for producerContinuation in producerContinuations { + producerContinuation(.failure(AlreadyFinishedError())) } - @inlinable - func enqueueProducer( - callbackToken: Source.WriteResult.CallbackToken, - onProduceMore: @escaping @Sendable (Result) -> Void - ) { - let action = self._stateMachine.withCriticalRegion { - $0.enqueueProducer(callbackToken: callbackToken, onProduceMore: onProduceMore) - } + case .none: + break + } + } - switch action { - case .resumeProducer(let onProduceMore): - onProduceMore(Result.success(())) + @inlinable + func next() async throws -> Element? { + let action = self._stateMachine.withCriticalRegion { + $0.next() + } - case .resumeProducerWithError(let onProduceMore, let error): - onProduceMore(Result.failure(error)) + switch action { + case .returnElement(let element): + return element - case .none: - break - } + case .returnElementAndResumeProducers(let element, let producerContinuations): + for producerContinuation in producerContinuations { + producerContinuation(Result.success(())) } - @inlinable - func cancelProducer(callbackToken: Source.WriteResult.CallbackToken) { - let action = self._stateMachine.withCriticalRegion { - $0.cancelProducer(callbackToken: callbackToken) - } + return element - switch action { - case .resumeProducerWithCancellationError(let onProduceMore): - onProduceMore(Result.failure(CancellationError())) + case .returnErrorAndCallOnTermination(let failure, let onTermination): + onTermination?() + switch failure { + case .some(let error): + throw error - case .none: - break - } + case .none: + return nil } - @inlinable - func finish(_ failure: (any Error)?) { - let action = self._stateMachine.withCriticalRegion { - $0.finish(failure) - } + case .returnNil: + return nil - switch action { - case .callOnTermination(let onTermination): - onTermination?() - - case .resumeConsumerAndCallOnTermination( - let consumerContinuation, - let failure, - let onTermination - ): - switch failure { - case .some(let error): - consumerContinuation.resume(throwing: error) - case .none: - consumerContinuation.resume(returning: nil) - } + case .suspendTask: + return try await self.suspendNext() + } + } - onTermination?() + @inlinable + func suspendNext() async throws -> Element? { + return try await withTaskCancellationHandler { + return try await withCheckedThrowingContinuation { continuation in + let action = self._stateMachine.withCriticalRegion { + $0.suspendNext(continuation: continuation) + } + + switch action { + case .resumeConsumerWithElement(let continuation, let element): + continuation.resume(returning: element) + + case .resumeConsumerWithElementAndProducers( + let continuation, + let element, + let producerContinuations + ): + continuation.resume(returning: element) + for producerContinuation in producerContinuations { + producerContinuation(Result.success(())) + } - case .resumeProducers(let producerContinuations): - for producerContinuation in producerContinuations { - producerContinuation(.failure(AlreadyFinishedError())) - } + case .resumeConsumerWithErrorAndCallOnTermination( + let continuation, + let failure, + let onTermination + ): + switch failure { + case .some(let error): + continuation.resume(throwing: error) case .none: - break - } - } - - @inlinable - func next() async throws -> Element? { - let action = self._stateMachine.withCriticalRegion { - $0.next() + continuation.resume(returning: nil) } + onTermination?() - switch action { - case .returnElement(let element): - return element + case .resumeConsumerWithNil(let continuation): + continuation.resume(returning: nil) - case .returnElementAndResumeProducers(let element, let producerContinuations): - for producerContinuation in producerContinuations { - producerContinuation(Result.success(())) - } - - return element - - case .returnErrorAndCallOnTermination(let failure, let onTermination): - onTermination?() - switch failure { - case .some(let error): - throw error - - case .none: - return nil - } - - case .returnNil: - return nil - - case .suspendTask: - return try await self.suspendNext() - } + case .none: + break + } + } + } onCancel: { + let action = self._stateMachine.withCriticalRegion { + $0.cancelNext() } - @inlinable - func suspendNext() async throws -> Element? { - return try await withTaskCancellationHandler { - return try await withCheckedThrowingContinuation { continuation in - let action = self._stateMachine.withCriticalRegion { - $0.suspendNext(continuation: continuation) - } - - switch action { - case .resumeConsumerWithElement(let continuation, let element): - continuation.resume(returning: element) - - case .resumeConsumerWithElementAndProducers( - let continuation, - let element, - let producerContinuations - ): - continuation.resume(returning: element) - for producerContinuation in producerContinuations { - producerContinuation(Result.success(())) - } - - case .resumeConsumerWithErrorAndCallOnTermination( - let continuation, - let failure, - let onTermination - ): - switch failure { - case .some(let error): - continuation.resume(throwing: error) - - case .none: - continuation.resume(returning: nil) - } - onTermination?() - - case .resumeConsumerWithNil(let continuation): - continuation.resume(returning: nil) - - case .none: - break - } - } - } onCancel: { - let action = self._stateMachine.withCriticalRegion { - $0.cancelNext() - } + switch action { + case .resumeConsumerWithCancellationErrorAndCallOnTermination( + let continuation, + let onTermination + ): + continuation.resume(throwing: CancellationError()) + onTermination?() - switch action { - case .resumeConsumerWithCancellationErrorAndCallOnTermination( - let continuation, - let onTermination - ): - continuation.resume(throwing: CancellationError()) - onTermination?() - - case .failProducersAndCallOnTermination( - let producerContinuations, - let onTermination - ): - for producerContinuation in producerContinuations { - producerContinuation(.failure(AlreadyFinishedError())) - } - onTermination?() - - case .none: - break - } - } + case .failProducersAndCallOnTermination( + let producerContinuations, + let onTermination + ): + for producerContinuation in producerContinuations { + producerContinuation(.failure(AlreadyFinishedError())) + } + onTermination?() + + case .none: + break } + } } + } } @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *) extension BufferedStream { - /// The state machine of the backpressured async stream. + /// The state machine of the backpressured async stream. + @usableFromInline + struct _StateMachine { @usableFromInline - struct _StateMachine { + enum _State { + @usableFromInline + struct Initial { + /// The backpressure strategy. @usableFromInline - enum _State { - @usableFromInline - struct Initial { - /// The backpressure strategy. - @usableFromInline - var backPressureStrategy: _InternalBackPressureStrategy - /// Indicates if the iterator was initialized. - @usableFromInline - var iteratorInitialized: Bool - /// The onTermination callback. - @usableFromInline - var onTermination: (@Sendable () -> Void)? - - @usableFromInline - init( - backPressureStrategy: _InternalBackPressureStrategy, - iteratorInitialized: Bool, - onTermination: (@Sendable () -> Void)? = nil - ) { - self.backPressureStrategy = backPressureStrategy - self.iteratorInitialized = iteratorInitialized - self.onTermination = onTermination - } - } - - @usableFromInline - struct Streaming { - /// The backpressure strategy. - @usableFromInline - var backPressureStrategy: _InternalBackPressureStrategy - /// Indicates if the iterator was initialized. - @usableFromInline - var iteratorInitialized: Bool - /// The onTermination callback. - @usableFromInline - var onTermination: (@Sendable () -> Void)? - /// The buffer of elements. - @usableFromInline - var buffer: Deque - /// The optional consumer continuation. - @usableFromInline - var consumerContinuation: CheckedContinuation? - /// The producer continuations. - @usableFromInline - var producerContinuations: Deque<(UInt, (Result) -> Void)> - /// The producers that have been cancelled. - @usableFromInline - var cancelledAsyncProducers: Deque - /// Indicates if we currently have outstanding demand. - @usableFromInline - var hasOutstandingDemand: Bool - - @usableFromInline - init( - backPressureStrategy: _InternalBackPressureStrategy, - iteratorInitialized: Bool, - onTermination: (@Sendable () -> Void)? = nil, - buffer: Deque, - consumerContinuation: CheckedContinuation? = nil, - producerContinuations: Deque<(UInt, (Result) -> Void)>, - cancelledAsyncProducers: Deque, - hasOutstandingDemand: Bool - ) { - self.backPressureStrategy = backPressureStrategy - self.iteratorInitialized = iteratorInitialized - self.onTermination = onTermination - self.buffer = buffer - self.consumerContinuation = consumerContinuation - self.producerContinuations = producerContinuations - self.cancelledAsyncProducers = cancelledAsyncProducers - self.hasOutstandingDemand = hasOutstandingDemand - } - } - - @usableFromInline - struct SourceFinished { - /// Indicates if the iterator was initialized. - @usableFromInline - var iteratorInitialized: Bool - /// The buffer of elements. - @usableFromInline - var buffer: Deque - /// The failure that should be thrown after the last element has been consumed. - @usableFromInline - var failure: (any Error)? - /// The onTermination callback. - @usableFromInline - var onTermination: (@Sendable () -> Void)? - - @usableFromInline - init( - iteratorInitialized: Bool, - buffer: Deque, - failure: (any Error)? = nil, - onTermination: (@Sendable () -> Void)? - ) { - self.iteratorInitialized = iteratorInitialized - self.buffer = buffer - self.failure = failure - self.onTermination = onTermination - } - } - - case initial(Initial) - /// The state once either any element was yielded or `next()` was called. - case streaming(Streaming) - /// The state once the underlying source signalled that it is finished. - case sourceFinished(SourceFinished) - - /// The state once there can be no outstanding demand. This can happen if: - /// 1. The iterator was deinited - /// 2. The underlying source finished and all buffered elements have been consumed - case finished(iteratorInitialized: Bool) - - /// An intermediate state to avoid CoWs. - case modify - } - - /// The state machine's current state. + var backPressureStrategy: _InternalBackPressureStrategy + /// Indicates if the iterator was initialized. @usableFromInline - var _state: _State - - // The ID used for the next CallbackToken. + var iteratorInitialized: Bool + /// The onTermination callback. @usableFromInline - var nextCallbackTokenID: UInt = 0 + var onTermination: (@Sendable () -> Void)? - @inlinable - var _onTermination: (@Sendable () -> Void)? { - set { - switch self._state { - case .initial(var initial): - initial.onTermination = newValue - self._state = .initial(initial) - - case .streaming(var streaming): - streaming.onTermination = newValue - self._state = .streaming(streaming) - - case .sourceFinished(var sourceFinished): - sourceFinished.onTermination = newValue - self._state = .sourceFinished(sourceFinished) - - case .finished: - break - - case .modify: - fatalError("AsyncStream internal inconsistency") - } - } - get { - switch self._state { - case .initial(let initial): - return initial.onTermination - - case .streaming(let streaming): - return streaming.onTermination - - case .sourceFinished(let sourceFinished): - return sourceFinished.onTermination - - case .finished: - return nil - - case .modify: - fatalError("AsyncStream internal inconsistency") - } - } - } - - /// Initializes a new `StateMachine`. - /// - /// We are passing and holding the back-pressure strategy here because - /// it is a customizable extension of the state machine. - /// - /// - Parameter backPressureStrategy: The back-pressure strategy. @usableFromInline init( - backPressureStrategy: _InternalBackPressureStrategy + backPressureStrategy: _InternalBackPressureStrategy, + iteratorInitialized: Bool, + onTermination: (@Sendable () -> Void)? = nil ) { - self._state = .initial( - .init( - backPressureStrategy: backPressureStrategy, - iteratorInitialized: false - ) - ) + self.backPressureStrategy = backPressureStrategy + self.iteratorInitialized = iteratorInitialized + self.onTermination = onTermination } + } - /// Generates the next callback token. - @inlinable - mutating func nextCallbackToken() -> Source.WriteResult.CallbackToken { - let id = self.nextCallbackTokenID - self.nextCallbackTokenID += 1 - return .init(id: id) - } - - /// Actions returned by `sequenceDeinitialized()`. + @usableFromInline + struct Streaming { + /// The backpressure strategy. @usableFromInline - enum SequenceDeinitializedAction { - /// Indicates that `onTermination` should be called. - case callOnTermination((@Sendable () -> Void)?) - /// Indicates that all producers should be failed and `onTermination` should be called. - case failProducersAndCallOnTermination( - [(Result) -> Void], - (@Sendable () -> Void)? - ) - } + var backPressureStrategy: _InternalBackPressureStrategy + /// Indicates if the iterator was initialized. + @usableFromInline + var iteratorInitialized: Bool + /// The onTermination callback. + @usableFromInline + var onTermination: (@Sendable () -> Void)? + /// The buffer of elements. + @usableFromInline + var buffer: Deque + /// The optional consumer continuation. + @usableFromInline + var consumerContinuation: CheckedContinuation? + /// The producer continuations. + @usableFromInline + var producerContinuations: Deque<(UInt, (Result) -> Void)> + /// The producers that have been cancelled. + @usableFromInline + var cancelledAsyncProducers: Deque + /// Indicates if we currently have outstanding demand. + @usableFromInline + var hasOutstandingDemand: Bool - @inlinable - mutating func sequenceDeinitialized() -> SequenceDeinitializedAction? { - switch self._state { - case .initial(let initial): - if initial.iteratorInitialized { - // An iterator was created and we deinited the sequence. - // This is an expected pattern and we just continue on normal. - return .none - } else { - // No iterator was created so we can transition to finished right away. - self._state = .finished(iteratorInitialized: false) - - return .callOnTermination(initial.onTermination) - } + @usableFromInline + init( + backPressureStrategy: _InternalBackPressureStrategy, + iteratorInitialized: Bool, + onTermination: (@Sendable () -> Void)? = nil, + buffer: Deque, + consumerContinuation: CheckedContinuation? = nil, + producerContinuations: Deque<(UInt, (Result) -> Void)>, + cancelledAsyncProducers: Deque, + hasOutstandingDemand: Bool + ) { + self.backPressureStrategy = backPressureStrategy + self.iteratorInitialized = iteratorInitialized + self.onTermination = onTermination + self.buffer = buffer + self.consumerContinuation = consumerContinuation + self.producerContinuations = producerContinuations + self.cancelledAsyncProducers = cancelledAsyncProducers + self.hasOutstandingDemand = hasOutstandingDemand + } + } + + @usableFromInline + struct SourceFinished { + /// Indicates if the iterator was initialized. + @usableFromInline + var iteratorInitialized: Bool + /// The buffer of elements. + @usableFromInline + var buffer: Deque + /// The failure that should be thrown after the last element has been consumed. + @usableFromInline + var failure: (any Error)? + /// The onTermination callback. + @usableFromInline + var onTermination: (@Sendable () -> Void)? - case .streaming(let streaming): - if streaming.iteratorInitialized { - // An iterator was created and we deinited the sequence. - // This is an expected pattern and we just continue on normal. - return .none - } else { - // No iterator was created so we can transition to finished right away. - self._state = .finished(iteratorInitialized: false) - - return .failProducersAndCallOnTermination( - Array(streaming.producerContinuations.map { $0.1 }), - streaming.onTermination - ) - } + @usableFromInline + init( + iteratorInitialized: Bool, + buffer: Deque, + failure: (any Error)? = nil, + onTermination: (@Sendable () -> Void)? + ) { + self.iteratorInitialized = iteratorInitialized + self.buffer = buffer + self.failure = failure + self.onTermination = onTermination + } + } + + case initial(Initial) + /// The state once either any element was yielded or `next()` was called. + case streaming(Streaming) + /// The state once the underlying source signalled that it is finished. + case sourceFinished(SourceFinished) + + /// The state once there can be no outstanding demand. This can happen if: + /// 1. The iterator was deinited + /// 2. The underlying source finished and all buffered elements have been consumed + case finished(iteratorInitialized: Bool) + + /// An intermediate state to avoid CoWs. + case modify + } - case .sourceFinished(let sourceFinished): - if sourceFinished.iteratorInitialized { - // An iterator was created and we deinited the sequence. - // This is an expected pattern and we just continue on normal. - return .none - } else { - // No iterator was created so we can transition to finished right away. - self._state = .finished(iteratorInitialized: false) + /// The state machine's current state. + @usableFromInline + var _state: _State - return .callOnTermination(sourceFinished.onTermination) - } + // The ID used for the next CallbackToken. + @usableFromInline + var nextCallbackTokenID: UInt = 0 - case .finished: - // We are already finished so there is nothing left to clean up. - // This is just the references dropping afterwards. - return .none + @inlinable + var _onTermination: (@Sendable () -> Void)? { + set { + switch self._state { + case .initial(var initial): + initial.onTermination = newValue + self._state = .initial(initial) - case .modify: - fatalError("AsyncStream internal inconsistency") - } - } + case .streaming(var streaming): + streaming.onTermination = newValue + self._state = .streaming(streaming) - @inlinable - mutating func iteratorInitialized() { - switch self._state { - case .initial(var initial): - if initial.iteratorInitialized { - // Our sequence is a unicast sequence and does not support multiple AsyncIterator's - fatalError("Only a single AsyncIterator can be created") - } else { - // The first and only iterator was initialized. - initial.iteratorInitialized = true - self._state = .initial(initial) - } + case .sourceFinished(var sourceFinished): + sourceFinished.onTermination = newValue + self._state = .sourceFinished(sourceFinished) - case .streaming(var streaming): - if streaming.iteratorInitialized { - // Our sequence is a unicast sequence and does not support multiple AsyncIterator's - fatalError("Only a single AsyncIterator can be created") - } else { - // The first and only iterator was initialized. - streaming.iteratorInitialized = true - self._state = .streaming(streaming) - } + case .finished: + break - case .sourceFinished(var sourceFinished): - if sourceFinished.iteratorInitialized { - // Our sequence is a unicast sequence and does not support multiple AsyncIterator's - fatalError("Only a single AsyncIterator can be created") - } else { - // The first and only iterator was initialized. - sourceFinished.iteratorInitialized = true - self._state = .sourceFinished(sourceFinished) - } + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } + get { + switch self._state { + case .initial(let initial): + return initial.onTermination - case .finished(iteratorInitialized: true): - // Our sequence is a unicast sequence and does not support multiple AsyncIterator's - fatalError("Only a single AsyncIterator can be created") + case .streaming(let streaming): + return streaming.onTermination - case .finished(iteratorInitialized: false): - // It is strange that an iterator is created after we are finished - // but it can definitely happen, e.g. - // Sequence.init -> source.finish -> sequence.makeAsyncIterator - self._state = .finished(iteratorInitialized: true) + case .sourceFinished(let sourceFinished): + return sourceFinished.onTermination - case .modify: - fatalError("AsyncStream internal inconsistency") - } - } + case .finished: + return nil - /// Actions returned by `iteratorDeinitialized()`. - @usableFromInline - enum IteratorDeinitializedAction { - /// Indicates that `onTermination` should be called. - case callOnTermination((@Sendable () -> Void)?) - /// Indicates that all producers should be failed and `onTermination` should be called. - case failProducersAndCallOnTermination( - [(Result) -> Void], - (@Sendable () -> Void)? - ) + case .modify: + fatalError("AsyncStream internal inconsistency") } + } + } - @inlinable - mutating func iteratorDeinitialized() -> IteratorDeinitializedAction? { - switch self._state { - case .initial(let initial): - if initial.iteratorInitialized { - // An iterator was created and deinited. Since we only support - // a single iterator we can now transition to finish. - self._state = .finished(iteratorInitialized: true) - return .callOnTermination(initial.onTermination) - } else { - // An iterator needs to be initialized before it can be deinitialized. - fatalError("AsyncStream internal inconsistency") - } - - case .streaming(let streaming): - if streaming.iteratorInitialized { - // An iterator was created and deinited. Since we only support - // a single iterator we can now transition to finish. - self._state = .finished(iteratorInitialized: true) - - return .failProducersAndCallOnTermination( - Array(streaming.producerContinuations.map { $0.1 }), - streaming.onTermination - ) - } else { - // An iterator needs to be initialized before it can be deinitialized. - fatalError("AsyncStream internal inconsistency") - } + /// Initializes a new `StateMachine`. + /// + /// We are passing and holding the back-pressure strategy here because + /// it is a customizable extension of the state machine. + /// + /// - Parameter backPressureStrategy: The back-pressure strategy. + @usableFromInline + init( + backPressureStrategy: _InternalBackPressureStrategy + ) { + self._state = .initial( + .init( + backPressureStrategy: backPressureStrategy, + iteratorInitialized: false + ) + ) + } - case .sourceFinished(let sourceFinished): - if sourceFinished.iteratorInitialized { - // An iterator was created and deinited. Since we only support - // a single iterator we can now transition to finish. - self._state = .finished(iteratorInitialized: true) - return .callOnTermination(sourceFinished.onTermination) - } else { - // An iterator needs to be initialized before it can be deinitialized. - fatalError("AsyncStream internal inconsistency") - } + /// Generates the next callback token. + @inlinable + mutating func nextCallbackToken() -> Source.WriteResult.CallbackToken { + let id = self.nextCallbackTokenID + self.nextCallbackTokenID += 1 + return .init(id: id) + } - case .finished: - // We are already finished so there is nothing left to clean up. - // This is just the references dropping afterwards. - return .none + /// Actions returned by `sequenceDeinitialized()`. + @usableFromInline + enum SequenceDeinitializedAction { + /// Indicates that `onTermination` should be called. + case callOnTermination((@Sendable () -> Void)?) + /// Indicates that all producers should be failed and `onTermination` should be called. + case failProducersAndCallOnTermination( + [(Result) -> Void], + (@Sendable () -> Void)? + ) + } - case .modify: - fatalError("AsyncStream internal inconsistency") - } - } + @inlinable + mutating func sequenceDeinitialized() -> SequenceDeinitializedAction? { + switch self._state { + case .initial(let initial): + if initial.iteratorInitialized { + // An iterator was created and we deinited the sequence. + // This is an expected pattern and we just continue on normal. + return .none + } else { + // No iterator was created so we can transition to finished right away. + self._state = .finished(iteratorInitialized: false) + + return .callOnTermination(initial.onTermination) + } + + case .streaming(let streaming): + if streaming.iteratorInitialized { + // An iterator was created and we deinited the sequence. + // This is an expected pattern and we just continue on normal. + return .none + } else { + // No iterator was created so we can transition to finished right away. + self._state = .finished(iteratorInitialized: false) + + return .failProducersAndCallOnTermination( + Array(streaming.producerContinuations.map { $0.1 }), + streaming.onTermination + ) + } + + case .sourceFinished(let sourceFinished): + if sourceFinished.iteratorInitialized { + // An iterator was created and we deinited the sequence. + // This is an expected pattern and we just continue on normal. + return .none + } else { + // No iterator was created so we can transition to finished right away. + self._state = .finished(iteratorInitialized: false) + + return .callOnTermination(sourceFinished.onTermination) + } + + case .finished: + // We are already finished so there is nothing left to clean up. + // This is just the references dropping afterwards. + return .none + + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } - /// Actions returned by `sourceDeinitialized()`. - @usableFromInline - enum SourceDeinitializedAction { - /// Indicates that `onTermination` should be called. - case callOnTermination((() -> Void)?) - /// Indicates that all producers should be failed and `onTermination` should be called. - case failProducersAndCallOnTermination( - CheckedContinuation?, - [(Result) -> Void], - (@Sendable () -> Void)? - ) - /// Indicates that all producers should be failed. - case failProducers([(Result) -> Void]) - } + @inlinable + mutating func iteratorInitialized() { + switch self._state { + case .initial(var initial): + if initial.iteratorInitialized { + // Our sequence is a unicast sequence and does not support multiple AsyncIterator's + fatalError("Only a single AsyncIterator can be created") + } else { + // The first and only iterator was initialized. + initial.iteratorInitialized = true + self._state = .initial(initial) + } + + case .streaming(var streaming): + if streaming.iteratorInitialized { + // Our sequence is a unicast sequence and does not support multiple AsyncIterator's + fatalError("Only a single AsyncIterator can be created") + } else { + // The first and only iterator was initialized. + streaming.iteratorInitialized = true + self._state = .streaming(streaming) + } + + case .sourceFinished(var sourceFinished): + if sourceFinished.iteratorInitialized { + // Our sequence is a unicast sequence and does not support multiple AsyncIterator's + fatalError("Only a single AsyncIterator can be created") + } else { + // The first and only iterator was initialized. + sourceFinished.iteratorInitialized = true + self._state = .sourceFinished(sourceFinished) + } + + case .finished(iteratorInitialized: true): + // Our sequence is a unicast sequence and does not support multiple AsyncIterator's + fatalError("Only a single AsyncIterator can be created") + + case .finished(iteratorInitialized: false): + // It is strange that an iterator is created after we are finished + // but it can definitely happen, e.g. + // Sequence.init -> source.finish -> sequence.makeAsyncIterator + self._state = .finished(iteratorInitialized: true) + + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } - @inlinable - mutating func sourceDeinitialized() -> SourceDeinitializedAction? { - switch self._state { - case .initial(let initial): - // The source got deinited before anything was written - self._state = .finished(iteratorInitialized: initial.iteratorInitialized) - return .callOnTermination(initial.onTermination) - - case .streaming(let streaming): - if streaming.buffer.isEmpty { - // We can transition to finished right away since the buffer is empty now - self._state = .finished(iteratorInitialized: streaming.iteratorInitialized) - - return .failProducersAndCallOnTermination( - streaming.consumerContinuation, - Array(streaming.producerContinuations.map { $0.1 }), - streaming.onTermination - ) - } else { - // The continuation must be `nil` if the buffer has elements - precondition(streaming.consumerContinuation == nil) - - self._state = .sourceFinished( - .init( - iteratorInitialized: streaming.iteratorInitialized, - buffer: streaming.buffer, - failure: nil, - onTermination: streaming.onTermination - ) - ) - - return .failProducers( - Array(streaming.producerContinuations.map { $0.1 }) - ) - } + /// Actions returned by `iteratorDeinitialized()`. + @usableFromInline + enum IteratorDeinitializedAction { + /// Indicates that `onTermination` should be called. + case callOnTermination((@Sendable () -> Void)?) + /// Indicates that all producers should be failed and `onTermination` should be called. + case failProducersAndCallOnTermination( + [(Result) -> Void], + (@Sendable () -> Void)? + ) + } - case .sourceFinished, .finished: - // This is normal and we just have to tolerate it - return .none + @inlinable + mutating func iteratorDeinitialized() -> IteratorDeinitializedAction? { + switch self._state { + case .initial(let initial): + if initial.iteratorInitialized { + // An iterator was created and deinited. Since we only support + // a single iterator we can now transition to finish. + self._state = .finished(iteratorInitialized: true) + return .callOnTermination(initial.onTermination) + } else { + // An iterator needs to be initialized before it can be deinitialized. + fatalError("AsyncStream internal inconsistency") + } + + case .streaming(let streaming): + if streaming.iteratorInitialized { + // An iterator was created and deinited. Since we only support + // a single iterator we can now transition to finish. + self._state = .finished(iteratorInitialized: true) + + return .failProducersAndCallOnTermination( + Array(streaming.producerContinuations.map { $0.1 }), + streaming.onTermination + ) + } else { + // An iterator needs to be initialized before it can be deinitialized. + fatalError("AsyncStream internal inconsistency") + } + + case .sourceFinished(let sourceFinished): + if sourceFinished.iteratorInitialized { + // An iterator was created and deinited. Since we only support + // a single iterator we can now transition to finish. + self._state = .finished(iteratorInitialized: true) + return .callOnTermination(sourceFinished.onTermination) + } else { + // An iterator needs to be initialized before it can be deinitialized. + fatalError("AsyncStream internal inconsistency") + } + + case .finished: + // We are already finished so there is nothing left to clean up. + // This is just the references dropping afterwards. + return .none + + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } - case .modify: - fatalError("AsyncStream internal inconsistency") - } - } + /// Actions returned by `sourceDeinitialized()`. + @usableFromInline + enum SourceDeinitializedAction { + /// Indicates that `onTermination` should be called. + case callOnTermination((() -> Void)?) + /// Indicates that all producers should be failed and `onTermination` should be called. + case failProducersAndCallOnTermination( + CheckedContinuation?, + [(Result) -> Void], + (@Sendable () -> Void)? + ) + /// Indicates that all producers should be failed. + case failProducers([(Result) -> Void]) + } - /// Actions returned by `write()`. - @usableFromInline - enum WriteAction { - /// Indicates that the producer should be notified to produce more. - case returnProduceMore - /// Indicates that the producer should be suspended to stop producing. - case returnEnqueue( - callbackToken: Source.WriteResult.CallbackToken - ) - /// Indicates that the consumer should be resumed and the producer should be notified to produce more. - case resumeConsumerAndReturnProduceMore( - continuation: CheckedContinuation, - element: Element - ) - /// Indicates that the consumer should be resumed and the producer should be suspended. - case resumeConsumerAndReturnEnqueue( - continuation: CheckedContinuation, - element: Element, - callbackToken: Source.WriteResult.CallbackToken + @inlinable + mutating func sourceDeinitialized() -> SourceDeinitializedAction? { + switch self._state { + case .initial(let initial): + // The source got deinited before anything was written + self._state = .finished(iteratorInitialized: initial.iteratorInitialized) + return .callOnTermination(initial.onTermination) + + case .streaming(let streaming): + if streaming.buffer.isEmpty { + // We can transition to finished right away since the buffer is empty now + self._state = .finished(iteratorInitialized: streaming.iteratorInitialized) + + return .failProducersAndCallOnTermination( + streaming.consumerContinuation, + Array(streaming.producerContinuations.map { $0.1 }), + streaming.onTermination + ) + } else { + // The continuation must be `nil` if the buffer has elements + precondition(streaming.consumerContinuation == nil) + + self._state = .sourceFinished( + .init( + iteratorInitialized: streaming.iteratorInitialized, + buffer: streaming.buffer, + failure: nil, + onTermination: streaming.onTermination ) - /// Indicates that the producer has been finished. - case throwFinishedError - - @inlinable - init( - callbackToken: Source.WriteResult.CallbackToken?, - continuationAndElement: (CheckedContinuation, Element)? = nil - ) { - switch (callbackToken, continuationAndElement) { - case (.none, .none): - self = .returnProduceMore - - case (.some(let callbackToken), .none): - self = .returnEnqueue(callbackToken: callbackToken) - - case (.none, .some((let continuation, let element))): - self = .resumeConsumerAndReturnProduceMore( - continuation: continuation, - element: element - ) - - case (.some(let callbackToken), .some((let continuation, let element))): - self = .resumeConsumerAndReturnEnqueue( - continuation: continuation, - element: element, - callbackToken: callbackToken - ) - } - } - } - - @inlinable - mutating func write(_ sequence: some Sequence) -> WriteAction { - switch self._state { - case .initial(var initial): - var buffer = Deque() - buffer.append(contentsOf: sequence) - - let shouldProduceMore = initial.backPressureStrategy.didYield(elements: buffer[...]) - let callbackToken = shouldProduceMore ? nil : self.nextCallbackToken() - - self._state = .streaming( - .init( - backPressureStrategy: initial.backPressureStrategy, - iteratorInitialized: initial.iteratorInitialized, - onTermination: initial.onTermination, - buffer: buffer, - consumerContinuation: nil, - producerContinuations: .init(), - cancelledAsyncProducers: .init(), - hasOutstandingDemand: shouldProduceMore - ) - ) - - return .init(callbackToken: callbackToken) - - case .streaming(var streaming): - self._state = .modify - - let bufferEndIndexBeforeAppend = streaming.buffer.endIndex - streaming.buffer.append(contentsOf: sequence) - - // We have an element and can resume the continuation - streaming.hasOutstandingDemand = streaming.backPressureStrategy.didYield( - elements: streaming.buffer[bufferEndIndexBeforeAppend...] - ) - - if let consumerContinuation = streaming.consumerContinuation { - guard let element = streaming.buffer.popFirst() else { - // We got a yield of an empty sequence. We just tolerate this. - self._state = .streaming(streaming) - - return .init(callbackToken: streaming.hasOutstandingDemand ? nil : self.nextCallbackToken()) - } - streaming.hasOutstandingDemand = streaming.backPressureStrategy.didConsume(element: element) - - // We got a consumer continuation and an element. We can resume the consumer now - streaming.consumerContinuation = nil - self._state = .streaming(streaming) - return .init( - callbackToken: streaming.hasOutstandingDemand ? nil : self.nextCallbackToken(), - continuationAndElement: (consumerContinuation, element) - ) - } else { - // We don't have a suspended consumer so we just buffer the elements - self._state = .streaming(streaming) - return .init( - callbackToken: streaming.hasOutstandingDemand ? nil : self.nextCallbackToken() - ) - } - - case .sourceFinished, .finished: - // If the source has finished we are dropping the elements. - return .throwFinishedError + ) - case .modify: - fatalError("AsyncStream internal inconsistency") - } + return .failProducers( + Array(streaming.producerContinuations.map { $0.1 }) + ) } - /// Actions returned by `enqueueProducer()`. - @usableFromInline - enum EnqueueProducerAction { - /// Indicates that the producer should be notified to produce more. - case resumeProducer((Result) -> Void) - /// Indicates that the producer should be notified about an error. - case resumeProducerWithError((Result) -> Void, any Error) - } + case .sourceFinished, .finished: + // This is normal and we just have to tolerate it + return .none - @inlinable - mutating func enqueueProducer( - callbackToken: Source.WriteResult.CallbackToken, - onProduceMore: @Sendable @escaping (Result) -> Void - ) -> EnqueueProducerAction? { - switch self._state { - case .initial: - // We need to transition to streaming before we can suspend - // This is enforced because the CallbackToken has no internal init so - // one must create it by calling `write` first. - fatalError("AsyncStream internal inconsistency") - - case .streaming(var streaming): - if let index = streaming.cancelledAsyncProducers.firstIndex(of: callbackToken.id) { - // Our producer got marked as cancelled. - self._state = .modify - streaming.cancelledAsyncProducers.remove(at: index) - self._state = .streaming(streaming) - - return .resumeProducerWithError(onProduceMore, CancellationError()) - } else if streaming.hasOutstandingDemand { - // We hit an edge case here where we wrote but the consuming thread got interleaved - return .resumeProducer(onProduceMore) - } else { - self._state = .modify - streaming.producerContinuations.append((callbackToken.id, onProduceMore)) - - self._state = .streaming(streaming) - return .none - } + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } - case .sourceFinished, .finished: - // Since we are unlocking between yielding and suspending the yield - // It can happen that the source got finished or the consumption fully finishes. - return .resumeProducerWithError(onProduceMore, AlreadyFinishedError()) + /// Actions returned by `write()`. + @usableFromInline + enum WriteAction { + /// Indicates that the producer should be notified to produce more. + case returnProduceMore + /// Indicates that the producer should be suspended to stop producing. + case returnEnqueue( + callbackToken: Source.WriteResult.CallbackToken + ) + /// Indicates that the consumer should be resumed and the producer should be notified to produce more. + case resumeConsumerAndReturnProduceMore( + continuation: CheckedContinuation, + element: Element + ) + /// Indicates that the consumer should be resumed and the producer should be suspended. + case resumeConsumerAndReturnEnqueue( + continuation: CheckedContinuation, + element: Element, + callbackToken: Source.WriteResult.CallbackToken + ) + /// Indicates that the producer has been finished. + case throwFinishedError + + @inlinable + init( + callbackToken: Source.WriteResult.CallbackToken?, + continuationAndElement: (CheckedContinuation, Element)? = nil + ) { + switch (callbackToken, continuationAndElement) { + case (.none, .none): + self = .returnProduceMore + + case (.some(let callbackToken), .none): + self = .returnEnqueue(callbackToken: callbackToken) + + case (.none, .some((let continuation, let element))): + self = .resumeConsumerAndReturnProduceMore( + continuation: continuation, + element: element + ) + + case (.some(let callbackToken), .some((let continuation, let element))): + self = .resumeConsumerAndReturnEnqueue( + continuation: continuation, + element: element, + callbackToken: callbackToken + ) + } + } + } - case .modify: - fatalError("AsyncStream internal inconsistency") - } - } + @inlinable + mutating func write(_ sequence: some Sequence) -> WriteAction { + switch self._state { + case .initial(var initial): + var buffer = Deque() + buffer.append(contentsOf: sequence) + + let shouldProduceMore = initial.backPressureStrategy.didYield(elements: buffer[...]) + let callbackToken = shouldProduceMore ? nil : self.nextCallbackToken() + + self._state = .streaming( + .init( + backPressureStrategy: initial.backPressureStrategy, + iteratorInitialized: initial.iteratorInitialized, + onTermination: initial.onTermination, + buffer: buffer, + consumerContinuation: nil, + producerContinuations: .init(), + cancelledAsyncProducers: .init(), + hasOutstandingDemand: shouldProduceMore + ) + ) - /// Actions returned by `cancelProducer()`. - @usableFromInline - enum CancelProducerAction { - /// Indicates that the producer should be notified about cancellation. - case resumeProducerWithCancellationError((Result) -> Void) - } + return .init(callbackToken: callbackToken) - @inlinable - mutating func cancelProducer( - callbackToken: Source.WriteResult.CallbackToken - ) -> CancelProducerAction? { - switch self._state { - case .initial: - // We need to transition to streaming before we can suspend - fatalError("AsyncStream internal inconsistency") - - case .streaming(var streaming): - if let index = streaming.producerContinuations.firstIndex(where: { - $0.0 == callbackToken.id - }) { - // We have an enqueued producer that we need to resume now - self._state = .modify - let continuation = streaming.producerContinuations.remove(at: index).1 - self._state = .streaming(streaming) - - return .resumeProducerWithCancellationError(continuation) - } else { - // The task that yields was cancelled before yielding so the cancellation handler - // got invoked right away - self._state = .modify - streaming.cancelledAsyncProducers.append(callbackToken.id) - self._state = .streaming(streaming) - - return .none - } + case .streaming(var streaming): + self._state = .modify - case .sourceFinished, .finished: - // Since we are unlocking between yielding and suspending the yield - // It can happen that the source got finished or the consumption fully finishes. - return .none + let bufferEndIndexBeforeAppend = streaming.buffer.endIndex + streaming.buffer.append(contentsOf: sequence) - case .modify: - fatalError("AsyncStream internal inconsistency") - } - } + // We have an element and can resume the continuation + streaming.hasOutstandingDemand = streaming.backPressureStrategy.didYield( + elements: streaming.buffer[bufferEndIndexBeforeAppend...] + ) - /// Actions returned by `finish()`. - @usableFromInline - enum FinishAction { - /// Indicates that `onTermination` should be called. - case callOnTermination((() -> Void)?) - /// Indicates that the consumer should be resumed with the failure, the producers - /// should be resumed with an error and `onTermination` should be called. - case resumeConsumerAndCallOnTermination( - consumerContinuation: CheckedContinuation, - failure: (any Error)?, - onTermination: (() -> Void)? - ) - /// Indicates that the producers should be resumed with an error. - case resumeProducers( - producerContinuations: [(Result) -> Void] - ) - } + if let consumerContinuation = streaming.consumerContinuation { + guard let element = streaming.buffer.popFirst() else { + // We got a yield of an empty sequence. We just tolerate this. + self._state = .streaming(streaming) - @inlinable - mutating func finish(_ failure: (any Error)?) -> FinishAction? { - switch self._state { - case .initial(let initial): - // Nothing was yielded nor did anybody call next - // This means we can transition to sourceFinished and store the failure - self._state = .sourceFinished( - .init( - iteratorInitialized: initial.iteratorInitialized, - buffer: .init(), - failure: failure, - onTermination: initial.onTermination - ) - ) - - return .callOnTermination(initial.onTermination) - - case .streaming(let streaming): - if let consumerContinuation = streaming.consumerContinuation { - // We have a continuation, this means our buffer must be empty - // Furthermore, we can now transition to finished - // and resume the continuation with the failure - precondition(streaming.buffer.isEmpty, "Expected an empty buffer") - precondition( - streaming.producerContinuations.isEmpty, - "Expected no suspended producers" - ) - - self._state = .finished(iteratorInitialized: streaming.iteratorInitialized) - - return .resumeConsumerAndCallOnTermination( - consumerContinuation: consumerContinuation, - failure: failure, - onTermination: streaming.onTermination - ) - } else { - self._state = .sourceFinished( - .init( - iteratorInitialized: streaming.iteratorInitialized, - buffer: streaming.buffer, - failure: failure, - onTermination: streaming.onTermination - ) - ) - - return .resumeProducers( - producerContinuations: Array(streaming.producerContinuations.map { $0.1 }) - ) - } + return .init( + callbackToken: streaming.hasOutstandingDemand ? nil : self.nextCallbackToken() + ) + } + streaming.hasOutstandingDemand = streaming.backPressureStrategy.didConsume( + element: element + ) + + // We got a consumer continuation and an element. We can resume the consumer now + streaming.consumerContinuation = nil + self._state = .streaming(streaming) + return .init( + callbackToken: streaming.hasOutstandingDemand ? nil : self.nextCallbackToken(), + continuationAndElement: (consumerContinuation, element) + ) + } else { + // We don't have a suspended consumer so we just buffer the elements + self._state = .streaming(streaming) + return .init( + callbackToken: streaming.hasOutstandingDemand ? nil : self.nextCallbackToken() + ) + } + + case .sourceFinished, .finished: + // If the source has finished we are dropping the elements. + return .throwFinishedError + + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } - case .sourceFinished, .finished: - // If the source has finished, finishing again has no effect. - return .none + /// Actions returned by `enqueueProducer()`. + @usableFromInline + enum EnqueueProducerAction { + /// Indicates that the producer should be notified to produce more. + case resumeProducer((Result) -> Void) + /// Indicates that the producer should be notified about an error. + case resumeProducerWithError((Result) -> Void, any Error) + } - case .modify: - fatalError("AsyncStream internal inconsistency") - } - } + @inlinable + mutating func enqueueProducer( + callbackToken: Source.WriteResult.CallbackToken, + onProduceMore: @Sendable @escaping (Result) -> Void + ) -> EnqueueProducerAction? { + switch self._state { + case .initial: + // We need to transition to streaming before we can suspend + // This is enforced because the CallbackToken has no internal init so + // one must create it by calling `write` first. + fatalError("AsyncStream internal inconsistency") + + case .streaming(var streaming): + if let index = streaming.cancelledAsyncProducers.firstIndex(of: callbackToken.id) { + // Our producer got marked as cancelled. + self._state = .modify + streaming.cancelledAsyncProducers.remove(at: index) + self._state = .streaming(streaming) + + return .resumeProducerWithError(onProduceMore, CancellationError()) + } else if streaming.hasOutstandingDemand { + // We hit an edge case here where we wrote but the consuming thread got interleaved + return .resumeProducer(onProduceMore) + } else { + self._state = .modify + streaming.producerContinuations.append((callbackToken.id, onProduceMore)) + + self._state = .streaming(streaming) + return .none + } + + case .sourceFinished, .finished: + // Since we are unlocking between yielding and suspending the yield + // It can happen that the source got finished or the consumption fully finishes. + return .resumeProducerWithError(onProduceMore, AlreadyFinishedError()) + + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } - /// Actions returned by `next()`. - @usableFromInline - enum NextAction { - /// Indicates that the element should be returned to the caller. - case returnElement(Element) - /// Indicates that the element should be returned to the caller and that all producers should be called. - case returnElementAndResumeProducers(Element, [(Result) -> Void]) - /// Indicates that the `Error` should be returned to the caller and that `onTermination` should be called. - case returnErrorAndCallOnTermination((any Error)?, (() -> Void)?) - /// Indicates that the `nil` should be returned to the caller. - case returnNil - /// Indicates that the `Task` of the caller should be suspended. - case suspendTask - } + /// Actions returned by `cancelProducer()`. + @usableFromInline + enum CancelProducerAction { + /// Indicates that the producer should be notified about cancellation. + case resumeProducerWithCancellationError((Result) -> Void) + } - @inlinable - mutating func next() -> NextAction { - switch self._state { - case .initial(let initial): - // We are not interacting with the back-pressure strategy here because - // we are doing this inside `next(:)` - self._state = .streaming( - .init( - backPressureStrategy: initial.backPressureStrategy, - iteratorInitialized: initial.iteratorInitialized, - onTermination: initial.onTermination, - buffer: Deque(), - consumerContinuation: nil, - producerContinuations: .init(), - cancelledAsyncProducers: .init(), - hasOutstandingDemand: false - ) - ) - - return .suspendTask - case .streaming(var streaming): - guard streaming.consumerContinuation == nil else { - // We have multiple AsyncIterators iterating the sequence - fatalError("AsyncStream internal inconsistency") - } + @inlinable + mutating func cancelProducer( + callbackToken: Source.WriteResult.CallbackToken + ) -> CancelProducerAction? { + switch self._state { + case .initial: + // We need to transition to streaming before we can suspend + fatalError("AsyncStream internal inconsistency") + + case .streaming(var streaming): + if let index = streaming.producerContinuations.firstIndex(where: { + $0.0 == callbackToken.id + }) { + // We have an enqueued producer that we need to resume now + self._state = .modify + let continuation = streaming.producerContinuations.remove(at: index).1 + self._state = .streaming(streaming) + + return .resumeProducerWithCancellationError(continuation) + } else { + // The task that yields was cancelled before yielding so the cancellation handler + // got invoked right away + self._state = .modify + streaming.cancelledAsyncProducers.append(callbackToken.id) + self._state = .streaming(streaming) + + return .none + } + + case .sourceFinished, .finished: + // Since we are unlocking between yielding and suspending the yield + // It can happen that the source got finished or the consumption fully finishes. + return .none + + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } - self._state = .modify - - if let element = streaming.buffer.popFirst() { - // We have an element to fulfil the demand right away. - streaming.hasOutstandingDemand = streaming.backPressureStrategy.didConsume(element: element) - - if streaming.hasOutstandingDemand { - // There is demand and we have to resume our producers - let producers = Array(streaming.producerContinuations.map { $0.1 }) - streaming.producerContinuations.removeAll() - self._state = .streaming(streaming) - return .returnElementAndResumeProducers(element, producers) - } else { - // We don't have any new demand, so we can just return the element. - self._state = .streaming(streaming) - return .returnElement(element) - } - } else { - // There is nothing in the buffer to fulfil the demand so we need to suspend. - // We are not interacting with the back-pressure strategy here because - // we are doing this inside `suspendNext` - self._state = .streaming(streaming) - - return .suspendTask - } + /// Actions returned by `finish()`. + @usableFromInline + enum FinishAction { + /// Indicates that `onTermination` should be called. + case callOnTermination((() -> Void)?) + /// Indicates that the consumer should be resumed with the failure, the producers + /// should be resumed with an error and `onTermination` should be called. + case resumeConsumerAndCallOnTermination( + consumerContinuation: CheckedContinuation, + failure: (any Error)?, + onTermination: (() -> Void)? + ) + /// Indicates that the producers should be resumed with an error. + case resumeProducers( + producerContinuations: [(Result) -> Void] + ) + } - case .sourceFinished(var sourceFinished): - // Check if we have an element left in the buffer and return it - self._state = .modify + @inlinable + mutating func finish(_ failure: (any Error)?) -> FinishAction? { + switch self._state { + case .initial(let initial): + // Nothing was yielded nor did anybody call next + // This means we can transition to sourceFinished and store the failure + self._state = .sourceFinished( + .init( + iteratorInitialized: initial.iteratorInitialized, + buffer: .init(), + failure: failure, + onTermination: initial.onTermination + ) + ) - if let element = sourceFinished.buffer.popFirst() { - self._state = .sourceFinished(sourceFinished) + return .callOnTermination(initial.onTermination) + + case .streaming(let streaming): + if let consumerContinuation = streaming.consumerContinuation { + // We have a continuation, this means our buffer must be empty + // Furthermore, we can now transition to finished + // and resume the continuation with the failure + precondition(streaming.buffer.isEmpty, "Expected an empty buffer") + precondition( + streaming.producerContinuations.isEmpty, + "Expected no suspended producers" + ) + + self._state = .finished(iteratorInitialized: streaming.iteratorInitialized) + + return .resumeConsumerAndCallOnTermination( + consumerContinuation: consumerContinuation, + failure: failure, + onTermination: streaming.onTermination + ) + } else { + self._state = .sourceFinished( + .init( + iteratorInitialized: streaming.iteratorInitialized, + buffer: streaming.buffer, + failure: failure, + onTermination: streaming.onTermination + ) + ) - return .returnElement(element) - } else { - // We are returning the queued failure now and can transition to finished - self._state = .finished(iteratorInitialized: sourceFinished.iteratorInitialized) + return .resumeProducers( + producerContinuations: Array(streaming.producerContinuations.map { $0.1 }) + ) + } - return .returnErrorAndCallOnTermination( - sourceFinished.failure, - sourceFinished.onTermination - ) - } + case .sourceFinished, .finished: + // If the source has finished, finishing again has no effect. + return .none - case .finished: - return .returnNil + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } - case .modify: - fatalError("AsyncStream internal inconsistency") - } - } + /// Actions returned by `next()`. + @usableFromInline + enum NextAction { + /// Indicates that the element should be returned to the caller. + case returnElement(Element) + /// Indicates that the element should be returned to the caller and that all producers should be called. + case returnElementAndResumeProducers(Element, [(Result) -> Void]) + /// Indicates that the `Error` should be returned to the caller and that `onTermination` should be called. + case returnErrorAndCallOnTermination((any Error)?, (() -> Void)?) + /// Indicates that the `nil` should be returned to the caller. + case returnNil + /// Indicates that the `Task` of the caller should be suspended. + case suspendTask + } - /// Actions returned by `suspendNext()`. - @usableFromInline - enum SuspendNextAction { - /// Indicates that the consumer should be resumed. - case resumeConsumerWithElement(CheckedContinuation, Element) - /// Indicates that the consumer and all producers should be resumed. - case resumeConsumerWithElementAndProducers( - CheckedContinuation, - Element, - [(Result) -> Void] - ) - /// Indicates that the consumer should be resumed with the failure and that `onTermination` should be called. - case resumeConsumerWithErrorAndCallOnTermination( - CheckedContinuation, - (any Error)?, - (() -> Void)? - ) - /// Indicates that the consumer should be resumed with `nil`. - case resumeConsumerWithNil(CheckedContinuation) - } + @inlinable + mutating func next() -> NextAction { + switch self._state { + case .initial(let initial): + // We are not interacting with the back-pressure strategy here because + // we are doing this inside `next(:)` + self._state = .streaming( + .init( + backPressureStrategy: initial.backPressureStrategy, + iteratorInitialized: initial.iteratorInitialized, + onTermination: initial.onTermination, + buffer: Deque(), + consumerContinuation: nil, + producerContinuations: .init(), + cancelledAsyncProducers: .init(), + hasOutstandingDemand: false + ) + ) - @inlinable - mutating func suspendNext( - continuation: CheckedContinuation - ) -> SuspendNextAction? { - switch self._state { - case .initial: - // We need to transition to streaming before we can suspend - preconditionFailure("AsyncStream internal inconsistency") - - case .streaming(var streaming): - guard streaming.consumerContinuation == nil else { - // We have multiple AsyncIterators iterating the sequence - fatalError( - "This should never happen since we only allow a single Iterator to be created" - ) - } + return .suspendTask + case .streaming(var streaming): + guard streaming.consumerContinuation == nil else { + // We have multiple AsyncIterators iterating the sequence + fatalError("AsyncStream internal inconsistency") + } + + self._state = .modify + + if let element = streaming.buffer.popFirst() { + // We have an element to fulfil the demand right away. + streaming.hasOutstandingDemand = streaming.backPressureStrategy.didConsume( + element: element + ) + + if streaming.hasOutstandingDemand { + // There is demand and we have to resume our producers + let producers = Array(streaming.producerContinuations.map { $0.1 }) + streaming.producerContinuations.removeAll() + self._state = .streaming(streaming) + return .returnElementAndResumeProducers(element, producers) + } else { + // We don't have any new demand, so we can just return the element. + self._state = .streaming(streaming) + return .returnElement(element) + } + } else { + // There is nothing in the buffer to fulfil the demand so we need to suspend. + // We are not interacting with the back-pressure strategy here because + // we are doing this inside `suspendNext` + self._state = .streaming(streaming) + + return .suspendTask + } + + case .sourceFinished(var sourceFinished): + // Check if we have an element left in the buffer and return it + self._state = .modify + + if let element = sourceFinished.buffer.popFirst() { + self._state = .sourceFinished(sourceFinished) + + return .returnElement(element) + } else { + // We are returning the queued failure now and can transition to finished + self._state = .finished(iteratorInitialized: sourceFinished.iteratorInitialized) + + return .returnErrorAndCallOnTermination( + sourceFinished.failure, + sourceFinished.onTermination + ) + } + + case .finished: + return .returnNil + + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } - self._state = .modify - - // We have to check here again since we might have a producer interleave next and suspendNext - if let element = streaming.buffer.popFirst() { - // We have an element to fulfil the demand right away. - - streaming.hasOutstandingDemand = streaming.backPressureStrategy.didConsume(element: element) - - if streaming.hasOutstandingDemand { - // There is demand and we have to resume our producers - let producers = Array(streaming.producerContinuations.map { $0.1 }) - streaming.producerContinuations.removeAll() - self._state = .streaming(streaming) - return .resumeConsumerWithElementAndProducers( - continuation, - element, - producers - ) - } else { - // We don't have any new demand, so we can just return the element. - self._state = .streaming(streaming) - return .resumeConsumerWithElement(continuation, element) - } - } else { - // There is nothing in the buffer to fulfil the demand so we to store the continuation. - streaming.consumerContinuation = continuation - self._state = .streaming(streaming) - - return .none - } + /// Actions returned by `suspendNext()`. + @usableFromInline + enum SuspendNextAction { + /// Indicates that the consumer should be resumed. + case resumeConsumerWithElement(CheckedContinuation, Element) + /// Indicates that the consumer and all producers should be resumed. + case resumeConsumerWithElementAndProducers( + CheckedContinuation, + Element, + [(Result) -> Void] + ) + /// Indicates that the consumer should be resumed with the failure and that `onTermination` should be called. + case resumeConsumerWithErrorAndCallOnTermination( + CheckedContinuation, + (any Error)?, + (() -> Void)? + ) + /// Indicates that the consumer should be resumed with `nil`. + case resumeConsumerWithNil(CheckedContinuation) + } - case .sourceFinished(var sourceFinished): - // Check if we have an element left in the buffer and return it - self._state = .modify + @inlinable + mutating func suspendNext( + continuation: CheckedContinuation + ) -> SuspendNextAction? { + switch self._state { + case .initial: + // We need to transition to streaming before we can suspend + preconditionFailure("AsyncStream internal inconsistency") + + case .streaming(var streaming): + guard streaming.consumerContinuation == nil else { + // We have multiple AsyncIterators iterating the sequence + fatalError( + "This should never happen since we only allow a single Iterator to be created" + ) + } + + self._state = .modify + + // We have to check here again since we might have a producer interleave next and suspendNext + if let element = streaming.buffer.popFirst() { + // We have an element to fulfil the demand right away. + + streaming.hasOutstandingDemand = streaming.backPressureStrategy.didConsume( + element: element + ) + + if streaming.hasOutstandingDemand { + // There is demand and we have to resume our producers + let producers = Array(streaming.producerContinuations.map { $0.1 }) + streaming.producerContinuations.removeAll() + self._state = .streaming(streaming) + return .resumeConsumerWithElementAndProducers( + continuation, + element, + producers + ) + } else { + // We don't have any new demand, so we can just return the element. + self._state = .streaming(streaming) + return .resumeConsumerWithElement(continuation, element) + } + } else { + // There is nothing in the buffer to fulfil the demand so we to store the continuation. + streaming.consumerContinuation = continuation + self._state = .streaming(streaming) - if let element = sourceFinished.buffer.popFirst() { - self._state = .sourceFinished(sourceFinished) + return .none + } - return .resumeConsumerWithElement(continuation, element) - } else { - // We are returning the queued failure now and can transition to finished - self._state = .finished(iteratorInitialized: sourceFinished.iteratorInitialized) + case .sourceFinished(var sourceFinished): + // Check if we have an element left in the buffer and return it + self._state = .modify - return .resumeConsumerWithErrorAndCallOnTermination( - continuation, - sourceFinished.failure, - sourceFinished.onTermination - ) - } + if let element = sourceFinished.buffer.popFirst() { + self._state = .sourceFinished(sourceFinished) - case .finished: - return .resumeConsumerWithNil(continuation) + return .resumeConsumerWithElement(continuation, element) + } else { + // We are returning the queued failure now and can transition to finished + self._state = .finished(iteratorInitialized: sourceFinished.iteratorInitialized) - case .modify: - fatalError("AsyncStream internal inconsistency") - } + return .resumeConsumerWithErrorAndCallOnTermination( + continuation, + sourceFinished.failure, + sourceFinished.onTermination + ) } - /// Actions returned by `cancelNext()`. - @usableFromInline - enum CancelNextAction { - /// Indicates that the continuation should be resumed with a cancellation error, the producers should be finished and call onTermination. - case resumeConsumerWithCancellationErrorAndCallOnTermination( - CheckedContinuation, - (() -> Void)? - ) - /// Indicates that the producers should be finished and call onTermination. - case failProducersAndCallOnTermination([(Result) -> Void], (() -> Void)?) - } + case .finished: + return .resumeConsumerWithNil(continuation) - @inlinable - mutating func cancelNext() -> CancelNextAction? { - switch self._state { - case .initial: - // We need to transition to streaming before we can suspend - fatalError("AsyncStream internal inconsistency") - - case .streaming(let streaming): - self._state = .finished(iteratorInitialized: streaming.iteratorInitialized) - - if let consumerContinuation = streaming.consumerContinuation { - precondition( - streaming.producerContinuations.isEmpty, - "Internal inconsistency. Unexpected producer continuations." - ) - return .resumeConsumerWithCancellationErrorAndCallOnTermination( - consumerContinuation, - streaming.onTermination - ) - } else { - return .failProducersAndCallOnTermination( - Array(streaming.producerContinuations.map { $0.1 }), - streaming.onTermination - ) - } + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } - case .sourceFinished, .finished: - return .none + /// Actions returned by `cancelNext()`. + @usableFromInline + enum CancelNextAction { + /// Indicates that the continuation should be resumed with a cancellation error, the producers should be finished and call onTermination. + case resumeConsumerWithCancellationErrorAndCallOnTermination( + CheckedContinuation, + (() -> Void)? + ) + /// Indicates that the producers should be finished and call onTermination. + case failProducersAndCallOnTermination([(Result) -> Void], (() -> Void)?) + } - case .modify: - fatalError("AsyncStream internal inconsistency") - } - } + @inlinable + mutating func cancelNext() -> CancelNextAction? { + switch self._state { + case .initial: + // We need to transition to streaming before we can suspend + fatalError("AsyncStream internal inconsistency") + + case .streaming(let streaming): + self._state = .finished(iteratorInitialized: streaming.iteratorInitialized) + + if let consumerContinuation = streaming.consumerContinuation { + precondition( + streaming.producerContinuations.isEmpty, + "Internal inconsistency. Unexpected producer continuations." + ) + return .resumeConsumerWithCancellationErrorAndCallOnTermination( + consumerContinuation, + streaming.onTermination + ) + } else { + return .failProducersAndCallOnTermination( + Array(streaming.producerContinuations.map { $0.1 }), + streaming.onTermination + ) + } + + case .sourceFinished, .finished: + return .none + + case .modify: + fatalError("AsyncStream internal inconsistency") + } } + } } diff --git a/Sources/HTTPClient/HTTPClientFoundation/BufferedStream/Lock.swift b/Sources/HTTPClient/HTTPClientFoundation/BufferedStream/Lock.swift index 444d7c3..f4a62c9 100644 --- a/Sources/HTTPClient/HTTPClientFoundation/BufferedStream/Lock.swift +++ b/Sources/HTTPClient/HTTPClientFoundation/BufferedStream/Lock.swift @@ -11,7 +11,7 @@ // SPDX-License-Identifier: Apache-2.0 // //===----------------------------------------------------------------------===// -// swift-format-ignore-file + //===----------------------------------------------------------------------===// // // This source file is part of the SwiftNIO open source project @@ -27,23 +27,50 @@ //===----------------------------------------------------------------------===// #if canImport(Darwin) -import Darwin -#elseif canImport(Glibc) -import Glibc + import Darwin #elseif os(Windows) -import WinSDK + import ucrt + import WinSDK +#elseif canImport(Glibc) + @preconcurrency import Glibc +#elseif canImport(Musl) + @preconcurrency import Musl +#elseif canImport(Bionic) + @preconcurrency import Bionic +#elseif canImport(WASILibc) + @preconcurrency import WASILibc + #if canImport(wasi_pthread) + import wasi_pthread + #endif +#else + #error("The concurrency Lock module was unable to identify your C library.") #endif #if os(Windows) -@usableFromInline -typealias LockPrimitive = SRWLOCK + @usableFromInline + typealias LockPrimitive = SRWLOCK #else -@usableFromInline -typealias LockPrimitive = pthread_mutex_t + @usableFromInline + typealias LockPrimitive = pthread_mutex_t #endif +/// A utility function that runs the body code only in debug builds, without +/// emitting compiler warnings. +/// +/// This is currently the only way to do this in Swift: see +/// https://forums.swift.org/t/support-debug-only-code/11037 for a discussion. +@inlinable +internal func debugOnly(_ body: () -> Void) { + assert( + { + body() + return true + }() + ) +} + @usableFromInline -enum LockOperations {} +enum LockOperations: Sendable {} extension LockOperations { @inlinable @@ -51,13 +78,16 @@ extension LockOperations { mutex.assertValidAlignment() #if os(Windows) - InitializeSRWLock(mutex) - #else - var attr = pthread_mutexattr_t() - pthread_mutexattr_init(&attr) - - let err = pthread_mutex_init(mutex, &attr) - precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)") + InitializeSRWLock(mutex) + #elseif (compiler(<6.1) && !os(WASI)) || (compiler(>=6.1) && _runtime(_multithreaded)) + var attr = pthread_mutexattr_t() + pthread_mutexattr_init(&attr) + debugOnly { + pthread_mutexattr_settype(&attr, .init(PTHREAD_MUTEX_ERRORCHECK)) + } + + let err = pthread_mutex_init(mutex, &attr) + precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)") #endif } @@ -66,10 +96,10 @@ extension LockOperations { mutex.assertValidAlignment() #if os(Windows) - // SRWLOCK does not need to be freed - #else - let err = pthread_mutex_destroy(mutex) - precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)") + // SRWLOCK does not need to be free'd + #elseif (compiler(<6.1) && !os(WASI)) || (compiler(>=6.1) && _runtime(_multithreaded)) + let err = pthread_mutex_destroy(mutex) + precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)") #endif } @@ -78,10 +108,10 @@ extension LockOperations { mutex.assertValidAlignment() #if os(Windows) - AcquireSRWLockExclusive(mutex) - #else - let err = pthread_mutex_lock(mutex) - precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)") + AcquireSRWLockExclusive(mutex) + #elseif (compiler(<6.1) && !os(WASI)) || (compiler(>=6.1) && _runtime(_multithreaded)) + let err = pthread_mutex_lock(mutex) + precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)") #endif } @@ -90,10 +120,10 @@ extension LockOperations { mutex.assertValidAlignment() #if os(Windows) - ReleaseSRWLockExclusive(mutex) - #else - let err = pthread_mutex_unlock(mutex) - precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)") + ReleaseSRWLockExclusive(mutex) + #elseif (compiler(<6.1) && !os(WASI)) || (compiler(>=6.1) && _runtime(_multithreaded)) + let err = pthread_mutex_unlock(mutex) + precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)") #endif } } @@ -132,9 +162,11 @@ final class LockStorage: ManagedBuffer { @inlinable static func create(value: Value) -> Self { let buffer = Self.create(minimumCapacity: 1) { _ in - return value + value } - // Avoid 'unsafeDowncast' as there is a miscompilation on 5.10. + // Intentionally using a force cast here to avoid a miss compiliation in 5.10. + // This is as fast as an unsafeDownCast since ManagedBuffer is inlined and the optimizer + // can eliminate the upcast/downcast pair let storage = buffer as! Self storage.withUnsafeMutablePointers { _, lockPtr in @@ -158,7 +190,7 @@ final class LockStorage: ManagedBuffer { } } - @usableFromInline + @inlinable deinit { self.withUnsafeMutablePointerToElements { lockPtr in LockOperations.destroy(lockPtr) @@ -166,11 +198,10 @@ final class LockStorage: ManagedBuffer { } @inlinable - func withLockPrimitive( - _ body: (UnsafeMutablePointer) throws -> T - ) rethrows -> T { + func withLockPrimitive(_ body: (UnsafeMutablePointer) throws -> T) rethrows -> T + { try self.withUnsafeMutablePointerToElements { lockPtr in - return try body(lockPtr) + try body(lockPtr) } } @@ -184,11 +215,16 @@ final class LockStorage: ManagedBuffer { } } -extension LockStorage: @unchecked Sendable {} +// This compiler guard is here becaue `ManagedBuffer` is already declaring +// Sendable unavailability after 6.1, which `LockStorage` inherits. +#if compiler(<6.2) + @available(*, unavailable) + extension LockStorage: Sendable {} +#endif /// A threading lock based on `libpthread` instead of `libdispatch`. /// -/// - note: ``Lock`` has reference semantics. +/// - Note: ``Lock`` has reference semantics. /// /// This object provides a lock on top of a single `pthread_mutex_t`. This kind /// of lock is safe to use with `libpthread`-based threading models, such as the @@ -200,7 +236,7 @@ struct Lock { internal let _storage: LockStorage /// Create a new lock. - @usableFromInline + @inlinable init() { self._storage = .create(value: ()) } @@ -224,10 +260,10 @@ struct Lock { } @inlinable - internal func withLockPrimitive( - _ body: (UnsafeMutablePointer) throws -> T - ) rethrows -> T { - return try self._storage.withLockPrimitive(body) + internal func withLockPrimitive(_ body: (UnsafeMutablePointer) throws -> T) + rethrows -> T + { + try self._storage.withLockPrimitive(body) } } @@ -248,9 +284,14 @@ extension Lock { } return try body() } + + @inlinable + func withLockVoid(_ body: () throws -> Void) rethrows { + try self.withLock(body) + } } -extension Lock: Sendable {} +extension Lock: @unchecked Sendable {} extension UnsafeMutablePointer { @inlinable @@ -259,20 +300,74 @@ extension UnsafeMutablePointer { } } +/// Provides locked access to `Value`. +/// +/// - Note: ``LockedValueBox`` has reference semantics and holds the `Value` +/// alongside a lock behind a reference. +/// +/// This is no different than creating a ``Lock`` and protecting all +/// accesses to a value using the lock. But it's easy to forget to actually +/// acquire/release the lock in the correct place. ``LockedValueBox`` makes +/// that much easier. @usableFromInline struct LockedValueBox { - @usableFromInline - let storage: LockStorage @usableFromInline + internal let _storage: LockStorage + + /// Initialize the `Value`. + @inlinable init(_ value: Value) { - self.storage = .create(value: value) + self._storage = .create(value: value) } + /// Access the `Value`, allowing mutation of it. @inlinable func withLockedValue(_ mutate: (inout Value) throws -> T) rethrows -> T { - return try self.storage.withLockedValue(mutate) + try self._storage.withLockedValue(mutate) + } + + /// Provides an unsafe view over the lock and its value. + /// + /// This can be beneficial when you require fine grained control over the lock in some + /// situations but don't want lose the benefits of ``withLockedValue(_:)`` in others by + /// switching to ``NIOLock``. + var unsafe: Unsafe { + Unsafe(_storage: self._storage) + } + + /// Provides an unsafe view over the lock and its value. + struct Unsafe { + @usableFromInline + let _storage: LockStorage + + /// Manually acquire the lock. + @inlinable + func lock() { + self._storage.lock() + } + + /// Manually release the lock. + @inlinable + func unlock() { + self._storage.unlock() + } + + /// Mutate the value, assuming the lock has been acquired manually. + /// + /// - Parameter mutate: A closure with scoped access to the value. + /// - Returns: The result of the `mutate` closure. + @inlinable + func withValueAssumingLockIsAcquired( + _ mutate: (_ value: inout Value) throws -> Result + ) rethrows -> Result { + try self._storage.withUnsafeMutablePointerToHeader { value in + try mutate(&value.pointee) + } + } } } -extension LockedValueBox: Sendable where Value: Sendable {} \ No newline at end of file +extension LockedValueBox: @unchecked Sendable where Value: Sendable {} + +extension LockedValueBox.Unsafe: @unchecked Sendable where Value: Sendable {} diff --git a/Sources/HTTPClient/HTTPClientFoundation/URLSessionBidirectionalStreaming/BidirectionalStreamingURLSessionDelegate.swift b/Sources/HTTPClient/HTTPClientFoundation/URLSessionBidirectionalStreaming/BidirectionalStreamingURLSessionDelegate.swift index 7377c54..a17a96b 100644 --- a/Sources/HTTPClient/HTTPClientFoundation/URLSessionBidirectionalStreaming/BidirectionalStreamingURLSessionDelegate.swift +++ b/Sources/HTTPClient/HTTPClientFoundation/URLSessionBidirectionalStreaming/BidirectionalStreamingURLSessionDelegate.swift @@ -83,21 +83,21 @@ import HTTPTypes /// Use `bidirectionalStreamingRequest(for:baseURL:requestBody:requestStreamBufferSize:responseStreamWatermarks:)`. init( - requestBody: HTTPBody?, requestStreamBufferSize: Int, + requestBody: HTTPBody?, + requestStreamBufferSize: Int, responseStreamWatermarks: (low: Int, high: Int) ) { self.requestBody = requestBody self.hasAlreadyIteratedRequestBody = false self.hasSuspendedURLSessionTask = LockedValueBox(false) self.requestStreamBufferSize = requestStreamBufferSize - (self.responseBodyStream, self.responseBodyStreamSource) = - ResponseBodyStream.makeStream( - backPressureStrategy: .customWatermark( - low: responseStreamWatermarks.low, - high: responseStreamWatermarks.high, - waterLevelForElement: { $0.count } - ) + (self.responseBodyStream, self.responseBodyStreamSource) = ResponseBodyStream.makeStream( + backPressureStrategy: .customWatermark( + low: responseStreamWatermarks.low, + high: responseStreamWatermarks.high, + waterLevelForElement: { $0.count } ) + ) } func urlSession(_ session: URLSession, needNewBodyStreamForTask task: URLSessionTask) async @@ -117,8 +117,7 @@ import HTTPTypes hasAlreadyIteratedRequestBody = true // Create a fresh pair of streams. - let (inputStream, outputStream) = createStreamPair( - withBufferSize: requestStreamBufferSize) + let (inputStream, outputStream) = createStreamPair(withBufferSize: requestStreamBufferSize) // Bridge the output stream to the request body (which opens the output stream). requestStream = HTTPBodyOutputStreamBridge(outputStream, requestBody!) @@ -132,37 +131,28 @@ import HTTPTypes callbackLock.withLock { debug("Task delegate: didReceive data (numBytes: \(data.count))") do { - switch try responseBodyStreamSource.write( - contentsOf: CollectionOfOne(ArraySlice(data))) - { + switch try responseBodyStreamSource.write(contentsOf: CollectionOfOne(ArraySlice(data))) { case .produceMore: break case .enqueueCallback(let callbackToken): - let shouldActuallyEnqueueCallback = - hasSuspendedURLSessionTask.withLockedValue { - hasSuspendedURLSessionTask in - if hasSuspendedURLSessionTask { - debug( - "Task delegate: already suspended task, not enqueing another writer callback" - ) - return false - } - debug( - "Task delegate: response stream backpressure, suspending task and enqueing callback" - ) - dataTask.suspend() - hasSuspendedURLSessionTask = true - return true + let shouldActuallyEnqueueCallback = hasSuspendedURLSessionTask.withLockedValue { + hasSuspendedURLSessionTask in + if hasSuspendedURLSessionTask { + debug("Task delegate: already suspended task, not enqueing another writer callback") + return false } + debug( + "Task delegate: response stream backpressure, suspending task and enqueing callback" + ) + dataTask.suspend() + hasSuspendedURLSessionTask = true + return true + } if shouldActuallyEnqueueCallback { - responseBodyStreamSource.enqueueCallback(callbackToken: callbackToken) { - result in - self.hasSuspendedURLSessionTask.withLockedValue { - hasSuspendedURLSessionTask in + responseBodyStreamSource.enqueueCallback(callbackToken: callbackToken) { result in + self.hasSuspendedURLSessionTask.withLockedValue { hasSuspendedURLSessionTask in switch result { case .success: - debug( - "Task delegate: response stream callback, resuming task" - ) + debug("Task delegate: response stream callback, resuming task") dataTask.resume() hasSuspendedURLSessionTask = false case .failure(let error): @@ -183,7 +173,9 @@ import HTTPTypes } func urlSession( - _ session: URLSession, dataTask: URLSessionDataTask, didReceive response: URLResponse + _ session: URLSession, + dataTask: URLSessionDataTask, + didReceive response: URLResponse ) async -> URLSession.ResponseDisposition { @@ -196,7 +188,9 @@ import HTTPTypes } func urlSession( - _ session: URLSession, task: URLSessionTask, didCompleteWithError error: (any Error)? + _ session: URLSession, + task: URLSessionTask, + didCompleteWithError error: (any Error)? ) { callbackLock.withLock { debug("Task delegate: didCompleteWithError (error: \(String(describing: error)))") @@ -215,7 +209,10 @@ import HTTPTypes var inputStream: InputStream? var outputStream: OutputStream? Stream.getBoundStreams( - withBufferSize: bufferSize, inputStream: &inputStream, outputStream: &outputStream) + withBufferSize: bufferSize, + inputStream: &inputStream, + outputStream: &outputStream + ) guard let inputStream, let outputStream else { fatalError("getBoundStreams did not return non-nil streams") } diff --git a/Sources/HTTPClient/HTTPClientFoundation/URLSessionBidirectionalStreaming/HTTPBodyOutputStreamBridge.swift b/Sources/HTTPClient/HTTPClientFoundation/URLSessionBidirectionalStreaming/HTTPBodyOutputStreamBridge.swift index 5838f37..7e0c776 100644 --- a/Sources/HTTPClient/HTTPClientFoundation/URLSessionBidirectionalStreaming/HTTPBodyOutputStreamBridge.swift +++ b/Sources/HTTPClient/HTTPClientFoundation/URLSessionBidirectionalStreaming/HTTPBodyOutputStreamBridge.swift @@ -19,7 +19,9 @@ import HTTPTypes final class HTTPBodyOutputStreamBridge: NSObject, StreamDelegate { static let streamQueue = DispatchQueue( - label: "HTTPBodyStreamDelegate", autoreleaseFrequency: .workItem) + label: "HTTPBodyStreamDelegate", + autoreleaseFrequency: .workItem + ) let httpBody: HTTPBody let outputStream: OutputStream @@ -69,8 +71,7 @@ import HTTPTypes try await withCheckedThrowingContinuation { continuation in Self.streamQueue.async { debug("Output stream delegate produced chunk and suspended producer.") - self.performAction( - self.state.producedChunkAndSuspendedProducer(chunk, continuation)) + self.performAction(self.state.producedChunkAndSuspendedProducer(chunk, continuation)) } } } @@ -87,8 +88,7 @@ import HTTPTypes precondition(!bytesToWrite.isEmpty, "\(#function) must be called with non-empty bytes") guard outputStream.streamStatus == .open else { debug("Output stream closed unexpectedly.") - performAction( - state.wroteBytes(numBytesWritten: 0, streamStillHasSpaceAvailable: false)) + performAction(state.wroteBytes(numBytesWritten: 0, streamStillHasSpaceAvailable: false)) return } switch bytesToWrite.withUnsafeBytes({ @@ -107,10 +107,10 @@ import HTTPTypes performAction( state.wroteBytes( numBytesWritten: written, - streamStillHasSpaceAvailable: outputStream.hasSpaceAvailable) + streamStillHasSpaceAvailable: outputStream.hasSpaceAvailable + ) ) - default: - preconditionFailure("OutputStream.write(_:maxLength:) returned undocumented value") + default: preconditionFailure("OutputStream.write(_:maxLength:) returned undocumented value") } } @@ -157,7 +157,8 @@ import HTTPTypes } mutating func producedChunkAndSuspendedProducer( - _ chunk: Chunk, _ producerContinuation: ProducerContinuation + _ chunk: Chunk, + _ producerContinuation: ProducerContinuation ) -> Action { @@ -172,24 +173,21 @@ import HTTPTypes } } - mutating func wroteBytes(numBytesWritten: Int, streamStillHasSpaceAvailable: Bool) - -> Action - { + mutating func wroteBytes(numBytesWritten: Int, streamStillHasSpaceAvailable: Bool) -> Action { switch self { case .haveBytes(let spaceAvailable, let chunk, let producerContinuation): - guard spaceAvailable, numBytesWritten <= chunk.count else { - preconditionFailure() - } + guard spaceAvailable, numBytesWritten <= chunk.count else { preconditionFailure() } let remaining = chunk.dropFirst(numBytesWritten) guard remaining.isEmpty else { self = .haveBytes( - spaceAvailable: streamStillHasSpaceAvailable, remaining, - producerContinuation) + spaceAvailable: streamStillHasSpaceAvailable, + remaining, + producerContinuation + ) guard streamStillHasSpaceAvailable else { return .none } return .writeBytes(remaining) } - self = .needBytes( - spaceAvailable: streamStillHasSpaceAvailable, producerContinuation) + self = .needBytes(spaceAvailable: streamStillHasSpaceAvailable, producerContinuation) return .resumeProducer(producerContinuation) case .initial, .needBytes, .waitingForBytes, .closed: preconditionFailure("\(#function) called in invalid state: \(self)") @@ -245,8 +243,7 @@ import HTTPTypes case .needBytes(_, let producerContinuation): self = .closed(nil) return .cancelProducerAndCloseStream(producerContinuation) - case .initial, .closed: - preconditionFailure("\(#function) called in invalid state: \(self)") + case .initial, .closed: preconditionFailure("\(#function) called in invalid state: \(self)") } } diff --git a/Sources/HTTPClient/HTTPClientFoundation/URLSessionTransport.swift b/Sources/HTTPClient/HTTPClientFoundation/URLSessionTransport.swift index 04fd297..8ca2927 100644 --- a/Sources/HTTPClient/HTTPClientFoundation/URLSessionTransport.swift +++ b/Sources/HTTPClient/HTTPClientFoundation/URLSessionTransport.swift @@ -300,12 +300,11 @@ extension URLSessionTransportError: CustomStringConvertible { /// A textual representation of this instance. var description: String { switch self { - case let .invalidRequestURL(path: path, method: method, baseURL: baseURL): + case .invalidRequestURL(let path, let method, let baseURL): return "Invalid request URL from request path: \(path), method: \(method), relative to base URL: \(baseURL.absoluteString)" case .notHTTPResponse(let response): - return - "Received a non-HTTP response, of type: \(String(describing: type(of: response)))" + return "Received a non-HTTP response, of type: \(String(describing: type(of: response)))" case .invalidResponseStatusCode(let response): return "Received an HTTP response with invalid status code: \(response.statusCode))" case .noResponse(let url): @@ -315,12 +314,12 @@ extension URLSessionTransportError: CustomStringConvertible { } } -private let _debugLoggingEnabled = LockStorage.create(value: false) +private let _debugLoggingEnabled = LockedValueBox(false) var debugLoggingEnabled: Bool { get { _debugLoggingEnabled.withLockedValue { $0 } } set { _debugLoggingEnabled.withLockedValue { $0 = newValue } } } -private let _standardErrorLock = LockStorage.create(value: FileHandle.standardError) +private let _standardErrorLock = LockedValueBox(FileHandle.standardError) func debug( _ message: @autoclosure () -> String, function: String = #function, @@ -331,8 +330,7 @@ func debug( { if debugLoggingEnabled { _standardErrorLock.withLockedValue { - let logLine = - "[\(function) \(file.split(separator: "/").last!):\(line)] \(message())\n" + let logLine = "[\(function) \(file.split(separator: "/").last!):\(line)] \(message())\n" $0.write(Data((logLine).utf8)) } } @@ -349,9 +347,7 @@ extension URLSession { ) async throws -> (HTTPResponse, HTTPBody?) { try Task.checkCancellation() var urlRequest = try URLRequest(request, baseURL: baseURL) - if let requestBody { - urlRequest.httpBody = try await Data(collecting: requestBody, upTo: .max) - } + if let requestBody { urlRequest.httpBody = try await Data(collecting: requestBody, upTo: .max) } try Task.checkCancellation() /// Use `dataTask(with:completionHandler:)` here because `data(for:[delegate:]) async` is only available on @@ -361,8 +357,7 @@ extension URLSession { let (response, maybeResponseBodyData): (URLResponse, Data?) = try await withCheckedThrowingContinuation { continuation in - let task = self.dataTask(with: urlRequest) { - [urlRequest] data, response, error in + let task = self.dataTask(with: urlRequest) { [urlRequest] data, response, error in if let error { continuation.resume(throwing: error) return @@ -389,11 +384,7 @@ extension URLSession { } let maybeResponseBody = maybeResponseBodyData.map { data in - HTTPBody( - data, - length: HTTPBody.Length(from: response), - iterationBehavior: .multiple - ) + HTTPBody(data, length: HTTPBody.Length(from: response), iterationBehavior: .multiple) } return (try HTTPResponse(response), maybeResponseBody) } onCancel: {