diff --git a/Sources/OpenAPIURLSession/AsyncBackpressuredStream/AsyncBackpressuredStream.swift b/Sources/OpenAPIURLSession/AsyncBackpressuredStream/AsyncBackpressuredStream.swift deleted file mode 100644 index 2c792b4..0000000 --- a/Sources/OpenAPIURLSession/AsyncBackpressuredStream/AsyncBackpressuredStream.swift +++ /dev/null @@ -1,1384 +0,0 @@ -//===----------------------------------------------------------------------===// -// -// This source file is part of the SwiftOpenAPIGenerator open source project -// -// Copyright (c) 2023 Apple Inc. and the SwiftOpenAPIGenerator project authors -// Licensed under Apache License v2.0 -// -// See LICENSE.txt for license information -// See CONTRIBUTORS.txt for the list of SwiftOpenAPIGenerator project authors -// -// SPDX-License-Identifier: Apache-2.0 -// -//===----------------------------------------------------------------------===// -// swift-format-ignore-file -//===----------------------------------------------------------------------===// -// -// This source file is part of the Swift.org open source project -// -// Copyright (c) 2020-2021 Apple Inc. and the Swift project authors -// Licensed under Apache License v2.0 with Runtime Library Exception -// -// See https://swift.org/LICENSE.txt for license information -// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors -// -//===----------------------------------------------------------------------===// -import DequeModule - -struct AsyncBackpressuredStream: Sendable { - /// 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. - struct Source: Sendable { - /// A strategy that handles the back pressure of the asynchronous stream. - struct BackPressureStrategy: Sendable { - var internalBackPressureStrategy: InternalBackPressureStrategy - - /// When the high water mark is reached producers will be suspended. All producers will be resumed again once - /// the low water mark is reached. - static func highLowWatermark(lowWatermark: Int, highWatermark: Int) -> BackPressureStrategy { - .init( - internalBackPressureStrategy: .highLowWatermark( - .init(lowWatermark: lowWatermark, highWatermark: highWatermark) - ) - ) - } - - /// When the high water mark is reached producers will be suspended. All producers will be resumed again once - /// the low water mark is reached. When `usingElementCounts` is true, the counts of the element types will - /// be used to compute the watermark. - static func highLowWatermarkWithElementCounts(lowWatermark: Int, highWatermark: Int) - -> BackPressureStrategy where Element: RandomAccessCollection - { - .init( - internalBackPressureStrategy: .highLowWatermark( - .init( - lowWatermark: lowWatermark, - highWatermark: highWatermark, - waterLevelForElement: { $0.count } - ) - ) - ) - } - } - - /// A type that indicates the result of writing elements to the source. - enum WriteResult: Sendable { - /// A token that is returned when the asynchronous stream's back pressure strategy indicated that any producer should - /// be suspended. Use this token to enqueue a callback by calling the ``enqueueCallback(_:)`` method. - struct WriteToken: Sendable { - let id: UInt - - 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(WriteToken) - } - - private var storage: Storage - - init(storage: Storage) { self.storage = storage } - - /// 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. - /// - /// - Parameter sequence: The elements to write to the asynchronous stream. - /// - Returns: The result that indicates if more elements should be produced at this time. - func write(contentsOf sequence: S) throws -> WriteResult where S.Element == Element { - try self.storage.write(contentsOf: sequence) - } - - /// Enqueues a callback that will be invoked once more elements should be produced. - /// - /// Call this method after ``write(contentsOf:)`` returned a ``WriteResult/enqueueCallback(_:)``. - /// - /// - Parameters: - /// - writeToken: The write token produced by ``write(contentsOf:)``. - /// - onProduceMore: The callback which gets invoked once more elements should be produced. - func enqueueCallback( - writeToken: WriteResult.WriteToken, - onProduceMore: @escaping @Sendable (Result) -> Void - ) { self.storage.enqueueProducer(writeToken: writeToken, onProduceMore: onProduceMore) } - - /// Cancel an enqueued callback. - /// - /// Call this method to cancel a callback enqueued by the ``enqueueCallback(writeToken:onProduceMore:)`` method. - /// - /// > Note: This methods supports being called before ``enqueueCallback(writeToken:onProduceMore:)`` is called and - /// will mark the passed `writeToken` as cancelled. - /// - Parameter writeToken: The write token produced by ``write(contentsOf:)``. - func cancelCallback(writeToken: WriteResult.WriteToken) { - self.storage.cancelProducer(writeToken: writeToken) - } - - /// Write new elements to the asynchronous stream and provide a callback which will be invoked once more elements should be produced. - /// - /// - 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:)``. - func write( - contentsOf sequence: S, - onProduceMore: @escaping @Sendable (Result) -> Void - ) where S.Element == Element { - do { - let writeResult = try self.write(contentsOf: sequence) - - switch writeResult { - case .produceMore: onProduceMore(.success(())) - - case .enqueueCallback(let writeToken): - self.enqueueCallback(writeToken: writeToken, onProduceMore: onProduceMore) - } - } catch { onProduceMore(.failure(error)) } - } - - /// Write new elements to the asynchronous stream. - /// - /// This method returns once more elements should be produced. - /// - /// - Parameters: - /// - sequence: The elements to write to the asynchronous stream. - func asyncWrite(contentsOf sequence: S) async throws where S.Element == Element { - let writeResult = try self.write(contentsOf: sequence) - - switch writeResult { - case .produceMore: return - - case .enqueueCallback(let writeToken): - try await withTaskCancellationHandler { - try await withCheckedThrowingContinuation { continuation in - self.enqueueCallback( - writeToken: writeToken, - onProduceMore: { result in - switch result { - case .success(): continuation.resume(returning: ()) - case .failure(let error): continuation.resume(throwing: error) - } - } - ) - } - } onCancel: { - self.cancelCallback(writeToken: writeToken) - } - - } - } - - func finish(throwing failure: Failure?) { self.storage.finish(failure) } - } - - private var storage: Storage - - init(storage: Storage) { self.storage = storage } - - static func makeStream( - of elementType: Element.Type = Element.self, - backPressureStrategy: Source.BackPressureStrategy, - onTermination: (@Sendable () -> Void)? = nil - ) -> (Self, Source) where Failure == any Error { - let storage = Storage( - backPressureStrategy: backPressureStrategy.internalBackPressureStrategy, - onTerminate: onTermination - ) - let source = Source(storage: storage) - - return (.init(storage: storage), source) - } -} - -extension AsyncBackpressuredStream: AsyncSequence { - struct AsyncIterator: AsyncIteratorProtocol { - private var storage: Storage - - init(storage: Storage) { self.storage = storage } - - mutating func next() async throws -> Element? { return try await storage.next() } - } - - func makeAsyncIterator() -> AsyncIterator { return AsyncIterator(storage: self.storage) } -} - -extension AsyncBackpressuredStream { - struct HighLowWatermarkBackPressureStrategy { - private let lowWatermark: Int - private let highWatermark: Int - private(set) var currentWatermark: Int - - typealias CustomWaterLevelForElement = @Sendable (Element) -> Int - private let waterLevelForElement: CustomWaterLevelForElement? - - /// Initializes a new ``HighLowWatermarkBackPressureStrategy``. - /// - /// - Parameters: - /// - lowWatermark: The low watermark where demand should start. - /// - highWatermark: The high watermark where demand should be stopped. - init(lowWatermark: Int, highWatermark: Int, waterLevelForElement: CustomWaterLevelForElement? = nil) { - precondition(lowWatermark <= highWatermark, "Low watermark must be <= high watermark") - self.lowWatermark = lowWatermark - self.highWatermark = highWatermark - self.currentWatermark = 0 - self.waterLevelForElement = waterLevelForElement - } - - mutating func didYield(elements: Deque.SubSequence) -> Bool { - if let waterLevelForElement { - self.currentWatermark += elements.reduce(0) { $0 + waterLevelForElement($1) } - } else { - self.currentWatermark += elements.count - } - precondition(self.currentWatermark >= 0, "Watermark below zero") - // We are demanding more until we reach the high watermark - return self.currentWatermark < self.highWatermark - } - - mutating func didConsume(elements: Deque.SubSequence) -> Bool { - if let waterLevelForElement { - self.currentWatermark -= elements.reduce(0) { $0 + waterLevelForElement($1) } - } else { - self.currentWatermark -= elements.count - } - precondition(self.currentWatermark >= 0, "Watermark below zero") - // We start demanding again once we are below the low watermark - return self.currentWatermark < self.lowWatermark - } - - mutating func didConsume(element: Element) -> Bool { - if let waterLevelForElement { - self.currentWatermark -= waterLevelForElement(element) - } else { - self.currentWatermark -= 1 - } - precondition(self.currentWatermark >= 0, "Watermark below zero") - // We start demanding again once we are below the low watermark - return self.currentWatermark < self.lowWatermark - } - } - - enum InternalBackPressureStrategy { - case highLowWatermark(HighLowWatermarkBackPressureStrategy) - - mutating func didYield(elements: Deque.SubSequence) -> Bool { - switch self { - case .highLowWatermark(var strategy): - let result = strategy.didYield(elements: elements) - self = .highLowWatermark(strategy) - return result - } - } - - mutating func didConsume(elements: Deque.SubSequence) -> Bool { - switch self { - case .highLowWatermark(var strategy): - let result = strategy.didConsume(elements: elements) - self = .highLowWatermark(strategy) - return result - } - } - - mutating func didConsume(element: Element) -> Bool { - switch self { - case .highLowWatermark(var strategy): - let result = strategy.didConsume(element: element) - self = .highLowWatermark(strategy) - return result - } - } - } -} - -extension AsyncBackpressuredStream { - final class Storage: @unchecked Sendable { - /// The lock that protects the state machine and the nextProducerID. - let lock = NIOLock() - - /// The state machine. - var stateMachine: StateMachine - - /// The next producer's id. - var nextProducerID: UInt = 0 - - init(backPressureStrategy: InternalBackPressureStrategy, onTerminate: (() -> Void)?) { - self.stateMachine = .init(backPressureStrategy: backPressureStrategy, onTerminate: onTerminate) - } - - func sequenceDeinitialized() { - let onTerminate = self.lock.withLock { - let action = self.stateMachine.sequenceDeinitialized() - - switch action { - case .callOnTerminate(let onTerminate): - // We have to call onTerminate without the lock to avoid potential deadlocks - return onTerminate - - case .none: return nil - } - } - - onTerminate?() - } - - func iteratorInitialized() { self.lock.withLock { self.stateMachine.iteratorInitialized() } } - - func iteratorDeinitialized() { - let onTerminate = self.lock.withLock { - let action = self.stateMachine.iteratorDeinitialized() - - switch action { - case .callOnTerminate(let onTerminate): - // We have to call onTerminate without the lock to avoid potential deadlocks - return onTerminate - - case .none: return nil - } - } - - onTerminate?() - } - - func write(contentsOf sequence: S) throws -> Source.WriteResult where S.Element == Element { - let action = self.lock.withLock { - let stateBefore = self.stateMachine.state - let action = self.stateMachine.write(sequence) - let stateAfter = self.stateMachine.state - debug(""" - --- - event: write - state before: \(stateBefore) - state after: \(stateAfter) - action: \(action) - --- - """) - return action - } - - switch action { - case .returnProduceMore: return .produceMore - - case .returnEnqueue: - // TODO: Move the id into the state machine or use an atomic - let id = self.lock.withLock { - let id = self.nextProducerID - self.nextProducerID += 1 - return id - } - return .enqueueCallback(.init(id: id)) - - case .resumeConsumerContinuationAndReturnProduceMore(let continuation, let element): - continuation.resume(returning: element) - return .produceMore - - case .resumeConsumerContinuationAndReturnEnqueue(let continuation, let element): - continuation.resume(returning: element) - // TODO: Move the id into the state machine or use an atomic - let id = self.lock.withLock { - let id = self.nextProducerID - self.nextProducerID += 1 - return id - } - return .enqueueCallback(.init(id: id)) - - case .throwFinishedError: - // TODO: Introduce new Error - throw CancellationError() - } - } - - func enqueueProducer( - writeToken: Source.WriteResult.WriteToken, - onProduceMore: @escaping @Sendable (Result) -> Void - ) { - let action = self.lock.withLock { - let stateBefore = self.stateMachine.state - let action = self.stateMachine.enqueueProducer(writeToken: writeToken, onProduceMore: onProduceMore) - let stateAfter = self.stateMachine.state - debug(""" - --- - event: \(#function) - state before: \(stateBefore) - state after: \(stateAfter) - action: \(action) - --- - """) - return action - } - - switch action { - case .resumeProducer(let onProduceMore): onProduceMore(.success(())) - - case .resumeProducerWithCancellationError(let onProduceMore): onProduceMore(.failure(CancellationError())) - - case .none: break - } - } - - func cancelProducer(writeToken: Source.WriteResult.WriteToken) { - let action = self.lock.withLock { return self.stateMachine.cancelProducer(writeToken: writeToken) } - - switch action { - case .resumeProducerWithCancellationError(let onProduceMore): onProduceMore(.failure(CancellationError())) - - case .none: break - } - } - - func finish(_ failure: Failure?) { - let onTerminate = self.lock.withLock { - let action = self.stateMachine.finish(failure) - - switch action { - case .resumeAllContinuationsAndCallOnTerminate( - let consumerContinuation, - let failure, - let producerContinuations, - let onTerminate - ): - // It is safe to resume the continuation while holding the lock - // since the task will get enqueued on its executor and the resume method - // is returning immediately - switch failure { - case .some(let error): consumerContinuation.resume(throwing: error) - case .none: consumerContinuation.resume(returning: nil) - } - - for producerContinuation in producerContinuations { - // TODO: Throw a new cancelled error - producerContinuation(.failure(CancellationError())) - } - - return onTerminate - - case .resumeProducerContinuations(let producerContinuations): - for producerContinuation in producerContinuations { - // TODO: Throw a new cancelled error - producerContinuation(.failure(CancellationError())) - } - - return nil - - case .none: return nil - } - } - - onTerminate?() - } - - func next() async throws -> Element? { - let action = self.lock.withLock { - let stateBefore = self.stateMachine.state - let action = self.stateMachine.next() - let stateAfter = self.stateMachine.state - debug(""" - --- - event: next - state before: \(stateBefore) - state after: \(stateAfter) - action: \(action) - --- - """) - return action - } - - switch action { - case .returnElement(let element): return element - - case .returnElementAndResumeProducers(let element, let producerContinuations): - for producerContinuation in producerContinuations { producerContinuation(.success(())) } - - return element - - case .returnFailureAndCallOnTerminate(let failure, let onTerminate): - onTerminate?() - switch failure { - case .some(let error): throw error - - case .none: return nil - } - - case .returnNil: return nil - - case .suspendTask: return try await suspendNext() - } - } - - func suspendNext() async throws -> Element? { - return try await withTaskCancellationHandler { - return try await withCheckedThrowingContinuation { continuation in - let action = self.lock.withLock { - let stateBefore = self.stateMachine.state - let action = self.stateMachine.suspendNext(continuation: continuation) - let stateAfter = self.stateMachine.state - debug(""" - --- - event: \(#function) - state before: \(stateBefore) - state after: \(stateAfter) - action: \(action) - --- - """) - return action - } - - switch action { - case .resumeContinuationWithElement(let continuation, let element): - continuation.resume(returning: element) - - case .resumeContinuationWithElementAndProducers( - let continuation, - let element, - let producerContinuations - ): - continuation.resume(returning: element) - for producerContinuation in producerContinuations { producerContinuation(.success(())) } - - case .resumeContinuationWithFailureAndCallOnTerminate( - let continuation, - let failure, - let onTerminate - ): - onTerminate?() - switch failure { - case .some(let error): continuation.resume(throwing: error) - - case .none: continuation.resume(returning: nil) - } - - case .resumeContinuationWithNil(let continuation): continuation.resume(returning: nil) - - case .none: break - } - } - } onCancel: { - self.lock.withLockVoid { - let action = self.stateMachine.cancelNext() - - switch action { - case .resumeContinuationWithCancellationErrorAndFinishProducersAndCallOnTerminate( - let continuation, - let producerContinuations, - let onTerminate - ): - onTerminate?() - continuation.resume(throwing: CancellationError()) - for producerContinuation in producerContinuations { - // TODO: Throw a new cancelled error - producerContinuation(.failure(CancellationError())) - } - - case .finishProducersAndCallOnTerminate(let producerContinuations, let onTerminate): - onTerminate?() - for producerContinuation in producerContinuations { - // TODO: Throw a new cancelled error - producerContinuation(.failure(CancellationError())) - } - - case .none: break - } - } - } - } - } -} - -extension AsyncBackpressuredStream { - struct StateMachine { - enum State { - case initial( - backPressureStrategy: InternalBackPressureStrategy, - iteratorInitialized: Bool, - onTerminate: (() -> Void)? - ) - - /// The state once either any element was yielded or `next()` was called. - case streaming( - backPressureStrategy: InternalBackPressureStrategy, - buffer: Deque, - consumerContinuation: CheckedContinuation?, - producerContinuations: Deque<(UInt, (Result) -> Void)>, - cancelledAsyncProducers: Deque, - hasOutstandingDemand: Bool, - iteratorInitialized: Bool, - onTerminate: (() -> Void)? - ) - - /// The state once the underlying source signalled that it is finished. - case sourceFinished( - buffer: Deque, - iteratorInitialized: Bool, - failure: Failure?, - onTerminate: (() -> Void)? - ) - - /// 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) - } - - /// The state machine's current state. - var state: State - - var producerContinuationCounter: UInt = 0 - - /// 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. - init(backPressureStrategy: InternalBackPressureStrategy, onTerminate: (() -> Void)?) { - self.state = .initial( - backPressureStrategy: backPressureStrategy, - iteratorInitialized: false, - onTerminate: onTerminate - ) - } - - /// Actions returned by `sequenceDeinitialized()`. - enum SequenceDeinitializedAction { - /// Indicates that `onTerminate` should be called. - case callOnTerminate((() -> Void)?) - /// Indicates that nothing should be done. - case none - } - - mutating func sequenceDeinitialized() -> SequenceDeinitializedAction { - switch self.state { - case .initial(_, iteratorInitialized: false, let onTerminate), - .streaming(_, _, _, _, _, _, iteratorInitialized: false, let onTerminate), - .sourceFinished(_, iteratorInitialized: false, _, let onTerminate): - // No iterator was created so we can transition to finished right away. - self.state = .finished(iteratorInitialized: false) - - return .callOnTerminate(onTerminate) - - case .initial(_, iteratorInitialized: true, _), .streaming(_, _, _, _, _, _, iteratorInitialized: true, _), - .sourceFinished(_, iteratorInitialized: true, _, _): - // An iterator was created and we deinited the sequence. - // This is an expected pattern and we just continue on normal. - return .none - - case .finished: - // We are already finished so there is nothing left to clean up. - // This is just the references dropping afterwards. - return .none - } - } - - mutating func iteratorInitialized() { - switch self.state { - case .initial(_, iteratorInitialized: true, _), .streaming(_, _, _, _, _, _, iteratorInitialized: true, _), - .sourceFinished(_, iteratorInitialized: true, _, _), .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 .initial(let backPressureStrategy, iteratorInitialized: false, let onTerminate): - // The first and only iterator was initialized. - self.state = .initial( - backPressureStrategy: backPressureStrategy, - iteratorInitialized: true, - onTerminate: onTerminate - ) - - case .streaming( - let backPressureStrategy, - let buffer, - let consumerContinuation, - let producerContinuations, - let cancelledAsyncProducers, - let hasOutstandingDemand, - false, - let onTerminate - ): - // The first and only iterator was initialized. - self.state = .streaming( - backPressureStrategy: backPressureStrategy, - buffer: buffer, - consumerContinuation: consumerContinuation, - producerContinuations: producerContinuations, - cancelledAsyncProducers: cancelledAsyncProducers, - hasOutstandingDemand: hasOutstandingDemand, - iteratorInitialized: true, - onTerminate: onTerminate - ) - - case .sourceFinished(let buffer, false, let failure, let onTerminate): - // The first and only iterator was initialized. - self.state = .sourceFinished( - buffer: buffer, - iteratorInitialized: true, - failure: failure, - onTerminate: onTerminate - ) - - 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) - } - } - - /// Actions returned by `iteratorDeinitialized()`. - enum IteratorDeinitializedAction { - /// Indicates that `onTerminate` should be called. - case callOnTerminate((() -> Void)?) - /// Indicates that nothing should be done. - case none - } - - mutating func iteratorDeinitialized() -> IteratorDeinitializedAction { - switch self.state { - case .initial(_, iteratorInitialized: false, _), - .streaming(_, _, _, _, _, _, iteratorInitialized: false, _), - .sourceFinished(_, iteratorInitialized: false, _, _): - // An iterator needs to be initialized before it can be deinitialized. - preconditionFailure("Internal inconsistency") - - case .initial(_, iteratorInitialized: true, let onTerminate), - .streaming(_, _, _, _, _, _, iteratorInitialized: true, let onTerminate), - .sourceFinished(_, iteratorInitialized: true, _, let onTerminate): - // An iterator was created and deinited. Since we only support - // a single iterator we can now transition to finish and inform the delegate. - self.state = .finished(iteratorInitialized: true) - - return .callOnTerminate(onTerminate) - - 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 `yield()`. - 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 - /// Indicates that the consumer continuation should be resumed and the producer should be notified to produce more. - case resumeConsumerContinuationAndReturnProduceMore( - continuation: CheckedContinuation, - element: Element - ) - /// Indicates that the consumer continuation should be resumed and the producer should be suspended. - case resumeConsumerContinuationAndReturnEnqueue( - continuation: CheckedContinuation, - element: Element - ) - /// Indicates that the producer has been finished. - case throwFinishedError - - init( - shouldProduceMore: Bool, - continuationAndElement: (CheckedContinuation, Element)? = nil - ) { - switch (shouldProduceMore, continuationAndElement) { - case (true, .none): self = .returnProduceMore - - case (false, .none): self = .returnEnqueue - - case (true, .some((let continuation, let element))): - self = .resumeConsumerContinuationAndReturnProduceMore(continuation: continuation, element: element) - - case (false, .some((let continuation, let element))): - self = .resumeConsumerContinuationAndReturnEnqueue(continuation: continuation, element: element) - } - } - } - - mutating func write(_ sequence: S) -> WriteAction where S.Element == Element { - switch self.state { - case .initial(var backPressureStrategy, let iteratorInitialized, let onTerminate): - let buffer = Deque(sequence) - let shouldProduceMore = backPressureStrategy.didYield(elements: buffer[...]) - - self.state = .streaming( - backPressureStrategy: backPressureStrategy, - buffer: buffer, - consumerContinuation: nil, - producerContinuations: .init(), - cancelledAsyncProducers: .init(), - hasOutstandingDemand: shouldProduceMore, - iteratorInitialized: iteratorInitialized, - onTerminate: onTerminate - ) - - return .init(shouldProduceMore: shouldProduceMore) - - case .streaming( - var backPressureStrategy, - var buffer, - .some(let consumerContinuation), - let producerContinuations, - let cancelledAsyncProducers, - let hasOutstandingDemand, - let iteratorInitialized, - let onTerminate - ): - // The buffer should always be empty if we hold a continuation - precondition(buffer.isEmpty, "Expected an empty buffer") - - let bufferEndIndexBeforeAppend = buffer.endIndex - buffer.append(contentsOf: sequence) - _ = backPressureStrategy.didYield(elements: buffer[bufferEndIndexBeforeAppend...]) - - guard let element = buffer.popFirst() else { - // We got a yield of an empty sequence. We just tolerate this. - self.state = .streaming( - backPressureStrategy: backPressureStrategy, - buffer: buffer, - consumerContinuation: consumerContinuation, - producerContinuations: producerContinuations, - cancelledAsyncProducers: cancelledAsyncProducers, - hasOutstandingDemand: hasOutstandingDemand, - iteratorInitialized: iteratorInitialized, - onTerminate: onTerminate - ) - return .init(shouldProduceMore: hasOutstandingDemand) - } - - // We have an element and can resume the continuation - - let shouldProduceMore = backPressureStrategy.didConsume(element: element) - self.state = .streaming( - backPressureStrategy: backPressureStrategy, - buffer: buffer, - consumerContinuation: nil, // Setting this to nil since we are resuming the continuation - producerContinuations: producerContinuations, - cancelledAsyncProducers: cancelledAsyncProducers, - hasOutstandingDemand: shouldProduceMore, - iteratorInitialized: iteratorInitialized, - onTerminate: onTerminate - ) - - return .init( - shouldProduceMore: shouldProduceMore, - continuationAndElement: (consumerContinuation, element) - ) - - case .streaming( - var backPressureStrategy, - var buffer, - consumerContinuation: .none, - let producerContinuations, - let cancelledAsyncProducers, - _, - let iteratorInitialized, - let onTerminate - ): - let bufferEndIndexBeforeAppend = buffer.endIndex - buffer.append(contentsOf: sequence) - let shouldProduceMore = backPressureStrategy.didYield(elements: buffer[bufferEndIndexBeforeAppend...]) - - self.state = .streaming( - backPressureStrategy: backPressureStrategy, - buffer: buffer, - consumerContinuation: nil, - producerContinuations: producerContinuations, - cancelledAsyncProducers: cancelledAsyncProducers, - hasOutstandingDemand: shouldProduceMore, - iteratorInitialized: iteratorInitialized, - onTerminate: onTerminate - ) - - return .init(shouldProduceMore: shouldProduceMore) - - case .sourceFinished, .finished: - // If the source has finished we are dropping the elements. - return .throwFinishedError - } - } - - /// Actions returned by `suspendYield()`. - @usableFromInline enum EnqueueProducerAction { - case resumeProducer((Result) -> Void) - case resumeProducerWithCancellationError((Result) -> Void) - case none - } - - @inlinable mutating func enqueueProducer( - writeToken: Source.WriteResult.WriteToken, - onProduceMore: @escaping (Result) -> Void - ) -> EnqueueProducerAction { - switch self.state { - case .initial: - // We need to transition to streaming before we can suspend - preconditionFailure("Internal inconsistency") - - case .streaming( - let backPressureStrategy, - let buffer, - let consumerContinuation, - var producerContinuations, - var cancelledAsyncProducers, - let hasOutstandingDemand, - let iteratorInitialized, - let onTerminate - ): - if let index = cancelledAsyncProducers.firstIndex(of: writeToken.id) { - cancelledAsyncProducers.remove(at: index) - self.state = .streaming( - backPressureStrategy: backPressureStrategy, - buffer: buffer, - consumerContinuation: consumerContinuation, - producerContinuations: producerContinuations, - cancelledAsyncProducers: cancelledAsyncProducers, - hasOutstandingDemand: hasOutstandingDemand, - iteratorInitialized: iteratorInitialized, - onTerminate: onTerminate - ) - - return .resumeProducerWithCancellationError(onProduceMore) - } else if hasOutstandingDemand { - // We hit an edge case here where we yielded but got suspended afterwards - // and in-between yielding and suspending the yield we got consumption which lead us - // to produce more again. - return .resumeProducer(onProduceMore) - } else { - producerContinuations.append((writeToken.id, onProduceMore)) - - self.state = .streaming( - backPressureStrategy: backPressureStrategy, - buffer: buffer, - consumerContinuation: consumerContinuation, - producerContinuations: producerContinuations, - cancelledAsyncProducers: cancelledAsyncProducers, - hasOutstandingDemand: hasOutstandingDemand, - iteratorInitialized: iteratorInitialized, - onTerminate: onTerminate - ) - - 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 - } - } - - /// Actions returned by `cancelYield()`. - enum CancelYieldAction { - case resumeProducerWithCancellationError((Result) -> Void) - case none - } - - mutating func cancelProducer(writeToken: Source.WriteResult.WriteToken) -> CancelYieldAction { - switch self.state { - case .initial: - // We need to transition to streaming before we can suspend - preconditionFailure("Internal inconsistency") - - case .streaming( - let backPressureStrategy, - let buffer, - let consumerContinuation, - var producerContinuations, - var cancelledAsyncProducers, - let hasOutstandingDemand, - let iteratorInitialized, - let onTerminate - ): - guard let index = producerContinuations.firstIndex(where: { $0.0 == writeToken.id }) else { - // The task that yields was cancelled before yielding so the cancellation handler - // got invoked right away - cancelledAsyncProducers.append(writeToken.id) - self.state = .streaming( - backPressureStrategy: backPressureStrategy, - buffer: buffer, - consumerContinuation: consumerContinuation, - producerContinuations: producerContinuations, - cancelledAsyncProducers: cancelledAsyncProducers, - hasOutstandingDemand: hasOutstandingDemand, - iteratorInitialized: iteratorInitialized, - onTerminate: onTerminate - ) - - return .none - } - let continuation = producerContinuations.remove(at: index).1 - self.state = .streaming( - backPressureStrategy: backPressureStrategy, - buffer: buffer, - consumerContinuation: consumerContinuation, - producerContinuations: producerContinuations, - cancelledAsyncProducers: cancelledAsyncProducers, - hasOutstandingDemand: hasOutstandingDemand, - iteratorInitialized: iteratorInitialized, - onTerminate: onTerminate - ) - - return .resumeProducerWithCancellationError(continuation) - - 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 - } - } - - /// Actions returned by `finish()`. - @usableFromInline enum FinishAction { - /// Indicates that the consumer continuation should be resumed with the failure, the producer continuations - /// should be resumed with an error and `onTerminate` should be called. - case resumeAllContinuationsAndCallOnTerminate( - consumerContinuation: CheckedContinuation, - failure: Failure?, - producerContinuations: [(Result) -> Void], - onTerminate: (() -> Void)? - ) - /// Indicates that the producer continuations should be resumed with an error. - case resumeProducerContinuations(producerContinuations: [(Result) -> Void]) - /// Indicates that nothing should be done. - case none - } - - @inlinable mutating func finish(_ failure: Failure?) -> FinishAction { - switch self.state { - case .initial(_, let iteratorInitialized, let onTerminate): - // TODO: Should we call onTerminate here - // Nothing was yielded nor did anybody call next - // This means we can transition to sourceFinished and store the failure - self.state = .sourceFinished( - buffer: .init(), - iteratorInitialized: iteratorInitialized, - failure: failure, - onTerminate: onTerminate - ) - - return .none - - case .streaming( - _, - let buffer, - .some(let consumerContinuation), - let producerContinuations, - _, - _, - let iteratorInitialized, - let onTerminate - ): - // 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(buffer.isEmpty, "Expected an empty buffer") - - self.state = .finished(iteratorInitialized: iteratorInitialized) - - return .resumeAllContinuationsAndCallOnTerminate( - consumerContinuation: consumerContinuation, - failure: failure, - producerContinuations: Array(producerContinuations.map { $0.1 }), - onTerminate: onTerminate - ) - - case .streaming( - _, - let buffer, - consumerContinuation: .none, - let producerContinuations, - _, - _, - let iteratorInitialized, - let onTerminate - ): - self.state = .sourceFinished( - buffer: buffer, - iteratorInitialized: iteratorInitialized, - failure: failure, - onTerminate: onTerminate - ) - - return .resumeProducerContinuations(producerContinuations: Array(producerContinuations.map { $0.1 })) - - case .sourceFinished, .finished: - // If the source has finished, finishing again has no effect. - return .none - } - } - - /// Actions returned by `next()`. - 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 `Failure` should be returned to the caller and that `onTerminate` should be called. - case returnFailureAndCallOnTerminate(Failure?, (() -> 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 - } - - mutating func next() -> NextAction { - switch self.state { - case .initial(let backPressureStrategy, let iteratorInitialized, let onTerminate): - // We are not interacting with the back-pressure strategy here because - // we are doing this inside `next(:)` - self.state = .streaming( - backPressureStrategy: backPressureStrategy, - buffer: Deque(), - consumerContinuation: nil, - producerContinuations: .init(), - cancelledAsyncProducers: .init(), - hasOutstandingDemand: false, - iteratorInitialized: iteratorInitialized, - onTerminate: onTerminate - ) - - return .suspendTask - - case .streaming(_, _, .some, _, _, _, _, _): - // We have multiple AsyncIterators iterating the sequence - preconditionFailure("This should never happen since we only allow a single Iterator to be created") - - case .streaming( - var backPressureStrategy, - var buffer, - .none, - var producerContinuations, - let cancelledAsyncProducers, - let hasOutstandingDemand, - let iteratorInitialized, - let onTerminate - ): - guard let element = buffer.popFirst() 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( - backPressureStrategy: backPressureStrategy, - buffer: buffer, - consumerContinuation: nil, - producerContinuations: producerContinuations, - cancelledAsyncProducers: cancelledAsyncProducers, - hasOutstandingDemand: hasOutstandingDemand, - iteratorInitialized: iteratorInitialized, - onTerminate: onTerminate - ) - - return .suspendTask - } - // We have an element to fulfil the demand right away. - - let shouldProduceMore = backPressureStrategy.didConsume(element: element) - - guard shouldProduceMore else { - self.state = .streaming( - backPressureStrategy: backPressureStrategy, - buffer: buffer, - consumerContinuation: nil, - producerContinuations: producerContinuations, - cancelledAsyncProducers: cancelledAsyncProducers, - hasOutstandingDemand: shouldProduceMore, - iteratorInitialized: iteratorInitialized, - onTerminate: onTerminate - ) - // We don't have any new demand, so we can just return the element. - return .returnElement(element) - } - let producers = Array(producerContinuations.map { $0.1 }) - producerContinuations.removeAll() - self.state = .streaming( - backPressureStrategy: backPressureStrategy, - buffer: buffer, - consumerContinuation: nil, - producerContinuations: producerContinuations, - cancelledAsyncProducers: cancelledAsyncProducers, - hasOutstandingDemand: shouldProduceMore, - iteratorInitialized: iteratorInitialized, - onTerminate: onTerminate - ) - return .returnElementAndResumeProducers(element, producers) - - case .sourceFinished(var buffer, let iteratorInitialized, let failure, let onTerminate): - // Check if we have an element left in the buffer and return it - guard let element = buffer.popFirst() else { - // We are returning the queued failure now and can transition to finished - self.state = .finished(iteratorInitialized: iteratorInitialized) - - return .returnFailureAndCallOnTerminate(failure, onTerminate) - } - self.state = .sourceFinished( - buffer: buffer, - iteratorInitialized: iteratorInitialized, - failure: failure, - onTerminate: onTerminate - ) - - return .returnElement(element) - - case .finished: return .returnNil - } - } - - /// Actions returned by `suspendNext()`. - enum SuspendNextAction { - /// Indicates that the continuation should be resumed. - case resumeContinuationWithElement(CheckedContinuation, Element) - /// Indicates that the continuation and all producers should be resumed. - case resumeContinuationWithElementAndProducers( - CheckedContinuation, - Element, - [(Result) -> Void] - ) - /// Indicates that the continuation should be resumed with the failure and that `onTerminate` should be called. - case resumeContinuationWithFailureAndCallOnTerminate( - CheckedContinuation, - Failure?, - (() -> Void)? - ) - /// Indicates that the continuation should be resumed with `nil`. - case resumeContinuationWithNil(CheckedContinuation) - /// Indicates that nothing should be done. - case none - } - - mutating func suspendNext(continuation: CheckedContinuation) -> SuspendNextAction { - switch self.state { - case .initial: - // We need to transition to streaming before we can suspend - preconditionFailure("Internal inconsistency") - - case .streaming(_, _, .some, _, _, _, _, _): - // We have multiple AsyncIterators iterating the sequence - preconditionFailure("This should never happen since we only allow a single Iterator to be created") - - case .streaming( - var backPressureStrategy, - var buffer, - .none, - var producerContinuations, - let cancelledAsyncProducers, - let hasOutstandingDemand, - let iteratorInitialized, - let onTerminate - ): - // We have to check here again since we might have a producer interleave next and suspendNext - guard let element = buffer.popFirst() else { - // There is nothing in the buffer to fulfil the demand so we to store the continuation. - self.state = .streaming( - backPressureStrategy: backPressureStrategy, - buffer: buffer, - consumerContinuation: continuation, - producerContinuations: producerContinuations, - cancelledAsyncProducers: cancelledAsyncProducers, - hasOutstandingDemand: hasOutstandingDemand, - iteratorInitialized: iteratorInitialized, - onTerminate: onTerminate - ) - - return .none - } - // We have an element to fulfil the demand right away. - - let shouldProduceMore = backPressureStrategy.didConsume(element: element) - - guard shouldProduceMore else { - // We don't have any new demand, so we can just return the element. - self.state = .streaming( - backPressureStrategy: backPressureStrategy, - buffer: buffer, - consumerContinuation: nil, - producerContinuations: producerContinuations, - cancelledAsyncProducers: cancelledAsyncProducers, - hasOutstandingDemand: hasOutstandingDemand, - iteratorInitialized: iteratorInitialized, - onTerminate: onTerminate - ) - return .resumeContinuationWithElement(continuation, element) - } - let producers = Array(producerContinuations.map { $0.1 }) - producerContinuations.removeAll() - self.state = .streaming( - backPressureStrategy: backPressureStrategy, - buffer: buffer, - consumerContinuation: nil, - producerContinuations: producerContinuations, - cancelledAsyncProducers: cancelledAsyncProducers, - hasOutstandingDemand: shouldProduceMore, - iteratorInitialized: iteratorInitialized, - onTerminate: onTerminate - ) - return .resumeContinuationWithElementAndProducers(continuation, element, producers) - - case .sourceFinished(var buffer, let iteratorInitialized, let failure, let onTerminate): - // Check if we have an element left in the buffer and return it - guard let element = buffer.popFirst() else { - // We are returning the queued failure now and can transition to finished - self.state = .finished(iteratorInitialized: iteratorInitialized) - - return .resumeContinuationWithFailureAndCallOnTerminate(continuation, failure, onTerminate) - } - self.state = .sourceFinished( - buffer: buffer, - iteratorInitialized: iteratorInitialized, - failure: failure, - onTerminate: onTerminate - ) - - return .resumeContinuationWithElement(continuation, element) - - case .finished: return .resumeContinuationWithNil(continuation) - } - } - - /// Actions returned by `cancelNext()`. - enum CancelNextAction { - /// Indicates that the continuation should be resumed with a cancellation error, the producers should be finished and call onTerminate. - case resumeContinuationWithCancellationErrorAndFinishProducersAndCallOnTerminate( - CheckedContinuation, - [(Result) -> Void], - (() -> Void)? - ) - /// Indicates that the producers should be finished and call onTerminate. - case finishProducersAndCallOnTerminate([(Result) -> Void], (() -> Void)?) - /// Indicates that nothing should be done. - case none - } - - mutating func cancelNext() -> CancelNextAction { - switch self.state { - case .initial: - // We need to transition to streaming before we can suspend - preconditionFailure("Internal inconsistency") - - case .streaming( - _, - _, - let consumerContinuation, - let producerContinuations, - _, - _, - let iteratorInitialized, - let onTerminate - ): - self.state = .finished(iteratorInitialized: iteratorInitialized) - - guard let consumerContinuation = consumerContinuation else { - return .finishProducersAndCallOnTerminate(Array(producerContinuations.map { $0.1 }), onTerminate) - } - return .resumeContinuationWithCancellationErrorAndFinishProducersAndCallOnTerminate( - consumerContinuation, - Array(producerContinuations.map { $0.1 }), - onTerminate - ) - - case .sourceFinished, .finished: return .none - } - } - } -} diff --git a/Sources/OpenAPIURLSession/AsyncBackpressuredStream/NIOLock.swift b/Sources/OpenAPIURLSession/AsyncBackpressuredStream/NIOLock.swift deleted file mode 100644 index 783d37e..0000000 --- a/Sources/OpenAPIURLSession/AsyncBackpressuredStream/NIOLock.swift +++ /dev/null @@ -1,213 +0,0 @@ -//===----------------------------------------------------------------------===// -// -// This source file is part of the SwiftOpenAPIGenerator open source project -// -// Copyright (c) 2023 Apple Inc. and the SwiftOpenAPIGenerator project authors -// Licensed under Apache License v2.0 -// -// See LICENSE.txt for license information -// See CONTRIBUTORS.txt for the list of SwiftOpenAPIGenerator project authors -// -// SPDX-License-Identifier: Apache-2.0 -// -//===----------------------------------------------------------------------===// -// swift-format-ignore-file -//===----------------------------------------------------------------------===// -// -// This source file is part of the SwiftNIO open source project -// -// Copyright (c) 2017-2022 Apple Inc. and the SwiftNIO project authors -// Licensed under Apache License v2.0 -// -// See LICENSE.txt for license information -// See CONTRIBUTORS.txt for the list of SwiftNIO project authors -// -// SPDX-License-Identifier: Apache-2.0 -// -//===----------------------------------------------------------------------===// - -#if canImport(Darwin) -import Darwin -#elseif os(Windows) -import ucrt -import WinSDK -#elseif canImport(Glibc) -import Glibc -#elseif canImport(Musl) -import Musl -#else -#error("The concurrency NIOLock module was unable to identify your C library.") -#endif - -#if os(Windows) -@usableFromInline typealias LockPrimitive = SRWLOCK -#else -@usableFromInline typealias LockPrimitive = pthread_mutex_t -#endif - -@usableFromInline enum LockOperations {} - -extension LockOperations { - @inlinable static func create(_ mutex: UnsafeMutablePointer) { - 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)") - #endif - } - - @inlinable static func destroy(_ mutex: UnsafeMutablePointer) { - mutex.assertValidAlignment() - - #if os(Windows) - // SRWLOCK does not need to be free'd - #else - let err = pthread_mutex_destroy(mutex) - precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)") - #endif - } - - @inlinable static func lock(_ mutex: UnsafeMutablePointer) { - 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)") - #endif - } - - @inlinable static func unlock(_ mutex: UnsafeMutablePointer) { - 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)") - #endif - } -} - -// Tail allocate both the mutex and a generic value using ManagedBuffer. -// Both the header pointer and the elements pointer are stable for -// the class's entire lifetime. -// -// However, for safety reasons, we elect to place the lock in the "elements" -// section of the buffer instead of the head. The reasoning here is subtle, -// so buckle in. -// -// _As a practical matter_, the implementation of ManagedBuffer ensures that -// the pointer to the header is stable across the lifetime of the class, and so -// each time you call `withUnsafeMutablePointers` or `withUnsafeMutablePointerToHeader` -// the value of the header pointer will be the same. This is because ManagedBuffer uses -// `Builtin.addressOf` to load the value of the header, and that does ~magic~ to ensure -// that it does not invoke any weird Swift accessors that might copy the value. -// -// _However_, the header is also available via the `.header` field on the ManagedBuffer. -// This presents a problem! The reason there's an issue is that `Builtin.addressOf` and friends -// do not interact with Swift's exclusivity model. That is, the various `with` functions do not -// conceptually trigger a mutating access to `.header`. For elements this isn't a concern because -// there's literally no other way to perform the access, but for `.header` it's entirely possible -// to accidentally recursively read it. -// -// Our implementation is free from these issues, so we don't _really_ need to worry about it. -// However, out of an abundance of caution, we store the Value in the header, and the LockPrimitive -// in the trailing elements. We still don't use `.header`, but it's better to be safe than sorry, -// and future maintainers will be happier that we were cautious. -// -// See also: https://github.com/apple/swift/pull/40000 -@usableFromInline final class LockStorage: ManagedBuffer { - - @inlinable static func create(value: Value) -> Self { - let buffer = Self.create(minimumCapacity: 1) { _ in return value } - let storage = unsafeDowncast(buffer, to: Self.self) - - storage.withUnsafeMutablePointers { _, lockPtr in LockOperations.create(lockPtr) } - - return storage - } - - @inlinable func lock() { self.withUnsafeMutablePointerToElements { lockPtr in LockOperations.lock(lockPtr) } } - - @inlinable func unlock() { self.withUnsafeMutablePointerToElements { lockPtr in LockOperations.unlock(lockPtr) } } - - @inlinable deinit { self.withUnsafeMutablePointerToElements { lockPtr in LockOperations.destroy(lockPtr) } } - - @inlinable func withLockPrimitive(_ body: (UnsafeMutablePointer) throws -> T) rethrows -> T { - try self.withUnsafeMutablePointerToElements { lockPtr in return try body(lockPtr) } - } - - @inlinable func withLockedValue(_ mutate: (inout Value) throws -> T) rethrows -> T { - try self.withUnsafeMutablePointers { valuePtr, lockPtr in LockOperations.lock(lockPtr) - defer { LockOperations.unlock(lockPtr) } - return try mutate(&valuePtr.pointee) - } - } -} - -extension LockStorage: @unchecked Sendable {} - -/// A threading lock based on `libpthread` instead of `libdispatch`. -/// -/// - note: ``NIOLock`` 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 -/// one used by NIO. On Windows, the lock is based on the substantially similar -/// `SRWLOCK` type. -public struct NIOLock { - @usableFromInline internal let _storage: LockStorage - - /// Create a new lock. - @inlinable public init() { self._storage = .create(value: ()) } - - /// Acquire the lock. - /// - /// Whenever possible, consider using `withLock` instead of this method and - /// `unlock`, to simplify lock handling. - @inlinable public func lock() { self._storage.lock() } - - /// Release the lock. - /// - /// Whenever possible, consider using `withLock` instead of this method and - /// `lock`, to simplify lock handling. - @inlinable public func unlock() { self._storage.unlock() } - - @inlinable internal func withLockPrimitive(_ body: (UnsafeMutablePointer) throws -> T) rethrows - -> T - { return try self._storage.withLockPrimitive(body) } -} - -extension NIOLock { - /// Acquire the lock for the duration of the given block. - /// - /// This convenience method should be preferred to `lock` and `unlock` in - /// most situations, as it ensures that the lock will be released regardless - /// of how `body` exits. - /// - /// - Parameter body: The block to execute while holding the lock. - /// - Returns: The value returned by the block. - @inlinable public func withLock(_ body: () throws -> T) rethrows -> T { - self.lock() - defer { self.unlock() } - return try body() - } - - @inlinable public func withLockVoid(_ body: () throws -> Void) rethrows { try self.withLock(body) } -} - -extension NIOLock: Sendable {} - -extension UnsafeMutablePointer { - @inlinable func assertValidAlignment() { - assert(UInt(bitPattern: self) % UInt(MemoryLayout.alignment) == 0) - } -} diff --git a/Sources/OpenAPIURLSession/BufferedStream/BufferedStream.swift b/Sources/OpenAPIURLSession/BufferedStream/BufferedStream.swift new file mode 100644 index 0000000..797a86e --- /dev/null +++ b/Sources/OpenAPIURLSession/BufferedStream/BufferedStream.swift @@ -0,0 +1,1973 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the SwiftOpenAPIGenerator open source project +// +// Copyright (c) 2023 Apple Inc. and the SwiftOpenAPIGenerator project authors +// Licensed under Apache License v2.0 +// +// See LICENSE.txt for license information +// See CONTRIBUTORS.txt for the list of SwiftOpenAPIGenerator project authors +// +// SPDX-License-Identifier: Apache-2.0 +// +//===----------------------------------------------------------------------===// +// swift-format-ignore-file +//===----------------------------------------------------------------------===// +// +// This source file is part of the Swift.org open source project +// +// Copyright (c) 2020-2021 Apple Inc. and the Swift project authors +// Licensed under Apache License v2.0 with Runtime Library Exception +// +// See https://swift.org/LICENSE.txt for license information +// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors +// +//===----------------------------------------------------------------------===// + +import DequeModule + +/// An asynchronous sequence generated from an error-throwing closure that +/// calls a continuation to produce new elements. +/// +/// `BufferedStream` conforms to `AsyncSequence`, providing a convenient +/// way to create an asynchronous sequence without manually implementing an +/// asynchronous iterator. In particular, an asynchronous stream is well-suited +/// to adapt callback- or delegation-based APIs to participate with +/// `async`-`await`. +/// +/// In contrast to `AsyncStream`, this type can throw an error from the awaited +/// `next()`, which terminates the stream with the thrown error. +/// +/// You initialize an `BufferedStream` with a closure that receives an +/// `BufferedStream.Continuation`. Produce elements in this closure, then +/// provide them to the stream by calling the continuation's `yield(_:)` method. +/// When there are no further elements to produce, call the continuation's +/// `finish()` method. This causes the sequence iterator to produce a `nil`, +/// which terminates the sequence. If an error occurs, call the continuation's +/// `finish(throwing:)` method, which causes the iterator's `next()` method to +/// throw the error to the awaiting call point. The continuation is `Sendable`, +/// which permits calling it from concurrent contexts external to the iteration +/// of the `BufferedStream`. +/// +/// An arbitrary source of elements can produce elements faster than they are +/// consumed by a caller iterating over them. Because of this, `BufferedStream` +/// defines a buffering behavior, allowing the stream to buffer a specific +/// number of oldest or newest elements. By default, the buffer limit is +/// `Int.max`, which means it's unbounded. +/// +/// ### Adapting Existing Code to Use Streams +/// +/// To adapt existing callback code to use `async`-`await`, use the callbacks +/// to provide values to the stream, by using the continuation's `yield(_:)` +/// method. +/// +/// Consider a hypothetical `QuakeMonitor` type that provides callers with +/// `Quake` instances every time it detects an earthquake. To receive callbacks, +/// callers set a custom closure as the value of the monitor's +/// `quakeHandler` property, which the monitor calls back as necessary. Callers +/// can also set an `errorHandler` to receive asynchronous error notifications, +/// such as the monitor service suddenly becoming unavailable. +/// +/// class QuakeMonitor { +/// var quakeHandler: ((Quake) -> Void)? +/// var errorHandler: ((Error) -> Void)? +/// +/// func startMonitoring() {…} +/// func stopMonitoring() {…} +/// } +/// +/// To adapt this to use `async`-`await`, extend the `QuakeMonitor` to add a +/// `quakes` property, of type `BufferedStream`. In the getter for +/// this property, return an `BufferedStream`, whose `build` closure -- +/// called at runtime to create the stream -- uses the continuation to +/// perform the following steps: +/// +/// 1. Creates a `QuakeMonitor` instance. +/// 2. Sets the monitor's `quakeHandler` property to a closure that receives +/// each `Quake` instance and forwards it to the stream by calling the +/// continuation's `yield(_:)` method. +/// 3. Sets the monitor's `errorHandler` property to a closure that receives +/// any error from the monitor and forwards it to the stream by calling the +/// continuation's `finish(throwing:)` method. This causes the stream's +/// iterator to throw the error and terminate the stream. +/// 4. Sets the continuation's `onTermination` property to a closure that +/// calls `stopMonitoring()` on the monitor. +/// 5. Calls `startMonitoring` on the `QuakeMonitor`. +/// +/// ``` +/// extension QuakeMonitor { +/// +/// static var throwingQuakes: BufferedStream { +/// BufferedStream { continuation in +/// let monitor = QuakeMonitor() +/// monitor.quakeHandler = { quake in +/// continuation.yield(quake) +/// } +/// monitor.errorHandler = { error in +/// continuation.finish(throwing: error) +/// } +/// continuation.onTermination = { @Sendable _ in +/// monitor.stopMonitoring() +/// } +/// monitor.startMonitoring() +/// } +/// } +/// } +/// ``` +/// +/// +/// Because the stream is an `AsyncSequence`, the call point uses the +/// `for`-`await`-`in` syntax to process each `Quake` instance as produced by the stream: +/// +/// do { +/// for try await quake in quakeStream { +/// print("Quake: \(quake.date)") +/// } +/// print("Stream done.") +/// } catch { +/// print("Error: \(error)") +/// } +/// +@available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *) +@usableFromInline +internal struct BufferedStream { + @usableFromInline + final class _Backing: Sendable { + @usableFromInline + let storage: _BackPressuredStorage + + @inlinable + init(storage: _BackPressuredStorage) { + self.storage = storage + } + + deinit { + self.storage.sequenceDeinitialized() + } + } + + @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()`. + @usableFromInline + internal struct Iterator: AsyncIteratorProtocol { + @usableFromInline + final class _Backing { + @usableFromInline + let storage: _BackPressuredStorage + + @inlinable + 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) + } + + @usableFromInline + var implementation: _Implementation + + @inlinable + init(implementation: _Implementation) { + self.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() + } + } + } + + /// 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, *) +extension BufferedStream: Sendable where Element: Sendable {} + +@usableFromInline +internal struct _ManagedCriticalState: @unchecked Sendable { + @usableFromInline + let lock: LockedValueBox + + @inlinable + 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 { + @inlinable + 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. + @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)) + ) + } + + @inlinable + init(internalBackPressureStrategy: _InternalBackPressureStrategy) { + self._internalBackPressureStrategy = internalBackPressureStrategy + } + + @usableFromInline + let _internalBackPressureStrategy: _InternalBackPressureStrategy + } + + /// 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 + 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. + @usableFromInline + final class _Backing: Sendable { + @usableFromInline + let storage: _BackPressuredStorage + + @inlinable + init(storage: _BackPressuredStorage) { + self.storage = storage + } + + deinit { + self.storage.sourceDeinitialized() + } + } + + /// 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 + } + } + + @usableFromInline + var _backing: _Backing + + @inlinable + internal init(storage: _BackPressuredStorage) { + self._backing = .init(storage: storage) + } + + /// 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 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)) + } + + /// 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 + ) + } + + /// 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 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)) + } + } + + /// 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) + } + + /// 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) + } + } + } + + /// 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 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)) + } + } + + /// 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) + } + + @inlinable + 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 + 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. + @inlinable + 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 + } + + @inlinable + 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 + } + + @inlinable + 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 + } + + @inlinable + 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 + 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 + } + } + + @inlinable + mutating func didConsume(elements: Deque.SubSequence) -> Bool { + switch self { + case .watermark(var strategy): + let result = strategy.didConsume(elements: elements) + self = .watermark(strategy) + return result + } + } + + @inlinable + 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. + @usableFromInline + final class _BackPressuredStorage: Sendable { + /// The state machine + @usableFromInline + let _stateMachine: _ManagedCriticalState<_StateMachine> + + @usableFromInline + var onTermination: (@Sendable () -> Void)? { + set { + self._stateMachine.withCriticalRegion { + $0._onTermination = newValue + } + } + get { + self._stateMachine.withCriticalRegion { + $0._onTermination + } + } + } + + @inlinable + init( + backPressureStrategy: _InternalBackPressureStrategy + ) { + self._stateMachine = .init(.init(backPressureStrategy: backPressureStrategy)) + } + + @inlinable + func sequenceDeinitialized() { + let action = self._stateMachine.withCriticalRegion { + $0.sequenceDeinitialized() + } + + switch action { + case .callOnTermination(let onTermination): + onTermination?() + + case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): + for producerContinuation in producerContinuations { + producerContinuation(.failure(AlreadyFinishedError())) + } + onTermination?() + + case .none: + break + } + } + + @inlinable + func iteratorInitialized() { + self._stateMachine.withCriticalRegion { + $0.iteratorInitialized() + } + } + + @inlinable + func iteratorDeinitialized() { + let action = self._stateMachine.withCriticalRegion { + $0.iteratorDeinitialized() + } + + switch action { + case .callOnTermination(let onTermination): + onTermination?() + + case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): + for producerContinuation in producerContinuations { + producerContinuation(.failure(AlreadyFinishedError())) + } + onTermination?() + + case .none: + break + } + } + + @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 + } + } + + @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() + } + } + + @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 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 + } + } + + @inlinable + func finish(_ failure: (any Error)?) { + let action = self._stateMachine.withCriticalRegion { + $0.finish(failure) + } + + 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) + } + + onTermination?() + + case .resumeProducers(let producerContinuations): + for producerContinuation in producerContinuations { + producerContinuation(.failure(AlreadyFinishedError())) + } + + case .none: + break + } + } + + @inlinable + func next() async throws -> Element? { + let action = self._stateMachine.withCriticalRegion { + $0.next() + } + + switch action { + case .returnElement(let element): + return element + + 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() + } + } + + @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?() + + 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. + @usableFromInline + struct _StateMachine { + @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)? + + @inlinable + 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 + + @inlinable + 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)? + + @inlinable + 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. + @usableFromInline + var _state: _State + + // The ID used for the next CallbackToken. + @usableFromInline + var nextCallbackTokenID: UInt = 0 + + @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. + @inlinable + init( + backPressureStrategy: _InternalBackPressureStrategy + ) { + self._state = .initial( + .init( + backPressureStrategy: backPressureStrategy, + iteratorInitialized: false + ) + ) + } + + /// 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 + 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)? + ) + } + + @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") + } + } + + @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") + } + } + + /// 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)? + ) + } + + @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") + } + } + + /// 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 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 }) + ) + } + + case .sourceFinished, .finished: + // This is normal and we just have to tolerate it + return .none + + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } + + /// 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 + ) + } + } + } + + @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") + } + } + + /// 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) + } + + @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 `cancelProducer()`. + @usableFromInline + enum CancelProducerAction { + /// Indicates that the producer should be notified about cancellation. + case resumeProducerWithCancellationError((Result) -> Void) + } + + @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") + } + } + + /// 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] + ) + } + + @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 }) + ) + } + + case .sourceFinished, .finished: + // If the source has finished, finishing again has no effect. + return .none + + 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 + } + + @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") + } + + 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") + } + } + + /// 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 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) + + return .none + } + + 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 .resumeConsumerWithElement(continuation, element) + } else { + // We are returning the queued failure now and can transition to finished + self._state = .finished(iteratorInitialized: sourceFinished.iteratorInitialized) + + return .resumeConsumerWithErrorAndCallOnTermination( + continuation, + sourceFinished.failure, + sourceFinished.onTermination + ) + } + + case .finished: + return .resumeConsumerWithNil(continuation) + + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } + + /// 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)?) + } + + @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/OpenAPIURLSession/BufferedStream/Lock.swift b/Sources/OpenAPIURLSession/BufferedStream/Lock.swift new file mode 100644 index 0000000..38f7890 --- /dev/null +++ b/Sources/OpenAPIURLSession/BufferedStream/Lock.swift @@ -0,0 +1,254 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the SwiftOpenAPIGenerator open source project +// +// Copyright (c) 2023 Apple Inc. and the SwiftOpenAPIGenerator project authors +// Licensed under Apache License v2.0 +// +// See LICENSE.txt for license information +// See CONTRIBUTORS.txt for the list of SwiftOpenAPIGenerator project authors +// +// SPDX-License-Identifier: Apache-2.0 +// +//===----------------------------------------------------------------------===// +// swift-format-ignore-file +//===----------------------------------------------------------------------===// +// +// This source file is part of the SwiftNIO open source project +// +// Copyright (c) 2017-2022 Apple Inc. and the SwiftNIO project authors +// Licensed under Apache License v2.0 +// +// See LICENSE.txt for license information +// See CONTRIBUTORS.txt for the list of SwiftNIO project authors +// +// SPDX-License-Identifier: Apache-2.0 +// +//===----------------------------------------------------------------------===// + +#if canImport(Darwin) +import Darwin +#elseif canImport(Glibc) +import Glibc +#endif + +@usableFromInline +typealias LockPrimitive = pthread_mutex_t + +@usableFromInline +enum LockOperations {} + +extension LockOperations { + @inlinable + static func create(_ mutex: UnsafeMutablePointer) { + mutex.assertValidAlignment() + + 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)") + } + + @inlinable + static func destroy(_ mutex: UnsafeMutablePointer) { + mutex.assertValidAlignment() + + let err = pthread_mutex_destroy(mutex) + precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)") + } + + @inlinable + static func lock(_ mutex: UnsafeMutablePointer) { + mutex.assertValidAlignment() + + let err = pthread_mutex_lock(mutex) + precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)") + } + + @inlinable + static func unlock(_ mutex: UnsafeMutablePointer) { + mutex.assertValidAlignment() + + let err = pthread_mutex_unlock(mutex) + precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)") + } +} + +// Tail allocate both the mutex and a generic value using ManagedBuffer. +// Both the header pointer and the elements pointer are stable for +// the class's entire lifetime. +// +// However, for safety reasons, we elect to place the lock in the "elements" +// section of the buffer instead of the head. The reasoning here is subtle, +// so buckle in. +// +// _As a practical matter_, the implementation of ManagedBuffer ensures that +// the pointer to the header is stable across the lifetime of the class, and so +// each time you call `withUnsafeMutablePointers` or `withUnsafeMutablePointerToHeader` +// the value of the header pointer will be the same. This is because ManagedBuffer uses +// `Builtin.addressOf` to load the value of the header, and that does ~magic~ to ensure +// that it does not invoke any weird Swift accessors that might copy the value. +// +// _However_, the header is also available via the `.header` field on the ManagedBuffer. +// This presents a problem! The reason there's an issue is that `Builtin.addressOf` and friends +// do not interact with Swift's exclusivity model. That is, the various `with` functions do not +// conceptually trigger a mutating access to `.header`. For elements this isn't a concern because +// there's literally no other way to perform the access, but for `.header` it's entirely possible +// to accidentally recursively read it. +// +// Our implementation is free from these issues, so we don't _really_ need to worry about it. +// However, out of an abundance of caution, we store the Value in the header, and the LockPrimitive +// in the trailing elements. We still don't use `.header`, but it's better to be safe than sorry, +// and future maintainers will be happier that we were cautious. +// +// See also: https://github.com/apple/swift/pull/40000 +@usableFromInline +final class LockStorage: ManagedBuffer { + + @inlinable + static func create(value: Value) -> Self { + let buffer = Self.create(minimumCapacity: 1) { _ in + return value + } + let storage = unsafeDowncast(buffer, to: Self.self) + + storage.withUnsafeMutablePointers { _, lockPtr in + LockOperations.create(lockPtr) + } + + return storage + } + + @inlinable + func lock() { + self.withUnsafeMutablePointerToElements { lockPtr in + LockOperations.lock(lockPtr) + } + } + + @inlinable + func unlock() { + self.withUnsafeMutablePointerToElements { lockPtr in + LockOperations.unlock(lockPtr) + } + } + + @inlinable + deinit { + self.withUnsafeMutablePointerToElements { lockPtr in + LockOperations.destroy(lockPtr) + } + } + + @inlinable + func withLockPrimitive( + _ body: (UnsafeMutablePointer) throws -> T + ) rethrows -> T { + try self.withUnsafeMutablePointerToElements { lockPtr in + return try body(lockPtr) + } + } + + @inlinable + func withLockedValue(_ mutate: (inout Value) throws -> T) rethrows -> T { + try self.withUnsafeMutablePointers { valuePtr, lockPtr in + LockOperations.lock(lockPtr) + defer { LockOperations.unlock(lockPtr) } + return try mutate(&valuePtr.pointee) + } + } +} + +extension LockStorage: @unchecked Sendable {} + +/// A threading lock based on `libpthread` instead of `libdispatch`. +/// +/// - 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 +/// one used by NIO. On Windows, the lock is based on the substantially similar +/// `SRWLOCK` type. +@usableFromInline +struct Lock { + @usableFromInline + internal let _storage: LockStorage + + /// Create a new lock. + @inlinable + init() { + self._storage = .create(value: ()) + } + + /// Acquire the lock. + /// + /// Whenever possible, consider using `withLock` instead of this method and + /// `unlock`, to simplify lock handling. + @inlinable + func lock() { + self._storage.lock() + } + + /// Release the lock. + /// + /// Whenever possible, consider using `withLock` instead of this method and + /// `lock`, to simplify lock handling. + @inlinable + func unlock() { + self._storage.unlock() + } + + @inlinable + internal func withLockPrimitive( + _ body: (UnsafeMutablePointer) throws -> T + ) rethrows -> T { + return try self._storage.withLockPrimitive(body) + } +} + +extension Lock { + /// Acquire the lock for the duration of the given block. + /// + /// This convenience method should be preferred to `lock` and `unlock` in + /// most situations, as it ensures that the lock will be released regardless + /// of how `body` exits. + /// + /// - Parameter body: The block to execute while holding the lock. + /// - Returns: The value returned by the block. + @inlinable + func withLock(_ body: () throws -> T) rethrows -> T { + self.lock() + defer { + self.unlock() + } + return try body() + } +} + +extension Lock: Sendable {} + +extension UnsafeMutablePointer { + @inlinable + func assertValidAlignment() { + assert(UInt(bitPattern: self) % UInt(MemoryLayout.alignment) == 0) + } +} + +@usableFromInline +struct LockedValueBox { + @usableFromInline + let storage: LockStorage + + @inlinable + init(_ value: Value) { + self.storage = .create(value: value) + } + + @inlinable + func withLockedValue(_ mutate: (inout Value) throws -> T) rethrows -> T { + return try self.storage.withLockedValue(mutate) + } +} + +extension LockedValueBox: Sendable where Value: Sendable {} diff --git a/Sources/OpenAPIURLSession/URLSessionBidirectionalStreaming/BidirectionalStreamingURLSessionDelegate.swift b/Sources/OpenAPIURLSession/URLSessionBidirectionalStreaming/BidirectionalStreamingURLSessionDelegate.swift index f35e380..f457582 100644 --- a/Sources/OpenAPIURLSession/URLSessionBidirectionalStreaming/BidirectionalStreamingURLSessionDelegate.swift +++ b/Sources/OpenAPIURLSession/URLSessionBidirectionalStreaming/BidirectionalStreamingURLSessionDelegate.swift @@ -53,14 +53,16 @@ final class BidirectionalStreamingURLSessionDelegate: NSObject, URLSessionTaskDe let requestBody: HTTPBody? var hasAlreadyIteratedRequestBody: Bool - var hasSuspendedURLSessionTask: Bool + /// In addition to the callback lock, there is one point of rentrancy, where the response stream callback gets fired + /// immediately, for this we have a different lock, which protects `hasSuspendedURLSessionTask`. + var hasSuspendedURLSessionTask: LockedValueBox let requestStreamBufferSize: Int var requestStream: HTTPBodyOutputStreamBridge? typealias ResponseContinuation = CheckedContinuation var responseContinuation: ResponseContinuation? - typealias ResponseBodyStream = AsyncBackpressuredStream + typealias ResponseBodyStream = BufferedStream var responseBodyStream: ResponseBodyStream var responseBodyStreamSource: ResponseBodyStream.Source @@ -74,22 +76,19 @@ final class BidirectionalStreamingURLSessionDelegate: NSObject, URLSessionTaskDe /// /// Therefore, even though the `suspend()`, `resume()`, and `cancel()` URLSession methods are thread-safe, we need /// to protect any mutable state within the delegate itself. - let callbackLock = NIOLock() - - /// In addition to the callback lock, there is one point of rentrancy, where the response stream callback gets fired - /// immediately, for this we have a different lock, which protects `hasSuspendedURLSessionTask`. - let hasSuspendedURLSessionTaskLock = NIOLock() + let callbackLock = Lock() /// Use `bidirectionalStreamingRequest(for:baseURL:requestBody:requestStreamBufferSize:responseStreamWatermarks:)`. init(requestBody: HTTPBody?, requestStreamBufferSize: Int, responseStreamWatermarks: (low: Int, high: Int)) { self.requestBody = requestBody self.hasAlreadyIteratedRequestBody = false - self.hasSuspendedURLSessionTask = false + self.hasSuspendedURLSessionTask = LockedValueBox(false) self.requestStreamBufferSize = requestStreamBufferSize - (self.responseBodyStream, self.responseBodyStreamSource) = AsyncBackpressuredStream.makeStream( - backPressureStrategy: .highLowWatermarkWithElementCounts( - lowWatermark: responseStreamWatermarks.low, - highWatermark: responseStreamWatermarks.high + (self.responseBodyStream, self.responseBodyStreamSource) = ResponseBodyStream.makeStream( + backPressureStrategy: .customWatermark( + low: responseStreamWatermarks.low, + high: responseStreamWatermarks.high, + waterLevelForElement: { $0.count } ) ) } @@ -125,8 +124,9 @@ final class BidirectionalStreamingURLSessionDelegate: NSObject, URLSessionTaskDe do { switch try responseBodyStreamSource.write(contentsOf: CollectionOfOne(ArraySlice(data))) { case .produceMore: break - case .enqueueCallback(let writeToken): - let shouldActuallyEnqueueCallback = hasSuspendedURLSessionTaskLock.withLock { + 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 @@ -137,13 +137,13 @@ final class BidirectionalStreamingURLSessionDelegate: NSObject, URLSessionTaskDe return true } if shouldActuallyEnqueueCallback { - responseBodyStreamSource.enqueueCallback(writeToken: writeToken) { result in - self.hasSuspendedURLSessionTaskLock.withLock { + responseBodyStreamSource.enqueueCallback(callbackToken: callbackToken) { result in + self.hasSuspendedURLSessionTask.withLockedValue { hasSuspendedURLSessionTask in switch result { case .success: debug("Task delegate: response stream callback, resuming task") dataTask.resume() - self.hasSuspendedURLSessionTask = false + hasSuspendedURLSessionTask = false case .failure(let error): debug("Task delegate: response stream callback, cancelling task, error: \(error)") dataTask.cancel() diff --git a/Tests/OpenAPIURLSessionTests/AsyncBackpressuredStreamTests/AsyncBackpressuredStreamTests.swift b/Tests/OpenAPIURLSessionTests/AsyncBackpressuredStreamTests/AsyncBackpressuredStreamTests.swift deleted file mode 100644 index 72ca741..0000000 --- a/Tests/OpenAPIURLSessionTests/AsyncBackpressuredStreamTests/AsyncBackpressuredStreamTests.swift +++ /dev/null @@ -1,359 +0,0 @@ -//===----------------------------------------------------------------------===// -// -// This source file is part of the SwiftOpenAPIGenerator open source project -// -// Copyright (c) 2023 Apple Inc. and the SwiftOpenAPIGenerator project authors -// Licensed under Apache License v2.0 -// -// See LICENSE.txt for license information -// See CONTRIBUTORS.txt for the list of SwiftOpenAPIGenerator project authors -// -// SPDX-License-Identifier: Apache-2.0 -// -//===----------------------------------------------------------------------===// -//===----------------------------------------------------------------------===// -// -// This source file is part of the Swift.org open source project -// -// Copyright (c) 2020-2021 Apple Inc. and the Swift project authors -// Licensed under Apache License v2.0 with Runtime Library Exception -// -// See https://swift.org/LICENSE.txt for license information -// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors -// -//===----------------------------------------------------------------------===// -import XCTest -@testable import OpenAPIURLSession - -final class AsyncBackpressuredStreamTests: XCTestCase { - func testYield() async throws { - let (stream, source) = AsyncBackpressuredStream.makeStream( - of: Int.self, - backPressureStrategy: .highLowWatermark(lowWatermark: 5, highWatermark: 10) - ) - - try await source.asyncWrite(contentsOf: [1, 2, 3, 4, 5, 6]) - source.finish(throwing: nil) - - let result = try await stream.collect() - XCTAssertEqual(result, [1, 2, 3, 4, 5, 6]) - } - - func testBackPressure() async throws { - let (stream, source) = AsyncBackpressuredStream.makeStream( - of: Int.self, - backPressureStrategy: .highLowWatermark(lowWatermark: 2, highWatermark: 4) - ) - - let (backPressureEventStream, backPressureEventContinuation) = AsyncStream.makeStream(of: Void.self) - - try await withThrowingTaskGroup(of: Void.self) { group in - group.addTask { - while true { - backPressureEventContinuation.yield(()) - debug("Yielding") - try await source.asyncWrite(contentsOf: [1]) - } - } - - var backPressureEventIterator = backPressureEventStream.makeAsyncIterator() - var iterator = stream.makeAsyncIterator() - - await backPressureEventIterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() - - debug("Waited 4 times") - - _ = try await iterator.next() - _ = try await iterator.next() - _ = try await iterator.next() - debug("Consumed three") - - await backPressureEventIterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() - - group.cancelAll() - } - } - - func testBackPressureSync() async throws { - let (stream, source) = AsyncBackpressuredStream.makeStream( - of: Int.self, - backPressureStrategy: .highLowWatermark(lowWatermark: 2, highWatermark: 4) - ) - - let (backPressureEventStream, backPressureEventContinuation) = AsyncStream.makeStream(of: Void.self) - - try await withThrowingTaskGroup(of: Void.self) { group in - group.addTask { - @Sendable func yield() { - backPressureEventContinuation.yield(()) - debug("Yielding") - source.write(contentsOf: [1]) { result in - switch result { - case .success: yield() - - case .failure: debug("Stopping to yield") - } - } - } - - yield() - } - - var backPressureEventIterator = backPressureEventStream.makeAsyncIterator() - var iterator = stream.makeAsyncIterator() - - await backPressureEventIterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() - - debug("Waited 4 times") - - _ = try await iterator.next() - _ = try await iterator.next() - _ = try await iterator.next() - debug("Consumed three") - - await backPressureEventIterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() - - group.cancelAll() - } - } - - func testWatermarkBackPressureStrategy() async throws { - typealias Strategy = AsyncBackpressuredStream.HighLowWatermarkBackPressureStrategy - var strategy = Strategy(lowWatermark: 2, highWatermark: 3) - - XCTAssertEqual(strategy.currentWatermark, 0) - XCTAssertEqual(strategy.didYield(elements: Slice([])), true) - XCTAssertEqual(strategy.currentWatermark, 0) - XCTAssertEqual(strategy.didYield(elements: Slice(["*", "*"])), true) - XCTAssertEqual(strategy.currentWatermark, 2) - XCTAssertEqual(strategy.didYield(elements: Slice(["*"])), false) - XCTAssertEqual(strategy.currentWatermark, 3) - XCTAssertEqual(strategy.didYield(elements: Slice(["*"])), false) - XCTAssertEqual(strategy.currentWatermark, 4) - - XCTAssertEqual(strategy.currentWatermark, 4) - XCTAssertEqual(strategy.didConsume(elements: Slice([])), false) - XCTAssertEqual(strategy.currentWatermark, 4) - XCTAssertEqual(strategy.didConsume(elements: Slice(["*", "*"])), false) - XCTAssertEqual(strategy.currentWatermark, 2) - XCTAssertEqual(strategy.didConsume(elements: Slice(["*"])), true) - XCTAssertEqual(strategy.currentWatermark, 1) - XCTAssertEqual(strategy.didConsume(elements: Slice(["*"])), true) - XCTAssertEqual(strategy.currentWatermark, 0) - XCTAssertEqual(strategy.didConsume(elements: Slice([])), true) - XCTAssertEqual(strategy.currentWatermark, 0) - } - - func testWatermarkWithoutElementCountsBackPressureStrategy() async throws { - typealias Strategy = AsyncBackpressuredStream<[String], any Error>.HighLowWatermarkBackPressureStrategy - var strategy = Strategy(lowWatermark: 2, highWatermark: 3) - - XCTAssertEqual(strategy.currentWatermark, 0) - XCTAssertEqual(strategy.didYield(elements: Slice([])), true) - XCTAssertEqual(strategy.currentWatermark, 0) - XCTAssertEqual(strategy.didYield(elements: Slice([["*", "*"]])), true) - XCTAssertEqual(strategy.currentWatermark, 1) - XCTAssertEqual(strategy.didYield(elements: Slice([["*", "*"]])), true) - XCTAssertEqual(strategy.currentWatermark, 2) - - XCTAssertEqual(strategy.currentWatermark, 2) - XCTAssertEqual(strategy.didConsume(elements: Slice([])), false) - XCTAssertEqual(strategy.currentWatermark, 2) - XCTAssertEqual(strategy.didConsume(elements: Slice([["*", "*"]])), true) - XCTAssertEqual(strategy.currentWatermark, 1) - XCTAssertEqual(strategy.didConsume(elements: Slice([["*", "*"]])), true) - XCTAssertEqual(strategy.currentWatermark, 0) - XCTAssertEqual(strategy.didConsume(elements: Slice([])), true) - XCTAssertEqual(strategy.currentWatermark, 0) - } - - func testWatermarkWithElementCountsBackPressureStrategy() async throws { - typealias Strategy = AsyncBackpressuredStream<[String], any Error>.HighLowWatermarkBackPressureStrategy - var strategy = Strategy(lowWatermark: 2, highWatermark: 3, waterLevelForElement: { $0.count }) - XCTAssertEqual(strategy.currentWatermark, 0) - XCTAssertEqual(strategy.didYield(elements: Slice([])), true) - XCTAssertEqual(strategy.currentWatermark, 0) - XCTAssertEqual(strategy.didYield(elements: Slice([["*", "*"]])), true) - XCTAssertEqual(strategy.currentWatermark, 2) - XCTAssertEqual(strategy.didYield(elements: Slice([["*", "*"]])), false) - XCTAssertEqual(strategy.currentWatermark, 4) - - XCTAssertEqual(strategy.currentWatermark, 4) - XCTAssertEqual(strategy.didConsume(elements: Slice([])), false) - XCTAssertEqual(strategy.currentWatermark, 4) - XCTAssertEqual(strategy.didConsume(elements: Slice([["*", "*"]])), false) - XCTAssertEqual(strategy.currentWatermark, 2) - XCTAssertEqual(strategy.didConsume(elements: Slice([["*", "*"]])), true) - XCTAssertEqual(strategy.currentWatermark, 0) - XCTAssertEqual(strategy.didConsume(elements: Slice([])), true) - XCTAssertEqual(strategy.currentWatermark, 0) - } - - func testWritingOverWatermark() async throws { - try await withThrowingTaskGroup(of: Void.self) { group in - let (stream, continuation) = AsyncBackpressuredStream - .makeStream(backPressureStrategy: .highLowWatermark(lowWatermark: 1, highWatermark: 1)) - - group.addTask { - for i in 1...10 { - debug("Producer writing element \(i)...") - let writeResult = try continuation.write(contentsOf: CollectionOfOne(i)) - debug("Producer wrote element \(i), result = \(writeResult)") - // ignore backpressure result and write again anyway - } - debug("Producer finished") - continuation.finish(throwing: nil) - } - - var iterator = stream.makeAsyncIterator() - var numElementsConsumed = 0 - var expectedNextValue = 1 - while true { - debug("Consumer reading element...") - guard let element = try await iterator.next() else { break } - XCTAssertEqual(element, expectedNextValue) - debug("Consumer read element: \(element), expected: \(expectedNextValue)") - numElementsConsumed += 1 - expectedNextValue += 1 - } - XCTAssertEqual(numElementsConsumed, 10) - - group.cancelAll() - } - } - - func testStateMachineSuspendNext() async throws { - typealias Stream = AsyncBackpressuredStream - - var strategy = Stream.InternalBackPressureStrategy.highLowWatermark(.init(lowWatermark: 1, highWatermark: 1)) - _ = strategy.didYield(elements: Slice([1, 2, 3])) - var stateMachine = Stream.StateMachine(backPressureStrategy: strategy, onTerminate: nil) - stateMachine.state = .streaming( - backPressureStrategy: strategy, - buffer: [1, 2, 3], - consumerContinuation: nil, - producerContinuations: [], - cancelledAsyncProducers: [], - hasOutstandingDemand: false, - iteratorInitialized: true, - onTerminate: nil - ) - - guard case .streaming(_, let buffer, let consumerContinuation, _, _, _, _, _) = stateMachine.state else { - XCTFail("Unexpected state: \(stateMachine.state)") - return - } - XCTAssertEqual(buffer, [1, 2, 3]) - XCTAssertNil(consumerContinuation) - - _ = try await withCheckedThrowingContinuation { continuation in - let action = stateMachine.suspendNext(continuation: continuation) - - guard case .resumeContinuationWithElement(_, let element) = action else { - XCTFail("Unexpected action: \(action)") - return - } - XCTAssertEqual(element, 1) - - guard case .streaming(_, let buffer, let consumerContinuation, _, _, _, _, _) = stateMachine.state else { - XCTFail("Unexpected state: \(stateMachine.state)") - return - } - XCTAssertEqual(buffer, [2, 3]) - XCTAssertNil(consumerContinuation) - - continuation.resume(returning: element) - } - } -} - -extension AsyncBackpressuredStream.Source.WriteResult: CustomStringConvertible { - // swift-format-ignore: AllPublicDeclarationsHaveDocumentation - public var description: String { - switch self { - case .enqueueCallback: return "enqueueCallBack" - case .produceMore: return "produceMore" - } - } -} - -extension AsyncBackpressuredStream.StateMachine.SuspendNextAction: CustomStringConvertible { - // swift-format-ignore: AllPublicDeclarationsHaveDocumentation - public var description: String { - switch self { - case .none: return "none" - case .resumeContinuationWithElement: return "resumeContinuationWithElement" - case .resumeContinuationWithElementAndProducers: return "resumeContinuationWithElementAndProducers" - case .resumeContinuationWithFailureAndCallOnTerminate: return "resumeContinuationWithFailureAndCallOnTerminate" - case .resumeContinuationWithNil: return "resumeContinuationWithNil" - } - } -} - -extension AsyncBackpressuredStream.StateMachine.State: CustomStringConvertible { - // swift-format-ignore: AllPublicDeclarationsHaveDocumentation - public var description: String { - switch self { - case .initial: return "initial" - case .streaming(_, let buffer, let consumer, let producers, _, let demand, _, _): - return - "streaming(buffer.count: \(buffer.count), consumer: \(consumer != nil ? "yes" : "no"), producers: \(producers), demand: \(demand))" - case .finished: return "finished" - case .sourceFinished: return "sourceFinished" - } - } -} - -extension AsyncSequence { - /// Collect all elements in the sequence into an array. - fileprivate func collect() async rethrows -> [Element] { - try await self.reduce(into: []) { accumulated, next in accumulated.append(next) } - } -} - -extension AsyncBackpressuredStream.StateMachine.NextAction: CustomStringConvertible { - // swift-format-ignore: AllPublicDeclarationsHaveDocumentation - public var description: String { - switch self { - case .returnNil: return "returnNil" - case .returnElementAndResumeProducers: return "returnElementAndResumeProducers" - case .returnFailureAndCallOnTerminate: return "returnFailureAndCallOnTerminate" - case .returnElement: return "returnElement" - case .suspendTask: return "suspendTask" - } - } -} - -extension AsyncBackpressuredStream.StateMachine.WriteAction: CustomStringConvertible { - // swift-format-ignore: AllPublicDeclarationsHaveDocumentation - public var description: String { - switch self { - case .returnProduceMore: return "returnProduceMore" - case .returnEnqueue: return "returnEnqueue" - case .resumeConsumerContinuationAndReturnProduceMore: return "resumeConsumerContinuationAndReturnProduceMore" - case .resumeConsumerContinuationAndReturnEnqueue: return "resumeConsumerContinuationAndReturnEnqueue" - case .throwFinishedError: return "throwFinishedError" - } - } -} - -extension AsyncBackpressuredStream.StateMachine.EnqueueProducerAction: CustomStringConvertible { - // swift-format-ignore: AllPublicDeclarationsHaveDocumentation - public var description: String { - switch self { - case .resumeProducer: return "resumeProducer" - case .resumeProducerWithCancellationError: return "resumeProducerWithCancellationError" - case .none: return "none" - } - } -} diff --git a/Tests/OpenAPIURLSessionTests/BufferedStreamTests/BufferedStreamTests.swift b/Tests/OpenAPIURLSessionTests/BufferedStreamTests/BufferedStreamTests.swift new file mode 100644 index 0000000..a2ff290 --- /dev/null +++ b/Tests/OpenAPIURLSessionTests/BufferedStreamTests/BufferedStreamTests.swift @@ -0,0 +1,1185 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the SwiftOpenAPIGenerator open source project +// +// Copyright (c) 2023 Apple Inc. and the SwiftOpenAPIGenerator project authors +// Licensed under Apache License v2.0 +// +// See LICENSE.txt for license information +// See CONTRIBUTORS.txt for the list of SwiftOpenAPIGenerator project authors +// +// SPDX-License-Identifier: Apache-2.0 +// +//===----------------------------------------------------------------------===// +// swift-format-ignore-file +//===----------------------------------------------------------------------===// +// +// This source file is part of the Swift.org open source project +// +// Copyright (c) 2020-2021 Apple Inc. and the Swift project authors +// Licensed under Apache License v2.0 with Runtime Library Exception +// +// See https://swift.org/LICENSE.txt for license information +// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors +// +//===----------------------------------------------------------------------===// + +import XCTest + +@testable import OpenAPIURLSession + +@available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *) +final class BufferedStreamTests: XCTestCase { + // MARK: - sequenceDeinitialized + + func testSequenceDeinitialized_whenNoIterator() async throws { + var (stream, source): (BufferedStream?, BufferedStream.Source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 5, high: 10) + ) + + let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() + source.onTermination = { + onTerminationContinuation.finish() + } + + await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + while !Task.isCancelled { + onTerminationContinuation.yield() + try await Task.sleep(nanoseconds: 200_000_000) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + withExtendedLifetime(stream) {} + stream = nil + + let terminationResult: Void? = await onTerminationIterator.next() + XCTAssertNil(terminationResult) + + do { + _ = try { try source.write(2) }() + XCTFail("Expected an error to be thrown") + } catch { + XCTAssertTrue(error is AlreadyFinishedError) + } + + group.cancelAll() + } + } + + func testSequenceDeinitialized_whenIterator() async throws { + var (stream, source): (BufferedStream?, BufferedStream.Source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 5, high: 10) + ) + + var iterator = stream?.makeAsyncIterator() + + let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() + source.onTermination = { + onTerminationContinuation.finish() + } + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + while !Task.isCancelled { + onTerminationContinuation.yield() + try await Task.sleep(nanoseconds: 200_000_000) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + try withExtendedLifetime(stream) { + let writeResult = try source.write(1) + writeResult.assertIsProducerMore() + } + + stream = nil + + do { + let writeResult = try { try source.write(2) }() + writeResult.assertIsProducerMore() + } catch { + XCTFail("Expected no error to be thrown") + } + + let element1 = try await iterator?.next() + XCTAssertEqual(element1, 1) + let element2 = try await iterator?.next() + XCTAssertEqual(element2, 2) + + group.cancelAll() + } + } + + func testSequenceDeinitialized_whenFinished() async throws { + var (stream, source): (BufferedStream?, BufferedStream.Source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 5, high: 10) + ) + + let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() + source.onTermination = { + onTerminationContinuation.finish() + } + + await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + while !Task.isCancelled { + onTerminationContinuation.yield() + try await Task.sleep(nanoseconds: 200_000_000) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + withExtendedLifetime(stream) { + source.finish(throwing: nil) + } + + stream = nil + + let terminationResult: Void? = await onTerminationIterator.next() + XCTAssertNil(terminationResult) + + do { + _ = try { try source.write(1) }() + XCTFail("Expected an error to be thrown") + } catch { + XCTAssertTrue(error is AlreadyFinishedError) + } + + group.cancelAll() + } + } + + func testSequenceDeinitialized_whenStreaming_andSuspendedProducer() async throws { + var (stream, source): (BufferedStream?, BufferedStream.Source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 1, high: 2) + ) + + _ = try { try source.write(1) }() + + do { + try await withCheckedThrowingContinuation { continuation in + source.write(1) { result in + continuation.resume(with: result) + } + + stream = nil + _ = stream?.makeAsyncIterator() + } + } catch { + XCTAssertTrue(error is AlreadyFinishedError) + } + } + + // MARK: - iteratorInitialized + + func testIteratorInitialized_whenInitial() async throws { + let (stream, _) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 5, high: 10) + ) + + _ = stream.makeAsyncIterator() + } + + func testIteratorInitialized_whenStreaming() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 5, high: 10) + ) + + try await source.write(1) + + var iterator = stream.makeAsyncIterator() + let element = try await iterator.next() + XCTAssertEqual(element, 1) + } + + func testIteratorInitialized_whenSourceFinished() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 5, high: 10) + ) + + try await source.write(1) + source.finish(throwing: nil) + + var iterator = stream.makeAsyncIterator() + let element1 = try await iterator.next() + XCTAssertEqual(element1, 1) + let element2 = try await iterator.next() + XCTAssertNil(element2) + } + + func testIteratorInitialized_whenFinished() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 5, high: 10) + ) + + source.finish(throwing: nil) + + var iterator = stream.makeAsyncIterator() + let element = try await iterator.next() + XCTAssertNil(element) + } + + // MARK: - iteratorDeinitialized + + func testIteratorDeinitialized_whenInitial() async throws { + var (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 5, high: 10) + ) + + let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() + source.onTermination = { + onTerminationContinuation.finish() + } + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + while !Task.isCancelled { + onTerminationContinuation.yield() + try await Task.sleep(nanoseconds: 200_000_000) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + var iterator: BufferedStream.AsyncIterator? = stream.makeAsyncIterator() + iterator = nil + _ = try await iterator?.next() + + let terminationResult: Void? = await onTerminationIterator.next() + XCTAssertNil(terminationResult) + + group.cancelAll() + } + } + + func testIteratorDeinitialized_whenStreaming() async throws { + var (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 5, high: 10) + ) + + let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() + source.onTermination = { + onTerminationContinuation.finish() + } + + try await source.write(1) + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + while !Task.isCancelled { + onTerminationContinuation.yield() + try await Task.sleep(nanoseconds: 200_000_000) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + var iterator: BufferedStream.AsyncIterator? = stream.makeAsyncIterator() + iterator = nil + _ = try await iterator?.next() + + let terminationResult: Void? = await onTerminationIterator.next() + XCTAssertNil(terminationResult) + + group.cancelAll() + } + } + + func testIteratorDeinitialized_whenSourceFinished() async throws { + var (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 5, high: 10) + ) + + let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() + source.onTermination = { + onTerminationContinuation.finish() + } + + try await source.write(1) + source.finish(throwing: nil) + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + while !Task.isCancelled { + onTerminationContinuation.yield() + try await Task.sleep(nanoseconds: 200_000_000) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + var iterator: BufferedStream.AsyncIterator? = stream.makeAsyncIterator() + iterator = nil + _ = try await iterator?.next() + + let terminationResult: Void? = await onTerminationIterator.next() + XCTAssertNil(terminationResult) + + group.cancelAll() + } + } + + func testIteratorDeinitialized_whenFinished() async throws { + var (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 5, high: 10) + ) + + let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() + source.onTermination = { + onTerminationContinuation.finish() + } + + source.finish(throwing: nil) + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + while !Task.isCancelled { + onTerminationContinuation.yield() + try await Task.sleep(nanoseconds: 200_000_000) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + var iterator: BufferedStream.AsyncIterator? = stream.makeAsyncIterator() + iterator = nil + _ = try await iterator?.next() + + let terminationResult: Void? = await onTerminationIterator.next() + XCTAssertNil(terminationResult) + + group.cancelAll() + } + } + + func testIteratorDeinitialized_whenStreaming_andSuspendedProducer() async throws { + var (stream, source): (BufferedStream?, BufferedStream.Source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 1, high: 2) + ) + + var iterator: BufferedStream.AsyncIterator? = stream?.makeAsyncIterator() + stream = nil + + _ = try { try source.write(1) }() + + do { + try await withCheckedThrowingContinuation { continuation in + source.write(1) { result in + continuation.resume(with: result) + } + + iterator = nil + } + } catch { + XCTAssertTrue(error is AlreadyFinishedError) + } + + _ = try await iterator?.next() + } + + // MARK: - sourceDeinitialized + + func testSourceDeinitialized_whenInitial() async throws { + var (stream, source): (BufferedStream, BufferedStream.Source?) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 5, high: 10) + ) + + let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() + source?.onTermination = { + onTerminationContinuation.finish() + } + + await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + while !Task.isCancelled { + onTerminationContinuation.yield() + try await Task.sleep(nanoseconds: 200_000_000) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + source = nil + + let terminationResult: Void? = await onTerminationIterator.next() + XCTAssertNil(terminationResult) + + group.cancelAll() + } + + withExtendedLifetime(stream) {} + } + + func testSourceDeinitialized_whenStreaming_andEmptyBuffer() async throws { + var (stream, source): (BufferedStream, BufferedStream.Source?) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 5, high: 10) + ) + + let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() + source?.onTermination = { + onTerminationContinuation.finish() + } + + try await source?.write(1) + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + while !Task.isCancelled { + onTerminationContinuation.yield() + try await Task.sleep(nanoseconds: 200_000_000) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + var iterator: BufferedStream.AsyncIterator? = stream.makeAsyncIterator() + _ = try await iterator?.next() + + source = nil + + let terminationResult: Void? = await onTerminationIterator.next() + XCTAssertNil(terminationResult) + + group.cancelAll() + } + } + + func testSourceDeinitialized_whenStreaming_andNotEmptyBuffer() async throws { + var (stream, source): (BufferedStream, BufferedStream.Source?) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 5, high: 10) + ) + + let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() + source?.onTermination = { + onTerminationContinuation.finish() + } + + try await source?.write(1) + try await source?.write(2) + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + while !Task.isCancelled { + onTerminationContinuation.yield() + try await Task.sleep(nanoseconds: 200_000_000) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + var iterator: BufferedStream.AsyncIterator? = stream.makeAsyncIterator() + _ = try await iterator?.next() + + source = nil + + _ = await onTerminationIterator.next() + + _ = try await iterator?.next() + _ = try await iterator?.next() + + let terminationResult: Void? = await onTerminationIterator.next() + XCTAssertNil(terminationResult) + + group.cancelAll() + } + } + + func testSourceDeinitialized_whenSourceFinished() async throws { + var (stream, source): (BufferedStream, BufferedStream.Source?) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 5, high: 10) + ) + + let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() + source?.onTermination = { + onTerminationContinuation.finish() + } + + try await source?.write(1) + try await source?.write(2) + source?.finish(throwing: nil) + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + while !Task.isCancelled { + onTerminationContinuation.yield() + try await Task.sleep(nanoseconds: 200_000_000) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + var iterator: BufferedStream.AsyncIterator? = stream.makeAsyncIterator() + _ = try await iterator?.next() + + source = nil + + _ = await onTerminationIterator.next() + + _ = try await iterator?.next() + _ = try await iterator?.next() + + let terminationResult: Void? = await onTerminationIterator.next() + XCTAssertNil(terminationResult) + + group.cancelAll() + } + } + + func testSourceDeinitialized_whenFinished() async throws { + var (stream, source): (BufferedStream, BufferedStream.Source?) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 5, high: 10) + ) + + let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() + source?.onTermination = { + onTerminationContinuation.finish() + } + + source?.finish(throwing: nil) + + await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + while !Task.isCancelled { + onTerminationContinuation.yield() + try await Task.sleep(nanoseconds: 200_000_000) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + _ = stream.makeAsyncIterator() + + source = nil + + _ = await onTerminationIterator.next() + + let terminationResult: Void? = await onTerminationIterator.next() + XCTAssertNil(terminationResult) + + group.cancelAll() + } + } + + func testSourceDeinitialized_whenStreaming_andSuspendedProducer() async throws { + var (stream, source): (BufferedStream, BufferedStream.Source?) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 0, high: 0) + ) + let (producerStream, producerContinuation) = AsyncThrowingStream.makeStream() + var iterator = stream.makeAsyncIterator() + + source?.write(1) { + producerContinuation.yield(with: $0) + } + + _ = try await iterator.next() + source = nil + + do { + try await producerStream.first { _ in true } + XCTFail("We expected to throw here") + } catch { + XCTAssertTrue(error is AlreadyFinishedError) + } + } + + // MARK: - write + + func testWrite_whenInitial() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 2, high: 5) + ) + + try await source.write(1) + + var iterator = stream.makeAsyncIterator() + let element = try await iterator.next() + XCTAssertEqual(element, 1) + } + + func testWrite_whenStreaming_andNoConsumer() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 2, high: 5) + ) + + try await source.write(1) + try await source.write(2) + + var iterator = stream.makeAsyncIterator() + let element1 = try await iterator.next() + XCTAssertEqual(element1, 1) + let element2 = try await iterator.next() + XCTAssertEqual(element2, 2) + } + + func testWrite_whenStreaming_andSuspendedConsumer() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 2, high: 5) + ) + + try await withThrowingTaskGroup(of: Int?.self) { group in + group.addTask { + return try await stream.first { _ in true } + } + + // This is always going to be a bit racy since we need the call to next() suspend + try await Task.sleep(nanoseconds: 500_000_000) + + try await source.write(1) + let element = try await group.next() + XCTAssertEqual(element, 1) + } + } + + func testWrite_whenStreaming_andSuspendedConsumer_andEmptySequence() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 2, high: 5) + ) + + try await withThrowingTaskGroup(of: Int?.self) { group in + group.addTask { + return try await stream.first { _ in true } + } + + // This is always going to be a bit racy since we need the call to next() suspend + try await Task.sleep(nanoseconds: 500_000_000) + + try await source.write(contentsOf: []) + try await source.write(contentsOf: [1]) + let element = try await group.next() + XCTAssertEqual(element, 1) + } + } + + // MARK: - enqueueProducer + + func testEnqueueProducer_whenStreaming_andAndCancelled() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 1, high: 2) + ) + + let (producerStream, producerSource) = AsyncThrowingStream.makeStream() + + try await source.write(1) + + let writeResult = try { try source.write(2) }() + + switch writeResult { + case .produceMore: + preconditionFailure() + case .enqueueCallback(let callbackToken): + source.cancelCallback(callbackToken: callbackToken) + + source.enqueueCallback(callbackToken: callbackToken) { result in + producerSource.yield(with: result) + } + } + + do { + _ = try await producerStream.first { _ in true } + XCTFail("Expected an error to be thrown") + } catch { + XCTAssertTrue(error is CancellationError) + } + + let element = try await stream.first { _ in true } + XCTAssertEqual(element, 1) + } + + func testEnqueueProducer_whenStreaming_andAndCancelled_andAsync() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 1, high: 2) + ) + + try await source.write(1) + + await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + try await source.write(2) + } + + group.cancelAll() + do { + try await group.next() + XCTFail("Expected an error to be thrown") + } catch { + XCTAssertTrue(error is CancellationError) + } + } + + let element = try await stream.first { _ in true } + XCTAssertEqual(element, 1) + } + + func testEnqueueProducer_whenStreaming_andInterleaving() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 1, high: 1) + ) + var iterator = stream.makeAsyncIterator() + + let (producerStream, producerSource) = AsyncThrowingStream.makeStream() + + let writeResult = try { try source.write(1) }() + + switch writeResult { + case .produceMore: + preconditionFailure() + case .enqueueCallback(let callbackToken): + let element = try await iterator.next() + XCTAssertEqual(element, 1) + + source.enqueueCallback(callbackToken: callbackToken) { result in + producerSource.yield(with: result) + } + } + + do { + _ = try await producerStream.first { _ in true } + } catch { + XCTFail("Expected no error to be thrown") + } + } + + func testEnqueueProducer_whenStreaming_andSuspending() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 1, high: 1) + ) + var iterator = stream.makeAsyncIterator() + + let (producerStream, producerSource) = AsyncThrowingStream.makeStream() + + let writeResult = try { try source.write(1) }() + + switch writeResult { + case .produceMore: + preconditionFailure() + case .enqueueCallback(let callbackToken): + source.enqueueCallback(callbackToken: callbackToken) { result in + producerSource.yield(with: result) + } + } + + let element = try await iterator.next() + XCTAssertEqual(element, 1) + + do { + _ = try await producerStream.first { _ in true } + } catch { + XCTFail("Expected no error to be thrown") + } + } + + func testEnqueueProducer_whenFinished() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 1, high: 1) + ) + var iterator = stream.makeAsyncIterator() + + let (producerStream, producerSource) = AsyncThrowingStream.makeStream() + + let writeResult = try { try source.write(1) }() + + switch writeResult { + case .produceMore: + preconditionFailure() + case .enqueueCallback(let callbackToken): + source.finish(throwing: nil) + + source.enqueueCallback(callbackToken: callbackToken) { result in + producerSource.yield(with: result) + } + } + + let element = try await iterator.next() + XCTAssertEqual(element, 1) + + do { + _ = try await producerStream.first { _ in true } + XCTFail("Expected an error to be thrown") + } catch { + XCTAssertTrue(error is AlreadyFinishedError) + } + } + + // MARK: - cancelProducer + + func testCancelProducer_whenStreaming() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 1, high: 2) + ) + + let (producerStream, producerSource) = AsyncThrowingStream.makeStream() + + try await source.write(1) + + let writeResult = try { try source.write(2) }() + + switch writeResult { + case .produceMore: + preconditionFailure() + case .enqueueCallback(let callbackToken): + source.enqueueCallback(callbackToken: callbackToken) { result in + producerSource.yield(with: result) + } + + source.cancelCallback(callbackToken: callbackToken) + } + + do { + _ = try await producerStream.first { _ in true } + XCTFail("Expected an error to be thrown") + } catch { + XCTAssertTrue(error is CancellationError) + } + + let element = try await stream.first { _ in true } + XCTAssertEqual(element, 1) + } + + func testCancelProducer_whenSourceFinished() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 1, high: 2) + ) + + let (producerStream, producerSource) = AsyncThrowingStream.makeStream() + + try await source.write(1) + + let writeResult = try { try source.write(2) }() + + switch writeResult { + case .produceMore: + preconditionFailure() + case .enqueueCallback(let callbackToken): + source.enqueueCallback(callbackToken: callbackToken) { result in + producerSource.yield(with: result) + } + + source.finish(throwing: nil) + + source.cancelCallback(callbackToken: callbackToken) + } + + do { + _ = try await producerStream.first { _ in true } + XCTFail("Expected an error to be thrown") + } catch { + XCTAssertTrue(error is AlreadyFinishedError) + } + + let element = try await stream.first { _ in true } + XCTAssertEqual(element, 1) + } + + // MARK: - finish + + func testFinish_whenStreaming_andConsumerSuspended() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 1, high: 1) + ) + + try await withThrowingTaskGroup(of: Int?.self) { group in + group.addTask { + return try await stream.first { $0 == 2 } + } + + // This is always going to be a bit racy since we need the call to next() suspend + try await Task.sleep(nanoseconds: 500_000_000) + + source.finish(throwing: nil) + let element = try await group.next() + XCTAssertEqual(element, .some(nil)) + } + } + + func testFinish_whenInitial() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 1, high: 1) + ) + + source.finish(throwing: CancellationError()) + + do { + for try await _ in stream {} + XCTFail("Expected an error to be thrown") + } catch { + XCTAssertTrue(error is CancellationError) + } + + } + + // MARK: - Backpressure + + func testBackPressure() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 2, high: 4) + ) + + let (backPressureEventStream, backPressureEventContinuation) = AsyncStream.makeStream( + of: Void.self + ) + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + while true { + backPressureEventContinuation.yield(()) + try await source.write(contentsOf: [1]) + } + } + + var backPressureEventIterator = backPressureEventStream.makeAsyncIterator() + var iterator = stream.makeAsyncIterator() + + await backPressureEventIterator.next() + await backPressureEventIterator.next() + await backPressureEventIterator.next() + await backPressureEventIterator.next() + + _ = try await iterator.next() + _ = try await iterator.next() + _ = try await iterator.next() + + await backPressureEventIterator.next() + await backPressureEventIterator.next() + await backPressureEventIterator.next() + + group.cancelAll() + } + } + + func testBackPressureSync() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 2, high: 4) + ) + + let (backPressureEventStream, backPressureEventContinuation) = AsyncStream.makeStream( + of: Void.self + ) + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + @Sendable func yield() { + backPressureEventContinuation.yield(()) + source.write(contentsOf: [1]) { result in + switch result { + case .success: + yield() + + case .failure: + break + } + } + } + + yield() + } + + var backPressureEventIterator = backPressureEventStream.makeAsyncIterator() + var iterator = stream.makeAsyncIterator() + + await backPressureEventIterator.next() + await backPressureEventIterator.next() + await backPressureEventIterator.next() + await backPressureEventIterator.next() + + _ = try await iterator.next() + _ = try await iterator.next() + _ = try await iterator.next() + + await backPressureEventIterator.next() + await backPressureEventIterator.next() + await backPressureEventIterator.next() + + group.cancelAll() + } + } + + func testThrowsError() async throws { + let (stream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 2, high: 4) + ) + + try await source.write(1) + try await source.write(2) + source.finish(throwing: CancellationError()) + + var elements = [Int]() + var iterator = stream.makeAsyncIterator() + + do { + while let element = try await iterator.next() { + elements.append(element) + } + XCTFail("Expected an error to be thrown") + } catch { + XCTAssertTrue(error is CancellationError) + XCTAssertEqual(elements, [1, 2]) + } + + let element = try await iterator.next() + XCTAssertNil(element) + } + + func testAsyncSequenceWrite() async throws { + let (stream, continuation) = AsyncStream.makeStream() + let (backpressuredStream, source) = BufferedStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 2, high: 4) + ) + + continuation.yield(1) + continuation.yield(2) + continuation.finish() + + try await source.write(contentsOf: stream) + source.finish(throwing: nil) + + let elements = try await backpressuredStream.collect() + XCTAssertEqual(elements, [1, 2]) + } + + func testWatermarkBackPressureStrategy() async throws { + typealias Strategy = BufferedStream._WatermarkBackPressureStrategy + var strategy = Strategy(low: 2, high: 3) + + XCTAssertEqual(strategy._current, 0) + XCTAssertEqual(strategy.didYield(elements: Slice([])), true) + XCTAssertEqual(strategy._current, 0) + XCTAssertEqual(strategy.didYield(elements: Slice(["*", "*"])), true) + XCTAssertEqual(strategy._current, 2) + XCTAssertEqual(strategy.didYield(elements: Slice(["*"])), false) + XCTAssertEqual(strategy._current, 3) + XCTAssertEqual(strategy.didYield(elements: Slice(["*"])), false) + XCTAssertEqual(strategy._current, 4) + + XCTAssertEqual(strategy._current, 4) + XCTAssertEqual(strategy.didConsume(elements: Slice([])), false) + XCTAssertEqual(strategy._current, 4) + XCTAssertEqual(strategy.didConsume(elements: Slice(["*", "*"])), false) + XCTAssertEqual(strategy._current, 2) + XCTAssertEqual(strategy.didConsume(elements: Slice(["*"])), true) + XCTAssertEqual(strategy._current, 1) + XCTAssertEqual(strategy.didConsume(elements: Slice(["*"])), true) + XCTAssertEqual(strategy._current, 0) + XCTAssertEqual(strategy.didConsume(elements: Slice([])), true) + XCTAssertEqual(strategy._current, 0) + } + + func testWatermarkWithoutElementCountsBackPressureStrategy() async throws { + typealias Strategy = BufferedStream<[String]>._WatermarkBackPressureStrategy + var strategy = Strategy(low: 2, high: 3) + + XCTAssertEqual(strategy._current, 0) + XCTAssertEqual(strategy.didYield(elements: Slice([])), true) + XCTAssertEqual(strategy._current, 0) + XCTAssertEqual(strategy.didYield(elements: Slice([["*", "*"]])), true) + XCTAssertEqual(strategy._current, 1) + XCTAssertEqual(strategy.didYield(elements: Slice([["*", "*"]])), true) + XCTAssertEqual(strategy._current, 2) + + XCTAssertEqual(strategy._current, 2) + XCTAssertEqual(strategy.didConsume(elements: Slice([])), false) + XCTAssertEqual(strategy._current, 2) + XCTAssertEqual(strategy.didConsume(elements: Slice([["*", "*"]])), true) + XCTAssertEqual(strategy._current, 1) + XCTAssertEqual(strategy.didConsume(elements: Slice([["*", "*"]])), true) + XCTAssertEqual(strategy._current, 0) + XCTAssertEqual(strategy.didConsume(elements: Slice([])), true) + XCTAssertEqual(strategy._current, 0) + } + + func testWatermarkWithElementCountsBackPressureStrategy() async throws { + typealias Strategy = BufferedStream<[String]>._WatermarkBackPressureStrategy + var strategy = Strategy(low: 2, high: 3, waterLevelForElement: { $0.count }) + XCTAssertEqual(strategy._current, 0) + XCTAssertEqual(strategy.didYield(elements: Slice([])), true) + XCTAssertEqual(strategy._current, 0) + XCTAssertEqual(strategy.didYield(elements: Slice([["*", "*"]])), true) + XCTAssertEqual(strategy._current, 2) + XCTAssertEqual(strategy.didYield(elements: Slice([["*", "*"]])), false) + XCTAssertEqual(strategy._current, 4) + + XCTAssertEqual(strategy._current, 4) + XCTAssertEqual(strategy.didConsume(elements: Slice([])), false) + XCTAssertEqual(strategy._current, 4) + XCTAssertEqual(strategy.didConsume(elements: Slice([["*", "*"]])), false) + XCTAssertEqual(strategy._current, 2) + XCTAssertEqual(strategy.didConsume(elements: Slice([["*", "*"]])), true) + XCTAssertEqual(strategy._current, 0) + XCTAssertEqual(strategy.didConsume(elements: Slice([])), true) + XCTAssertEqual(strategy._current, 0) + } +} + +extension BufferedStream.Source.WriteResult { + func assertIsProducerMore() { + switch self { + case .produceMore: + return + + case .enqueueCallback: + XCTFail("Expected produceMore") + } + } + + func assertIsEnqueueCallback() { + switch self { + case .produceMore: + XCTFail("Expected enqueueCallback") + + case .enqueueCallback: + return + } + } +} diff --git a/Tests/OpenAPIURLSessionTests/TestUtils.swift b/Tests/OpenAPIURLSessionTests/TestUtils.swift index 7798e77..a6dc085 100644 --- a/Tests/OpenAPIURLSessionTests/TestUtils.swift +++ b/Tests/OpenAPIURLSessionTests/TestUtils.swift @@ -153,14 +153,34 @@ final class LockedValueBox: @unchecked Sendable where Value: Sendable { } } +#if swift(<5.9) extension AsyncStream { - // We have this here until we drop 5.8, since it's in the standard library in Swift 5.9+. static func makeStream( of elementType: Element.Type = Element.self, bufferingPolicy limit: Self.Continuation.BufferingPolicy = .unbounded ) -> (stream: Self, continuation: Self.Continuation) { var continuation: Self.Continuation! - let stream = Self(elementType, bufferingPolicy: limit) { continuation = $0 } + let stream = Self(bufferingPolicy: limit) { continuation = $0 } return (stream, continuation) } } + +extension AsyncThrowingStream { + static func makeStream( + of elementType: Element.Type = Element.self, + throwing failureType: Failure.Type = Failure.self, + bufferingPolicy limit: Self.Continuation.BufferingPolicy = .unbounded + ) -> (stream: Self, continuation: Self.Continuation) where Failure == Error { + var continuation: Self.Continuation! + let stream = Self(bufferingPolicy: limit) { continuation = $0 } + return (stream, continuation!) + } +} +#endif // #if swift(<5.9) + +extension AsyncSequence { + /// Collect all elements in the sequence into an array. + func collect() async throws -> [Element] { + try await self.reduce(into: []) { accumulated, next in accumulated.append(next) } + } +}