From 65e8957854abff207f6b132af89f17c9dfc498e9 Mon Sep 17 00:00:00 2001 From: Franz Busch Date: Tue, 19 Dec 2023 11:59:17 +0100 Subject: [PATCH 1/7] Add `AsyncBackpressuredStream` proposal and implementation # Motivation The pitch to add external backpressure support to the standard libraries `AsyncStream` got returned for revision since there are larger open questions around `AsyncSequence`. However, having external backpressure in a source asynchronous sequence is becoming more and more important. # Modification This PR adds a modified proposal and implementation that brings the Swift Evolution proposal over to Swift Async Algorithms. --- .../{NNNN-channel.md => 0012-channel.md} | 0 Evolution/{NNNN-chunk.md => 0013-chunk.md} | 0 ...NNN-rate-limits.md => 0014-rate-limits.md} | 0 ...{NNNN-reductions.md => 0015-reductions.md} | 0 Evolution/0016-backpressured-stream.md | 731 ++++++++++ .../AsyncBackPressuredStream.swift | 425 ++++++ .../AsyncBackPressuredStream+Internal.swift | 1192 +++++++++++++++++ .../AsyncNonThrowingBackPressuredStream.swift | 429 ++++++ .../BackPressuredStreamTests.swift | 1166 ++++++++++++++++ 9 files changed, 3943 insertions(+) rename Evolution/{NNNN-channel.md => 0012-channel.md} (100%) rename Evolution/{NNNN-chunk.md => 0013-chunk.md} (100%) rename Evolution/{NNNN-rate-limits.md => 0014-rate-limits.md} (100%) rename Evolution/{NNNN-reductions.md => 0015-reductions.md} (100%) create mode 100644 Evolution/0016-backpressured-stream.md create mode 100644 Sources/AsyncAlgorithms/BackPressuredStream/AsyncBackPressuredStream.swift create mode 100644 Sources/AsyncAlgorithms/BackpressuredStream/AsyncBackPressuredStream+Internal.swift create mode 100644 Sources/AsyncAlgorithms/BackpressuredStream/AsyncNonThrowingBackPressuredStream.swift create mode 100644 Tests/AsyncAlgorithmsTests/BackpressuredStream/BackPressuredStreamTests.swift diff --git a/Evolution/NNNN-channel.md b/Evolution/0012-channel.md similarity index 100% rename from Evolution/NNNN-channel.md rename to Evolution/0012-channel.md diff --git a/Evolution/NNNN-chunk.md b/Evolution/0013-chunk.md similarity index 100% rename from Evolution/NNNN-chunk.md rename to Evolution/0013-chunk.md diff --git a/Evolution/NNNN-rate-limits.md b/Evolution/0014-rate-limits.md similarity index 100% rename from Evolution/NNNN-rate-limits.md rename to Evolution/0014-rate-limits.md diff --git a/Evolution/NNNN-reductions.md b/Evolution/0015-reductions.md similarity index 100% rename from Evolution/NNNN-reductions.md rename to Evolution/0015-reductions.md diff --git a/Evolution/0016-backpressured-stream.md b/Evolution/0016-backpressured-stream.md new file mode 100644 index 00000000..2d8f76a8 --- /dev/null +++ b/Evolution/0016-backpressured-stream.md @@ -0,0 +1,731 @@ +# Externally backpressured support for AsyncStream + +* Proposal: [SAA-0016](0016-backpressured-stream.md) +* Authors: [Franz Busch](https://github.com/FranzBusch) +* Review Manager: TBD +* Status: **Implemented** + +## Revision +- 2023/12/18: Migrate proposal from Swift Evolution to Swift Async Algorithms. +- 2023/12/19: Add element size dependent strategy + +## Introduction + +[SE-0314](https://github.com/apple/swift-evolution/blob/main/proposals/0314-async-stream.md) +introduced new `Async[Throwing]Stream` types which act as root asynchronous +sequences. These two types allow bridging from synchronous callbacks such as +delegates to an asynchronous sequence. This proposal adds a new way of +constructing asynchronous streams with the goal to bridge backpressured systems +into an asynchronous sequence. Furthermore, this proposal aims to clarify the +cancellation behaviour both when the consuming task is cancelled and when +the production side indicates termination. + +## Motivation + +After using the `AsyncSequence` protocol and the `Async[Throwing]Stream` types +extensively over the past years, we learned that there are a few important +behavioral details that any `AsyncSequence` implementation needs to support. +These behaviors are: + +1. Backpressure +2. Multi/single consumer support +3. Downstream consumer termination +4. Upstream producer termination + +In general, `AsyncSequence` implementations can be divided into two kinds: Root +asynchronous sequences that are the source of values such as +`Async[Throwing]Stream` and transformational asynchronous sequences such as +`AsyncMapSequence`. Most transformational asynchronous sequences implicitly +fulfill the above behaviors since they forward any demand to a base asynchronous +sequence that should implement the behaviors. On the other hand, root +asynchronous sequences need to make sure that all of the above behaviors are +correctly implemented. Let's look at the current behavior of +`Async[Throwing]Stream` to see if and how it achieves these behaviors. + +### Backpressure + +Root asynchronous sequences need to relay the backpressure to the producing +system. `Async[Throwing]Stream` aims to support backpressure by providing a +configurable buffer and returning +`Async[Throwing]Stream.Continuation.YieldResult` which contains the current +buffer depth from the `yield()` method. However, only providing the current +buffer depth on `yield()` is not enough to bridge a backpressured system into +an asynchronous sequence since this can only be used as a "stop" signal but we +are missing a signal to indicate resuming the production. The only viable +backpressure strategy that can be implemented with the current API is a timed +backoff where we stop producing for some period of time and then speculatively +produce again. This is a very inefficient pattern that produces high latencies +and inefficient use of resources. + +### Multi/single consumer support + +The `AsyncSequence` protocol itself makes no assumptions about whether the +implementation supports multiple consumers or not. This allows the creation of +unicast and multicast asynchronous sequences. The difference between a unicast +and multicast asynchronous sequence is if they allow multiple iterators to be +created. `AsyncStream` does support the creation of multiple iterators and it +does handle multiple consumers correctly. On the other hand, +`AsyncThrowingStream` also supports multiple iterators but does `fatalError` +when more than one iterator has to suspend. The original proposal states: + +> As with any sequence, iterating over an AsyncStream multiple times, or +creating multiple iterators and iterating over them separately, may produce an +unexpected series of values. + +While that statement leaves room for any behavior we learned that a clear distinction +of behavior for root asynchronous sequences is beneficial; especially, when it comes to +how transformation algorithms are applied on top. + +### Downstream consumer termination + +Downstream consumer termination allows the producer to notify the consumer that +no more values are going to be produced. `Async[Throwing]Stream` does support +this by calling the `finish()` or `finish(throwing:)` methods of the +`Async[Throwing]Stream.Continuation`. However, `Async[Throwing]Stream` does not +handle the case that the `Continuation` may be `deinit`ed before one of the +finish methods is called. This currently leads to async streams that never +terminate. The behavior could be changed but it could result in semantically +breaking code. + +### Upstream producer termination + +Upstream producer termination is the inverse of downstream consumer termination +where the producer is notified once the consumption has terminated. Currently, +`Async[Throwing]Stream` does expose the `onTermination` property on the +`Continuation`. The `onTermination` closure is invoked once the consumer has +terminated. The consumer can terminate in four separate cases: + +1. The asynchronous sequence was `deinit`ed and no iterator was created +2. The iterator was `deinit`ed and the asynchronous sequence is unicast +3. The consuming task is canceled +4. The asynchronous sequence returned `nil` or threw + +`Async[Throwing]Stream` currently invokes `onTermination` in all cases; however, +since `Async[Throwing]Stream` supports multiple consumers (as discussed in the +`Multi/single consumer support` section), a single consumer task being canceled +leads to the termination of all consumers. This is not expected from multicast +asynchronous sequences in general. + +## Proposed solution + +The above motivation lays out the expected behaviors from a root asynchronous +sequence and compares them to the behaviors of `Async[Throwing]Stream`. These +are the behaviors where `Async[Throwing]Stream` diverges from the expectations. + +- Backpressure: Doesn't expose a "resumption" signal to the producer +- Multi/single consumer: + - Divergent implementation between throwing and non-throwing variant + - Supports multiple consumers even though proposal positions it as a unicast + asynchronous sequence +- Consumer termination: Doesn't handle the `Continuation` being `deinit`ed +- Producer termination: Happens on first consumer termination + +This section proposes a new type called `AsyncBackpressuredStream` that implement all of +the above-mentioned behaviors. + +### Creating an AsyncBackpressuredStream + +You can create an `AsyncBackpressuredStream` instance using the new +`makeStream(of: backpressureStrategy:)` method. This method returns you the +stream and the source. The source can be used to write new values to the +asynchronous stream. The new API specifically provides a +multi-producer/single-consumer pattern. + +```swift +let (stream, source) = AsyncBackpressuredStream.makeStream( + of: Int.self, + backpressureStrategy: .watermark(low: 2, high: 4) +) +``` + +The new proposed APIs offer three different ways to bridge a backpressured +system. The foundation is the multi-step synchronous interface. Below is an +example of how it can be used: + +```swift +do { + let writeResult = try source.write(contentsOf: sequence) + + switch writeResult { + case .produceMore: + // Trigger more production + + case .enqueueCallback(let callbackToken): + source.enqueueCallback(token: callbackToken, onProduceMore: { result in + switch result { + case .success: + // Trigger more production + case .failure(let error): + // Terminate the underlying producer + } + }) + } +} catch { + // `write(contentsOf:)` throws if the asynchronous stream already terminated +} +``` + +The above API offers the most control and highest performance when bridging a +synchronous producer to an asynchronous sequence. First, you have to write +values using the `write(contentsOf:)` which returns a `WriteResult`. The result +either indicates that more values should be produced or that a callback should +be enqueued by calling the `enqueueCallback(callbackToken: onProduceMore:)` +method. This callback is invoked once the backpressure strategy decided that +more values should be produced. This API aims to offer the most flexibility with +the greatest performance. The callback only has to be allocated in the case +where the producer needs to be suspended. + +Additionally, the above API is the building block for some higher-level and +easier-to-use APIs to write values to the asynchronous stream. Below is an +example of the two higher-level APIs. + +```swift +// Writing new values and providing a callback when to produce more +try source.write(contentsOf: sequence, onProduceMore: { result in + switch result { + case .success: + // Trigger more production + case .failure(let error): + // Terminate the underlying producer + } +}) + +// This method suspends until more values should be produced +try await source.write(contentsOf: sequence) +``` + +With the above APIs, we should be able to effectively bridge any system into an +asynchronous stream regardless if the system is callback-based, blocking or +asynchronous. + +### Downstream consumer termination + +> When reading the next two examples around termination behaviour keep in mind +that the newly proposed APIs are providing a strict unicast asynchronous sequence. + +Calling `finish()` terminates the downstream consumer. Below is an example of +this: + +```swift +// Termination through calling finish +let (stream, source) = AsyncBackpressuredStream.makeStream( + of: Int.self, + backpressureStrategy: .watermark(low: 2, high: 4) +) + +_ = try await source.write(1) +source.finish() + +for try await element in stream { + print(element) +} +print("Finished") + +// Prints +// 1 +// Finished +``` + +The other way to terminate the consumer is by deiniting the source. This has the +same effect as calling `finish()` and makes sure that no consumer is stuck +indefinitely. + +```swift +// Termination through deiniting the source +let (stream, _) = AsyncStream.makeStream( + of: Int.self, + backpressureStrategy: .watermark(low: 2, high: 4) +) + +for await element in stream { + print(element) +} +print("Finished") + +// Prints +// Finished +``` + +Trying to write more elements after the source has been finish will result in an +error thrown from the write methods. + +### Upstream producer termination + +The producer will get notified about termination through the `onTerminate` +callback. Termination of the producer happens in the following scenarios: + +```swift +// Termination through task cancellation +let (stream, source) = AsyncStream.makeStream( + of: Int.self, + backpressureStrategy: .watermark(low: 2, high: 4) +) + +let task = Task { + for await element in stream { + + } +} +task.cancel() +``` + +```swift +// Termination through deiniting the sequence +let (_, source) = AsyncStream.makeStream( + of: Int.self, + backpressureStrategy: .watermark(low: 2, high: 4) +) +``` + +```swift +// Termination through deiniting the iterator +let (stream, source) = AsyncStream.makeStream( + of: Int.self, + backpressureStrategy: .watermark(low: 2, high: 4) +) +_ = stream.makeAsyncIterator() +``` + +```swift +// Termination through calling finish +let (stream, source) = AsyncStream.makeStream( + of: Int.self, + backpressureStrategy: .watermark(low: 2, high: 4) +) + +_ = try source.write(1) +source.finish() + +for await element in stream {} + +// onTerminate will be called after all elements have been consumed +``` + +Similar to the downstream consumer termination, trying to write more elements after the +producer has been terminated will result in an error thrown from the write methods. + +## Detailed design + +All new APIs on `AsyncStream` and `AsyncThrowingStream` are as follows: + +```swift +/// Error that is thrown from the various `write` methods of the +/// ``AsyncBackpressuredStream/Source``. +/// +/// This error is thrown when the asynchronous stream is already finished when +/// trying to write new elements. +public struct AsyncBackpressuredStreamAlreadyFinishedError : Error { +} + +/// A struct that acts as a source asynchronous sequence. +/// +/// The ``AsyncBackpressuredStream`` provides a ``AsyncBackpressuredStream/Source`` to +/// write values to the stream. The source exposes the internal backpressure of the asynchronous sequence to the +/// external producer. This allows to bridge both synchronous and asynchronous producers into an asynchronous sequence. +/// +/// ## Using an AsyncBackpressuredStream +/// +/// To use an ``AsyncBackpressuredStream`` you have to create a new stream with it's source first by calling +/// the ``AsyncBackpressuredStream/makeStream(of:throwing:backpressureStrategy:)`` method. +/// Afterwards, you can pass the source to the producer and the stream to the consumer. +/// +/// ``` +/// let (stream, source) = AsyncBackpressuredStream.makeStream( +/// backpressureStrategy: .watermark(low: 2, high: 4) +/// ) +/// +/// try await withThrowingTaskGroup(of: Void.self) { group in +/// group.addTask { +/// try await source.write(1) +/// try await source.write(2) +/// try await source.write(3) +/// } +/// +/// for await element in stream { +/// print(element) +/// } +/// } +/// ``` +/// +/// The source also exposes synchronous write methods that communicate the backpressure via callbacks. +public struct AsyncBackpressuredStream: AsyncSequence { + /// Initializes a new ``AsyncBackpressuredStream`` and an ``AsyncBackpressuredStream/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. + public static func makeStream( + of elementType: Element.Type = Element.self, + throwing failureType: Failure.Type = Failure.self, + backPressureStrategy: Source.BackPressureStrategy + ) -> (`Self`, Source) where Failure == any Error +} + +extension AsyncBackpressuredStream { + /// A struct 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. + /// + /// - Important: You must terminate the source by calling one of the `finish` methods otherwise the stream's iterator + /// will never terminate. + public struct Source: Sendable { + /// A strategy that handles the backpressure of the asynchronous stream. + public struct BackPressureStrategy: Sendable { + /// A backpressure strategy using a high and low watermark to suspend and resume production respectively. + /// + /// - Parameters: + /// - low: When the number of buffered elements drops below the low watermark, producers will be resumed. + /// - high: When the number of buffered elements rises above the high watermark, producers will be suspended. + public static func watermark(low: Int, high: Int) -> BackPressureStrategy + + /// A backpressure strategy using a high and low watermark to suspend and resume production respectively. + /// + /// - Parameters: + /// - low: When the number of buffered elements drops below the low watermark, producers will be resumed. + /// - high: When the number of buffered elements rises above the high watermark, producers will be suspended. + /// - waterLevelForElement: A closure used to compute the contribution of each buffered element to the current water level. + /// + /// - Note, `waterLevelForElement` will be called on each element when it is written into the source and when + /// it is consumed from the stream, so it is recommended to provide an function that runs in constant time. + public static func watermark(low: Int, high: Int, waterLevelForElement: @escaping @Sendable (Element) -> Int) -> BackPressureStrategy + } + + /// A type that indicates the result of writing elements to the source. + public 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. + public struct CallbackToken: Sendable {} + + /// 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) + } + + /// 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 + public var onTermination: (@Sendable () -> Void)? { get set } + + /// 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. + public func write(contentsOf sequence: S) throws -> WriteResult where Element == S.Element, S : 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. + public func write(_ element: Element) throws -> WriteResult + + /// 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. + public func enqueueCallback(callbackToken: WriteResult.CallbackToken, onProduceMore: @escaping @Sendable (Result) -> Void) + + /// 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. + public func cancelCallback(callbackToken: WriteResult.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:)``. + public func write(contentsOf sequence: S, onProduceMore: @escaping @Sendable (Result) -> Void) where Element == S.Element, S : Sequence + + /// 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:)``. + public func write(_ element: Element, onProduceMore: @escaping @Sendable (Result) -> Void) + + /// 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. + public func write(contentsOf sequence: S) async throws where Element == S.Element, S : Sequence + + /// 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. + public func write(_ element: Element) async throws + + /// 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. + public func write(contentsOf sequence: S) async throws where Element == S.Element, S : AsyncSequence + + /// 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. + public func finish(throwing error: Failure?) + } +} + +extension AsyncBackpressuredStream { + public struct Iterator : AsyncIteratorProtocol { + public mutating func next() async throws -> Element? + } + public func makeAsyncIterator() -> Iterator +} + +extension AsyncBackpressuredStream: Sendable where Element: Sendable {} + +@available(*, unavailable) +extension AsyncBackpressuredStream.Iterator: Sendable {} +``` + +Additionally, this proposal adds a new `AsyncNonThrowingBackpressuredStream` +which is identical to the above except that it has only one generic parameter +and that the `Source` doesn't offer a `finish(throwing:)` method. + +```swift +public struct AsyncNonThrowingBackpressuredStream: AsyncSequence { ... } +``` + +## Comparison to other root asynchronous sequences + +### swift-async-algorithm: AsyncChannel + +The `AsyncChannel` is a multi-consumer/multi-producer root asynchronous sequence +which can be used to communicate between two tasks. It only offers asynchronous +production APIs and has no internal buffer. This means that any producer will be +suspended until its value has been consumed. `AsyncChannel` can handle multiple +consumers and resumes them in FIFO order. + +### swift-nio: NIOAsyncSequenceProducer + +The NIO team have created their own root asynchronous sequence with the goal to +provide a high performance sequence that can be used to bridge a NIO `Channel` +inbound stream into Concurrency. The `NIOAsyncSequenceProducer` is a highly +generic and fully inlinable type and quite unwiedly to use. This proposal is +heavily inspired by the learnings from this type but tries to create a more +flexible and easier to use API that fits into the standard library. + +## Source compatibility + +This change is additive and does not affect source compatibility. + +## ABI compatibility + +This change is additive and does not affect ABI compatibility. All new methods +are non-inlineable leaving us flexiblity to change the implementation in the +future. + +## Future directions + +### Adaptive backpressure strategy + +The high/low watermark strategy is common in networking code; however, there are +other strategies such as an adaptive strategy that we could offer in the future. +An adaptive strategy regulates the backpressure based on the rate of +consumption and production. With the proposed new APIs we can easily add further +strategies. + +### Element size dependent strategy + +When the stream's element is a collection type then the proposed high/low +watermark backpressure strategy might lead to unexpected results since each +element can vary in actual memory size. In the future, we could provide a new +backpressure strategy that supports inspecting the size of the collection. + +### Deprecate `Async[Throwing]Stream.Continuation` + +In the future, we could deprecate the current continuation based APIs since the +new proposed APIs are also capable of bridging non-backpressured producers by +just discarding the `WriteResult`. The only use-case that the new APIs do not +cover is the _anycast_ behaviour of the current `AsyncStream` where one can +create multiple iterators to the stream as long as no two iterators are +consuming the stream at the same time. This can be solved via additional +algorithms such as `broadcast` in the `swift-async-algorithms` package. + +To give developers more time to adopt the new APIs the deprecation of the +current APIs should be deferred to a future version. Especially since those new +APIs are not backdeployed like the current Concurrency runtime. + +### Introduce a `Writer` and an `AsyncWriter` protocol + +The newly introduced `Source` type offers a bunch of different write methods. We +have seen similar types used in other places such as file abstraction or +networking APIs. We could introduce a new `Writer` and `AsyncWriter` protocol in +the future to enable writing generic algorithms on top of writers. The `Source` +type could then conform to these new protocols. + +## Alternatives considered + +### Providing an `Async[Throwing]Stream.Continuation.onConsume` + +We could add a new closure property to the `Async[Throwing]Stream.Continuation` +which is invoked once an element has been consumed to implement a backpressure +strategy; however, this requires the usage of a synchronization mechanism since +the consumption and production often happen on separate threads. The +added complexity and performance impact led to avoiding this approach. + +### Provide a getter for the current buffer depth + +We could provide a getter for the current buffer depth on the +`Async[Throwing]Stream.Continuation`. This could be used to query the buffer +depth at an arbitrary time; however, it wouldn't allow us to implement +backpressure strategies such as high/low watermarks without continuously asking +what the buffer depth is. That would result in a very inefficient +implementation. + +### Extending `Async[Throwing]Stream.Continuation` + +Extending the current APIs to support all expected behaviors is problematic +since it would change the semantics and might lead to currently working code +misbehaving. Furthermore, extending the current APIs to support backpressure +turns out to be problematic without compromising performance or usability. + +### Introducing a new type + +We could introduce a new type such as `AsyncBackpressured[Throwing]Stream`; +however, one of the original intentions of `Async[Throwing]Stream` was to be +able to bridge backpressured systems. Furthermore, `Async[Throwing]Stream` is +the best name. Therefore, this proposal decided to provide new interfaces to +`Async[Throwing]Stream`. + +### Stick with the current `Continuation` and `yield` naming + +The proposal decided against sticking to the current names since the existing +names caused confusion to them being used in multiple places. Continuation was +both used by the `AsyncStream` but also by Swift Concurrency via +`CheckedContinuation` and `UnsafeContinuation`. Similarly, yield was used by +both `AsyncStream.Continuation.yield()`, `Task.yield()` and the `yield` keyword. +Having different names for these different concepts makes it easier to explain +their usage. The currently proposed `write` names were choosen to align with the +future direction of adding an `AsyncWriter` protocol. `Source` is a common name +in flow based systems such as Akka. Other names that were considered: + +- `enqueue` +- `send` + +### Provide the `onTermination` callback to the factory method + +During development of the new APIs, I first tried to provide the `onTermination` +callback in the `makeStream` method. However, that showed significant usability +problems in scenarios where one wants to store the source in a type and +reference `self` in the `onTermination` closure at the same time; hence, I kept +the current pattern of setting the `onTermination` closure on the source. + +### Provide a `onConsumerCancellation` callback + +During the pitch phase, it was raised that we should provide a +`onConsumerCancellation` callback which gets invoked once the asynchronous +stream notices that the consuming task got cancelled. This callback could be +used to customize how cancellation is handled by the stream e.g. one could +imagine writing a few more elements to the stream before finishing it. Right now +the stream immediately returns `nil` or throws a `CancellationError` when it +notices cancellation. This proposal decided to not provide this customization +because it opens up the possiblity that asynchronous streams are not terminating +when implemented incorrectly. Additionally, asynchronous sequences are not the +only place where task cancellation leads to an immediate error being thrown i.e. +`Task.sleep()` does the same. Hence, the value of the asynchronous not +terminating immediately brings little value when the next call in the iterating +task might throw. However, the implementation is flexible enough to add this in +the future and we can just default it to the current behaviour. + +### Create a custom type for the `Result` of the `onProduceMore` callback + +The `onProducerMore` callback takes a `Result` which is used to +indicate if the producer should produce more or if the asynchronous stream +finished. We could introduce a new type for this but the proposal decided +against it since it effectively is a result type. + +### Use an initializer instead of factory methods + +Instead of providing a `makeStream` factory method we could use an initializer +approach that takes a closure which gets the `Source` passed into. A similar API +has been offered with the `Continuation` based approach and +[SE-0388](https://github.com/apple/swift-evolution/blob/main/proposals/0388-async-stream-factory.md) +introduced new factory methods to solve some of the usability ergonomics with +the initializer based APIs. + +### Follow the `AsyncStream` & `AsyncThrowingStream` naming + +All other types that offer throwing and non-throwing variants are currently +following the naming scheme where the throwing variant gets an extra `Throwing` +in its name. Now that Swift is gaining typed throws support this would make the +type with the `Failure` parameter capable to express both throwing and +non-throwing variants. However, the less flexible type has the better name. +Hence, this proposal uses the good name for the throwing variant with the +potential in the future to deprecate the `AsyncNonThrowingBackpressuredStream` +in favour of adopting typed throws. + +## Acknowledgements + +- [Johannes Weiss](https://github.com/weissi) - For making me aware how +important this problem is and providing great ideas on how to shape the API. +- [Philippe Hausler](https://github.com/phausler) - For helping me designing the +APIs and continuously providing feedback +- [George Barnett](https://github.com/glbrntt) - For providing extensive code +reviews and testing the implementation. +- [Si Beaumont](https://github.com/simonjbeaumont) - For implementing the element size dependent strategy diff --git a/Sources/AsyncAlgorithms/BackPressuredStream/AsyncBackPressuredStream.swift b/Sources/AsyncAlgorithms/BackPressuredStream/AsyncBackPressuredStream.swift new file mode 100644 index 00000000..4b7d3e2a --- /dev/null +++ b/Sources/AsyncAlgorithms/BackPressuredStream/AsyncBackPressuredStream.swift @@ -0,0 +1,425 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the Swift Async Algorithms open source project +// +// Copyright (c) 2023 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 +// +//===----------------------------------------------------------------------===// + +/// Error that is thrown from the various `write` methods of the +/// ``AsyncBackPressuredStream/Source``. +/// +/// This error is thrown when the asynchronous stream is already finished when +/// trying to write new elements. +public struct AsyncBackPressuredStreamAlreadyFinishedError: Error {} + +/// A struct that acts as a source asynchronous sequence. +/// +/// The ``AsyncBackPressuredStream`` provides a ``AsyncBackPressuredStream/Source`` to +/// write values to the stream. The source exposes the internal backpressure of the asynchronous sequence to the +/// external producer. This allows to bridge both synchronous and asynchronous producers into an asynchronous sequence. +/// +/// ## Using an AsyncBackPressuredStream +/// +/// To use an ``AsyncBackPressuredStream`` you have to create a new stream with it's source first by calling +/// the ``AsyncBackPressuredStream/makeStream(of:throwing:backPressureStrategy:)`` method. +/// Afterwards, you can pass the source to the producer and the stream to the consumer. +/// +/// ``` +/// let (stream, source) = AsyncBackPressuredStream.makeStream( +/// backpressureStrategy: .watermark(low: 2, high: 4) +/// ) +/// +/// try await withThrowingTaskGroup(of: Void.self) { group in +/// group.addTask { +/// try await source.write(1) +/// try await source.write(2) +/// try await source.write(3) +/// } +/// +/// for await element in stream { +/// print(element) +/// } +/// } +/// ``` +/// +/// The source also exposes synchronous write methods that communicate the backpressure via callbacks. +public struct AsyncBackPressuredStream: AsyncSequence { + /// A private class to give the ``AsyncBackPressuredStream`` a deinit so we + /// can tell the producer when any potential consumer went away. + private final class _Backing: Sendable { + /// The underlying storage. + fileprivate let storage: _AsyncBackPressuredStreamBackPressuredStorage + + init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { + self.storage = storage + } + + deinit { + storage.sequenceDeinitialized() + } + } + + /// The backing storage. + private let backing: _Backing + + /// Initializes a new ``AsyncBackPressuredStream`` and an ``AsyncBackPressuredStream/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. + public static func makeStream( + of elementType: Element.Type = Element.self, + throwing failureType: Failure.Type = Failure.self, + backPressureStrategy: Source.BackPressureStrategy + ) -> (`Self`, Source) where Failure == Error { + let storage = _AsyncBackPressuredStreamBackPressuredStorage( + backPressureStrategy: backPressureStrategy.internalBackPressureStrategy + ) + let source = Source(storage: storage) + + return (.init(storage: storage), source) + } + + init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { + self.backing = .init(storage: storage) + } +} + +extension AsyncBackPressuredStream { + /// A struct 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. + /// + /// - Important: You must terminate the source by calling one of the `finish` methods otherwise the stream's iterator + /// will never terminate. + public struct Source: Sendable { + /// A strategy that handles the backpressure of the asynchronous stream. + public struct BackPressureStrategy: Sendable { + var internalBackPressureStrategy: _AsyncBackPressuredStreamInternalBackPressureStrategy + + /// A backpressure strategy using a high and low watermark to suspend and resume production respectively. + /// + /// - Parameters: + /// - low: When the number of buffered elements drops below the low watermark, producers will be resumed. + /// - high: When the number of buffered elements rises above the high watermark, producers will be suspended. + public static func watermark(low: Int, high: Int) -> BackPressureStrategy { + .init( + internalBackPressureStrategy: .watermark( + .init(low: low, high: high, waterLevelForElement: nil) + ) + ) + } + + /// A backpressure strategy using a high and low watermark to suspend and resume production respectively. + /// + /// - Parameters: + /// - low: When the number of buffered elements drops below the low watermark, producers will be resumed. + /// - high: When the number of buffered elements rises above the high watermark, producers will be suspended. + /// - waterLevelForElement: A closure used to compute the contribution of each buffered element to the current water level. + /// + /// - Note, `waterLevelForElement` will be called on each element when it is written into the source and when + /// it is consumed from the stream, so it is recommended to provide an function that runs in constant time. + public static func watermark( + low: Int, + high: Int, + waterLevelForElement: @escaping @Sendable (Element) -> Int + ) -> BackPressureStrategy { + .init( + internalBackPressureStrategy: .watermark( + .init(low: low, high: high, waterLevelForElement: waterLevelForElement) + ) + ) + } + } + + /// A type that indicates the result of writing elements to the source. + public 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. + public struct CallbackToken: Sendable { + let id: UInt + } + + /// 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. + final class _Backing: Sendable { + let storage: _AsyncBackPressuredStreamBackPressuredStorage + + init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { + self.storage = storage + } + + // TODO: Double check + 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 + public var onTermination: (@Sendable () -> Void)? { + set { + self._backing.storage.onTermination = newValue + } + get { + self._backing.storage.onTermination + } + } + + private var _backing: _Backing + + internal init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { + 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. + public 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. + public 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:)-4amlx`` or ``write(_:)-8e7el`` 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. + public 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. + public 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:)``. + public 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:)``. + public 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. + public 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. + public 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. + public 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. + public func finish(throwing error: Failure?) { + self._backing.storage.finish(error) + } + } +} + +extension AsyncBackPressuredStream { + /// 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()`. + public struct Iterator: AsyncIteratorProtocol { + private class _Backing { + let storage: _AsyncBackPressuredStreamBackPressuredStorage + + init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { + self.storage = storage + self.storage.iteratorInitialized() + } + + deinit { + self.storage.iteratorDeinitialized() + } + } + + private let backing: _Backing + + init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { + self.backing = .init(storage: storage) + } + + /// The next value from the asynchronous stream. + /// + /// When `next()` returns `nil`, this signifies the end of the + /// `AsyncThrowingStream`. + /// + /// 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 `AsyncThrowingStream` terminates. In this case, + /// `next()` may return `nil` immediately, or else return `nil` on + /// subsequent calls. + public mutating func next() async throws -> Element? { + try await self.backing.storage.next() + } + } + + /// Creates the asynchronous iterator that produces elements of this + /// asynchronous sequence. + public func makeAsyncIterator() -> Iterator { + Iterator(storage: self.backing.storage) + } +} + +extension AsyncBackPressuredStream: Sendable where Element: Sendable {} + +@available(*, unavailable) +extension AsyncBackPressuredStream.Iterator: Sendable {} diff --git a/Sources/AsyncAlgorithms/BackpressuredStream/AsyncBackPressuredStream+Internal.swift b/Sources/AsyncAlgorithms/BackpressuredStream/AsyncBackPressuredStream+Internal.swift new file mode 100644 index 00000000..50579730 --- /dev/null +++ b/Sources/AsyncAlgorithms/BackpressuredStream/AsyncBackPressuredStream+Internal.swift @@ -0,0 +1,1192 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the Swift Async Algorithms open source project +// +// Copyright (c) 2023 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 +// +//===----------------------------------------------------------------------===// + +import DequeModule + +struct _AsyncBackPressuredStreamWatermarkBackPressureStrategy { + /// The low watermark where demand should start. + private let low: Int + /// The high watermark where demand should be stopped. + private let high: Int + private var currentWatermark: Int = 0 + private 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. + init(low: Int, high: Int, waterLevelForElement: (@Sendable (Element) -> Int)?) { + precondition(low <= high) + self.low = low + self.high = high + 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) + // We are demanding more until we reach the high watermark + return self.currentWatermark < self.high + } + + mutating func didConsume(element: Element) -> Bool { + if let waterLevelForElement { + self.currentWatermark -= waterLevelForElement(element) + } else { + self.currentWatermark -= 1 + } + precondition(self.currentWatermark >= 0) + // We start demanding again once we are below the low watermark + return self.currentWatermark < self.low + } +} + +enum _AsyncBackPressuredStreamInternalBackPressureStrategy { + case watermark(_AsyncBackPressuredStreamWatermarkBackPressureStrategy) + + mutating func didYield(elements: Deque.SubSequence) -> Bool { + switch self { + case .watermark(var strategy): + let result = strategy.didYield(elements: elements) + self = .watermark(strategy) + return result + } + } + + mutating func didConsume(element: Element) -> Bool { + switch self { + case .watermark(var strategy): + let result = strategy.didConsume(element: element) + self = .watermark(strategy) + return result + } + } +} + +// We are unchecked Sendable since we are protecting our state with a lock. +final class _AsyncBackPressuredStreamBackPressuredStorage: @unchecked Sendable { + /// The state machine + var _stateMachine: ManagedCriticalState<_AsyncBackPressuredStateMachine> + + var onTermination: (@Sendable () -> Void)? { + set { + self._stateMachine.withCriticalRegion { + $0._onTermination = newValue + } + } + get { + self._stateMachine.withCriticalRegion { + $0._onTermination + } + } + } + + init( + backPressureStrategy: _AsyncBackPressuredStreamInternalBackPressureStrategy + ) { + self._stateMachine = .init(.init(backPressureStrategy: backPressureStrategy)) + } + + 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(AsyncBackPressuredStreamAlreadyFinishedError())) + } + onTermination?() + + case .none: + break + } + } + + func iteratorInitialized() { + self._stateMachine.withCriticalRegion { + $0.iteratorInitialized() + } + } + + 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(AsyncBackPressuredStreamAlreadyFinishedError())) + } + onTermination?() + + case .none: + break + } + } + + func sourceDeinitialized() { + let action = self._stateMachine.withCriticalRegion { + $0.sourceDeinitialized() + } + + switch action { + case .callOnTermination(let onTermination): + onTermination?() + + case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): + for producerContinuation in producerContinuations { + producerContinuation(.failure(AsyncBackPressuredStreamAlreadyFinishedError())) + } + onTermination?() + + case .failProducers(let producerContinuations): + for producerContinuation in producerContinuations { + producerContinuation(.failure(AsyncBackPressuredStreamAlreadyFinishedError())) + } + + case .none: + break + } + } + + func write( + contentsOf sequence: some Sequence + ) throws -> AsyncBackPressuredStream.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 AsyncBackPressuredStreamAlreadyFinishedError() + } + } + + func enqueueProducer( + callbackToken: AsyncBackPressuredStream.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 + } + } + + func cancelProducer( + callbackToken: AsyncBackPressuredStream.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 + } + } + + func finish(_ failure: Failure?) { + 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(AsyncBackPressuredStreamAlreadyFinishedError())) + } + + case .none: + break + } + } + + 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 .returnFailureAndCallOnTermination(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() + } + } + + 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 .resumeConsumerWithFailureAndCallOnTermination(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 .resumeConsumerWithNilAndCallOnTermination(let continuation, let onTermination): + continuation.resume(returning: nil) + onTermination?() + + case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): + for producerContinuation in producerContinuations { + producerContinuation(.failure(AsyncBackPressuredStreamAlreadyFinishedError())) + } + onTermination?() + + case .none: + break + } + } + } +} + +/// The state machine of the backpressured async stream. +struct _AsyncBackPressuredStateMachine: Sendable { + enum _State { + struct Initial { + /// The backpressure strategy. + var backPressureStrategy: _AsyncBackPressuredStreamInternalBackPressureStrategy + /// Indicates if the iterator was initialized. + var iteratorInitialized: Bool + /// The onTermination callback. + var onTermination: (@Sendable () -> Void)? + } + + struct Streaming { + /// The backpressure strategy. + var backPressureStrategy: _AsyncBackPressuredStreamInternalBackPressureStrategy + /// Indicates if the iterator was initialized. + var iteratorInitialized: Bool + /// The onTermination callback. + var onTermination: (@Sendable () -> Void)? + /// The buffer of elements. + var buffer: Deque + /// The optional consumer continuation. + var consumerContinuation: CheckedContinuation? + /// The producer continuations. + var producerContinuations: Deque<(UInt, (Result) -> Void)> + /// The producers that have been cancelled. + var cancelledAsyncProducers: Deque + /// Indicates if we currently have outstanding demand. + var hasOutstandingDemand: Bool + } + + struct SourceFinished { + /// Indicates if the iterator was initialized. + var iteratorInitialized: Bool + /// The buffer of elements. + var buffer: Deque + /// The failure that should be thrown after the last element has been consumed. + var failure: Failure? + /// The onTermination callback. + var onTermination: (@Sendable () -> Void)? + } + + case initial(Initial) + /// The state once either any element was yielded or `next()` was called. + case streaming(Streaming) + /// The state once the underlying source signalled that it is finished. + case sourceFinished(SourceFinished) + + /// The state once there can be no outstanding demand. This can happen if: + /// 1. The iterator was deinited + /// 2. The underlying source finished and all buffered elements have been consumed + case finished(iteratorInitialized: Bool) + + /// An intermediate state to avoid CoWs. + case modify + } + + /// The state machine's current state. + var _state: _State + + // The ID used for the next CallbackToken. + var nextCallbackTokenID: UInt = 0 + + 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. + init( + backPressureStrategy: _AsyncBackPressuredStreamInternalBackPressureStrategy + ) { + self._state = .initial( + .init( + backPressureStrategy: backPressureStrategy, + iteratorInitialized: false, + onTermination: nil + ) + ) + } + + /// Generates the next callback token. + mutating func nextCallbackToken() -> AsyncBackPressuredStream.Source.WriteResult.CallbackToken { + let id = self.nextCallbackTokenID + self.nextCallbackTokenID += 1 + return .init(id: id) + } + + /// Actions returned by `sequenceDeinitialized()`. + 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)? + ) + } + + mutating func sequenceDeinitialized() -> SequenceDeinitializedAction? { + switch self._state { + case .initial(let initial): + guard initial.iteratorInitialized else { + // No iterator was created so we can transition to finished right away. + self._state = .finished(iteratorInitialized: false) + + return .callOnTermination(initial.onTermination) + } + // An iterator was created and we deinited the sequence. + // This is an expected pattern and we just continue on normal. + return .none + + case .streaming(let streaming): + guard streaming.iteratorInitialized 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 + ) + } + // An iterator was created and we deinited the sequence. + // This is an expected pattern and we just continue on normal. + return .none + + case .sourceFinished(let sourceFinished): + guard sourceFinished.iteratorInitialized else { + // No iterator was created so we can transition to finished right away. + self._state = .finished(iteratorInitialized: false) + + return .callOnTermination(sourceFinished.onTermination) + } + // 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 + + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } + + 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()`. + 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)? + ) + } + + 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()`. + 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( + [(Result) -> Void], + (@Sendable () -> Void)? + ) + /// Indicates that all producers should be failed. + case failProducers([(Result) -> Void]) + } + + 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): + guard streaming.buffer.isEmpty 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 }) + ) + } + // We can transition to finished right away since the buffer is empty now + self._state = .finished(iteratorInitialized: streaming.iteratorInitialized) + + return .failProducersAndCallOnTermination( + Array(streaming.producerContinuations.map { $0.1 }), + streaming.onTermination + ) + + 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()`. + 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: AsyncBackPressuredStream.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: AsyncBackPressuredStream.Source.WriteResult.CallbackToken + ) + /// Indicates that the producer has been finished. + case throwFinishedError + + init( + callbackToken: AsyncBackPressuredStream.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 + ) + } + } + } + + 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 + + // We have an element and can resume the continuation + let bufferEndIndexBeforeAppend = streaming.buffer.endIndex + streaming.buffer.append(contentsOf: sequence) + let shouldProduceMore = streaming.backPressureStrategy.didYield( + elements: streaming.buffer[bufferEndIndexBeforeAppend...] + ) + streaming.hasOutstandingDemand = shouldProduceMore + let callbackToken = shouldProduceMore ? nil : self.nextCallbackToken() + + guard let consumerContinuation = streaming.consumerContinuation else { + // We don't have a suspended consumer so we just buffer the elements + self._state = .streaming(streaming) + return .init( + callbackToken: callbackToken + ) + } + 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: callbackToken) + } + + // We got a consumer continuation and an element. We can resume the consumer now + streaming.consumerContinuation = nil + self._state = .streaming(streaming) + return .init( + callbackToken: callbackToken, + continuationAndElement: (consumerContinuation, element) + ) + + 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()`. + 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, Error) + } + + mutating func enqueueProducer( + callbackToken: AsyncBackPressuredStream.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 public 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, AsyncBackPressuredStreamAlreadyFinishedError()) + + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } + + /// Actions returned by `cancelProducer()`. + enum CancelProducerAction { + /// Indicates that the producer should be notified about cancellation. + case resumeProducerWithCancellationError((Result) -> Void) + } + + mutating func cancelProducer( + callbackToken: AsyncBackPressuredStream.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): + guard let index = streaming.producerContinuations.firstIndex(where: { $0.0 == callbackToken.id }) 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 + } + // 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) + + 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()`. + 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: Failure?, + onTermination: (() -> Void)? + ) + /// Indicates that the producers should be resumed with an error. + case resumeProducers( + producerContinuations: [(Result) -> Void] + ) + } + + @inlinable + mutating func finish(_ failure: Failure?) -> 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): + guard let consumerContinuation = streaming.consumerContinuation 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 })) + } + // 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 + ) + + 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()`. + 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 `onTermination` should be called. + case returnFailureAndCallOnTermination(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 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 + + guard let element = streaming.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(streaming) + + return .suspendTask + } + // We have an element to fulfil the demand right away. + let shouldProduceMore = streaming.backPressureStrategy.didConsume(element: element) + streaming.hasOutstandingDemand = shouldProduceMore + + guard shouldProduceMore else { + // We don't have any new demand, so we can just return the element. + self._state = .streaming(streaming) + return .returnElement(element) + } + // 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) + + case .sourceFinished(var sourceFinished): + // Check if we have an element left in the buffer and return it + self._state = .modify + + guard let element = sourceFinished.buffer.popFirst() else { + // We are returning the queued failure now and can transition to finished + self._state = .finished(iteratorInitialized: sourceFinished.iteratorInitialized) + + return .returnFailureAndCallOnTermination(sourceFinished.failure, sourceFinished.onTermination) + } + self._state = .sourceFinished(sourceFinished) + + return .returnElement(element) + + case .finished: + return .returnNil + + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } + + /// Actions returned by `suspendNext()`. + 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 resumeConsumerWithFailureAndCallOnTermination( + CheckedContinuation, + Failure?, + (() -> Void)? + ) + /// Indicates that the consumer should be resumed with `nil`. + case resumeConsumerWithNil(CheckedContinuation) + } + + 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 + guard let element = streaming.buffer.popFirst() 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 + } + // We have an element to fulfil the demand right away. + + let shouldProduceMore = streaming.backPressureStrategy.didConsume(element: element) + streaming.hasOutstandingDemand = shouldProduceMore + + guard shouldProduceMore else { + // We don't have any new demand, so we can just return the element. + self._state = .streaming(streaming) + return .resumeConsumerWithElement(continuation, element) + } + // 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) + + case .sourceFinished(var sourceFinished): + // Check if we have an element left in the buffer and return it + self._state = .modify + + guard let element = sourceFinished.buffer.popFirst() else { + // We are returning the queued failure now and can transition to finished + self._state = .finished(iteratorInitialized: sourceFinished.iteratorInitialized) + + return .resumeConsumerWithFailureAndCallOnTermination( + continuation, + sourceFinished.failure, + sourceFinished.onTermination + ) + } + self._state = .sourceFinished(sourceFinished) + + return .resumeConsumerWithElement(continuation, element) + + case .finished: + return .resumeConsumerWithNil(continuation) + + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } + + /// Actions returned by `cancelNext()`. + enum CancelNextAction { + /// Indicates that the continuation should be resumed with nil, the producers should be finished and call onTermination. + case resumeConsumerWithNilAndCallOnTermination(CheckedContinuation, (() -> Void)?) + /// Indicates that the producers should be finished and call onTermination. + case failProducersAndCallOnTermination([(Result) -> Void], (() -> Void)?) + } + + 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) + + guard let consumerContinuation = streaming.consumerContinuation else { + return .failProducersAndCallOnTermination( + Array(streaming.producerContinuations.map { $0.1 }), + streaming.onTermination + ) + } + precondition( + streaming.producerContinuations.isEmpty, + "Internal inconsistency. Unexpected producer continuations." + ) + return .resumeConsumerWithNilAndCallOnTermination( + consumerContinuation, + streaming.onTermination + ) + + case .sourceFinished, .finished: + return .none + + case .modify: + fatalError("AsyncStream internal inconsistency") + } + } +} diff --git a/Sources/AsyncAlgorithms/BackpressuredStream/AsyncNonThrowingBackPressuredStream.swift b/Sources/AsyncAlgorithms/BackpressuredStream/AsyncNonThrowingBackPressuredStream.swift new file mode 100644 index 00000000..07910595 --- /dev/null +++ b/Sources/AsyncAlgorithms/BackpressuredStream/AsyncNonThrowingBackPressuredStream.swift @@ -0,0 +1,429 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the Swift Async Algorithms open source project +// +// Copyright (c) 2023 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 +// +//===----------------------------------------------------------------------===// + +/// A struct that acts as a source asynchronous sequence. +/// +/// The ``AsyncNonThrowingBackPressuredStream`` provides a ``AsyncNonThrowingBackPressuredStream/Source`` to +/// write values to the stream. The source exposes the internal backpressure of the asynchronous sequence to the +/// external producer. This allows to bridge both synchronous and asynchronous producers into an asynchronous sequence. +/// +/// ## Using an AsyncNonThrowingBackPressuredStream +/// +/// To use an ``AsyncNonThrowingBackPressuredStream`` you have to create a new stream with it's source first by calling +/// the ``AsyncNonThrowingBackPressuredStream/makeStream(of:backpressureStrategy:)`` method. +/// Afterwards, you can pass the source to the producer and the stream to the consumer. +/// +/// ``` +/// let (stream, source) = AsyncNonThrowingBackPressuredStream.makeStream( +/// backpressureStrategy: .watermark(low: 2, high: 4) +/// ) +/// +/// try await withThrowingTaskGroup(of: Void.self) { group in +/// group.addTask { +/// try await source.write(1) +/// try await source.write(2) +/// try await source.write(3) +/// } +/// +/// for await element in stream { +/// print(element) +/// } +/// } +/// ``` +/// +/// The source also exposes synchronous write methods that communicate the backpressure via callbacks. +public struct AsyncNonThrowingBackPressuredStream: AsyncSequence { + /// A private class to give the ``AsyncNonThrowingBackPressuredStream`` a deinit so we + /// can tell the producer when any potential consumer went away. + private final class _Backing: Sendable { + /// The underlying storage. + fileprivate let storage: _AsyncBackPressuredStreamBackPressuredStorage + + init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { + self.storage = storage + } + + deinit { + storage.sequenceDeinitialized() + } + } + + /// The backing storage. + private let backing: _Backing + + /// Initializes a new ``AsyncNonThrowingBackPressuredStream`` and an ``AsyncNonThrowingBackPressuredStream/Source``. + /// + /// - Parameters: + /// - elementType: The element 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. + public static func makeStream( + of elementType: Element.Type = Element.self, + backPressureStrategy: Source.BackPressureStrategy + ) -> (`Self`, Source) { + let storage = _AsyncBackPressuredStreamBackPressuredStorage( + backPressureStrategy: backPressureStrategy.internalBackPressureStrategy + ) + let source = Source(storage: storage) + + return (.init(storage: storage), source) + } + + init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { + self.backing = .init(storage: storage) + } +} + +extension AsyncNonThrowingBackPressuredStream { + /// A struct 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. + /// + /// - Important: You must terminate the source by calling one of the `finish` methods otherwise the stream's iterator + /// will never terminate. + public struct Source: Sendable { + /// A strategy that handles the backpressure of the asynchronous stream. + public struct BackPressureStrategy: Sendable { + var internalBackPressureStrategy: _AsyncBackPressuredStreamInternalBackPressureStrategy + + /// A backpressure strategy using a high and low watermark to suspend and resume production respectively. + /// + /// - Parameters: + /// - low: When the number of buffered elements drops below the low watermark, producers will be resumed. + /// - high: When the number of buffered elements rises above the high watermark, producers will be suspended. + public static func watermark(low: Int, high: Int) -> BackPressureStrategy { + .init( + internalBackPressureStrategy: .watermark( + .init(low: low, high: high, waterLevelForElement: nil) + ) + ) + } + + /// A backpressure strategy using a high and low watermark to suspend and resume production respectively. + /// + /// - Parameters: + /// - low: When the number of buffered elements drops below the low watermark, producers will be resumed. + /// - high: When the number of buffered elements rises above the high watermark, producers will be suspended. + /// - waterLevelForElement: A closure used to compute the contribution of each buffered element to the current water level. + /// + /// - Note, `waterLevelForElement` will be called on each element when it is written into the source and when + /// it is consumed from the stream, so it is recommended to provide an function that runs in constant time. + public static func watermark( + low: Int, + high: Int, + waterLevelForElement: @escaping @Sendable (Element) -> Int + ) -> BackPressureStrategy { + .init( + internalBackPressureStrategy: .watermark( + .init(low: low, high: high, waterLevelForElement: waterLevelForElement) + ) + ) + } + } + + /// A type that indicates the result of writing elements to the source. + public 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. + public struct CallbackToken: Sendable { + let id: UInt + } + + /// 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. + final class _Backing: Sendable { + let storage: _AsyncBackPressuredStreamBackPressuredStorage + + init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { + self.storage = storage + } + + // TODO: Double check + 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 + public var onTermination: (@Sendable () -> Void)? { + set { + self._backing.storage.onTermination = newValue + } + get { + self._backing.storage.onTermination + } + } + + private var _backing: _Backing + + internal init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { + 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. + public func write(contentsOf sequence: S) throws -> WriteResult where Element == S.Element, S: Sequence { + switch try self._backing.storage.write(contentsOf: sequence) { + case .produceMore: + return .produceMore + case .enqueueCallback(let callbackToken): + return .enqueueCallback(.init(id: callbackToken.id)) + } + } + + /// 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. + public func write(_ element: Element) throws -> WriteResult { + switch try self._backing.storage.write(contentsOf: CollectionOfOne(element)) { + case .produceMore: + return .produceMore + case .enqueueCallback(let callbackToken): + return .enqueueCallback(.init(id: callbackToken.id)) + } + } + + /// 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. + public func enqueueCallback( + callbackToken: WriteResult.CallbackToken, + onProduceMore: @escaping @Sendable (Result) -> Void + ) { + let callbackToken = AsyncBackPressuredStream.Source.WriteResult.CallbackToken( + id: callbackToken.id + ) + 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. + public func cancelCallback(callbackToken: WriteResult.CallbackToken) { + let callbackToken = AsyncBackPressuredStream.Source.WriteResult.CallbackToken( + id: callbackToken.id + ) + 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:)``. + public 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:)``. + public 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. + public 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. + public 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. + public 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`. + /// + /// Calling this function more than once has no effect. After calling finish, the stream enters a terminal state and doesn't accept + /// new elements. + public func finish() { + self._backing.storage.finish(nil) + } + } +} + +extension AsyncNonThrowingBackPressuredStream { + /// 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()`. + public struct Iterator: AsyncIteratorProtocol { + private class _Backing { + let storage: _AsyncBackPressuredStreamBackPressuredStorage + + init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { + self.storage = storage + self.storage.iteratorInitialized() + } + + deinit { + self.storage.iteratorDeinitialized() + } + } + + private let backing: _Backing + + init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { + self.backing = .init(storage: storage) + } + + /// The next value from the asynchronous stream. + /// + /// When `next()` returns `nil`, this signifies the end of the + /// `AsyncThrowingStream`. + /// + /// 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 `AsyncThrowingStream` terminates. In this case, + /// `next()` may return `nil` immediately, or else return `nil` on + /// subsequent calls. + public mutating func next() async -> Element? { + try! await self.backing.storage.next() + } + } + + /// Creates the asynchronous iterator that produces elements of this + /// asynchronous sequence. + public func makeAsyncIterator() -> Iterator { + Iterator(storage: self.backing.storage) + } +} + +extension AsyncNonThrowingBackPressuredStream: Sendable where Element: Sendable {} + +@available(*, unavailable) +extension AsyncNonThrowingBackPressuredStream.Iterator: Sendable {} diff --git a/Tests/AsyncAlgorithmsTests/BackpressuredStream/BackPressuredStreamTests.swift b/Tests/AsyncAlgorithmsTests/BackpressuredStream/BackPressuredStreamTests.swift new file mode 100644 index 00000000..ac74282b --- /dev/null +++ b/Tests/AsyncAlgorithmsTests/BackpressuredStream/BackPressuredStreamTests.swift @@ -0,0 +1,1166 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the Swift Async Algorithms open source project +// +// Copyright (c) 2023 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 +// +//===----------------------------------------------------------------------===// + +import AsyncAlgorithms +import XCTest + +final class BackPressuredStreamTests: XCTestCase { + // MARK: - sequenceDeinitialized + + func testSequenceDeinitialized_whenNoIterator() async throws { + var (stream, source): (AsyncBackPressuredStream?, AsyncBackPressuredStream.Source) = + AsyncBackPressuredStream.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(for: .seconds(0.2)) + } + } + + 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 AsyncBackPressuredStreamAlreadyFinishedError) + } + + group.cancelAll() + } + } + + func testSequenceDeinitialized_whenIterator() async throws { + var (stream, source): (AsyncBackPressuredStream?, AsyncBackPressuredStream.Source) = + AsyncBackPressuredStream.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(for: .seconds(0.2)) + } + } + + 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): (AsyncBackPressuredStream?, AsyncBackPressuredStream.Source) = + AsyncBackPressuredStream.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(for: .seconds(0.2)) + } + } + + 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 AsyncBackPressuredStreamAlreadyFinishedError) + } + + group.cancelAll() + } + } + + func testSequenceDeinitialized_whenStreaming_andSuspendedProducer() async throws { + var (stream, source): (AsyncBackPressuredStream?, AsyncBackPressuredStream.Source) = + AsyncBackPressuredStream.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 AsyncBackPressuredStreamAlreadyFinishedError) + } + } + + // MARK: - iteratorInitialized + + func testIteratorInitialized_whenInitial() async throws { + let (stream, _) = AsyncBackPressuredStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 5, high: 10) + ) + + _ = stream.makeAsyncIterator() + } + + func testIteratorInitialized_whenStreaming() async throws { + let (stream, source) = AsyncBackPressuredStream.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) = AsyncBackPressuredStream.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) = AsyncBackPressuredStream.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) = AsyncBackPressuredStream.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(for: .seconds(0.2)) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + var iterator: AsyncBackPressuredStream.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) = AsyncBackPressuredStream.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(for: .seconds(0.2)) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + var iterator: AsyncBackPressuredStream.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) = AsyncBackPressuredStream.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(for: .seconds(0.2)) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + var iterator: AsyncBackPressuredStream.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) = AsyncBackPressuredStream.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(for: .seconds(0.2)) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + var iterator: AsyncBackPressuredStream.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): (AsyncBackPressuredStream?, AsyncBackPressuredStream.Source) = + AsyncBackPressuredStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 1, high: 2) + ) + + var iterator: AsyncBackPressuredStream.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 AsyncBackPressuredStreamAlreadyFinishedError) + } + + _ = try await iterator?.next() + } + + // MARK: - sourceDeinitialized + + func testSourceDeinitialized_whenInitial() async throws { + var (stream, source): (AsyncBackPressuredStream, AsyncBackPressuredStream.Source?) = + AsyncBackPressuredStream.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(for: .seconds(0.2)) + } + } + + 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): (AsyncBackPressuredStream, AsyncBackPressuredStream.Source?) = + AsyncBackPressuredStream.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(for: .seconds(0.2)) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + var iterator: AsyncBackPressuredStream.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): (AsyncBackPressuredStream, AsyncBackPressuredStream.Source?) = + AsyncBackPressuredStream.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(for: .seconds(0.2)) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + var iterator: AsyncBackPressuredStream.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): (AsyncBackPressuredStream, AsyncBackPressuredStream.Source?) = + AsyncBackPressuredStream.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(for: .seconds(0.2)) + } + } + + var onTerminationIterator = onTerminationStream.makeAsyncIterator() + _ = await onTerminationIterator.next() + + var iterator: AsyncBackPressuredStream.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): (AsyncBackPressuredStream, AsyncBackPressuredStream.Source?) = + AsyncBackPressuredStream.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(for: .seconds(0.2)) + } + } + + 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): (AsyncBackPressuredStream, AsyncBackPressuredStream.Source?) = + AsyncBackPressuredStream.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 AsyncBackPressuredStreamAlreadyFinishedError) + } + } + + // MARK: - write + + func testWrite_whenInitial() async throws { + let (stream, source) = AsyncBackPressuredStream.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) = AsyncBackPressuredStream.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) = AsyncBackPressuredStream.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(for: .seconds(0.5)) + + try await source.write(1) + let element = try await group.next() + XCTAssertEqual(element, 1) + } + } + + func testWrite_whenStreaming_andSuspendedConsumer_andEmptySequence() async throws { + let (stream, source) = AsyncBackPressuredStream.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(for: .seconds(0.5)) + + 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) = AsyncBackPressuredStream.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) = AsyncBackPressuredStream.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) = AsyncBackPressuredStream.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) = AsyncBackPressuredStream.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) = AsyncBackPressuredStream.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 AsyncBackPressuredStreamAlreadyFinishedError) + } + } + + // MARK: - cancelProducer + + func testCancelProducer_whenStreaming() async throws { + let (stream, source) = AsyncBackPressuredStream.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) = AsyncBackPressuredStream.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 AsyncBackPressuredStreamAlreadyFinishedError) + } + + let element = try await stream.first { _ in true } + XCTAssertEqual(element, 1) + } + + // MARK: - finish + + func testFinish_whenStreaming_andConsumerSuspended() async throws { + let (stream, source) = AsyncBackPressuredStream.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(for: .seconds(0.5)) + + source.finish(throwing: nil) + let element = try await group.next() + XCTAssertEqual(element, .some(nil)) + } + } + + func testFinish_whenInitial() async throws { + let (stream, source) = AsyncBackPressuredStream.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) = AsyncBackPressuredStream.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) = AsyncBackPressuredStream.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 testWatermarkWithCustomCoount() async throws { + let (stream, source) = AsyncBackPressuredStream.makeStream( + of: [Int].self, + backPressureStrategy: .watermark(low: 2, high: 4, waterLevelForElement: { $0.count }) + ) + var iterator = stream.makeAsyncIterator() + + try await source.write([1, 1, 1]) + + _ = try await iterator.next() + + try await source.write([1, 1, 1]) + + _ = try await iterator.next() + } + + func testThrowsError() async throws { + let (stream, source) = AsyncBackPressuredStream.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) = AsyncBackPressuredStream.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]) + } + + // MARK: NonThrowing + + func testNonThrowing() async throws { + let (stream, source) = AsyncNonThrowingBackPressuredStream.makeStream( + of: Int.self, + backPressureStrategy: .watermark(low: 2, high: 4) + ) + + let (backPressureEventStream, backPressureEventContinuation) = AsyncStream.makeStream(of: Void.self) + + 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() + + _ = await iterator.next() + _ = await iterator.next() + _ = await iterator.next() + + await backPressureEventIterator.next() + await backPressureEventIterator.next() + await backPressureEventIterator.next() + + group.cancelAll() + } + } +} + +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.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 + } + } +} From f09d0ce2cb283d4a0a25953e3c1d028903a27261 Mon Sep 17 00:00:00 2001 From: Franz Busch Date: Tue, 14 May 2024 14:43:49 +0200 Subject: [PATCH 2/7] rename --- ...mutli-producer-single-consumer-channel.md} | 389 +++----- .../AsyncBackPressuredStream.swift | 425 -------- .../AsyncNonThrowingBackPressuredStream.swift | 429 --------- .../AsyncAlgorithms/Internal/_TinyArray.swift | 329 +++++++ Sources/AsyncAlgorithms/Locking.swift | 20 +- ...ducerSingleConsumerChannel+Internal.swift} | 910 ++++++++++++------ .../MultiProducerSingleConsumerChannel.swift | 608 ++++++++++++ ...iProducerSingleConsumerChannelTests.swift} | 694 +++++++------ 8 files changed, 2030 insertions(+), 1774 deletions(-) rename Evolution/{0016-backpressured-stream.md => 0016-mutli-producer-single-consumer-channel.md} (57%) delete mode 100644 Sources/AsyncAlgorithms/BackPressuredStream/AsyncBackPressuredStream.swift delete mode 100644 Sources/AsyncAlgorithms/BackpressuredStream/AsyncNonThrowingBackPressuredStream.swift create mode 100644 Sources/AsyncAlgorithms/Internal/_TinyArray.swift rename Sources/AsyncAlgorithms/{BackpressuredStream/AsyncBackPressuredStream+Internal.swift => MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel+Internal.swift} (50%) create mode 100644 Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel.swift rename Tests/AsyncAlgorithmsTests/{BackpressuredStream/BackPressuredStreamTests.swift => MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift} (53%) diff --git a/Evolution/0016-backpressured-stream.md b/Evolution/0016-mutli-producer-single-consumer-channel.md similarity index 57% rename from Evolution/0016-backpressured-stream.md rename to Evolution/0016-mutli-producer-single-consumer-channel.md index 2d8f76a8..596a1ccc 100644 --- a/Evolution/0016-backpressured-stream.md +++ b/Evolution/0016-mutli-producer-single-consumer-channel.md @@ -1,6 +1,6 @@ -# Externally backpressured support for AsyncStream +# MutliProducerSingleConsumerChannel -* Proposal: [SAA-0016](0016-backpressured-stream.md) +* Proposal: [SAA-0016](0016-multi-producer-single-consumer-channel.md) * Authors: [Franz Busch](https://github.com/FranzBusch) * Review Manager: TBD * Status: **Implemented** @@ -8,17 +8,17 @@ ## Revision - 2023/12/18: Migrate proposal from Swift Evolution to Swift Async Algorithms. - 2023/12/19: Add element size dependent strategy +- 2024/05/19: Rename to multi producer single consumer channel +- 2024/05/28: Add unbounded strategy ## Introduction [SE-0314](https://github.com/apple/swift-evolution/blob/main/proposals/0314-async-stream.md) introduced new `Async[Throwing]Stream` types which act as root asynchronous sequences. These two types allow bridging from synchronous callbacks such as -delegates to an asynchronous sequence. This proposal adds a new way of -constructing asynchronous streams with the goal to bridge backpressured systems -into an asynchronous sequence. Furthermore, this proposal aims to clarify the -cancellation behaviour both when the consuming task is cancelled and when -the production side indicates termination. +delegates to an asynchronous sequence. This proposal adds a new root +asynchronous sequence with the goal to bridge multi producer systems +into an asynchronous sequence. ## Motivation @@ -120,19 +120,19 @@ are the behaviors where `Async[Throwing]Stream` diverges from the expectations. - Consumer termination: Doesn't handle the `Continuation` being `deinit`ed - Producer termination: Happens on first consumer termination -This section proposes a new type called `AsyncBackpressuredStream` that implement all of +This section proposes a new type called `MutliProducerSingleConsumerChannel` that implement all of the above-mentioned behaviors. -### Creating an AsyncBackpressuredStream +### Creating an MutliProducerSingleConsumerChannel -You can create an `AsyncBackpressuredStream` instance using the new -`makeStream(of: backpressureStrategy:)` method. This method returns you the -stream and the source. The source can be used to write new values to the -asynchronous stream. The new API specifically provides a +You can create an `MutliProducerSingleConsumerChannel` instance using the new +`makeChannel(of: backpressureStrategy:)` method. This method returns you the +channel and the source. The source can be used to send new values to the +asynchronous channel. The new API specifically provides a multi-producer/single-consumer pattern. ```swift -let (stream, source) = AsyncBackpressuredStream.makeStream( +let (channel, source) = MutliProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 2, high: 4) ) @@ -144,9 +144,9 @@ example of how it can be used: ```swift do { - let writeResult = try source.write(contentsOf: sequence) + let sendResult = try source.send(contentsOf: sequence) - switch writeResult { + switch sendResult { case .produceMore: // Trigger more production @@ -161,13 +161,13 @@ do { }) } } catch { - // `write(contentsOf:)` throws if the asynchronous stream already terminated + // `send(contentsOf:)` throws if the asynchronous stream already terminated } ``` The above API offers the most control and highest performance when bridging a -synchronous producer to an asynchronous sequence. First, you have to write -values using the `write(contentsOf:)` which returns a `WriteResult`. The result +synchronous producer to an asynchronous sequence. First, you have to send +values using the `send(contentsOf:)` which returns a `SendResult`. The result either indicates that more values should be produced or that a callback should be enqueued by calling the `enqueueCallback(callbackToken: onProduceMore:)` method. This callback is invoked once the backpressure strategy decided that @@ -176,12 +176,12 @@ the greatest performance. The callback only has to be allocated in the case where the producer needs to be suspended. Additionally, the above API is the building block for some higher-level and -easier-to-use APIs to write values to the asynchronous stream. Below is an +easier-to-use APIs to send values to the channel. Below is an example of the two higher-level APIs. ```swift // Writing new values and providing a callback when to produce more -try source.write(contentsOf: sequence, onProduceMore: { result in +try source.send(contentsOf: sequence, onProduceMore: { result in switch result { case .success: // Trigger more production @@ -191,11 +191,11 @@ try source.write(contentsOf: sequence, onProduceMore: { result in }) // This method suspends until more values should be produced -try await source.write(contentsOf: sequence) +try await source.send(contentsOf: sequence) ``` With the above APIs, we should be able to effectively bridge any system into an -asynchronous stream regardless if the system is callback-based, blocking or +asynchronous sequence regardless if the system is callback-based, blocking or asynchronous. ### Downstream consumer termination @@ -208,15 +208,15 @@ this: ```swift // Termination through calling finish -let (stream, source) = AsyncBackpressuredStream.makeStream( +let (channel, source) = MutliProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 2, high: 4) ) -_ = try await source.write(1) +_ = try await source.send(1) source.finish() -for try await element in stream { +for try await element in channel { print(element) } print("Finished") @@ -232,12 +232,12 @@ indefinitely. ```swift // Termination through deiniting the source -let (stream, _) = AsyncStream.makeStream( +let (channel, _) = MutliProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 2, high: 4) ) -for await element in stream { +for await element in channel { print(element) } print("Finished") @@ -246,8 +246,8 @@ print("Finished") // Finished ``` -Trying to write more elements after the source has been finish will result in an -error thrown from the write methods. +Trying to send more elements after the source has been finish will result in an +error thrown from the send methods. ### Upstream producer termination @@ -256,13 +256,13 @@ callback. Termination of the producer happens in the following scenarios: ```swift // Termination through task cancellation -let (stream, source) = AsyncStream.makeStream( +let (channel source) = MutliProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 2, high: 4) ) let task = Task { - for await element in stream { + for await element in channel { } } @@ -271,7 +271,7 @@ task.cancel() ```swift // Termination through deiniting the sequence -let (_, source) = AsyncStream.makeStream( +let (_, source) = MutliProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 2, high: 4) ) @@ -279,109 +279,120 @@ let (_, source) = AsyncStream.makeStream( ```swift // Termination through deiniting the iterator -let (stream, source) = AsyncStream.makeStream( +let (channel, source) = MutliProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 2, high: 4) ) -_ = stream.makeAsyncIterator() +_ = channel.makeAsyncIterator() ``` ```swift // Termination through calling finish -let (stream, source) = AsyncStream.makeStream( +let (channel, source) = MutliProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 2, high: 4) ) -_ = try source.write(1) +_ = try source.send(1) source.finish() -for await element in stream {} +for await element in channel {} // onTerminate will be called after all elements have been consumed ``` -Similar to the downstream consumer termination, trying to write more elements after the -producer has been terminated will result in an error thrown from the write methods. +Similar to the downstream consumer termination, trying to send more elements after the +producer has been terminated will result in an error thrown from the send methods. ## Detailed design -All new APIs on `AsyncStream` and `AsyncThrowingStream` are as follows: - ```swift -/// Error that is thrown from the various `write` methods of the -/// ``AsyncBackpressuredStream/Source``. +/// An error that is thrown from the various `send` methods of the +/// ``MultiProducerSingleConsumerChannel/Source``. /// -/// This error is thrown when the asynchronous stream is already finished when -/// trying to write new elements. -public struct AsyncBackpressuredStreamAlreadyFinishedError : Error { +/// This error is thrown when the channel is already finished when +/// trying to send new elements to the source. +public struct MultiProducerSingleConsumerChannelAlreadyFinishedError : Error { + + @usableFromInline + internal init() } -/// A struct that acts as a source asynchronous sequence. +/// A multi producer single consumer channel. +/// +/// The ``MultiProducerSingleConsumerChannel`` provides a ``MultiProducerSingleConsumerChannel/Source`` to +/// send values to the channel. The source exposes the internal backpressure of the asynchronous sequence to the +/// producer. Additionally, the source can be used from synchronous and asynchronous contexts. /// -/// The ``AsyncBackpressuredStream`` provides a ``AsyncBackpressuredStream/Source`` to -/// write values to the stream. The source exposes the internal backpressure of the asynchronous sequence to the -/// external producer. This allows to bridge both synchronous and asynchronous producers into an asynchronous sequence. /// -/// ## Using an AsyncBackpressuredStream +/// ## Using a MultiProducerSingleConsumerChannel /// -/// To use an ``AsyncBackpressuredStream`` you have to create a new stream with it's source first by calling -/// the ``AsyncBackpressuredStream/makeStream(of:throwing:backpressureStrategy:)`` method. -/// Afterwards, you can pass the source to the producer and the stream to the consumer. +/// To use a ``MultiProducerSingleConsumerChannel`` you have to create a new channel with it's source first by calling +/// the ``MultiProducerSingleConsumerChannel/makeChannel(of:throwing:BackpressureStrategy:)`` method. +/// Afterwards, you can pass the source to the producer and the channel to the consumer. /// /// ``` -/// let (stream, source) = AsyncBackpressuredStream.makeStream( +/// let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( /// backpressureStrategy: .watermark(low: 2, high: 4) /// ) +/// ``` +/// +/// ### Asynchronous producers /// +/// Values can be send to the source from asynchronous contexts using ``MultiProducerSingleConsumerChannel/Source/send(_:)-9b5do`` +/// and ``MultiProducerSingleConsumerChannel/Source/send(contentsOf:)-4myrz``. Backpressure results in calls +/// to the `send` methods to be suspended. Once more elements should be produced the `send` methods will be resumed. +/// +/// ``` /// try await withThrowingTaskGroup(of: Void.self) { group in /// group.addTask { -/// try await source.write(1) -/// try await source.write(2) -/// try await source.write(3) +/// try await source.send(1) +/// try await source.send(2) +/// try await source.send(3) /// } /// -/// for await element in stream { +/// for await element in channel { /// print(element) /// } /// } /// ``` /// -/// The source also exposes synchronous write methods that communicate the backpressure via callbacks. -public struct AsyncBackpressuredStream: AsyncSequence { - /// Initializes a new ``AsyncBackpressuredStream`` and an ``AsyncBackpressuredStream/Source``. +/// ### Synchronous producers +/// +/// Values can also be send to the source from synchronous context. Backpressure is also exposed on the synchronous contexts; however, +/// it is up to the caller to decide how to properly translate the backpressure to underlying producer e.g. by blocking the thread. +/// +/// ## Finishing the source +/// +/// To properly notify the consumer if the production of values has been finished the source's ``MultiProducerSingleConsumerChannel/Source/finish(throwing:)`` **must** be called. +public struct MultiProducerSingleConsumerChannel: AsyncSequence { + /// Initializes a new ``MultiProducerSingleConsumerChannel`` and an ``MultiProducerSingleConsumerChannel/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. - public static func makeStream( - of elementType: Element.Type = Element.self, - throwing failureType: Failure.Type = Failure.self, - backPressureStrategy: Source.BackPressureStrategy - ) -> (`Self`, Source) where Failure == any Error + /// - elementType: The element type of the channel. + /// - failureType: The failure type of the channel. + /// - BackpressureStrategy: The backpressure strategy that the channel should use. + /// - Returns: A tuple containing the channel and its source. The source should be passed to the + /// producer while the channel should be passed to the consumer. + public static func makeChannel(of elementType: Element.Type = Element.self, throwing failureType: Failure.Type = Never.self, backpressureStrategy: Source.BackpressureStrategy) -> (`Self`, Source) } -extension AsyncBackpressuredStream { - /// A struct to interface between producer code and an asynchronous stream. +extension MultiProducerSingleConsumerChannel { + /// A struct to send values to the channel. /// - /// 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. + /// Use this source to provide elements to the channel by calling one of the `send` methods. /// - /// - Important: You must terminate the source by calling one of the `finish` methods otherwise the stream's iterator - /// will never terminate. + /// - Important: You must terminate the source by calling ``finish(throwing:)``. public struct Source: Sendable { - /// A strategy that handles the backpressure of the asynchronous stream. - public struct BackPressureStrategy: Sendable { + /// A strategy that handles the backpressure of the channel. + public struct BackpressureStrategy: Sendable { + /// A backpressure strategy using a high and low watermark to suspend and resume production respectively. /// /// - Parameters: /// - low: When the number of buffered elements drops below the low watermark, producers will be resumed. /// - high: When the number of buffered elements rises above the high watermark, producers will be suspended. - public static func watermark(low: Int, high: Int) -> BackPressureStrategy + public static func watermark(low: Int, high: Int) -> BackpressureStrategy /// A backpressure strategy using a high and low watermark to suspend and resume production respectively. /// @@ -391,15 +402,15 @@ extension AsyncBackpressuredStream { /// - waterLevelForElement: A closure used to compute the contribution of each buffered element to the current water level. /// /// - Note, `waterLevelForElement` will be called on each element when it is written into the source and when - /// it is consumed from the stream, so it is recommended to provide an function that runs in constant time. - public static func watermark(low: Int, high: Int, waterLevelForElement: @escaping @Sendable (Element) -> Int) -> BackPressureStrategy + /// it is consumed from the channel, so it is recommended to provide an function that runs in constant time. + public static func watermark(low: Int, high: Int, waterLevelForElement: @escaping @Sendable (Element) -> Int) -> BackpressureStrategy } - /// A type that indicates the result of writing elements to the source. - public enum WriteResult: Sendable { - /// A token that is returned when the asynchronous stream's backpressure strategy indicated that production should + /// A type that indicates the result of sending elements to the source. + public enum SendResult: Sendable { + /// A token that is returned when the channel's backpressure strategy indicated that production should /// be suspended. Use this token to enqueue a callback by calling the ``enqueueCallback(_:)`` method. - public struct CallbackToken: Sendable {} + public struct CallbackToken: Sendable { } /// Indicates that more elements should be produced and written to the source. case produceMore @@ -410,44 +421,44 @@ extension AsyncBackpressuredStream { case enqueueCallback(CallbackToken) } - /// A callback to invoke when the stream finished. + /// A callback to invoke when the channel finished. /// - /// The stream finishes and calls this closure in the following cases: + /// The channel 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 public var onTermination: (@Sendable () -> Void)? { get set } - /// Writes new elements to the asynchronous stream. + /// Sends new elements to the channel. /// - /// 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 + /// If there is a task consuming the channel and awaiting the next element then the task will get resumed with the + /// first element of the provided sequence. If the channel already terminated then this method will throw an error /// indicating the failure. /// - /// - Parameter sequence: The elements to write to the asynchronous stream. + /// - Parameter sequence: The elements to send to the channel. /// - Returns: The result that indicates if more elements should be produced at this time. - public func write(contentsOf sequence: S) throws -> WriteResult where Element == S.Element, S : Sequence + public func send(contentsOf sequence: S) throws -> SendResult where Element == S.Element, S : Sequence - /// Write the element to the asynchronous stream. + /// Send the element to the channel. /// - /// 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 + /// If there is a task consuming the channel and awaiting the next element then the task will get resumed with the + /// provided element. If the channel already terminated then this method will throw an error /// indicating the failure. /// - /// - Parameter element: The element to write to the asynchronous stream. + /// - Parameter element: The element to send to the channel. /// - Returns: The result that indicates if more elements should be produced at this time. - public func write(_ element: Element) throws -> WriteResult + public func send(_ element: Element) throws -> SendResult /// Enqueues a callback that will be invoked once more elements should be produced. /// - /// Call this method after ``write(contentsOf:)`` or ``write(:)`` returned ``WriteResult/enqueueCallback(_:)``. + /// Call this method after ``send(contentsOf:)-5honm`` or ``send(_:)-3jxzb`` returned ``SendResult/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. - public func enqueueCallback(callbackToken: WriteResult.CallbackToken, onProduceMore: @escaping @Sendable (Result) -> Void) + public func enqueueCallback(callbackToken: consuming SendResult.CallbackToken, onProduceMore: @escaping @Sendable (Result) -> Void) /// Cancel an enqueued callback. /// @@ -457,98 +468,94 @@ extension AsyncBackpressuredStream { /// will mark the passed `callbackToken` as cancelled. /// /// - Parameter callbackToken: The callback token. - public func cancelCallback(callbackToken: WriteResult.CallbackToken) + public func cancelCallback(callbackToken: consuming SendResult.CallbackToken) - /// Write new elements to the asynchronous stream and provide a callback which will be invoked once more elements should be produced. + /// Send new elements to the channel 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 + /// If there is a task consuming the channel and awaiting the next element then the task will get resumed with the + /// first element of the provided sequence. If the channel already terminated then `onProduceMore` will be invoked with /// a `Result.failure`. /// /// - Parameters: - /// - sequence: The elements to write to the asynchronous stream. + /// - sequence: The elements to send to the channel. /// - onProduceMore: The callback which gets invoked once more elements should be produced. This callback might be - /// invoked during the call to ``write(contentsOf:onProduceMore:)``. - public func write(contentsOf sequence: S, onProduceMore: @escaping @Sendable (Result) -> Void) where Element == S.Element, S : Sequence + /// invoked during the call to ``send(contentsOf:onProduceMore:)``. + public func send(contentsOf sequence: S, onProduceMore: @escaping @Sendable (Result) -> Void) where Element == S.Element, S : Sequence - /// Writes the element to the asynchronous stream. + /// Sends the element to the channel. /// - /// 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 + /// If there is a task consuming the channel and awaiting the next element then the task will get resumed with the + /// provided element. If the channel already terminated then `onProduceMore` will be invoked with /// a `Result.failure`. /// /// - Parameters: - /// - sequence: The element to write to the asynchronous stream. + /// - element: The element to send to the channel. /// - onProduceMore: The callback which gets invoked once more elements should be produced. This callback might be - /// invoked during the call to ``write(_:onProduceMore:)``. - public func write(_ element: Element, onProduceMore: @escaping @Sendable (Result) -> Void) + /// invoked during the call to ``send(_:onProduceMore:)``. + public func send(_ element: Element, onProduceMore: @escaping @Sendable (Result) -> Void) - /// Write new elements to the asynchronous stream. + /// Send new elements to the channel. /// - /// 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 + /// If there is a task consuming the channel and awaiting the next element then the task will get resumed with the + /// first element of the provided sequence. If the channel 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. - public func write(contentsOf sequence: S) async throws where Element == S.Element, S : Sequence + /// - sequence: The elements to send to the channel. + public func send(contentsOf sequence: S) async throws where Element == S.Element, S : Sequence - /// Write new element to the asynchronous stream. + /// Send new element to the channel. /// - /// 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 + /// If there is a task consuming the channel and awaiting the next element then the task will get resumed with the + /// provided element. If the channel 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. - public func write(_ element: Element) async throws + /// - element: The element to send to the channel. + public func send(_ element: Element) async throws - /// Write the elements of the asynchronous sequence to the asynchronous stream. + /// Send the elements of the asynchronous sequence to the channel. /// - /// This method returns once the provided asynchronous sequence or the the asynchronous stream finished. + /// This method returns once the provided asynchronous sequence or the channel 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. - public func write(contentsOf sequence: S) async throws where Element == S.Element, S : AsyncSequence + /// - sequence: The elements to send to the channel. + public func send(contentsOf sequence: S) async throws where Element == S.Element, S : AsyncSequence /// 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 + /// Calling this function more than once has no effect. After calling finish, the channel enters a terminal state and doesn't accept /// new elements. /// /// - Parameters: /// - error: The error to throw, or `nil`, to finish normally. - public func finish(throwing error: Failure?) + public func finish(throwing error: Failure? = nil) } } -extension AsyncBackpressuredStream { - public struct Iterator : AsyncIteratorProtocol { - public mutating func next() async throws -> Element? - } +extension MultiProducerSingleConsumerChannel { + /// The asynchronous iterator for iterating the channel. + /// + /// 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()`. + public struct Iterator: AsyncIteratorProtocol {} + + /// Creates the asynchronous iterator that produces elements of this + /// asynchronous sequence. public func makeAsyncIterator() -> Iterator } -extension AsyncBackpressuredStream: Sendable where Element: Sendable {} - -@available(*, unavailable) -extension AsyncBackpressuredStream.Iterator: Sendable {} -``` - -Additionally, this proposal adds a new `AsyncNonThrowingBackpressuredStream` -which is identical to the above except that it has only one generic parameter -and that the `Source` doesn't offer a `finish(throwing:)` method. - -```swift -public struct AsyncNonThrowingBackpressuredStream: AsyncSequence { ... } +extension MultiProducerSingleConsumerChannel: Sendable where Element : Sendable {} ``` ## Comparison to other root asynchronous sequences @@ -570,16 +577,6 @@ generic and fully inlinable type and quite unwiedly to use. This proposal is heavily inspired by the learnings from this type but tries to create a more flexible and easier to use API that fits into the standard library. -## Source compatibility - -This change is additive and does not affect source compatibility. - -## ABI compatibility - -This change is additive and does not affect ABI compatibility. All new methods -are non-inlineable leaving us flexiblity to change the implementation in the -future. - ## Future directions ### Adaptive backpressure strategy @@ -590,84 +587,8 @@ An adaptive strategy regulates the backpressure based on the rate of consumption and production. With the proposed new APIs we can easily add further strategies. -### Element size dependent strategy - -When the stream's element is a collection type then the proposed high/low -watermark backpressure strategy might lead to unexpected results since each -element can vary in actual memory size. In the future, we could provide a new -backpressure strategy that supports inspecting the size of the collection. - -### Deprecate `Async[Throwing]Stream.Continuation` - -In the future, we could deprecate the current continuation based APIs since the -new proposed APIs are also capable of bridging non-backpressured producers by -just discarding the `WriteResult`. The only use-case that the new APIs do not -cover is the _anycast_ behaviour of the current `AsyncStream` where one can -create multiple iterators to the stream as long as no two iterators are -consuming the stream at the same time. This can be solved via additional -algorithms such as `broadcast` in the `swift-async-algorithms` package. - -To give developers more time to adopt the new APIs the deprecation of the -current APIs should be deferred to a future version. Especially since those new -APIs are not backdeployed like the current Concurrency runtime. - -### Introduce a `Writer` and an `AsyncWriter` protocol - -The newly introduced `Source` type offers a bunch of different write methods. We -have seen similar types used in other places such as file abstraction or -networking APIs. We could introduce a new `Writer` and `AsyncWriter` protocol in -the future to enable writing generic algorithms on top of writers. The `Source` -type could then conform to these new protocols. - ## Alternatives considered -### Providing an `Async[Throwing]Stream.Continuation.onConsume` - -We could add a new closure property to the `Async[Throwing]Stream.Continuation` -which is invoked once an element has been consumed to implement a backpressure -strategy; however, this requires the usage of a synchronization mechanism since -the consumption and production often happen on separate threads. The -added complexity and performance impact led to avoiding this approach. - -### Provide a getter for the current buffer depth - -We could provide a getter for the current buffer depth on the -`Async[Throwing]Stream.Continuation`. This could be used to query the buffer -depth at an arbitrary time; however, it wouldn't allow us to implement -backpressure strategies such as high/low watermarks without continuously asking -what the buffer depth is. That would result in a very inefficient -implementation. - -### Extending `Async[Throwing]Stream.Continuation` - -Extending the current APIs to support all expected behaviors is problematic -since it would change the semantics and might lead to currently working code -misbehaving. Furthermore, extending the current APIs to support backpressure -turns out to be problematic without compromising performance or usability. - -### Introducing a new type - -We could introduce a new type such as `AsyncBackpressured[Throwing]Stream`; -however, one of the original intentions of `Async[Throwing]Stream` was to be -able to bridge backpressured systems. Furthermore, `Async[Throwing]Stream` is -the best name. Therefore, this proposal decided to provide new interfaces to -`Async[Throwing]Stream`. - -### Stick with the current `Continuation` and `yield` naming - -The proposal decided against sticking to the current names since the existing -names caused confusion to them being used in multiple places. Continuation was -both used by the `AsyncStream` but also by Swift Concurrency via -`CheckedContinuation` and `UnsafeContinuation`. Similarly, yield was used by -both `AsyncStream.Continuation.yield()`, `Task.yield()` and the `yield` keyword. -Having different names for these different concepts makes it easier to explain -their usage. The currently proposed `write` names were choosen to align with the -future direction of adding an `AsyncWriter` protocol. `Source` is a common name -in flow based systems such as Akka. Other names that were considered: - -- `enqueue` -- `send` - ### Provide the `onTermination` callback to the factory method During development of the new APIs, I first tried to provide the `onTermination` diff --git a/Sources/AsyncAlgorithms/BackPressuredStream/AsyncBackPressuredStream.swift b/Sources/AsyncAlgorithms/BackPressuredStream/AsyncBackPressuredStream.swift deleted file mode 100644 index 4b7d3e2a..00000000 --- a/Sources/AsyncAlgorithms/BackPressuredStream/AsyncBackPressuredStream.swift +++ /dev/null @@ -1,425 +0,0 @@ -//===----------------------------------------------------------------------===// -// -// This source file is part of the Swift Async Algorithms open source project -// -// Copyright (c) 2023 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 -// -//===----------------------------------------------------------------------===// - -/// Error that is thrown from the various `write` methods of the -/// ``AsyncBackPressuredStream/Source``. -/// -/// This error is thrown when the asynchronous stream is already finished when -/// trying to write new elements. -public struct AsyncBackPressuredStreamAlreadyFinishedError: Error {} - -/// A struct that acts as a source asynchronous sequence. -/// -/// The ``AsyncBackPressuredStream`` provides a ``AsyncBackPressuredStream/Source`` to -/// write values to the stream. The source exposes the internal backpressure of the asynchronous sequence to the -/// external producer. This allows to bridge both synchronous and asynchronous producers into an asynchronous sequence. -/// -/// ## Using an AsyncBackPressuredStream -/// -/// To use an ``AsyncBackPressuredStream`` you have to create a new stream with it's source first by calling -/// the ``AsyncBackPressuredStream/makeStream(of:throwing:backPressureStrategy:)`` method. -/// Afterwards, you can pass the source to the producer and the stream to the consumer. -/// -/// ``` -/// let (stream, source) = AsyncBackPressuredStream.makeStream( -/// backpressureStrategy: .watermark(low: 2, high: 4) -/// ) -/// -/// try await withThrowingTaskGroup(of: Void.self) { group in -/// group.addTask { -/// try await source.write(1) -/// try await source.write(2) -/// try await source.write(3) -/// } -/// -/// for await element in stream { -/// print(element) -/// } -/// } -/// ``` -/// -/// The source also exposes synchronous write methods that communicate the backpressure via callbacks. -public struct AsyncBackPressuredStream: AsyncSequence { - /// A private class to give the ``AsyncBackPressuredStream`` a deinit so we - /// can tell the producer when any potential consumer went away. - private final class _Backing: Sendable { - /// The underlying storage. - fileprivate let storage: _AsyncBackPressuredStreamBackPressuredStorage - - init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { - self.storage = storage - } - - deinit { - storage.sequenceDeinitialized() - } - } - - /// The backing storage. - private let backing: _Backing - - /// Initializes a new ``AsyncBackPressuredStream`` and an ``AsyncBackPressuredStream/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. - public static func makeStream( - of elementType: Element.Type = Element.self, - throwing failureType: Failure.Type = Failure.self, - backPressureStrategy: Source.BackPressureStrategy - ) -> (`Self`, Source) where Failure == Error { - let storage = _AsyncBackPressuredStreamBackPressuredStorage( - backPressureStrategy: backPressureStrategy.internalBackPressureStrategy - ) - let source = Source(storage: storage) - - return (.init(storage: storage), source) - } - - init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { - self.backing = .init(storage: storage) - } -} - -extension AsyncBackPressuredStream { - /// A struct 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. - /// - /// - Important: You must terminate the source by calling one of the `finish` methods otherwise the stream's iterator - /// will never terminate. - public struct Source: Sendable { - /// A strategy that handles the backpressure of the asynchronous stream. - public struct BackPressureStrategy: Sendable { - var internalBackPressureStrategy: _AsyncBackPressuredStreamInternalBackPressureStrategy - - /// A backpressure strategy using a high and low watermark to suspend and resume production respectively. - /// - /// - Parameters: - /// - low: When the number of buffered elements drops below the low watermark, producers will be resumed. - /// - high: When the number of buffered elements rises above the high watermark, producers will be suspended. - public static func watermark(low: Int, high: Int) -> BackPressureStrategy { - .init( - internalBackPressureStrategy: .watermark( - .init(low: low, high: high, waterLevelForElement: nil) - ) - ) - } - - /// A backpressure strategy using a high and low watermark to suspend and resume production respectively. - /// - /// - Parameters: - /// - low: When the number of buffered elements drops below the low watermark, producers will be resumed. - /// - high: When the number of buffered elements rises above the high watermark, producers will be suspended. - /// - waterLevelForElement: A closure used to compute the contribution of each buffered element to the current water level. - /// - /// - Note, `waterLevelForElement` will be called on each element when it is written into the source and when - /// it is consumed from the stream, so it is recommended to provide an function that runs in constant time. - public static func watermark( - low: Int, - high: Int, - waterLevelForElement: @escaping @Sendable (Element) -> Int - ) -> BackPressureStrategy { - .init( - internalBackPressureStrategy: .watermark( - .init(low: low, high: high, waterLevelForElement: waterLevelForElement) - ) - ) - } - } - - /// A type that indicates the result of writing elements to the source. - public 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. - public struct CallbackToken: Sendable { - let id: UInt - } - - /// 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. - final class _Backing: Sendable { - let storage: _AsyncBackPressuredStreamBackPressuredStorage - - init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { - self.storage = storage - } - - // TODO: Double check - 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 - public var onTermination: (@Sendable () -> Void)? { - set { - self._backing.storage.onTermination = newValue - } - get { - self._backing.storage.onTermination - } - } - - private var _backing: _Backing - - internal init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { - 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. - public 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. - public 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:)-4amlx`` or ``write(_:)-8e7el`` 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. - public 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. - public 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:)``. - public 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:)``. - public 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. - public 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. - public 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. - public 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. - public func finish(throwing error: Failure?) { - self._backing.storage.finish(error) - } - } -} - -extension AsyncBackPressuredStream { - /// 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()`. - public struct Iterator: AsyncIteratorProtocol { - private class _Backing { - let storage: _AsyncBackPressuredStreamBackPressuredStorage - - init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { - self.storage = storage - self.storage.iteratorInitialized() - } - - deinit { - self.storage.iteratorDeinitialized() - } - } - - private let backing: _Backing - - init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { - self.backing = .init(storage: storage) - } - - /// The next value from the asynchronous stream. - /// - /// When `next()` returns `nil`, this signifies the end of the - /// `AsyncThrowingStream`. - /// - /// 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 `AsyncThrowingStream` terminates. In this case, - /// `next()` may return `nil` immediately, or else return `nil` on - /// subsequent calls. - public mutating func next() async throws -> Element? { - try await self.backing.storage.next() - } - } - - /// Creates the asynchronous iterator that produces elements of this - /// asynchronous sequence. - public func makeAsyncIterator() -> Iterator { - Iterator(storage: self.backing.storage) - } -} - -extension AsyncBackPressuredStream: Sendable where Element: Sendable {} - -@available(*, unavailable) -extension AsyncBackPressuredStream.Iterator: Sendable {} diff --git a/Sources/AsyncAlgorithms/BackpressuredStream/AsyncNonThrowingBackPressuredStream.swift b/Sources/AsyncAlgorithms/BackpressuredStream/AsyncNonThrowingBackPressuredStream.swift deleted file mode 100644 index 07910595..00000000 --- a/Sources/AsyncAlgorithms/BackpressuredStream/AsyncNonThrowingBackPressuredStream.swift +++ /dev/null @@ -1,429 +0,0 @@ -//===----------------------------------------------------------------------===// -// -// This source file is part of the Swift Async Algorithms open source project -// -// Copyright (c) 2023 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 -// -//===----------------------------------------------------------------------===// - -/// A struct that acts as a source asynchronous sequence. -/// -/// The ``AsyncNonThrowingBackPressuredStream`` provides a ``AsyncNonThrowingBackPressuredStream/Source`` to -/// write values to the stream. The source exposes the internal backpressure of the asynchronous sequence to the -/// external producer. This allows to bridge both synchronous and asynchronous producers into an asynchronous sequence. -/// -/// ## Using an AsyncNonThrowingBackPressuredStream -/// -/// To use an ``AsyncNonThrowingBackPressuredStream`` you have to create a new stream with it's source first by calling -/// the ``AsyncNonThrowingBackPressuredStream/makeStream(of:backpressureStrategy:)`` method. -/// Afterwards, you can pass the source to the producer and the stream to the consumer. -/// -/// ``` -/// let (stream, source) = AsyncNonThrowingBackPressuredStream.makeStream( -/// backpressureStrategy: .watermark(low: 2, high: 4) -/// ) -/// -/// try await withThrowingTaskGroup(of: Void.self) { group in -/// group.addTask { -/// try await source.write(1) -/// try await source.write(2) -/// try await source.write(3) -/// } -/// -/// for await element in stream { -/// print(element) -/// } -/// } -/// ``` -/// -/// The source also exposes synchronous write methods that communicate the backpressure via callbacks. -public struct AsyncNonThrowingBackPressuredStream: AsyncSequence { - /// A private class to give the ``AsyncNonThrowingBackPressuredStream`` a deinit so we - /// can tell the producer when any potential consumer went away. - private final class _Backing: Sendable { - /// The underlying storage. - fileprivate let storage: _AsyncBackPressuredStreamBackPressuredStorage - - init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { - self.storage = storage - } - - deinit { - storage.sequenceDeinitialized() - } - } - - /// The backing storage. - private let backing: _Backing - - /// Initializes a new ``AsyncNonThrowingBackPressuredStream`` and an ``AsyncNonThrowingBackPressuredStream/Source``. - /// - /// - Parameters: - /// - elementType: The element 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. - public static func makeStream( - of elementType: Element.Type = Element.self, - backPressureStrategy: Source.BackPressureStrategy - ) -> (`Self`, Source) { - let storage = _AsyncBackPressuredStreamBackPressuredStorage( - backPressureStrategy: backPressureStrategy.internalBackPressureStrategy - ) - let source = Source(storage: storage) - - return (.init(storage: storage), source) - } - - init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { - self.backing = .init(storage: storage) - } -} - -extension AsyncNonThrowingBackPressuredStream { - /// A struct 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. - /// - /// - Important: You must terminate the source by calling one of the `finish` methods otherwise the stream's iterator - /// will never terminate. - public struct Source: Sendable { - /// A strategy that handles the backpressure of the asynchronous stream. - public struct BackPressureStrategy: Sendable { - var internalBackPressureStrategy: _AsyncBackPressuredStreamInternalBackPressureStrategy - - /// A backpressure strategy using a high and low watermark to suspend and resume production respectively. - /// - /// - Parameters: - /// - low: When the number of buffered elements drops below the low watermark, producers will be resumed. - /// - high: When the number of buffered elements rises above the high watermark, producers will be suspended. - public static func watermark(low: Int, high: Int) -> BackPressureStrategy { - .init( - internalBackPressureStrategy: .watermark( - .init(low: low, high: high, waterLevelForElement: nil) - ) - ) - } - - /// A backpressure strategy using a high and low watermark to suspend and resume production respectively. - /// - /// - Parameters: - /// - low: When the number of buffered elements drops below the low watermark, producers will be resumed. - /// - high: When the number of buffered elements rises above the high watermark, producers will be suspended. - /// - waterLevelForElement: A closure used to compute the contribution of each buffered element to the current water level. - /// - /// - Note, `waterLevelForElement` will be called on each element when it is written into the source and when - /// it is consumed from the stream, so it is recommended to provide an function that runs in constant time. - public static func watermark( - low: Int, - high: Int, - waterLevelForElement: @escaping @Sendable (Element) -> Int - ) -> BackPressureStrategy { - .init( - internalBackPressureStrategy: .watermark( - .init(low: low, high: high, waterLevelForElement: waterLevelForElement) - ) - ) - } - } - - /// A type that indicates the result of writing elements to the source. - public 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. - public struct CallbackToken: Sendable { - let id: UInt - } - - /// 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. - final class _Backing: Sendable { - let storage: _AsyncBackPressuredStreamBackPressuredStorage - - init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { - self.storage = storage - } - - // TODO: Double check - 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 - public var onTermination: (@Sendable () -> Void)? { - set { - self._backing.storage.onTermination = newValue - } - get { - self._backing.storage.onTermination - } - } - - private var _backing: _Backing - - internal init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { - 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. - public func write(contentsOf sequence: S) throws -> WriteResult where Element == S.Element, S: Sequence { - switch try self._backing.storage.write(contentsOf: sequence) { - case .produceMore: - return .produceMore - case .enqueueCallback(let callbackToken): - return .enqueueCallback(.init(id: callbackToken.id)) - } - } - - /// 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. - public func write(_ element: Element) throws -> WriteResult { - switch try self._backing.storage.write(contentsOf: CollectionOfOne(element)) { - case .produceMore: - return .produceMore - case .enqueueCallback(let callbackToken): - return .enqueueCallback(.init(id: callbackToken.id)) - } - } - - /// 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. - public func enqueueCallback( - callbackToken: WriteResult.CallbackToken, - onProduceMore: @escaping @Sendable (Result) -> Void - ) { - let callbackToken = AsyncBackPressuredStream.Source.WriteResult.CallbackToken( - id: callbackToken.id - ) - 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. - public func cancelCallback(callbackToken: WriteResult.CallbackToken) { - let callbackToken = AsyncBackPressuredStream.Source.WriteResult.CallbackToken( - id: callbackToken.id - ) - 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:)``. - public 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:)``. - public 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. - public 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. - public 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. - public 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`. - /// - /// Calling this function more than once has no effect. After calling finish, the stream enters a terminal state and doesn't accept - /// new elements. - public func finish() { - self._backing.storage.finish(nil) - } - } -} - -extension AsyncNonThrowingBackPressuredStream { - /// 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()`. - public struct Iterator: AsyncIteratorProtocol { - private class _Backing { - let storage: _AsyncBackPressuredStreamBackPressuredStorage - - init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { - self.storage = storage - self.storage.iteratorInitialized() - } - - deinit { - self.storage.iteratorDeinitialized() - } - } - - private let backing: _Backing - - init(storage: _AsyncBackPressuredStreamBackPressuredStorage) { - self.backing = .init(storage: storage) - } - - /// The next value from the asynchronous stream. - /// - /// When `next()` returns `nil`, this signifies the end of the - /// `AsyncThrowingStream`. - /// - /// 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 `AsyncThrowingStream` terminates. In this case, - /// `next()` may return `nil` immediately, or else return `nil` on - /// subsequent calls. - public mutating func next() async -> Element? { - try! await self.backing.storage.next() - } - } - - /// Creates the asynchronous iterator that produces elements of this - /// asynchronous sequence. - public func makeAsyncIterator() -> Iterator { - Iterator(storage: self.backing.storage) - } -} - -extension AsyncNonThrowingBackPressuredStream: Sendable where Element: Sendable {} - -@available(*, unavailable) -extension AsyncNonThrowingBackPressuredStream.Iterator: Sendable {} diff --git a/Sources/AsyncAlgorithms/Internal/_TinyArray.swift b/Sources/AsyncAlgorithms/Internal/_TinyArray.swift new file mode 100644 index 00000000..07357ccb --- /dev/null +++ b/Sources/AsyncAlgorithms/Internal/_TinyArray.swift @@ -0,0 +1,329 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the Swift Async Algorithms open source project +// +// Copyright (c) 2023 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 +// +//===----------------------------------------------------------------------===// +//===----------------------------------------------------------------------===// +// +// This source file is part of the SwiftCertificates open source project +// +// Copyright (c) 2023 Apple Inc. and the SwiftCertificates project authors +// Licensed under Apache License v2.0 +// +// See LICENSE.txt for license information +// See CONTRIBUTORS.txt for the list of SwiftCertificates project authors +// +// SPDX-License-Identifier: Apache-2.0 +// +//===----------------------------------------------------------------------===// + +/// ``_TinyArray`` is a ``RandomAccessCollection`` optimised to store zero or one ``Element``. +/// It supports arbitrary many elements but if only up to one ``Element`` is stored it does **not** allocate separate storage on the heap +/// and instead stores the ``Element`` inline. +@usableFromInline +struct _TinyArray { + @usableFromInline + enum Storage { + case one(Element) + case arbitrary([Element]) + } + + @usableFromInline + var storage: Storage +} + +// MARK: - TinyArray "public" interface + +extension _TinyArray: Equatable where Element: Equatable {} +extension _TinyArray: Hashable where Element: Hashable {} +extension _TinyArray: Sendable where Element: Sendable {} + +extension _TinyArray: RandomAccessCollection { + @usableFromInline + typealias Element = Element + + @usableFromInline + typealias Index = Int + + @inlinable + subscript(position: Int) -> Element { + get { + self.storage[position] + } + set { + self.storage[position] = newValue + } + } + + @inlinable + var startIndex: Int { + self.storage.startIndex + } + + @inlinable + var endIndex: Int { + self.storage.endIndex + } +} + +extension _TinyArray { + @inlinable + init(_ elements: some Sequence) { + self.storage = .init(elements) + } + + @inlinable + init() { + self.storage = .init() + } + + @inlinable + mutating func append(_ newElement: Element) { + self.storage.append(newElement) + } + + @inlinable + mutating func append(contentsOf newElements: some Sequence) { + self.storage.append(contentsOf: newElements) + } + + @discardableResult + @inlinable + mutating func remove(at index: Int) -> Element { + self.storage.remove(at: index) + } + + @inlinable + mutating func removeAll(where shouldBeRemoved: (Element) throws -> Bool) rethrows { + try self.storage.removeAll(where: shouldBeRemoved) + } + + @inlinable + mutating func sort(by areInIncreasingOrder: (Element, Element) throws -> Bool) rethrows { + try self.storage.sort(by: areInIncreasingOrder) + } +} + +// MARK: - TinyArray.Storage "private" implementation + +extension _TinyArray.Storage: Equatable where Element: Equatable { + @inlinable + static func == (lhs: Self, rhs: Self) -> Bool { + switch (lhs, rhs) { + case (.one(let lhs), .one(let rhs)): + return lhs == rhs + case (.arbitrary(let lhs), .arbitrary(let rhs)): + // we don't use lhs.elementsEqual(rhs) so we can hit the fast path from Array + // if both arrays share the same underlying storage: https://github.com/apple/swift/blob/b42019005988b2d13398025883e285a81d323efa/stdlib/public/core/Array.swift#L1775 + return lhs == rhs + + case (.one(let element), .arbitrary(let array)), + (.arbitrary(let array), .one(let element)): + guard array.count == 1 else { + return false + } + return element == array[0] + + } + } +} +extension _TinyArray.Storage: Hashable where Element: Hashable { + @inlinable + func hash(into hasher: inout Hasher) { + // same strategy as Array: https://github.com/apple/swift/blob/b42019005988b2d13398025883e285a81d323efa/stdlib/public/core/Array.swift#L1801 + hasher.combine(count) + for element in self { + hasher.combine(element) + } + } +} +extension _TinyArray.Storage: Sendable where Element: Sendable {} + +extension _TinyArray.Storage: RandomAccessCollection { + @inlinable + subscript(position: Int) -> Element { + get { + switch self { + case .one(let element): + guard position == 0 else { + fatalError("index \(position) out of bounds") + } + return element + case .arbitrary(let elements): + return elements[position] + } + } + set { + switch self { + case .one: + guard position == 0 else { + fatalError("index \(position) out of bounds") + } + self = .one(newValue) + case .arbitrary(var elements): + elements[position] = newValue + self = .arbitrary(elements) + } + } + } + + @inlinable + var startIndex: Int { + 0 + } + + @inlinable + var endIndex: Int { + switch self { + case .one: return 1 + case .arbitrary(let elements): return elements.endIndex + } + } +} + +extension _TinyArray.Storage { + @inlinable + init(_ elements: some Sequence) { + var iterator = elements.makeIterator() + guard let firstElement = iterator.next() else { + self = .arbitrary([]) + return + } + guard let secondElement = iterator.next() else { + // newElements just contains a single element + // and we hit the fast path + self = .one(firstElement) + return + } + + var elements: [Element] = [] + elements.reserveCapacity(elements.underestimatedCount) + elements.append(firstElement) + elements.append(secondElement) + while let nextElement = iterator.next() { + elements.append(nextElement) + } + self = .arbitrary(elements) + } + + @inlinable + init() { + self = .arbitrary([]) + } + + @inlinable + mutating func append(_ newElement: Element) { + self.append(contentsOf: CollectionOfOne(newElement)) + } + + @inlinable + mutating func append(contentsOf newElements: some Sequence) { + switch self { + case .one(let firstElement): + var iterator = newElements.makeIterator() + guard let secondElement = iterator.next() else { + // newElements is empty, nothing to do + return + } + var elements: [Element] = [] + elements.reserveCapacity(1 + newElements.underestimatedCount) + elements.append(firstElement) + elements.append(secondElement) + elements.appendRemainingElements(from: &iterator) + self = .arbitrary(elements) + + case .arbitrary(var elements): + if elements.isEmpty { + // if `self` is currently empty and `newElements` just contains a single + // element, we skip allocating an array and set `self` to `.one(firstElement)` + var iterator = newElements.makeIterator() + guard let firstElement = iterator.next() else { + // newElements is empty, nothing to do + return + } + guard let secondElement = iterator.next() else { + // newElements just contains a single element + // and we hit the fast path + self = .one(firstElement) + return + } + elements.reserveCapacity(elements.count + newElements.underestimatedCount) + elements.append(firstElement) + elements.append(secondElement) + elements.appendRemainingElements(from: &iterator) + self = .arbitrary(elements) + + } else { + elements.append(contentsOf: newElements) + self = .arbitrary(elements) + } + + } + } + + @discardableResult + @inlinable + mutating func remove(at index: Int) -> Element { + switch self { + case .one(let oldElement): + guard index == 0 else { + fatalError("index \(index) out of bounds") + } + self = .arbitrary([]) + return oldElement + + case .arbitrary(var elements): + defer { + self = .arbitrary(elements) + } + return elements.remove(at: index) + + } + } + + @inlinable + mutating func removeAll(where shouldBeRemoved: (Element) throws -> Bool) rethrows { + switch self { + case .one(let oldElement): + if try shouldBeRemoved(oldElement) { + self = .arbitrary([]) + } + + case .arbitrary(var elements): + defer { + self = .arbitrary(elements) + } + return try elements.removeAll(where: shouldBeRemoved) + + } + } + + @inlinable + mutating func sort(by areInIncreasingOrder: (Element, Element) throws -> Bool) rethrows { + switch self { + case .one: + // a collection of just one element is always sorted, nothing to do + break + case .arbitrary(var elements): + defer { + self = .arbitrary(elements) + } + + try elements.sort(by: areInIncreasingOrder) + } + } +} + +extension Array { + @inlinable + mutating func appendRemainingElements(from iterator: inout some IteratorProtocol) { + while let nextElement = iterator.next() { + append(nextElement) + } + } +} diff --git a/Sources/AsyncAlgorithms/Locking.swift b/Sources/AsyncAlgorithms/Locking.swift index 952b13c8..ac8d34a3 100644 --- a/Sources/AsyncAlgorithms/Locking.swift +++ b/Sources/AsyncAlgorithms/Locking.swift @@ -17,17 +17,23 @@ import Glibc import WinSDK #endif +@usableFromInline internal struct Lock { #if canImport(Darwin) + @usableFromInline typealias Primitive = os_unfair_lock #elseif canImport(Glibc) + @usableFromInline typealias Primitive = pthread_mutex_t #elseif canImport(WinSDK) + @usableFromInline typealias Primitive = SRWLOCK #else + @usableFromInline typealias Primitive = Int #endif + @usableFromInline typealias PlatformLock = UnsafeMutablePointer let platformLock: PlatformLock @@ -54,7 +60,8 @@ internal struct Lock { platformLock.deinitialize(count: 1) } - fileprivate static func lock(_ platformLock: PlatformLock) { + @usableFromInline + static func lock(_ platformLock: PlatformLock) { #if canImport(Darwin) os_unfair_lock_lock(platformLock) #elseif canImport(Glibc) @@ -64,7 +71,8 @@ internal struct Lock { #endif } - fileprivate static func unlock(_ platformLock: PlatformLock) { + @usableFromInline + static func unlock(_ platformLock: PlatformLock) { #if canImport(Darwin) os_unfair_lock_unlock(platformLock) #elseif canImport(Glibc) @@ -115,14 +123,17 @@ internal struct Lock { } } +@usableFromInline struct ManagedCriticalState { - private final class LockedBuffer: ManagedBuffer { + @usableFromInline + final class LockedBuffer: ManagedBuffer { deinit { withUnsafeMutablePointerToElements { Lock.deinitialize($0) } } } - private let buffer: ManagedBuffer + @usableFromInline + let buffer: ManagedBuffer init(_ initial: State) { buffer = LockedBuffer.create(minimumCapacity: 1) { buffer in @@ -131,6 +142,7 @@ struct ManagedCriticalState { } } + @inlinable func withCriticalRegion(_ critical: (inout State) throws -> R) rethrows -> R { try buffer.withUnsafeMutablePointers { header, lock in Lock.lock(lock) diff --git a/Sources/AsyncAlgorithms/BackpressuredStream/AsyncBackPressuredStream+Internal.swift b/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel+Internal.swift similarity index 50% rename from Sources/AsyncAlgorithms/BackpressuredStream/AsyncBackPressuredStream+Internal.swift rename to Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel+Internal.swift index 50579730..0e4997df 100644 --- a/Sources/AsyncAlgorithms/BackpressuredStream/AsyncBackPressuredStream+Internal.swift +++ b/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel+Internal.swift @@ -11,75 +11,132 @@ import DequeModule -struct _AsyncBackPressuredStreamWatermarkBackPressureStrategy { +@usableFromInline +struct _MultiProducerSingleConsumerChannelWatermarkBackpressureStrategy: Sendable, CustomStringConvertible { /// The low watermark where demand should start. - private let low: Int + @usableFromInline + let _low: Int + /// The high watermark where demand should be stopped. - private let high: Int - private var currentWatermark: Int = 0 - private let waterLevelForElement: (@Sendable (Element) -> Int)? + @usableFromInline + let _high: Int + + /// The current watermark level. + @usableFromInline + var _currentWatermark: Int = 0 + + /// A closure that can be used to calculate the watermark impact of a single element + @usableFromInline + let _waterLevelForElement: (@Sendable (Element) -> Int)? + + @usableFromInline + var description: String { + "watermark(\(self._currentWatermark))" + } - /// Initializes a new ``_WatermarkBackPressureStrategy``. - /// - /// - Parameters: - /// - low: The low watermark where demand should start. - /// - high: The high watermark where demand should be stopped. init(low: Int, high: Int, waterLevelForElement: (@Sendable (Element) -> Int)?) { precondition(low <= high) - self.low = low - self.high = high - self.waterLevelForElement = waterLevelForElement + self._low = low + self._high = high + self._waterLevelForElement = waterLevelForElement } - mutating func didYield(elements: Deque.SubSequence) -> Bool { - if let waterLevelForElement { - self.currentWatermark += elements.reduce(0) { $0 + waterLevelForElement($1) } + @inlinable + mutating func didSend(elements: Deque.SubSequence) -> Bool { + if let waterLevelForElement = self._waterLevelForElement { + self._currentWatermark += elements.reduce(0) { $0 + waterLevelForElement($1) } } else { - self.currentWatermark += elements.count + self._currentWatermark += elements.count } - precondition(self.currentWatermark >= 0) + precondition(self._currentWatermark >= 0) // We are demanding more until we reach the high watermark - return self.currentWatermark < self.high + return self._currentWatermark < self._high } + @inlinable mutating func didConsume(element: Element) -> Bool { - if let waterLevelForElement { - self.currentWatermark -= waterLevelForElement(element) + if let waterLevelForElement = self._waterLevelForElement { + self._currentWatermark -= waterLevelForElement(element) } else { - self.currentWatermark -= 1 + self._currentWatermark -= 1 } - precondition(self.currentWatermark >= 0) + precondition(self._currentWatermark >= 0) // We start demanding again once we are below the low watermark - return self.currentWatermark < self.low + return self._currentWatermark < self._low + } +} + +@usableFromInline +struct _MultiProducerSingleConsumerChannelUnboundedBackpressureStrategy: Sendable, CustomStringConvertible { + @usableFromInline + var description: String { + return "unbounded" + } + + init() { } + + @inlinable + mutating func didSend(elements: Deque.SubSequence) -> Bool { + return true + } + + @inlinable + mutating func didConsume(element: Element) -> Bool { + return true } } -enum _AsyncBackPressuredStreamInternalBackPressureStrategy { - case watermark(_AsyncBackPressuredStreamWatermarkBackPressureStrategy) +@usableFromInline +enum _MultiProducerSingleConsumerChannelInternalBackpressureStrategy: Sendable, CustomStringConvertible { + /// A watermark based strategy. + case watermark(_MultiProducerSingleConsumerChannelWatermarkBackpressureStrategy) + /// An unbounded based strategy. + case unbounded(_MultiProducerSingleConsumerChannelUnboundedBackpressureStrategy) - mutating func didYield(elements: Deque.SubSequence) -> Bool { + @usableFromInline + var description: String { + switch self { + case .watermark(let strategy): + return strategy.description + case .unbounded(let unbounded): + return unbounded.description + } + } + + @inlinable + mutating func didSend(elements: Deque.SubSequence) -> Bool { switch self { case .watermark(var strategy): - let result = strategy.didYield(elements: elements) + let result = strategy.didSend(elements: elements) self = .watermark(strategy) return result + case .unbounded(var strategy): + let result = strategy.didSend(elements: elements) + self = .unbounded(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 + case .unbounded(var strategy): + let result = strategy.didConsume(element: element) + self = .unbounded(strategy) + return result } } } -// We are unchecked Sendable since we are protecting our state with a lock. -final class _AsyncBackPressuredStreamBackPressuredStorage: @unchecked Sendable { +@usableFromInline +final class _MultiProducerSingleConsumerChannelBackpressuredStorage { /// The state machine - var _stateMachine: ManagedCriticalState<_AsyncBackPressuredStateMachine> + @usableFromInline + var _stateMachine: ManagedCriticalState<_MultiProducerSingleConsumerStateMachine> var onTermination: (@Sendable () -> Void)? { set { @@ -95,9 +152,9 @@ final class _AsyncBackPressuredStreamBackPressuredStorage + backpressureStrategy: _MultiProducerSingleConsumerChannelInternalBackpressureStrategy ) { - self._stateMachine = .init(.init(backPressureStrategy: backPressureStrategy)) + self._stateMachine = .init(.init(backpressureStrategy: backpressureStrategy)) } func sequenceDeinitialized() { @@ -111,7 +168,7 @@ final class _AsyncBackPressuredStreamBackPressuredStorage - ) throws -> AsyncBackPressuredStream.Source.WriteResult { + ) throws -> MultiProducerSingleConsumerChannel.Source.SendResult { let action = self._stateMachine.withCriticalRegion { - return $0.write(sequence) + return $0.send(sequence) } switch action { @@ -183,7 +241,7 @@ final class _AsyncBackPressuredStreamBackPressuredStorage.Source.WriteResult.CallbackToken, + callbackToken: UInt64, onProduceMore: @escaping @Sendable (Result) -> Void ) { let action = self._stateMachine.withCriticalRegion { @@ -218,8 +277,9 @@ final class _AsyncBackPressuredStreamBackPressuredStorage.Source.WriteResult.CallbackToken + callbackToken: UInt64 ) { let action = self._stateMachine.withCriticalRegion { $0.cancelProducer(callbackToken: callbackToken) @@ -234,6 +294,7 @@ final class _AsyncBackPressuredStreamBackPressuredStorage=6.0) + @inlinable + func next(isolation actor: isolated (any Actor)?) 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 .returnFailureAndCallOnTermination(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(isolation: actor) + } + } + #else + @inlinable func next() async throws -> Element? { let action = self._stateMachine.withCriticalRegion { $0.next() @@ -296,10 +394,13 @@ final class _AsyncBackPressuredStreamBackPressuredStorage Element? { + #if compiler(>=6.0) + @inlinable + func suspendNext(isolation actor: isolated (any Actor)?) async throws -> Element? { return try await withTaskCancellationHandler { - return try await withCheckedThrowingContinuation { continuation in + return try await withUnsafeThrowingContinuation { continuation in let action = self._stateMachine.withCriticalRegion { $0.suspendNext(continuation: continuation) } @@ -343,7 +444,7 @@ final class _AsyncBackPressuredStreamBackPressuredStorage Element? { + return try await withTaskCancellationHandler { + return try await withUnsafeThrowingContinuation { continuation in + let action = self._stateMachine.withCriticalRegion { + $0.suspendNext(continuation: continuation) + } -/// The state machine of the backpressured async stream. -struct _AsyncBackPressuredStateMachine: Sendable { - enum _State { - struct Initial { - /// The backpressure strategy. - var backPressureStrategy: _AsyncBackPressuredStreamInternalBackPressureStrategy - /// Indicates if the iterator was initialized. - var iteratorInitialized: Bool - /// The onTermination callback. - var onTermination: (@Sendable () -> Void)? - } + switch action { + case .resumeConsumerWithElement(let continuation, let element): + continuation.resume(returning: element) - struct Streaming { - /// The backpressure strategy. - var backPressureStrategy: _AsyncBackPressuredStreamInternalBackPressureStrategy - /// Indicates if the iterator was initialized. - var iteratorInitialized: Bool - /// The onTermination callback. - var onTermination: (@Sendable () -> Void)? - /// The buffer of elements. - var buffer: Deque - /// The optional consumer continuation. - var consumerContinuation: CheckedContinuation? - /// The producer continuations. - var producerContinuations: Deque<(UInt, (Result) -> Void)> - /// The producers that have been cancelled. - var cancelledAsyncProducers: Deque - /// Indicates if we currently have outstanding demand. - var hasOutstandingDemand: Bool - } + case .resumeConsumerWithElementAndProducers(let continuation, let element, let producerContinuations): + continuation.resume(returning: element) + for producerContinuation in producerContinuations { + producerContinuation(Result.success(())) + } - struct SourceFinished { - /// Indicates if the iterator was initialized. - var iteratorInitialized: Bool - /// The buffer of elements. - var buffer: Deque - /// The failure that should be thrown after the last element has been consumed. - var failure: Failure? - /// The onTermination callback. - var onTermination: (@Sendable () -> Void)? - } + case .resumeConsumerWithFailureAndCallOnTermination(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() + } - 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) + switch action { + case .resumeConsumerWithNilAndCallOnTermination(let continuation, let onTermination): + continuation.resume(returning: nil) + onTermination?() - /// 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) + case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): + for producerContinuation in producerContinuations { + producerContinuation(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) + } + onTermination?() - /// An intermediate state to avoid CoWs. - case modify + case .none: + break + } + } } + #endif +} +/// The state machine of the channel. +@usableFromInline +struct _MultiProducerSingleConsumerStateMachine { /// The state machine's current state. - var _state: _State + @usableFromInline + var _state: _MultiProducerSingleConsumerState // The ID used for the next CallbackToken. - var nextCallbackTokenID: UInt = 0 + var _nextCallbackTokenID: UInt64 = 0 var _onTermination: (@Sendable () -> Void)? { set { @@ -424,9 +530,9 @@ struct _AsyncBackPressuredStateMachine: Sendable { initial.onTermination = newValue self._state = .initial(initial) - case .streaming(var streaming): - streaming.onTermination = newValue - self._state = .streaming(streaming) + case .channeling(var channeling): + channeling.onTermination = newValue + self._state = .channeling(channeling) case .sourceFinished(var sourceFinished): sourceFinished.onTermination = newValue @@ -436,7 +542,7 @@ struct _AsyncBackPressuredStateMachine: Sendable { break case .modify: - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") } } get { @@ -444,8 +550,8 @@ struct _AsyncBackPressuredStateMachine: Sendable { case .initial(let initial): return initial.onTermination - case .streaming(let streaming): - return streaming.onTermination + case .channeling(let channeling): + return channeling.onTermination case .sourceFinished(let sourceFinished): return sourceFinished.onTermination @@ -454,23 +560,23 @@ struct _AsyncBackPressuredStateMachine: Sendable { return nil case .modify: - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") } } } /// Initializes a new `StateMachine`. /// - /// We are passing and holding the back-pressure strategy here because + /// We are passing and holding the backpressure strategy here because /// it is a customizable extension of the state machine. /// - /// - Parameter backPressureStrategy: The back-pressure strategy. + /// - Parameter backpressureStrategy: The backpressure strategy. init( - backPressureStrategy: _AsyncBackPressuredStreamInternalBackPressureStrategy + backpressureStrategy: _MultiProducerSingleConsumerChannelInternalBackpressureStrategy ) { self._state = .initial( .init( - backPressureStrategy: backPressureStrategy, + backpressureStrategy: backpressureStrategy, iteratorInitialized: false, onTermination: nil ) @@ -478,29 +584,32 @@ struct _AsyncBackPressuredStateMachine: Sendable { } /// Generates the next callback token. - mutating func nextCallbackToken() -> AsyncBackPressuredStream.Source.WriteResult.CallbackToken { - let id = self.nextCallbackTokenID - self.nextCallbackTokenID += 1 - return .init(id: id) + @usableFromInline + mutating func nextCallbackToken() -> UInt64 { + let id = self._nextCallbackTokenID + self._nextCallbackTokenID += 1 + return 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. + /// Indicates that all producers should be failed and `onTermination` should be called. case failProducersAndCallOnTermination( - [(Result) -> Void], + _TinyArray<(Result) -> Void>, (@Sendable () -> Void)? ) } + @inlinable mutating func sequenceDeinitialized() -> SequenceDeinitializedAction? { switch self._state { case .initial(let initial): guard initial.iteratorInitialized else { // No iterator was created so we can transition to finished right away. - self._state = .finished(iteratorInitialized: false) + self._state = .finished(.init(iteratorInitialized: false, sourceFinished: false)) return .callOnTermination(initial.onTermination) } @@ -508,14 +617,14 @@ struct _AsyncBackPressuredStateMachine: Sendable { // This is an expected pattern and we just continue on normal. return .none - case .streaming(let streaming): - guard streaming.iteratorInitialized else { + case .channeling(let channeling): + guard channeling.iteratorInitialized else { // No iterator was created so we can transition to finished right away. - self._state = .finished(iteratorInitialized: false) + self._state = .finished(.init(iteratorInitialized: false, sourceFinished: false)) return .failProducersAndCallOnTermination( - Array(streaming.producerContinuations.map { $0.1 }), - streaming.onTermination + .init(channeling.producerContinuations.lazy.map { $0.1 }), + channeling.onTermination ) } // An iterator was created and we deinited the sequence. @@ -525,7 +634,7 @@ struct _AsyncBackPressuredStateMachine: Sendable { case .sourceFinished(let sourceFinished): guard sourceFinished.iteratorInitialized else { // No iterator was created so we can transition to finished right away. - self._state = .finished(iteratorInitialized: false) + self._state = .finished(.init(iteratorInitialized: false, sourceFinished: true)) return .callOnTermination(sourceFinished.onTermination) } @@ -539,10 +648,11 @@ struct _AsyncBackPressuredStateMachine: Sendable { return .none case .modify: - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") } } + @inlinable mutating func iteratorInitialized() { switch self._state { case .initial(var initial): @@ -555,14 +665,14 @@ struct _AsyncBackPressuredStateMachine: Sendable { self._state = .initial(initial) } - case .streaming(var streaming): - if streaming.iteratorInitialized { + case .channeling(var channeling): + if channeling.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) + channeling.iteratorInitialized = true + self._state = .channeling(channeling) } case .sourceFinished(var sourceFinished): @@ -575,69 +685,69 @@ struct _AsyncBackPressuredStateMachine: Sendable { 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 .finished(let finished): + if finished.iteratorInitialized { + // Our sequence is a unicast sequence and does not support multiple AsyncIterator's + fatalError("Only a single AsyncIterator can be created") + } else { + self._state = .finished(.init(iteratorInitialized: true, sourceFinished: finished.sourceFinished)) + } case .modify: - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel 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], + _TinyArray<(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) + self._state = .finished(.init(iteratorInitialized: true, sourceFinished: false)) return .callOnTermination(initial.onTermination) } else { // An iterator needs to be initialized before it can be deinitialized. - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") } - case .streaming(let streaming): - if streaming.iteratorInitialized { + case .channeling(let channeling): + if channeling.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) + self._state = .finished(.init(iteratorInitialized: true, sourceFinished: false)) return .failProducersAndCallOnTermination( - Array(streaming.producerContinuations.map { $0.1 }), - streaming.onTermination + .init(channeling.producerContinuations.lazy.map { $0.1 }), + channeling.onTermination ) } else { // An iterator needs to be initialized before it can be deinitialized. - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel 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) + self._state = .finished(.init(iteratorInitialized: true, sourceFinished: true)) return .callOnTermination(sourceFinished.onTermination) } else { // An iterator needs to be initialized before it can be deinitialized. - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") } case .finished: @@ -646,90 +756,74 @@ struct _AsyncBackPressuredStateMachine: Sendable { return .none case .modify: - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel 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. + /// Indicates that all producers should be failed and `onTermination` should be called. case failProducersAndCallOnTermination( - [(Result) -> Void], + _TinyArray<(Result) -> Void>, (@Sendable () -> Void)? ) /// Indicates that all producers should be failed. - case failProducers([(Result) -> Void]) + case failProducers(_TinyArray<(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 .initial, .channeling: + fatalError("The channel's source hasn't been finished but deinited") - case .streaming(let streaming): - guard streaming.buffer.isEmpty 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 - ) - ) + case .sourceFinished: + // This is the expected case where finish was called and then the source deinited + return .none - return .failProducers( - Array(streaming.producerContinuations.map { $0.1 }) - ) + case .finished(let finished): + if finished.sourceFinished { + // The source already got finished so this is fine. + return .none + } else { + fatalError("The channel's source hasn't been finished but deinited") } - // We can transition to finished right away since the buffer is empty now - self._state = .finished(iteratorInitialized: streaming.iteratorInitialized) - - return .failProducersAndCallOnTermination( - Array(streaming.producerContinuations.map { $0.1 }), - streaming.onTermination - ) - - case .sourceFinished, .finished: - // This is normal and we just have to tolerate it - return .none case .modify: - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") } } - /// Actions returned by `write()`. - enum WriteAction { + /// Actions returned by `send()`. + @usableFromInline + enum SendAction { /// 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: AsyncBackPressuredStream.Source.WriteResult.CallbackToken + callbackToken: UInt64 ) /// Indicates that the consumer should be resumed and the producer should be notified to produce more. case resumeConsumerAndReturnProduceMore( - continuation: CheckedContinuation, + continuation: UnsafeContinuation, element: Element ) /// Indicates that the consumer should be resumed and the producer should be suspended. case resumeConsumerAndReturnEnqueue( - continuation: CheckedContinuation, + continuation: UnsafeContinuation, element: Element, - callbackToken: AsyncBackPressuredStream.Source.WriteResult.CallbackToken + callbackToken: UInt64 ) /// Indicates that the producer has been finished. case throwFinishedError + @inlinable init( - callbackToken: AsyncBackPressuredStream.Source.WriteResult.CallbackToken?, - continuationAndElement: (CheckedContinuation, Element)? = nil + callbackToken: UInt64?, + continuationAndElement: (UnsafeContinuation, Element)? = nil ) { switch (callbackToken, continuationAndElement) { case (.none, .none): @@ -754,18 +848,19 @@ struct _AsyncBackPressuredStateMachine: Sendable { } } - mutating func write(_ sequence: some Sequence) -> WriteAction { + @inlinable + mutating func send(_ sequence: some Sequence) -> SendAction { switch self._state { case .initial(var initial): var buffer = Deque() buffer.append(contentsOf: sequence) - let shouldProduceMore = initial.backPressureStrategy.didYield(elements: buffer[...]) + let shouldProduceMore = initial.backpressureStrategy.didSend(elements: buffer[...]) let callbackToken = shouldProduceMore ? nil : self.nextCallbackToken() - self._state = .streaming( + self._state = .channeling( .init( - backPressureStrategy: initial.backPressureStrategy, + backpressureStrategy: initial.backpressureStrategy, iteratorInitialized: initial.iteratorInitialized, onTermination: initial.onTermination, buffer: buffer, @@ -778,37 +873,38 @@ struct _AsyncBackPressuredStateMachine: Sendable { return .init(callbackToken: callbackToken) - case .streaming(var streaming): + case .channeling(var channeling): self._state = .modify // We have an element and can resume the continuation - let bufferEndIndexBeforeAppend = streaming.buffer.endIndex - streaming.buffer.append(contentsOf: sequence) - let shouldProduceMore = streaming.backPressureStrategy.didYield( - elements: streaming.buffer[bufferEndIndexBeforeAppend...] + let bufferEndIndexBeforeAppend = channeling.buffer.endIndex + channeling.buffer.append(contentsOf: sequence) + var shouldProduceMore = channeling.backpressureStrategy.didSend( + elements: channeling.buffer[bufferEndIndexBeforeAppend...] ) - streaming.hasOutstandingDemand = shouldProduceMore - let callbackToken = shouldProduceMore ? nil : self.nextCallbackToken() + channeling.hasOutstandingDemand = shouldProduceMore - guard let consumerContinuation = streaming.consumerContinuation else { + guard let consumerContinuation = channeling.consumerContinuation else { // We don't have a suspended consumer so we just buffer the elements - self._state = .streaming(streaming) + self._state = .channeling(channeling) return .init( - callbackToken: callbackToken + callbackToken: shouldProduceMore ? nil : self.nextCallbackToken() ) } - 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: callbackToken) + guard let element = channeling.buffer.popFirst() else { + // We got a send of an empty sequence. We just tolerate this. + self._state = .channeling(channeling) + return .init(callbackToken: shouldProduceMore ? nil : self.nextCallbackToken()) } + // We need to tell the back pressure strategy that we consumed + shouldProduceMore = channeling.backpressureStrategy.didConsume(element: element) + channeling.hasOutstandingDemand = shouldProduceMore // We got a consumer continuation and an element. We can resume the consumer now - streaming.consumerContinuation = nil - self._state = .streaming(streaming) + channeling.consumerContinuation = nil + self._state = .channeling(channeling) return .init( - callbackToken: callbackToken, + callbackToken: shouldProduceMore ? nil : self.nextCallbackToken(), continuationAndElement: (consumerContinuation, element) ) @@ -817,11 +913,12 @@ struct _AsyncBackPressuredStateMachine: Sendable { return .throwFinishedError case .modify: - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") } } /// Actions returned by `enqueueProducer()`. + @usableFromInline enum EnqueueProducerAction { /// Indicates that the producer should be notified to produce more. case resumeProducer((Result) -> Void) @@ -829,101 +926,102 @@ struct _AsyncBackPressuredStateMachine: Sendable { case resumeProducerWithError((Result) -> Void, Error) } + @inlinable mutating func enqueueProducer( - callbackToken: AsyncBackPressuredStream.Source.WriteResult.CallbackToken, + callbackToken: UInt64, 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 public init so - // one must create it by calling `write` first. - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - case .streaming(var streaming): - if let index = streaming.cancelledAsyncProducers.firstIndex(of: callbackToken.id) { + case .channeling(var channeling): + if let index = channeling.cancelledAsyncProducers.firstIndex(of: callbackToken) { // Our producer got marked as cancelled. self._state = .modify - streaming.cancelledAsyncProducers.remove(at: index) - self._state = .streaming(streaming) + channeling.cancelledAsyncProducers.remove(at: index) + self._state = .channeling(channeling) return .resumeProducerWithError(onProduceMore, CancellationError()) - } else if streaming.hasOutstandingDemand { + } else if channeling.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)) + channeling.producerContinuations.append((callbackToken, onProduceMore)) - self._state = .streaming(streaming) + self._state = .channeling(channeling) return .none } case .sourceFinished, .finished: - // Since we are unlocking between yielding and suspending the yield + // Since we are unlocking between sending elements and suspending the send // It can happen that the source got finished or the consumption fully finishes. - return .resumeProducerWithError(onProduceMore, AsyncBackPressuredStreamAlreadyFinishedError()) + return .resumeProducerWithError(onProduceMore, MultiProducerSingleConsumerChannelAlreadyFinishedError()) case .modify: - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel 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: AsyncBackPressuredStream.Source.WriteResult.CallbackToken + callbackToken: UInt64 ) -> CancelProducerAction? { + //print(#function, self._state.description) switch self._state { case .initial: - // We need to transition to streaming before we can suspend - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - case .streaming(var streaming): - guard let index = streaming.producerContinuations.firstIndex(where: { $0.0 == callbackToken.id }) else { - // The task that yields was cancelled before yielding so the cancellation handler + case .channeling(var channeling): + guard let index = channeling.producerContinuations.firstIndex(where: { $0.0 == callbackToken }) else { + // The task that sends was cancelled before sending elements so the cancellation handler // got invoked right away self._state = .modify - streaming.cancelledAsyncProducers.append(callbackToken.id) - self._state = .streaming(streaming) + channeling.cancelledAsyncProducers.append(callbackToken) + self._state = .channeling(channeling) return .none } // 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) + let continuation = channeling.producerContinuations.remove(at: index).1 + self._state = .channeling(channeling) return .resumeProducerWithCancellationError(continuation) case .sourceFinished, .finished: - // Since we are unlocking between yielding and suspending the yield + // Since we are unlocking between sending elements and suspending the send // It can happen that the source got finished or the consumption fully finishes. return .none case .modify: - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel 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, + consumerContinuation: UnsafeContinuation, failure: Failure?, onTermination: (() -> Void)? ) /// Indicates that the producers should be resumed with an error. case resumeProducers( - producerContinuations: [(Result) -> Void] + producerContinuations: _TinyArray<(Result) -> Void> ) } @@ -931,7 +1029,7 @@ struct _AsyncBackPressuredStateMachine: Sendable { mutating func finish(_ failure: Failure?) -> FinishAction? { switch self._state { case .initial(let initial): - // Nothing was yielded nor did anybody call next + // Nothing was sent nor did anybody call next // This means we can transition to sourceFinished and store the failure self._state = .sourceFinished( .init( @@ -944,48 +1042,55 @@ struct _AsyncBackPressuredStateMachine: Sendable { return .callOnTermination(initial.onTermination) - case .streaming(let streaming): - guard let consumerContinuation = streaming.consumerContinuation else { + case .channeling(let channeling): + guard let consumerContinuation = channeling.consumerContinuation else { + // We don't have a suspended consumer so we are just going to mark + // the source as finished and terminate the current suspended producers. self._state = .sourceFinished( .init( - iteratorInitialized: streaming.iteratorInitialized, - buffer: streaming.buffer, + iteratorInitialized: channeling.iteratorInitialized, + buffer: channeling.buffer, failure: failure, - onTermination: streaming.onTermination + onTermination: channeling.onTermination ) ) - return .resumeProducers(producerContinuations: Array(streaming.producerContinuations.map { $0.1 })) + return .resumeProducers(producerContinuations: .init(channeling.producerContinuations.lazy.map { $0.1 })) } // 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") + precondition(channeling.buffer.isEmpty, "Expected an empty buffer") - self._state = .finished(iteratorInitialized: streaming.iteratorInitialized) + self._state = .finished(.init(iteratorInitialized: channeling.iteratorInitialized, sourceFinished: true)) return .resumeConsumerAndCallOnTermination( consumerContinuation: consumerContinuation, failure: failure, - onTermination: streaming.onTermination + onTermination: channeling.onTermination ) - case .sourceFinished, .finished: + case .sourceFinished: // If the source has finished, finishing again has no effect. return .none + case .finished(var finished): + finished.sourceFinished = true + self._state = .finished(finished) + return .none + case .modify: - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel 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]) + case returnElementAndResumeProducers(Element, _TinyArray<(Result) -> Void>) /// Indicates that the `Failure` should be returned to the caller and that `onTermination` should be called. case returnFailureAndCallOnTermination(Failure?, (() -> Void)?) /// Indicates that the `nil` should be returned to the caller. @@ -994,14 +1099,15 @@ struct _AsyncBackPressuredStateMachine: Sendable { 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( + // We are not interacting with the backpressure strategy here because + // we are doing this inside `suspendNext` + self._state = .channeling( .init( - backPressureStrategy: initial.backPressureStrategy, + backpressureStrategy: initial.backpressureStrategy, iteratorInitialized: initial.iteratorInitialized, onTermination: initial.onTermination, buffer: Deque(), @@ -1013,35 +1119,35 @@ struct _AsyncBackPressuredStateMachine: Sendable { ) return .suspendTask - case .streaming(var streaming): - guard streaming.consumerContinuation == nil else { + case .channeling(var channeling): + guard channeling.consumerContinuation == nil else { // We have multiple AsyncIterators iterating the sequence - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") } self._state = .modify - guard let element = streaming.buffer.popFirst() else { + guard let element = channeling.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 not interacting with the backpressure strategy here because // we are doing this inside `suspendNext` - self._state = .streaming(streaming) + self._state = .channeling(channeling) return .suspendTask } // We have an element to fulfil the demand right away. - let shouldProduceMore = streaming.backPressureStrategy.didConsume(element: element) - streaming.hasOutstandingDemand = shouldProduceMore + let shouldProduceMore = channeling.backpressureStrategy.didConsume(element: element) + channeling.hasOutstandingDemand = shouldProduceMore guard shouldProduceMore else { // We don't have any new demand, so we can just return the element. - self._state = .streaming(streaming) + self._state = .channeling(channeling) return .returnElement(element) } // 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) + let producers = _TinyArray(channeling.producerContinuations.lazy.map { $0.1 }) + channeling.producerContinuations.removeAll(keepingCapacity: true) + self._state = .channeling(channeling) return .returnElementAndResumeProducers(element, producers) case .sourceFinished(var sourceFinished): @@ -1050,7 +1156,7 @@ struct _AsyncBackPressuredStateMachine: Sendable { guard let element = sourceFinished.buffer.popFirst() else { // We are returning the queued failure now and can transition to finished - self._state = .finished(iteratorInitialized: sourceFinished.iteratorInitialized) + self._state = .finished(.init(iteratorInitialized: sourceFinished.iteratorInitialized, sourceFinished: true)) return .returnFailureAndCallOnTermination(sourceFinished.failure, sourceFinished.onTermination) } @@ -1062,66 +1168,67 @@ struct _AsyncBackPressuredStateMachine: Sendable { return .returnNil case .modify: - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") } } /// Actions returned by `suspendNext()`. + @usableFromInline enum SuspendNextAction { /// Indicates that the consumer should be resumed. - case resumeConsumerWithElement(CheckedContinuation, Element) + case resumeConsumerWithElement(UnsafeContinuation, Element) /// Indicates that the consumer and all producers should be resumed. case resumeConsumerWithElementAndProducers( - CheckedContinuation, + UnsafeContinuation, Element, - [(Result) -> Void] + _TinyArray<(Result) -> Void> ) /// Indicates that the consumer should be resumed with the failure and that `onTermination` should be called. case resumeConsumerWithFailureAndCallOnTermination( - CheckedContinuation, + UnsafeContinuation, Failure?, (() -> Void)? ) /// Indicates that the consumer should be resumed with `nil`. - case resumeConsumerWithNil(CheckedContinuation) + case resumeConsumerWithNil(UnsafeContinuation) } - mutating func suspendNext(continuation: CheckedContinuation) -> SuspendNextAction? { + @inlinable + mutating func suspendNext(continuation: UnsafeContinuation) -> SuspendNextAction? { switch self._state { case .initial: - // We need to transition to streaming before we can suspend - preconditionFailure("AsyncStream internal inconsistency") + preconditionFailure("MultiProducerSingleConsumerChannel internal inconsistency") - case .streaming(var streaming): - guard streaming.consumerContinuation == nil else { + case .channeling(var channeling): + guard channeling.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") + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") } self._state = .modify // We have to check here again since we might have a producer interleave next and suspendNext - guard let element = streaming.buffer.popFirst() else { + guard let element = channeling.buffer.popFirst() else { // There is nothing in the buffer to fulfil the demand so we to store the continuation. - streaming.consumerContinuation = continuation - self._state = .streaming(streaming) + channeling.consumerContinuation = continuation + self._state = .channeling(channeling) return .none } // We have an element to fulfil the demand right away. - let shouldProduceMore = streaming.backPressureStrategy.didConsume(element: element) - streaming.hasOutstandingDemand = shouldProduceMore + let shouldProduceMore = channeling.backpressureStrategy.didConsume(element: element) + channeling.hasOutstandingDemand = shouldProduceMore guard shouldProduceMore else { // We don't have any new demand, so we can just return the element. - self._state = .streaming(streaming) + self._state = .channeling(channeling) return .resumeConsumerWithElement(continuation, element) } // 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) + let producers = _TinyArray(channeling.producerContinuations.lazy.map { $0.1 }) + channeling.producerContinuations.removeAll(keepingCapacity: true) + self._state = .channeling(channeling) return .resumeConsumerWithElementAndProducers(continuation, element, producers) case .sourceFinished(var sourceFinished): @@ -1130,7 +1237,7 @@ struct _AsyncBackPressuredStateMachine: Sendable { guard let element = sourceFinished.buffer.popFirst() else { // We are returning the queued failure now and can transition to finished - self._state = .finished(iteratorInitialized: sourceFinished.iteratorInitialized) + self._state = .finished(.init(iteratorInitialized: sourceFinished.iteratorInitialized, sourceFinished: true)) return .resumeConsumerWithFailureAndCallOnTermination( continuation, @@ -1146,47 +1253,222 @@ struct _AsyncBackPressuredStateMachine: Sendable { return .resumeConsumerWithNil(continuation) case .modify: - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") } } /// Actions returned by `cancelNext()`. + @usableFromInline enum CancelNextAction { /// Indicates that the continuation should be resumed with nil, the producers should be finished and call onTermination. - case resumeConsumerWithNilAndCallOnTermination(CheckedContinuation, (() -> Void)?) + case resumeConsumerWithNilAndCallOnTermination(UnsafeContinuation, (() -> Void)?) /// Indicates that the producers should be finished and call onTermination. - case failProducersAndCallOnTermination([(Result) -> Void], (() -> Void)?) + case failProducersAndCallOnTermination(_TinyArray<(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") + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - case .streaming(let streaming): - self._state = .finished(iteratorInitialized: streaming.iteratorInitialized) + case .channeling(let channeling): + self._state = .finished(.init(iteratorInitialized: channeling.iteratorInitialized, sourceFinished: false)) - guard let consumerContinuation = streaming.consumerContinuation else { + guard let consumerContinuation = channeling.consumerContinuation else { return .failProducersAndCallOnTermination( - Array(streaming.producerContinuations.map { $0.1 }), - streaming.onTermination + .init(channeling.producerContinuations.lazy.map { $0.1 }), + channeling.onTermination ) } precondition( - streaming.producerContinuations.isEmpty, + channeling.producerContinuations.isEmpty, "Internal inconsistency. Unexpected producer continuations." ) return .resumeConsumerWithNilAndCallOnTermination( consumerContinuation, - streaming.onTermination + channeling.onTermination ) case .sourceFinished, .finished: return .none case .modify: - fatalError("AsyncStream internal inconsistency") + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + } + } +} + +@usableFromInline +enum _MultiProducerSingleConsumerState: CustomStringConvertible { + @usableFromInline + struct Initial: CustomStringConvertible { + /// The backpressure strategy. + @usableFromInline + var backpressureStrategy: _MultiProducerSingleConsumerChannelInternalBackpressureStrategy + + /// Indicates if the iterator was initialized. + @usableFromInline + var iteratorInitialized: Bool + + /// The onTermination callback. + @usableFromInline + var onTermination: (@Sendable () -> Void)? + + @usableFromInline + var description: String { + "backpressure:\(self.backpressureStrategy.description) iteratorInitialized:\(self.iteratorInitialized)" + } + } + + @usableFromInline + struct Channeling { + /// The backpressure strategy. + @usableFromInline + var backpressureStrategy: _MultiProducerSingleConsumerChannelInternalBackpressureStrategy + + /// 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: UnsafeContinuation? + + /// The producer continuations. + @usableFromInline + var producerContinuations: Deque<(UInt64, (Result) -> Void)> + + /// The producers that have been cancelled. + @usableFromInline + var cancelledAsyncProducers: Deque + + /// Indicates if we currently have outstanding demand. + @usableFromInline + var hasOutstandingDemand: Bool + + var description: String { + "backpressure:\(self.backpressureStrategy.description) iteratorInitialized:\(self.iteratorInitialized) buffer:\(self.buffer.count) consumerContinuation:\(self.consumerContinuation == nil) producerContinuations:\(self.producerContinuations.count) cancelledProducers:\(self.cancelledAsyncProducers.count) hasOutstandingDemand:\(self.hasOutstandingDemand)" + } + + @usableFromInline + init( + backpressureStrategy: _MultiProducerSingleConsumerChannelInternalBackpressureStrategy, iteratorInitialized: Bool, + onTermination: (@Sendable () -> Void)? = nil, + buffer: Deque, + consumerContinuation: UnsafeContinuation? = nil, + producerContinuations: Deque<(UInt64, (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: Failure? + + /// The onTermination callback. + @usableFromInline + var onTermination: (@Sendable () -> Void)? + + var description: String { + "iteratorInitialized:\(self.iteratorInitialized) buffer:\(self.buffer.count) failure:\(self.failure == nil)" + } + + @usableFromInline + init( + iteratorInitialized: Bool, + buffer: Deque, + failure: Failure? = nil, + onTermination: (@Sendable () -> Void)? = nil + ) { + self.iteratorInitialized = iteratorInitialized + self.buffer = buffer + self.failure = failure + self.onTermination = onTermination + } + } + + @usableFromInline + struct Finished { + /// Indicates if the iterator was initialized. + @usableFromInline + var iteratorInitialized: Bool + + /// Indicates if the source was finished. + @usableFromInline + var sourceFinished: Bool + + var description: String { + "iteratorInitialized:\(self.iteratorInitialized) sourceFinished:\(self.sourceFinished)" + } + + @inlinable + init( + iteratorInitialized: Bool, + sourceFinished: Bool + ) { + self.iteratorInitialized = iteratorInitialized + self.sourceFinished = sourceFinished + } + } + + /// The initial state. + case initial(Initial) + + /// The state once either any element was sent or `next()` was called. + case channeling(Channeling) + + /// 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(Finished) + + /// An intermediate state to avoid CoWs. + case modify + + @usableFromInline + var description: String { + switch self { + case .initial(let initial): + return "initial \(initial.description)" + case .channeling(let channeling): + return "channeling \(channeling.description)" + case .sourceFinished(let sourceFinished): + return "sourceFinished \(sourceFinished.description)" + case .finished(let finished): + return "finished \(finished.description)" + case .modify: + fatalError() } } } diff --git a/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel.swift b/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel.swift new file mode 100644 index 00000000..d2b1d83e --- /dev/null +++ b/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel.swift @@ -0,0 +1,608 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the Swift Async Algorithms open source project +// +// Copyright (c) 2023 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 +// +//===----------------------------------------------------------------------===// + +/// An error that is thrown from the various `send` methods of the +/// ``MultiProducerSingleConsumerChannel/Source``. +/// +/// This error is thrown when the channel is already finished when +/// trying to send new elements to the source. +public struct MultiProducerSingleConsumerChannelAlreadyFinishedError: Error { + @usableFromInline + init() {} +} + +/// A multi producer single consumer channel. +/// +/// The ``MultiProducerSingleConsumerChannel`` provides a ``MultiProducerSingleConsumerChannel/Source`` to +/// send values to the channel. The source exposes the internal backpressure of the asynchronous sequence to the +/// producer. Additionally, the source can be used from synchronous and asynchronous contexts. +/// +/// +/// ## Using a MultiProducerSingleConsumerChannel +/// +/// To use a ``MultiProducerSingleConsumerChannel`` you have to create a new channel with it's source first by calling +/// the ``MultiProducerSingleConsumerChannel/makeChannel(of:throwing:BackpressureStrategy:)`` method. +/// Afterwards, you can pass the source to the producer and the channel to the consumer. +/// +/// ``` +/// let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( +/// backpressureStrategy: .watermark(low: 2, high: 4) +/// ) +/// ``` +/// +/// ### Asynchronous producers +/// +/// Values can be send to the source from asynchronous contexts using ``MultiProducerSingleConsumerChannel/Source/send(_:)-9b5do`` +/// and ``MultiProducerSingleConsumerChannel/Source/send(contentsOf:)-4myrz``. Backpressure results in calls +/// to the `send` methods to be suspended. Once more elements should be produced the `send` methods will be resumed. +/// +/// ``` +/// try await withThrowingTaskGroup(of: Void.self) { group in +/// group.addTask { +/// try await source.send(1) +/// try await source.send(2) +/// try await source.send(3) +/// } +/// +/// for await element in channel { +/// print(element) +/// } +/// } +/// ``` +/// +/// ### Synchronous producers +/// +/// Values can also be send to the source from synchronous context. Backpressure is also exposed on the synchronous contexts; however, +/// it is up to the caller to decide how to properly translate the backpressure to underlying producer e.g. by blocking the thread. +/// +/// ## Finishing the source +/// +/// To properly notify the consumer if the production of values has been finished the source's ``MultiProducerSingleConsumerChannel/Source/finish(throwing:)`` **must** be called. +public struct MultiProducerSingleConsumerChannel: AsyncSequence { + /// A private class to give the ``MultiProducerSingleConsumerChannel`` a deinit so we + /// can tell the producer when any potential consumer went away. + private final class _Backing: Sendable { + /// The underlying storage. + fileprivate let storage: _MultiProducerSingleConsumerChannelBackpressuredStorage + + init(storage: _MultiProducerSingleConsumerChannelBackpressuredStorage) { + self.storage = storage + } + + deinit { + storage.sequenceDeinitialized() + } + } + + /// The backing storage. + private let backing: _Backing + + /// Initializes a new ``MultiProducerSingleConsumerChannel`` and an ``MultiProducerSingleConsumerChannel/Source``. + /// + /// - Parameters: + /// - elementType: The element type of the channel. + /// - failureType: The failure type of the channel. + /// - BackpressureStrategy: The backpressure strategy that the channel should use. + /// - Returns: A tuple containing the channel and its source. The source should be passed to the + /// producer while the channel should be passed to the consumer. + public static func makeChannel( + of elementType: Element.Type = Element.self, + throwing failureType: Failure.Type = Never.self, + backpressureStrategy: Source.BackpressureStrategy + ) -> (`Self`, Source) { + let storage = _MultiProducerSingleConsumerChannelBackpressuredStorage( + backpressureStrategy: backpressureStrategy.internalBackpressureStrategy + ) + let source = Source(storage: storage) + + return (.init(storage: storage), source) + } + + init(storage: _MultiProducerSingleConsumerChannelBackpressuredStorage) { + self.backing = .init(storage: storage) + } +} + +extension MultiProducerSingleConsumerChannel { + /// A struct to send values to the channel. + /// + /// Use this source to provide elements to the channel by calling one of the `send` methods. + /// + /// - Important: You must terminate the source by calling ``finish(throwing:)``. + public struct Source: Sendable { + /// A strategy that handles the backpressure of the channel. + public struct BackpressureStrategy: Sendable { + var internalBackpressureStrategy: _MultiProducerSingleConsumerChannelInternalBackpressureStrategy + + /// A backpressure strategy using a high and low watermark to suspend and resume production respectively. + /// + /// - Parameters: + /// - low: When the number of buffered elements drops below the low watermark, producers will be resumed. + /// - high: When the number of buffered elements rises above the high watermark, producers will be suspended. + public static func watermark(low: Int, high: Int) -> BackpressureStrategy { + .init( + internalBackpressureStrategy: .watermark( + .init(low: low, high: high, waterLevelForElement: nil) + ) + ) + } + + /// A backpressure strategy using a high and low watermark to suspend and resume production respectively. + /// + /// - Parameters: + /// - low: When the number of buffered elements drops below the low watermark, producers will be resumed. + /// - high: When the number of buffered elements rises above the high watermark, producers will be suspended. + /// - waterLevelForElement: A closure used to compute the contribution of each buffered element to the current water level. + /// + /// - Note, `waterLevelForElement` will be called on each element when it is written into the source and when + /// it is consumed from the channel, so it is recommended to provide an function that runs in constant time. + public static func watermark( + low: Int, + high: Int, + waterLevelForElement: @escaping @Sendable (Element) -> Int // TODO: In the future this should become sending + ) -> BackpressureStrategy { + .init( + internalBackpressureStrategy: .watermark( + .init(low: low, high: high, waterLevelForElement: waterLevelForElement) + ) + ) + } + + /// An unbounded backpressure strategy. + /// + /// - Important: Only use this strategy if the production of elements is limited through some other mean. Otherwise + /// an unbounded backpressure strategy can result in infinite memory usage and open your application to denial of service + /// attacks. + public static func unbounded() -> BackpressureStrategy { + .init( + internalBackpressureStrategy: .unbounded(.init()) + ) + } + } + + /// A type that indicates the result of sending elements to the source. + public enum SendResult: Sendable { + /// A token that is returned when the channel's backpressure strategy indicated that production should + /// be suspended. Use this token to enqueue a callback by calling the ``enqueueCallback(_:)`` method. + public struct CallbackToken: Sendable { + @usableFromInline + let _id: UInt64 + + @usableFromInline + init(id: UInt64) { + 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: _MultiProducerSingleConsumerChannelBackpressuredStorage + + init(storage: _MultiProducerSingleConsumerChannelBackpressuredStorage) { + self.storage = storage + } + + deinit { + self.storage.sourceDeinitialized() + } + } + + /// A callback to invoke when the channel finished. + /// + /// The channel 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 + public var onTermination: (@Sendable () -> Void)? { + set { + self._backing.storage.onTermination = newValue + } + get { + self._backing.storage.onTermination + } + } + + @usableFromInline + var _backing: _Backing + + internal init(storage: _MultiProducerSingleConsumerChannelBackpressuredStorage) { + self._backing = .init(storage: storage) + } + + /// Sends new elements to the channel. + /// + /// If there is a task consuming the channel and awaiting the next element then the task will get resumed with the + /// first element of the provided sequence. If the channel already terminated then this method will throw an error + /// indicating the failure. + /// + /// - Parameter sequence: The elements to send to the channel. + /// - Returns: The result that indicates if more elements should be produced at this time. + #if compiler(>=6.0) + @inlinable + public func send(contentsOf sequence: sending S) throws -> SendResult where Element == S.Element, S: Sequence { + try self._backing.storage.send(contentsOf: sequence) + } + #else + @inlinable + public func send(contentsOf sequence: S) throws -> SendResult where Element == S.Element, S: Sequence & Sendable, Element: Sendable { + try self._backing.storage.send(contentsOf: sequence) + } + #endif + + /// Send the element to the channel. + /// + /// If there is a task consuming the channel and awaiting the next element then the task will get resumed with the + /// provided element. If the channel already terminated then this method will throw an error + /// indicating the failure. + /// + /// - Parameter element: The element to send to the channel. + /// - Returns: The result that indicates if more elements should be produced at this time. + #if compiler(>=6.0) + @inlinable + public func send(_ element: sending Element) throws -> SendResult { + try self._backing.storage.send(contentsOf: CollectionOfOne(element)) + } + #else + @inlinable + public func send(_ element: Element) throws -> SendResult where Element: Sendable { + try self._backing.storage.send(contentsOf: CollectionOfOne(element)) + } + #endif + + /// Enqueues a callback that will be invoked once more elements should be produced. + /// + /// Call this method after ``send(contentsOf:)-5honm`` or ``send(_:)-3jxzb`` returned ``SendResult/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. + #if compiler(>=6.0) + @inlinable + public func enqueueCallback( + callbackToken: consuming SendResult.CallbackToken, + onProduceMore: sending @escaping (Result) -> Void + ) { + self._backing.storage.enqueueProducer(callbackToken: callbackToken._id, onProduceMore: onProduceMore) + } + #else + @inlinable + public func enqueueCallback( + callbackToken: consuming SendResult.CallbackToken, + onProduceMore: @escaping @Sendable (Result) -> Void + ) { + self._backing.storage.enqueueProducer(callbackToken: callbackToken._id, onProduceMore: onProduceMore) + } + #endif + + /// 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 + public func cancelCallback(callbackToken: consuming SendResult.CallbackToken) { + self._backing.storage.cancelProducer(callbackToken: callbackToken._id) + } + + /// Send new elements to the channel and provide a callback which will be invoked once more elements should be produced. + /// + /// If there is a task consuming the channel and awaiting the next element then the task will get resumed with the + /// first element of the provided sequence. If the channel already terminated then `onProduceMore` will be invoked with + /// a `Result.failure`. + /// + /// - Parameters: + /// - sequence: The elements to send to the channel. + /// - onProduceMore: The callback which gets invoked once more elements should be produced. This callback might be + /// invoked during the call to ``send(contentsOf:onProduceMore:)``. + #if compiler(>=6.0) + @inlinable + public func send( + contentsOf sequence: sending S, + onProduceMore: @escaping @Sendable (Result) -> Void + ) where Element == S.Element, S: Sequence { + do { + let sendResult = try self.send(contentsOf: sequence) + + switch sendResult { + case .produceMore: + onProduceMore(Result.success(())) + + case .enqueueCallback(let callbackToken): + self.enqueueCallback(callbackToken: callbackToken, onProduceMore: onProduceMore) + } + } catch { + onProduceMore(.failure(error)) + } + } + #else + @inlinable + public func send( + contentsOf sequence: S, + onProduceMore: @escaping @Sendable (Result) -> Void + ) where Element == S.Element, S: Sequence & Sendable, Element: Sendable { + do { + let sendResult = try self.send(contentsOf: sequence) + + switch sendResult { + case .produceMore: + onProduceMore(Result.success(())) + + case .enqueueCallback(let callbackToken): + self.enqueueCallback(callbackToken: callbackToken, onProduceMore: onProduceMore) + } + } catch { + onProduceMore(.failure(error)) + } + } + #endif + + /// Sends the element to the channel. + /// + /// If there is a task consuming the channel and awaiting the next element then the task will get resumed with the + /// provided element. If the channel already terminated then `onProduceMore` will be invoked with + /// a `Result.failure`. + /// + /// - Parameters: + /// - element: The element to send to the channel. + /// - onProduceMore: The callback which gets invoked once more elements should be produced. This callback might be + /// invoked during the call to ``send(_:onProduceMore:)``. + #if compiler(>=6.0) + @inlinable + public func send( + _ element: sending Element, + onProduceMore: @escaping @Sendable (Result) -> Void + ) { + self.send(contentsOf: CollectionOfOne(element), onProduceMore: onProduceMore) + } + #else + @inlinable + public func send( + _ element: Element, + onProduceMore: @escaping @Sendable (Result) -> Void + ) where Element: Sendable { + self.send(contentsOf: CollectionOfOne(element), onProduceMore: onProduceMore) + } + #endif + + /// Send new elements to the channel. + /// + /// If there is a task consuming the channel and awaiting the next element then the task will get resumed with the + /// first element of the provided sequence. If the channel 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 send to the channel. + #if compiler(>=6.0) + @inlinable + public func send(contentsOf sequence: sending S) async throws where Element == S.Element, S: Sequence { + let sendResult = try { try self.send(contentsOf: sequence) }() + + switch sendResult { + case .produceMore: + return + + case .enqueueCallback(let callbackToken): + try await withTaskCancellationHandler { + try await withUnsafeThrowingContinuation { 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) + } + } + } + #else + @inlinable + public func send(contentsOf sequence: S) async throws where Element == S.Element, S: Sequence & Sendable, Element: Sendable { + let sendResult = try { try self.send(contentsOf: sequence) }() + + switch sendResult { + case .produceMore: + return + + case .enqueueCallback(let callbackToken): + try await withTaskCancellationHandler { + try await withUnsafeThrowingContinuation { 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) + } + } + } + #endif + + /// Send new element to the channel. + /// + /// If there is a task consuming the channel and awaiting the next element then the task will get resumed with the + /// provided element. If the channel already terminated then this method will throw an error + /// indicating the failure. + /// + /// This method returns once more elements should be produced. + /// + /// - Parameters: + /// - element: The element to send to the channel. + #if compiler(>=6.0) + @inlinable + public func send(_ element: sending Element) async throws { + try await self.send(contentsOf: CollectionOfOne(element)) + } + #else + @inlinable + public func send(_ element: Element) async throws where Element: Sendable { + try await self.send(contentsOf: CollectionOfOne(element)) + } + #endif + + /// Send the elements of the asynchronous sequence to the channel. + /// + /// This method returns once the provided asynchronous sequence or the channel finished. + /// + /// - Important: This method does not finish the source if consuming the upstream sequence terminated. + /// + /// - Parameters: + /// - sequence: The elements to send to the channel. + #if compiler(>=6.0) + @inlinable + public func send(contentsOf sequence: sending S) async throws where Element == S.Element, S: AsyncSequence { + for try await element in sequence { + try await self.send(contentsOf: CollectionOfOne(element)) + } + } + #else + @inlinable + public func send(contentsOf sequence: S) async throws where Element == S.Element, S: AsyncSequence & Sendable, Element: Sendable{ + for try await element in sequence { + try await self.send(contentsOf: CollectionOfOne(element)) + } + } + #endif + + /// 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 channel enters a terminal state and doesn't accept + /// new elements. + /// + /// - Parameters: + /// - error: The error to throw, or `nil`, to finish normally. + @inlinable + public func finish(throwing error: Failure? = nil) { + self._backing.storage.finish(error) + } + } +} + +extension MultiProducerSingleConsumerChannel { + /// The asynchronous iterator for iterating the channel. + /// + /// 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()`. + public struct Iterator: AsyncIteratorProtocol { + @usableFromInline + final class _Backing { + @usableFromInline + let storage: _MultiProducerSingleConsumerChannelBackpressuredStorage + + init(storage: _MultiProducerSingleConsumerChannelBackpressuredStorage) { + self.storage = storage + self.storage.iteratorInitialized() + } + + deinit { + self.storage.iteratorDeinitialized() + } + } + + @usableFromInline + let _backing: _Backing + + init(storage: _MultiProducerSingleConsumerChannelBackpressuredStorage) { + self._backing = .init(storage: storage) + } + + /// The next value from the channel. + /// + /// When `next()` returns `nil`, this signifies the end of the channel. + /// + /// 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 channel terminates. In this case, + /// `next()` may return `nil` immediately, or else return `nil` on + /// subsequent calls. + @_disfavoredOverload + @inlinable + public mutating func next() async throws -> Element? { + try await self._backing.storage.next() + } + + #if compiler(>=6.0) + /// The next value from the channel. + /// + /// When `next()` returns `nil`, this signifies the end of the channel. + /// + /// 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 channel terminates. In this case, + /// `next()` may return `nil` immediately, or else return `nil` on + /// subsequent calls. + @available(macOS 15.0, iOS 18.0, watchOS 11.0, tvOS 18.0, visionOS 2.0, *) + @inlinable + public mutating func next( + isolation actor: isolated (any Actor)? = #isolation + ) async throws(Failure) -> Element? { + do { + return try await self._backing.storage.next(isolation: actor) + } catch { + throw error as! Failure + } + } + #endif + } + + /// Creates the asynchronous iterator that produces elements of this + /// asynchronous sequence. + public func makeAsyncIterator() -> Iterator { + Iterator(storage: self.backing.storage) + } +} + +extension MultiProducerSingleConsumerChannel: Sendable where Element: Sendable {} + +@available(*, unavailable) +extension MultiProducerSingleConsumerChannel.Iterator: Sendable {} diff --git a/Tests/AsyncAlgorithmsTests/BackpressuredStream/BackPressuredStreamTests.swift b/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift similarity index 53% rename from Tests/AsyncAlgorithmsTests/BackpressuredStream/BackPressuredStreamTests.swift rename to Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift index ac74282b..93e5c3e8 100644 --- a/Tests/AsyncAlgorithmsTests/BackpressuredStream/BackPressuredStreamTests.swift +++ b/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift @@ -12,15 +12,19 @@ import AsyncAlgorithms import XCTest -final class BackPressuredStreamTests: XCTestCase { +@available(macOS 15.0, iOS 18.0, watchOS 11.0, tvOS 18.0, visionOS 2.0, *) +final class MultiProducerSingleConsumerChannelTests: XCTestCase { // MARK: - sequenceDeinitialized func testSequenceDeinitialized_whenNoIterator() async throws { - var (stream, source): (AsyncBackPressuredStream?, AsyncBackPressuredStream.Source) = - AsyncBackPressuredStream.makeStream( + var (channel, source): (MultiProducerSingleConsumerChannel?, MultiProducerSingleConsumerChannel.Source) = + MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 5, high: 10) + backpressureStrategy: .watermark(low: 5, high: 10) ) + defer { + source.finish() + } let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() source.onTermination = { @@ -38,17 +42,17 @@ final class BackPressuredStreamTests: XCTestCase { var onTerminationIterator = onTerminationStream.makeAsyncIterator() _ = await onTerminationIterator.next() - withExtendedLifetime(stream) {} - stream = nil + withExtendedLifetime(channel) {} + channel = nil let terminationResult: Void? = await onTerminationIterator.next() XCTAssertNil(terminationResult) do { - _ = try { try source.write(2) }() + _ = try { try source.send(2) }() XCTFail("Expected an error to be thrown") } catch { - XCTAssertTrue(error is AsyncBackPressuredStreamAlreadyFinishedError) + XCTAssertTrue(error is MultiProducerSingleConsumerChannelAlreadyFinishedError) } group.cancelAll() @@ -56,13 +60,16 @@ final class BackPressuredStreamTests: XCTestCase { } func testSequenceDeinitialized_whenIterator() async throws { - var (stream, source): (AsyncBackPressuredStream?, AsyncBackPressuredStream.Source) = - AsyncBackPressuredStream.makeStream( + var (channel, source): (MultiProducerSingleConsumerChannel?, MultiProducerSingleConsumerChannel.Source) = + MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 5, high: 10) + backpressureStrategy: .watermark(low: 5, high: 10) ) + defer { + source.finish() + } - var iterator = stream?.makeAsyncIterator() + var iterator = channel?.makeAsyncIterator() let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() source.onTermination = { @@ -80,23 +87,23 @@ final class BackPressuredStreamTests: XCTestCase { var onTerminationIterator = onTerminationStream.makeAsyncIterator() _ = await onTerminationIterator.next() - try withExtendedLifetime(stream) { - let writeResult = try source.write(1) + try withExtendedLifetime(channel) { + let writeResult = try source.send(1) writeResult.assertIsProducerMore() } - stream = nil + channel = nil do { - let writeResult = try { try source.write(2) }() + let writeResult = try { try source.send(2) }() writeResult.assertIsProducerMore() } catch { XCTFail("Expected no error to be thrown") } - let element1 = try await iterator?.next() + let element1 = await iterator?.next() XCTAssertEqual(element1, 1) - let element2 = try await iterator?.next() + let element2 = await iterator?.next() XCTAssertEqual(element2, 2) group.cancelAll() @@ -104,10 +111,10 @@ final class BackPressuredStreamTests: XCTestCase { } func testSequenceDeinitialized_whenFinished() async throws { - var (stream, source): (AsyncBackPressuredStream?, AsyncBackPressuredStream.Source) = - AsyncBackPressuredStream.makeStream( + var (channel, source): (MultiProducerSingleConsumerChannel?, MultiProducerSingleConsumerChannel.Source) = + MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 5, high: 10) + backpressureStrategy: .watermark(low: 5, high: 10) ) let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() @@ -126,116 +133,128 @@ final class BackPressuredStreamTests: XCTestCase { var onTerminationIterator = onTerminationStream.makeAsyncIterator() _ = await onTerminationIterator.next() - withExtendedLifetime(stream) { + withExtendedLifetime(channel) { source.finish(throwing: nil) } - stream = nil + channel = nil let terminationResult: Void? = await onTerminationIterator.next() XCTAssertNil(terminationResult) do { - _ = try { try source.write(1) }() + _ = try { try source.send(1) }() XCTFail("Expected an error to be thrown") } catch { - XCTAssertTrue(error is AsyncBackPressuredStreamAlreadyFinishedError) + XCTAssertTrue(error is MultiProducerSingleConsumerChannelAlreadyFinishedError) } group.cancelAll() } } - func testSequenceDeinitialized_whenStreaming_andSuspendedProducer() async throws { - var (stream, source): (AsyncBackPressuredStream?, AsyncBackPressuredStream.Source) = - AsyncBackPressuredStream.makeStream( + func testSequenceDeinitialized_whenChanneling_andSuspendedProducer() async throws { + var (channel, source): (MultiProducerSingleConsumerChannel?, MultiProducerSingleConsumerChannel.Source) = + MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 1, high: 2) + backpressureStrategy: .watermark(low: 1, high: 2) ) + defer { + source.finish() + } - _ = try { try source.write(1) }() + _ = try { try source.send(1) }() do { try await withCheckedThrowingContinuation { continuation in - source.write(1) { result in + source.send(1) { result in continuation.resume(with: result) } - stream = nil - _ = stream?.makeAsyncIterator() + channel = nil + _ = channel?.makeAsyncIterator() } } catch { - XCTAssertTrue(error is AsyncBackPressuredStreamAlreadyFinishedError) + XCTAssertTrue(error is MultiProducerSingleConsumerChannelAlreadyFinishedError) } } // MARK: - iteratorInitialized func testIteratorInitialized_whenInitial() async throws { - let (stream, _) = AsyncBackPressuredStream.makeStream( + let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 5, high: 10) + backpressureStrategy: .watermark(low: 5, high: 10) ) + defer { + source.finish() + } - _ = stream.makeAsyncIterator() + _ = channel.makeAsyncIterator() } - func testIteratorInitialized_whenStreaming() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + func testIteratorInitialized_whenChanneling() async throws { + let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 5, high: 10) + backpressureStrategy: .watermark(low: 5, high: 10) ) + defer { + source.finish() + } - try await source.write(1) + try await source.send(1) - var iterator = stream.makeAsyncIterator() - let element = try await iterator.next() + var iterator = channel.makeAsyncIterator() + let element = await iterator.next() XCTAssertEqual(element, 1) } func testIteratorInitialized_whenSourceFinished() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 5, high: 10) + backpressureStrategy: .watermark(low: 5, high: 10) ) - try await source.write(1) + try await source.send(1) source.finish(throwing: nil) - var iterator = stream.makeAsyncIterator() - let element1 = try await iterator.next() + var iterator = channel.makeAsyncIterator() + let element1 = await iterator.next() XCTAssertEqual(element1, 1) - let element2 = try await iterator.next() + let element2 = await iterator.next() XCTAssertNil(element2) } func testIteratorInitialized_whenFinished() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 5, high: 10) + backpressureStrategy: .watermark(low: 5, high: 10) ) source.finish(throwing: nil) - var iterator = stream.makeAsyncIterator() - let element = try await iterator.next() + var iterator = channel.makeAsyncIterator() + let element = await iterator.next() XCTAssertNil(element) } // MARK: - iteratorDeinitialized func testIteratorDeinitialized_whenInitial() async throws { - var (stream, source) = AsyncBackPressuredStream.makeStream( + var (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 5, high: 10) + backpressureStrategy: .watermark(low: 5, high: 10) ) + defer { + source.finish() + } let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() source.onTermination = { onTerminationContinuation.finish() } - try await withThrowingTaskGroup(of: Void.self) { group in + await withThrowingTaskGroup(of: Void.self) { group in group.addTask { while !Task.isCancelled { onTerminationContinuation.yield() @@ -246,9 +265,9 @@ final class BackPressuredStreamTests: XCTestCase { var onTerminationIterator = onTerminationStream.makeAsyncIterator() _ = await onTerminationIterator.next() - var iterator: AsyncBackPressuredStream.AsyncIterator? = stream.makeAsyncIterator() + var iterator: MultiProducerSingleConsumerChannel.AsyncIterator? = channel.makeAsyncIterator() iterator = nil - _ = try await iterator?.next() + _ = await iterator?.next() let terminationResult: Void? = await onTerminationIterator.next() XCTAssertNil(terminationResult) @@ -257,20 +276,23 @@ final class BackPressuredStreamTests: XCTestCase { } } - func testIteratorDeinitialized_whenStreaming() async throws { - var (stream, source) = AsyncBackPressuredStream.makeStream( + func testIteratorDeinitialized_whenChanneling() async throws { + var (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 5, high: 10) + backpressureStrategy: .watermark(low: 5, high: 10) ) + defer { + source.finish() + } let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() source.onTermination = { onTerminationContinuation.finish() } - try await source.write(1) + try await source.send(1) - try await withThrowingTaskGroup(of: Void.self) { group in + await withThrowingTaskGroup(of: Void.self) { group in group.addTask { while !Task.isCancelled { onTerminationContinuation.yield() @@ -281,9 +303,9 @@ final class BackPressuredStreamTests: XCTestCase { var onTerminationIterator = onTerminationStream.makeAsyncIterator() _ = await onTerminationIterator.next() - var iterator: AsyncBackPressuredStream.AsyncIterator? = stream.makeAsyncIterator() + var iterator: MultiProducerSingleConsumerChannel.AsyncIterator? = channel.makeAsyncIterator() iterator = nil - _ = try await iterator?.next() + _ = await iterator?.next(isolation: nil) let terminationResult: Void? = await onTerminationIterator.next() XCTAssertNil(terminationResult) @@ -293,9 +315,9 @@ final class BackPressuredStreamTests: XCTestCase { } func testIteratorDeinitialized_whenSourceFinished() async throws { - var (stream, source) = AsyncBackPressuredStream.makeStream( + var (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 5, high: 10) + backpressureStrategy: .watermark(low: 5, high: 10) ) let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() @@ -303,10 +325,10 @@ final class BackPressuredStreamTests: XCTestCase { onTerminationContinuation.finish() } - try await source.write(1) + try await source.send(1) source.finish(throwing: nil) - try await withThrowingTaskGroup(of: Void.self) { group in + await withThrowingTaskGroup(of: Void.self) { group in group.addTask { while !Task.isCancelled { onTerminationContinuation.yield() @@ -317,9 +339,9 @@ final class BackPressuredStreamTests: XCTestCase { var onTerminationIterator = onTerminationStream.makeAsyncIterator() _ = await onTerminationIterator.next() - var iterator: AsyncBackPressuredStream.AsyncIterator? = stream.makeAsyncIterator() + var iterator: MultiProducerSingleConsumerChannel.AsyncIterator? = channel.makeAsyncIterator() iterator = nil - _ = try await iterator?.next() + _ = await iterator?.next() let terminationResult: Void? = await onTerminationIterator.next() XCTAssertNil(terminationResult) @@ -329,9 +351,10 @@ final class BackPressuredStreamTests: XCTestCase { } func testIteratorDeinitialized_whenFinished() async throws { - var (stream, source) = AsyncBackPressuredStream.makeStream( + var (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 5, high: 10) + throwing: Error.self, + backpressureStrategy: .watermark(low: 5, high: 10) ) let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() @@ -352,7 +375,7 @@ final class BackPressuredStreamTests: XCTestCase { var onTerminationIterator = onTerminationStream.makeAsyncIterator() _ = await onTerminationIterator.next() - var iterator: AsyncBackPressuredStream.AsyncIterator? = stream.makeAsyncIterator() + var iterator: MultiProducerSingleConsumerChannel.AsyncIterator? = channel.makeAsyncIterator() iterator = nil _ = try await iterator?.next() @@ -363,28 +386,32 @@ final class BackPressuredStreamTests: XCTestCase { } } - func testIteratorDeinitialized_whenStreaming_andSuspendedProducer() async throws { - var (stream, source): (AsyncBackPressuredStream?, AsyncBackPressuredStream.Source) = - AsyncBackPressuredStream.makeStream( + func testIteratorDeinitialized_whenChanneling_andSuspendedProducer() async throws { + var (channel, source): (MultiProducerSingleConsumerChannel?, MultiProducerSingleConsumerChannel.Source) = + MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 1, high: 2) + throwing: Error.self, + backpressureStrategy: .watermark(low: 1, high: 2) ) + defer { + source.finish() + } - var iterator: AsyncBackPressuredStream.AsyncIterator? = stream?.makeAsyncIterator() - stream = nil + var iterator: MultiProducerSingleConsumerChannel.AsyncIterator? = channel?.makeAsyncIterator() + channel = nil - _ = try { try source.write(1) }() + _ = try { try source.send(1) }() do { try await withCheckedThrowingContinuation { continuation in - source.write(1) { result in + source.send(1) { result in continuation.resume(with: result) } iterator = nil } } catch { - XCTAssertTrue(error is AsyncBackPressuredStreamAlreadyFinishedError) + XCTAssertTrue(error is MultiProducerSingleConsumerChannelAlreadyFinishedError) } _ = try await iterator?.next() @@ -392,125 +419,12 @@ final class BackPressuredStreamTests: XCTestCase { // MARK: - sourceDeinitialized - func testSourceDeinitialized_whenInitial() async throws { - var (stream, source): (AsyncBackPressuredStream, AsyncBackPressuredStream.Source?) = - AsyncBackPressuredStream.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(for: .seconds(0.2)) - } - } - - 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): (AsyncBackPressuredStream, AsyncBackPressuredStream.Source?) = - AsyncBackPressuredStream.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(for: .seconds(0.2)) - } - } - - var onTerminationIterator = onTerminationStream.makeAsyncIterator() - _ = await onTerminationIterator.next() - - var iterator: AsyncBackPressuredStream.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): (AsyncBackPressuredStream, AsyncBackPressuredStream.Source?) = - AsyncBackPressuredStream.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(for: .seconds(0.2)) - } - } - - var onTerminationIterator = onTerminationStream.makeAsyncIterator() - _ = await onTerminationIterator.next() - - var iterator: AsyncBackPressuredStream.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): (AsyncBackPressuredStream, AsyncBackPressuredStream.Source?) = - AsyncBackPressuredStream.makeStream( + var (channel, source): (MultiProducerSingleConsumerChannel, MultiProducerSingleConsumerChannel.Source?) = + MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 5, high: 10) + throwing: Error.self, + backpressureStrategy: .watermark(low: 5, high: 10) ) let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() @@ -518,8 +432,8 @@ final class BackPressuredStreamTests: XCTestCase { onTerminationContinuation.finish() } - try await source?.write(1) - try await source?.write(2) + try await source?.send(1) + try await source?.send(2) source?.finish(throwing: nil) try await withThrowingTaskGroup(of: Void.self) { group in @@ -533,7 +447,7 @@ final class BackPressuredStreamTests: XCTestCase { var onTerminationIterator = onTerminationStream.makeAsyncIterator() _ = await onTerminationIterator.next() - var iterator: AsyncBackPressuredStream.AsyncIterator? = stream.makeAsyncIterator() + var iterator: MultiProducerSingleConsumerChannel.AsyncIterator? = channel.makeAsyncIterator() _ = try await iterator?.next() source = nil @@ -551,10 +465,10 @@ final class BackPressuredStreamTests: XCTestCase { } func testSourceDeinitialized_whenFinished() async throws { - var (stream, source): (AsyncBackPressuredStream, AsyncBackPressuredStream.Source?) = - AsyncBackPressuredStream.makeStream( + var (channel, source): (MultiProducerSingleConsumerChannel, MultiProducerSingleConsumerChannel.Source?) = + MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 5, high: 10) + backpressureStrategy: .watermark(low: 5, high: 10) ) let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() @@ -575,7 +489,7 @@ final class BackPressuredStreamTests: XCTestCase { var onTerminationIterator = onTerminationStream.makeAsyncIterator() _ = await onTerminationIterator.next() - _ = stream.makeAsyncIterator() + _ = channel.makeAsyncIterator() source = nil @@ -588,97 +502,85 @@ final class BackPressuredStreamTests: XCTestCase { } } - func testSourceDeinitialized_whenStreaming_andSuspendedProducer() async throws { - var (stream, source): (AsyncBackPressuredStream, AsyncBackPressuredStream.Source?) = - AsyncBackPressuredStream.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 AsyncBackPressuredStreamAlreadyFinishedError) - } - } - // MARK: - write func testWrite_whenInitial() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 2, high: 5) + backpressureStrategy: .watermark(low: 2, high: 5) ) + defer { + source.finish() + } - try await source.write(1) + try await source.send(1) - var iterator = stream.makeAsyncIterator() - let element = try await iterator.next() + var iterator = channel.makeAsyncIterator() + let element = await iterator.next() XCTAssertEqual(element, 1) } - func testWrite_whenStreaming_andNoConsumer() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + func testWrite_whenChanneling_andNoConsumer() async throws { + let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 2, high: 5) + backpressureStrategy: .watermark(low: 2, high: 5) ) + defer { + source.finish() + } - try await source.write(1) - try await source.write(2) + try await source.send(1) + try await source.send(2) var iterator = stream.makeAsyncIterator() - let element1 = try await iterator.next() + let element1 = await iterator.next() XCTAssertEqual(element1, 1) - let element2 = try await iterator.next() + let element2 = await iterator.next() XCTAssertEqual(element2, 2) } - func testWrite_whenStreaming_andSuspendedConsumer() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + func testWrite_whenChanneling_andSuspendedConsumer() async throws { + let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 2, high: 5) + backpressureStrategy: .watermark(low: 2, high: 5) ) + defer { + source.finish() + } try await withThrowingTaskGroup(of: Int?.self) { group in group.addTask { - return try await stream.first { _ in true } + return 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(for: .seconds(0.5)) - try await source.write(1) + try await source.send(1) let element = try await group.next() XCTAssertEqual(element, 1) } } - func testWrite_whenStreaming_andSuspendedConsumer_andEmptySequence() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + func testWrite_whenChanneling_andSuspendedConsumer_andEmptySequence() async throws { + let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 2, high: 5) + backpressureStrategy: .watermark(low: 2, high: 5) ) + defer { + source.finish() + } try await withThrowingTaskGroup(of: Int?.self) { group in group.addTask { - return try await stream.first { _ in true } + return 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(for: .seconds(0.5)) - try await source.write(contentsOf: []) - try await source.write(contentsOf: [1]) + try await source.send(contentsOf: []) + try await source.send(contentsOf: [1]) let element = try await group.next() XCTAssertEqual(element, 1) } @@ -686,17 +588,20 @@ final class BackPressuredStreamTests: XCTestCase { // MARK: - enqueueProducer - func testEnqueueProducer_whenStreaming_andAndCancelled() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + func testEnqueueProducer_whenChanneling_andAndCancelled() async throws { + let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 1, high: 2) + backpressureStrategy: .watermark(low: 1, high: 2) ) + defer { + source.finish() + } let (producerStream, producerSource) = AsyncThrowingStream.makeStream() - try await source.write(1) + try await source.send(1) - let writeResult = try { try source.write(2) }() + let writeResult = try { try source.send(2) }() switch writeResult { case .produceMore: @@ -716,21 +621,24 @@ final class BackPressuredStreamTests: XCTestCase { XCTAssertTrue(error is CancellationError) } - let element = try await stream.first { _ in true } + let element = await stream.first { _ in true } XCTAssertEqual(element, 1) } - func testEnqueueProducer_whenStreaming_andAndCancelled_andAsync() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + func testEnqueueProducer_whenChanneling_andAndCancelled_andAsync() async throws { + let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 1, high: 2) + backpressureStrategy: .watermark(low: 1, high: 2) ) + defer { + source.finish() + } - try await source.write(1) + try await source.send(1) await withThrowingTaskGroup(of: Void.self) { group in group.addTask { - try await source.write(2) + try await source.send(2) } group.cancelAll() @@ -742,26 +650,29 @@ final class BackPressuredStreamTests: XCTestCase { } } - let element = try await stream.first { _ in true } + let element = await stream.first { _ in true } XCTAssertEqual(element, 1) } - func testEnqueueProducer_whenStreaming_andInterleaving() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + func testEnqueueProducer_whenChanneling_andInterleaving() async throws { + let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 1, high: 1) + backpressureStrategy: .watermark(low: 1, high: 1) ) + defer { + source.finish() + } var iterator = stream.makeAsyncIterator() let (producerStream, producerSource) = AsyncThrowingStream.makeStream() - let writeResult = try { try source.write(1) }() + let writeResult = try { try source.send(1) }() switch writeResult { case .produceMore: preconditionFailure() case .enqueueCallback(let callbackToken): - let element = try await iterator.next() + let element = await iterator.next() XCTAssertEqual(element, 1) source.enqueueCallback(callbackToken: callbackToken) { result in @@ -776,16 +687,19 @@ final class BackPressuredStreamTests: XCTestCase { } } - func testEnqueueProducer_whenStreaming_andSuspending() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + func testEnqueueProducer_whenChanneling_andSuspending() async throws { + let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 1, high: 1) + backpressureStrategy: .watermark(low: 1, high: 1) ) + defer { + source.finish() + } var iterator = stream.makeAsyncIterator() let (producerStream, producerSource) = AsyncThrowingStream.makeStream() - let writeResult = try { try source.write(1) }() + let writeResult = try { try source.send(1) }() switch writeResult { case .produceMore: @@ -796,7 +710,7 @@ final class BackPressuredStreamTests: XCTestCase { } } - let element = try await iterator.next() + let element = await iterator.next() XCTAssertEqual(element, 1) do { @@ -807,15 +721,18 @@ final class BackPressuredStreamTests: XCTestCase { } func testEnqueueProducer_whenFinished() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 1, high: 1) + backpressureStrategy: .watermark(low: 1, high: 1) ) + defer { + source.finish() + } var iterator = stream.makeAsyncIterator() let (producerStream, producerSource) = AsyncThrowingStream.makeStream() - let writeResult = try { try source.write(1) }() + let writeResult = try { try source.send(1) }() switch writeResult { case .produceMore: @@ -828,30 +745,33 @@ final class BackPressuredStreamTests: XCTestCase { } } - let element = try await iterator.next() + let element = await iterator.next() XCTAssertEqual(element, 1) do { _ = try await producerStream.first { _ in true } XCTFail("Expected an error to be thrown") } catch { - XCTAssertTrue(error is AsyncBackPressuredStreamAlreadyFinishedError) + XCTAssertTrue(error is MultiProducerSingleConsumerChannelAlreadyFinishedError) } } // MARK: - cancelProducer - func testCancelProducer_whenStreaming() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + func testCancelProducer_whenChanneling() async throws { + let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 1, high: 2) + backpressureStrategy: .watermark(low: 1, high: 2) ) + defer { + source.finish() + } let (producerStream, producerSource) = AsyncThrowingStream.makeStream() - try await source.write(1) + try await source.send(1) - let writeResult = try { try source.write(2) }() + let writeResult = try { try source.send(2) }() switch writeResult { case .produceMore: @@ -871,21 +791,21 @@ final class BackPressuredStreamTests: XCTestCase { XCTAssertTrue(error is CancellationError) } - let element = try await stream.first { _ in true } + let element = await stream.first { _ in true } XCTAssertEqual(element, 1) } func testCancelProducer_whenSourceFinished() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 1, high: 2) + backpressureStrategy: .watermark(low: 1, high: 2) ) let (producerStream, producerSource) = AsyncThrowingStream.makeStream() - try await source.write(1) + try await source.send(1) - let writeResult = try { try source.write(2) }() + let writeResult = try { try source.send(2) }() switch writeResult { case .produceMore: @@ -904,24 +824,24 @@ final class BackPressuredStreamTests: XCTestCase { _ = try await producerStream.first { _ in true } XCTFail("Expected an error to be thrown") } catch { - XCTAssertTrue(error is AsyncBackPressuredStreamAlreadyFinishedError) + XCTAssertTrue(error is MultiProducerSingleConsumerChannelAlreadyFinishedError) } - let element = try await stream.first { _ in true } + let element = await stream.first { _ in true } XCTAssertEqual(element, 1) } // MARK: - finish - func testFinish_whenStreaming_andConsumerSuspended() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + func testFinish_whenChanneling_andConsumerSuspended() async throws { + let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 1, high: 1) + backpressureStrategy: .watermark(low: 1, high: 1) ) try await withThrowingTaskGroup(of: Int?.self) { group in group.addTask { - return try await stream.first { $0 == 2 } + return await stream.first { $0 == 2 } } // This is always going to be a bit racy since we need the call to next() suspend @@ -934,9 +854,10 @@ final class BackPressuredStreamTests: XCTestCase { } func testFinish_whenInitial() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 1, high: 1) + throwing: Error.self, + backpressureStrategy: .watermark(low: 1, high: 1) ) source.finish(throwing: CancellationError()) @@ -952,55 +873,61 @@ final class BackPressuredStreamTests: XCTestCase { // MARK: - Backpressure - func testBackPressure() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + func testBackpressure() async throws { + let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 2, high: 4) + backpressureStrategy: .watermark(low: 2, high: 4) ) + defer { + source.finish() + } - let (backPressureEventStream, backPressureEventContinuation) = AsyncStream.makeStream(of: Void.self) + let (backpressureEventStream, backpressureEventContinuation) = AsyncStream.makeStream(of: Void.self) - try await withThrowingTaskGroup(of: Void.self) { group in + await withThrowingTaskGroup(of: Void.self) { group in group.addTask { while true { - backPressureEventContinuation.yield(()) - try await source.write(contentsOf: [1]) + backpressureEventContinuation.yield(()) + try await source.send(contentsOf: [1]) } } - var backPressureEventIterator = backPressureEventStream.makeAsyncIterator() - var iterator = stream.makeAsyncIterator() + var backpressureEventIterator = backpressureEventStream.makeAsyncIterator() + var iterator = channel.makeAsyncIterator() - await backPressureEventIterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() + 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 iterator.next() + _ = await iterator.next() + _ = await iterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() group.cancelAll() } } - func testBackPressureSync() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + func testBackpressureSync() async throws { + let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 2, high: 4) + backpressureStrategy: .watermark(low: 2, high: 4) ) + defer { + source.finish() + } - let (backPressureEventStream, backPressureEventContinuation) = AsyncStream.makeStream(of: Void.self) + let (backpressureEventStream, backpressureEventContinuation) = AsyncStream.makeStream(of: Void.self) - try await withThrowingTaskGroup(of: Void.self) { group in + await withThrowingTaskGroup(of: Void.self) { group in group.addTask { @Sendable func yield() { - backPressureEventContinuation.yield(()) - source.write(contentsOf: [1]) { result in + backpressureEventContinuation.yield(()) + source.send(contentsOf: [1]) { result in switch result { case .success: yield() @@ -1014,54 +941,82 @@ final class BackPressuredStreamTests: XCTestCase { yield() } - var backPressureEventIterator = backPressureEventStream.makeAsyncIterator() - var iterator = stream.makeAsyncIterator() + var backpressureEventIterator = backpressureEventStream.makeAsyncIterator() + var iterator = channel.makeAsyncIterator() - await backPressureEventIterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() + 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 iterator.next() + _ = await iterator.next() + _ = await iterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() group.cancelAll() } } func testWatermarkWithCustomCoount() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( of: [Int].self, - backPressureStrategy: .watermark(low: 2, high: 4, waterLevelForElement: { $0.count }) + backpressureStrategy: .watermark(low: 2, high: 4, waterLevelForElement: { $0.count }) ) - var iterator = stream.makeAsyncIterator() + defer { + source.finish() + } + var iterator = channel.makeAsyncIterator() - try await source.write([1, 1, 1]) + try await source.send([1, 1, 1]) - _ = try await iterator.next() + _ = await iterator.next() - try await source.write([1, 1, 1]) + try await source.send([1, 1, 1]) - _ = try await iterator.next() + _ = await iterator.next() + } + + func testWatermarWithLotsOfElements() async throws { + // This test should in the future use a custom task executor to schedule to avoid sending + // 1000 elements. + let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + backpressureStrategy: .watermark(low: 2, high: 4) + ) + await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + for i in 0...10000 { + try await source.send(i) + } + source.finish() + } + + group.addTask { + var sum = 0 + for try await element in channel { + sum += element + } + } + } } func testThrowsError() async throws { - let (stream, source) = AsyncBackPressuredStream.makeStream( + let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 2, high: 4) + throwing: Error.self, + backpressureStrategy: .watermark(low: 2, high: 4) ) - try await source.write(1) - try await source.write(2) + try await source.send(1) + try await source.send(2) source.finish(throwing: CancellationError()) var elements = [Int]() - var iterator = stream.makeAsyncIterator() + var iterator = channel.makeAsyncIterator() do { while let element = try await iterator.next() { @@ -1079,55 +1034,58 @@ final class BackPressuredStreamTests: XCTestCase { func testAsyncSequenceWrite() async throws { let (stream, continuation) = AsyncStream.makeStream() - let (backpressuredStream, source) = AsyncBackPressuredStream.makeStream( + let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 2, high: 4) + backpressureStrategy: .watermark(low: 2, high: 4) ) continuation.yield(1) continuation.yield(2) continuation.finish() - try await source.write(contentsOf: stream) + try await source.send(contentsOf: stream) source.finish(throwing: nil) - let elements = try await backpressuredStream.collect() + let elements = await channel.collect() XCTAssertEqual(elements, [1, 2]) } // MARK: NonThrowing func testNonThrowing() async throws { - let (stream, source) = AsyncNonThrowingBackPressuredStream.makeStream( + let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, - backPressureStrategy: .watermark(low: 2, high: 4) + backpressureStrategy: .watermark(low: 2, high: 4) ) + defer { + source.finish() + } - let (backPressureEventStream, backPressureEventContinuation) = AsyncStream.makeStream(of: Void.self) + let (backpressureEventStream, backpressureEventContinuation) = AsyncStream.makeStream(of: Void.self) await withThrowingTaskGroup(of: Void.self) { group in group.addTask { while true { - backPressureEventContinuation.yield(()) - try await source.write(contentsOf: [1]) + backpressureEventContinuation.yield(()) + try await source.send(contentsOf: [1]) } } - var backPressureEventIterator = backPressureEventStream.makeAsyncIterator() - var iterator = stream.makeAsyncIterator() + var backpressureEventIterator = backpressureEventStream.makeAsyncIterator() + var iterator = channel.makeAsyncIterator() - await backPressureEventIterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() _ = await iterator.next() _ = await iterator.next() _ = await iterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() - await backPressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() group.cancelAll() } @@ -1143,7 +1101,7 @@ extension AsyncSequence { } } -extension AsyncBackPressuredStream.Source.WriteResult { +extension MultiProducerSingleConsumerChannel.Source.SendResult { func assertIsProducerMore() { switch self { case .produceMore: From 293dc60f8c71b217e300e05108a6da3337e8c5e1 Mon Sep 17 00:00:00 2001 From: Franz Busch Date: Thu, 13 Jun 2024 10:38:45 +0200 Subject: [PATCH 3/7] ~Copyable source --- ...oducerSingleConsumerChannel+Internal.swift | 269 +++++++++--------- .../MultiProducerSingleConsumerChannel.swift | 247 +++++----------- 2 files changed, 207 insertions(+), 309 deletions(-) diff --git a/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel+Internal.swift b/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel+Internal.swift index 0e4997df..ba717e60 100644 --- a/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel+Internal.swift +++ b/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel+Internal.swift @@ -9,6 +9,7 @@ // //===----------------------------------------------------------------------===// +#if compiler(>=6.0) import DequeModule @usableFromInline @@ -168,7 +169,12 @@ final class _MultiProducerSingleConsumerChannelBackpressuredStorage) -> Void + continuation: UnsafeContinuation + ) { + let action = self._stateMachine.withCriticalRegion { + $0.enqueueContinuation(callbackToken: callbackToken, continuation: continuation) + } + + switch action { + case .resumeProducer(let continuation): + continuation.resume() + + case .resumeProducerWithError(let continuation, let error): + continuation.resume(throwing: error) + + case .none: + break + } + } + + @inlinable + func enqueueProducer( + callbackToken: UInt64, + onProduceMore: sending @escaping (Result) -> Void ) { let action = self._stateMachine.withCriticalRegion { $0.enqueueProducer(callbackToken: callbackToken, onProduceMore: onProduceMore) @@ -287,7 +319,12 @@ final class _MultiProducerSingleConsumerChannelBackpressuredStorage.failure(CancellationError())) + switch onProduceMore { + case .closure(let onProduceMore): + onProduceMore(.failure(CancellationError())) + case .continuation(let continuation): + continuation.resume(throwing: CancellationError()) + } case .none: break @@ -316,7 +353,12 @@ final class _MultiProducerSingleConsumerChannelBackpressuredStorage=6.0) @inlinable func next(isolation actor: isolated (any Actor)?) async throws -> Element? { let action = self._stateMachine.withCriticalRegion { @@ -337,7 +378,12 @@ final class _MultiProducerSingleConsumerChannelBackpressuredStorage.success(())) + switch producerContinuation { + case .closure(let onProduceMore): + onProduceMore(.success(())) + case .continuation(let continuation): + continuation.resume() + } } return element @@ -359,44 +405,7 @@ final class _MultiProducerSingleConsumerChannelBackpressuredStorage 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 .returnFailureAndCallOnTermination(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() - } - } - #endif - #if compiler(>=6.0) @inlinable func suspendNext(isolation actor: isolated (any Actor)?) async throws -> Element? { return try await withTaskCancellationHandler { @@ -412,7 +421,12 @@ final class _MultiProducerSingleConsumerChannelBackpressuredStorage.success(())) + switch producerContinuation { + case .closure(let onProduceMore): + onProduceMore(.failure(CancellationError())) + case .continuation(let continuation): + continuation.resume() + } } case .resumeConsumerWithFailureAndCallOnTermination(let continuation, let failure, let onTermination): @@ -444,64 +458,12 @@ final class _MultiProducerSingleConsumerChannelBackpressuredStorage Element? { - return try await withTaskCancellationHandler { - return try await withUnsafeThrowingContinuation { 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(())) + switch producerContinuation { + case .closure(let onProduceMore): + onProduceMore(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) + case .continuation(let continuation): + continuation.resume(throwing: MultiProducerSingleConsumerChannelAlreadyFinishedError()) } - - case .resumeConsumerWithFailureAndCallOnTermination(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 .resumeConsumerWithNilAndCallOnTermination(let continuation, let onTermination): - continuation.resume(returning: nil) - onTermination?() - - case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): - for producerContinuation in producerContinuations { - producerContinuation(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) } onTermination?() @@ -510,7 +472,6 @@ final class _MultiProducerSingleConsumerChannelBackpressuredStorage { case callOnTermination((@Sendable () -> Void)?) /// Indicates that all producers should be failed and `onTermination` should be called. case failProducersAndCallOnTermination( - _TinyArray<(Result) -> Void>, + _TinyArray<_MultiProducerSingleConsumerSuspendedProducer>, (@Sendable () -> Void)? ) } @@ -623,7 +584,7 @@ struct _MultiProducerSingleConsumerStateMachine { self._state = .finished(.init(iteratorInitialized: false, sourceFinished: false)) return .failProducersAndCallOnTermination( - .init(channeling.producerContinuations.lazy.map { $0.1 }), + .init(channeling.suspendedProducers.lazy.map { $0.1 }), channeling.onTermination ) } @@ -705,7 +666,7 @@ struct _MultiProducerSingleConsumerStateMachine { case callOnTermination((@Sendable () -> Void)?) /// Indicates that all producers should be failed and `onTermination` should be called. case failProducersAndCallOnTermination( - _TinyArray<(Result) -> Void>, + _TinyArray<_MultiProducerSingleConsumerSuspendedProducer>, (@Sendable () -> Void)? ) } @@ -731,7 +692,7 @@ struct _MultiProducerSingleConsumerStateMachine { self._state = .finished(.init(iteratorInitialized: true, sourceFinished: false)) return .failProducersAndCallOnTermination( - .init(channeling.producerContinuations.lazy.map { $0.1 }), + .init(channeling.suspendedProducers.lazy.map { $0.1 }), channeling.onTermination ) } else { @@ -929,7 +890,7 @@ struct _MultiProducerSingleConsumerStateMachine { @inlinable mutating func enqueueProducer( callbackToken: UInt64, - onProduceMore: @Sendable @escaping (Result) -> Void + onProduceMore: sending @escaping (Result) -> Void ) -> EnqueueProducerAction? { switch self._state { case .initial: @@ -948,7 +909,7 @@ struct _MultiProducerSingleConsumerStateMachine { return .resumeProducer(onProduceMore) } else { self._state = .modify - channeling.producerContinuations.append((callbackToken, onProduceMore)) + channeling.suspendedProducers.append((callbackToken, .closure(onProduceMore))) self._state = .channeling(channeling) return .none @@ -964,11 +925,58 @@ struct _MultiProducerSingleConsumerStateMachine { } } + /// Actions returned by `enqueueContinuation()`. + @usableFromInline + enum EnqueueContinuationAction { + /// Indicates that the producer should be notified to produce more. + case resumeProducer(UnsafeContinuation) + /// Indicates that the producer should be notified about an error. + case resumeProducerWithError(UnsafeContinuation, Error) + } + + @inlinable + mutating func enqueueContinuation( + callbackToken: UInt64, + continuation: UnsafeContinuation + ) -> EnqueueContinuationAction? { + switch self._state { + case .initial: + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + + case .channeling(var channeling): + if let index = channeling.cancelledAsyncProducers.firstIndex(of: callbackToken) { + // Our producer got marked as cancelled. + self._state = .modify + channeling.cancelledAsyncProducers.remove(at: index) + self._state = .channeling(channeling) + + return .resumeProducerWithError(continuation, CancellationError()) + } else if channeling.hasOutstandingDemand { + // We hit an edge case here where we wrote but the consuming thread got interleaved + return .resumeProducer(continuation) + } else { + self._state = .modify + channeling.suspendedProducers.append((callbackToken, .continuation(continuation))) + + self._state = .channeling(channeling) + return .none + } + + case .sourceFinished, .finished: + // Since we are unlocking between sending elements and suspending the send + // It can happen that the source got finished or the consumption fully finishes. + return .resumeProducerWithError(continuation, MultiProducerSingleConsumerChannelAlreadyFinishedError()) + + case .modify: + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + } + } + /// Actions returned by `cancelProducer()`. @usableFromInline enum CancelProducerAction { /// Indicates that the producer should be notified about cancellation. - case resumeProducerWithCancellationError((Result) -> Void) + case resumeProducerWithCancellationError(_MultiProducerSingleConsumerSuspendedProducer) } @inlinable @@ -981,7 +989,7 @@ struct _MultiProducerSingleConsumerStateMachine { fatalError("MultiProducerSingleConsumerChannel internal inconsistency") case .channeling(var channeling): - guard let index = channeling.producerContinuations.firstIndex(where: { $0.0 == callbackToken }) else { + guard let index = channeling.suspendedProducers.firstIndex(where: { $0.0 == callbackToken }) else { // The task that sends was cancelled before sending elements so the cancellation handler // got invoked right away self._state = .modify @@ -992,7 +1000,7 @@ struct _MultiProducerSingleConsumerStateMachine { } // We have an enqueued producer that we need to resume now self._state = .modify - let continuation = channeling.producerContinuations.remove(at: index).1 + let continuation = channeling.suspendedProducers.remove(at: index).1 self._state = .channeling(channeling) return .resumeProducerWithCancellationError(continuation) @@ -1021,7 +1029,7 @@ struct _MultiProducerSingleConsumerStateMachine { ) /// Indicates that the producers should be resumed with an error. case resumeProducers( - producerContinuations: _TinyArray<(Result) -> Void> + producerContinuations: _TinyArray<_MultiProducerSingleConsumerSuspendedProducer> ) } @@ -1055,7 +1063,7 @@ struct _MultiProducerSingleConsumerStateMachine { ) ) - return .resumeProducers(producerContinuations: .init(channeling.producerContinuations.lazy.map { $0.1 })) + return .resumeProducers(producerContinuations: .init(channeling.suspendedProducers.lazy.map { $0.1 })) } // We have a continuation, this means our buffer must be empty // Furthermore, we can now transition to finished @@ -1090,7 +1098,7 @@ struct _MultiProducerSingleConsumerStateMachine { /// 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, _TinyArray<(Result) -> Void>) + case returnElementAndResumeProducers(Element, _TinyArray<_MultiProducerSingleConsumerSuspendedProducer>) /// Indicates that the `Failure` should be returned to the caller and that `onTermination` should be called. case returnFailureAndCallOnTermination(Failure?, (() -> Void)?) /// Indicates that the `nil` should be returned to the caller. @@ -1145,8 +1153,8 @@ struct _MultiProducerSingleConsumerStateMachine { return .returnElement(element) } // There is demand and we have to resume our producers - let producers = _TinyArray(channeling.producerContinuations.lazy.map { $0.1 }) - channeling.producerContinuations.removeAll(keepingCapacity: true) + let producers = _TinyArray(channeling.suspendedProducers.lazy.map { $0.1 }) + channeling.suspendedProducers.removeAll(keepingCapacity: true) self._state = .channeling(channeling) return .returnElementAndResumeProducers(element, producers) @@ -1181,7 +1189,7 @@ struct _MultiProducerSingleConsumerStateMachine { case resumeConsumerWithElementAndProducers( UnsafeContinuation, Element, - _TinyArray<(Result) -> Void> + _TinyArray<_MultiProducerSingleConsumerSuspendedProducer> ) /// Indicates that the consumer should be resumed with the failure and that `onTermination` should be called. case resumeConsumerWithFailureAndCallOnTermination( @@ -1226,8 +1234,8 @@ struct _MultiProducerSingleConsumerStateMachine { return .resumeConsumerWithElement(continuation, element) } // There is demand and we have to resume our producers - let producers = _TinyArray(channeling.producerContinuations.lazy.map { $0.1 }) - channeling.producerContinuations.removeAll(keepingCapacity: true) + let producers = _TinyArray(channeling.suspendedProducers.lazy.map { $0.1 }) + channeling.suspendedProducers.removeAll(keepingCapacity: true) self._state = .channeling(channeling) return .resumeConsumerWithElementAndProducers(continuation, element, producers) @@ -1263,7 +1271,7 @@ struct _MultiProducerSingleConsumerStateMachine { /// Indicates that the continuation should be resumed with nil, the producers should be finished and call onTermination. case resumeConsumerWithNilAndCallOnTermination(UnsafeContinuation, (() -> Void)?) /// Indicates that the producers should be finished and call onTermination. - case failProducersAndCallOnTermination(_TinyArray<(Result) -> Void>, (() -> Void)?) + case failProducersAndCallOnTermination(_TinyArray<_MultiProducerSingleConsumerSuspendedProducer>, (() -> Void)?) } @inlinable @@ -1277,12 +1285,12 @@ struct _MultiProducerSingleConsumerStateMachine { guard let consumerContinuation = channeling.consumerContinuation else { return .failProducersAndCallOnTermination( - .init(channeling.producerContinuations.lazy.map { $0.1 }), + .init(channeling.suspendedProducers.lazy.map { $0.1 }), channeling.onTermination ) } precondition( - channeling.producerContinuations.isEmpty, + channeling.suspendedProducers.isEmpty, "Internal inconsistency. Unexpected producer continuations." ) return .resumeConsumerWithNilAndCallOnTermination( @@ -1345,7 +1353,7 @@ enum _MultiProducerSingleConsumerState: CustomStringCon /// The producer continuations. @usableFromInline - var producerContinuations: Deque<(UInt64, (Result) -> Void)> + var suspendedProducers: Deque<(UInt64, _MultiProducerSingleConsumerSuspendedProducer)> /// The producers that have been cancelled. @usableFromInline @@ -1356,7 +1364,7 @@ enum _MultiProducerSingleConsumerState: CustomStringCon var hasOutstandingDemand: Bool var description: String { - "backpressure:\(self.backpressureStrategy.description) iteratorInitialized:\(self.iteratorInitialized) buffer:\(self.buffer.count) consumerContinuation:\(self.consumerContinuation == nil) producerContinuations:\(self.producerContinuations.count) cancelledProducers:\(self.cancelledAsyncProducers.count) hasOutstandingDemand:\(self.hasOutstandingDemand)" + "backpressure:\(self.backpressureStrategy.description) iteratorInitialized:\(self.iteratorInitialized) buffer:\(self.buffer.count) consumerContinuation:\(self.consumerContinuation == nil) producerContinuations:\(self.suspendedProducers.count) cancelledProducers:\(self.cancelledAsyncProducers.count) hasOutstandingDemand:\(self.hasOutstandingDemand)" } @usableFromInline @@ -1365,7 +1373,7 @@ enum _MultiProducerSingleConsumerState: CustomStringCon onTermination: (@Sendable () -> Void)? = nil, buffer: Deque, consumerContinuation: UnsafeContinuation? = nil, - producerContinuations: Deque<(UInt64, (Result) -> Void)>, + producerContinuations: Deque<(UInt64, _MultiProducerSingleConsumerSuspendedProducer)>, cancelledAsyncProducers: Deque, hasOutstandingDemand: Bool) { self.backpressureStrategy = backpressureStrategy @@ -1373,7 +1381,7 @@ enum _MultiProducerSingleConsumerState: CustomStringCon self.onTermination = onTermination self.buffer = buffer self.consumerContinuation = consumerContinuation - self.producerContinuations = producerContinuations + self.suspendedProducers = producerContinuations self.cancelledAsyncProducers = cancelledAsyncProducers self.hasOutstandingDemand = hasOutstandingDemand } @@ -1472,3 +1480,10 @@ enum _MultiProducerSingleConsumerState: CustomStringCon } } } + +@usableFromInline +enum _MultiProducerSingleConsumerSuspendedProducer { + case closure((Result) -> Void) + case continuation(UnsafeContinuation) +} +#endif diff --git a/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel.swift b/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel.swift index d2b1d83e..e4fb7782 100644 --- a/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel.swift +++ b/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel.swift @@ -9,6 +9,7 @@ // //===----------------------------------------------------------------------===// +#if compiler(>=6.0) /// An error that is thrown from the various `send` methods of the /// ``MultiProducerSingleConsumerChannel/Source``. /// @@ -85,6 +86,20 @@ public struct MultiProducerSingleConsumerChannel: Async /// The backing storage. private let backing: _Backing + @frozen + public struct NewChannel: ~Copyable { + public var channel: MultiProducerSingleConsumerChannel + public var source: Source + + public init( + channel: MultiProducerSingleConsumerChannel, + source: consuming Source + ) { + self.channel = channel + self.source = source + } + } + /// Initializes a new ``MultiProducerSingleConsumerChannel`` and an ``MultiProducerSingleConsumerChannel/Source``. /// /// - Parameters: @@ -97,13 +112,13 @@ public struct MultiProducerSingleConsumerChannel: Async of elementType: Element.Type = Element.self, throwing failureType: Failure.Type = Never.self, backpressureStrategy: Source.BackpressureStrategy - ) -> (`Self`, Source) { + ) -> NewChannel { let storage = _MultiProducerSingleConsumerChannelBackpressuredStorage( backpressureStrategy: backpressureStrategy.internalBackpressureStrategy ) let source = Source(storage: storage) - return (.init(storage: storage), source) + return .init(channel: .init(storage: storage), source: source) } init(storage: _MultiProducerSingleConsumerChannelBackpressuredStorage) { @@ -117,7 +132,7 @@ extension MultiProducerSingleConsumerChannel { /// Use this source to provide elements to the channel by calling one of the `send` methods. /// /// - Important: You must terminate the source by calling ``finish(throwing:)``. - public struct Source: Sendable { + public struct Source: ~Copyable, Sendable { /// A strategy that handles the backpressure of the channel. public struct BackpressureStrategy: Sendable { var internalBackpressureStrategy: _MultiProducerSingleConsumerChannelInternalBackpressureStrategy @@ -169,10 +184,10 @@ extension MultiProducerSingleConsumerChannel { } /// A type that indicates the result of sending elements to the source. - public enum SendResult: Sendable { + public enum SendResult: ~Copyable, Sendable { /// A token that is returned when the channel's backpressure strategy indicated that production should /// be suspended. Use this token to enqueue a callback by calling the ``enqueueCallback(_:)`` method. - public struct CallbackToken: Sendable { + public struct CallbackToken: ~Copyable, Sendable { @usableFromInline let _id: UInt64 @@ -191,20 +206,6 @@ extension MultiProducerSingleConsumerChannel { case enqueueCallback(CallbackToken) } - /// Backing class for the source used to hook a deinit. - @usableFromInline - final class _Backing: Sendable { - @usableFromInline - let storage: _MultiProducerSingleConsumerChannelBackpressuredStorage - - init(storage: _MultiProducerSingleConsumerChannelBackpressuredStorage) { - self.storage = storage - } - - deinit { - self.storage.sourceDeinitialized() - } - } /// A callback to invoke when the channel finished. /// @@ -214,18 +215,33 @@ extension MultiProducerSingleConsumerChannel { /// - After ``finish(throwing:)`` was called and all elements have been consumed public var onTermination: (@Sendable () -> Void)? { set { - self._backing.storage.onTermination = newValue + self._storage.onTermination = newValue } get { - self._backing.storage.onTermination + self._storage.onTermination } } @usableFromInline - var _backing: _Backing + let _storage: _MultiProducerSingleConsumerChannelBackpressuredStorage internal init(storage: _MultiProducerSingleConsumerChannelBackpressuredStorage) { - self._backing = .init(storage: storage) + self._storage = storage + } + + deinit { + // TODO: We can't finish here. +// self.finish() + } + + + /// Creates a new source which can be used to send elements to the channel concurrently. + /// + /// The channel will only automatically be finished if all existing sources have been deinited. + /// + /// - Returns: A new source for sending elements to the channel. + public mutating func copy() -> Self { + .init(storage: self._storage) } /// Sends new elements to the channel. @@ -236,17 +252,10 @@ extension MultiProducerSingleConsumerChannel { /// /// - Parameter sequence: The elements to send to the channel. /// - Returns: The result that indicates if more elements should be produced at this time. - #if compiler(>=6.0) - @inlinable - public func send(contentsOf sequence: sending S) throws -> SendResult where Element == S.Element, S: Sequence { - try self._backing.storage.send(contentsOf: sequence) - } - #else @inlinable - public func send(contentsOf sequence: S) throws -> SendResult where Element == S.Element, S: Sequence & Sendable, Element: Sendable { - try self._backing.storage.send(contentsOf: sequence) + public mutating func send(contentsOf sequence: sending S) throws -> SendResult where Element == S.Element, S: Sequence { + try self._storage.send(contentsOf: sequence) } - #endif /// Send the element to the channel. /// @@ -256,17 +265,10 @@ extension MultiProducerSingleConsumerChannel { /// /// - Parameter element: The element to send to the channel. /// - Returns: The result that indicates if more elements should be produced at this time. - #if compiler(>=6.0) - @inlinable - public func send(_ element: sending Element) throws -> SendResult { - try self._backing.storage.send(contentsOf: CollectionOfOne(element)) - } - #else @inlinable - public func send(_ element: Element) throws -> SendResult where Element: Sendable { - try self._backing.storage.send(contentsOf: CollectionOfOne(element)) + public mutating func send(_ element: sending Element) throws -> SendResult { + try self._storage.send(contentsOf: CollectionOfOne(element)) } - #endif /// Enqueues a callback that will be invoked once more elements should be produced. /// @@ -277,23 +279,13 @@ extension MultiProducerSingleConsumerChannel { /// - Parameters: /// - callbackToken: The callback token. /// - onProduceMore: The callback which gets invoked once more elements should be produced. - #if compiler(>=6.0) @inlinable - public func enqueueCallback( + public mutating func enqueueCallback( callbackToken: consuming SendResult.CallbackToken, onProduceMore: sending @escaping (Result) -> Void ) { - self._backing.storage.enqueueProducer(callbackToken: callbackToken._id, onProduceMore: onProduceMore) - } - #else - @inlinable - public func enqueueCallback( - callbackToken: consuming SendResult.CallbackToken, - onProduceMore: @escaping @Sendable (Result) -> Void - ) { - self._backing.storage.enqueueProducer(callbackToken: callbackToken._id, onProduceMore: onProduceMore) + self._storage.enqueueProducer(callbackToken: callbackToken._id, onProduceMore: onProduceMore) } - #endif /// Cancel an enqueued callback. /// @@ -304,8 +296,8 @@ extension MultiProducerSingleConsumerChannel { /// /// - Parameter callbackToken: The callback token. @inlinable - public func cancelCallback(callbackToken: consuming SendResult.CallbackToken) { - self._backing.storage.cancelProducer(callbackToken: callbackToken._id) + public mutating func cancelCallback(callbackToken: consuming SendResult.CallbackToken) { + self._storage.cancelProducer(callbackToken: callbackToken._id) } /// Send new elements to the channel and provide a callback which will be invoked once more elements should be produced. @@ -318,16 +310,15 @@ extension MultiProducerSingleConsumerChannel { /// - sequence: The elements to send to the channel. /// - onProduceMore: The callback which gets invoked once more elements should be produced. This callback might be /// invoked during the call to ``send(contentsOf:onProduceMore:)``. - #if compiler(>=6.0) @inlinable - public func send( + public mutating func send( contentsOf sequence: sending S, onProduceMore: @escaping @Sendable (Result) -> Void ) where Element == S.Element, S: Sequence { do { let sendResult = try self.send(contentsOf: sequence) - switch sendResult { + switch consume sendResult { case .produceMore: onProduceMore(Result.success(())) @@ -338,27 +329,6 @@ extension MultiProducerSingleConsumerChannel { onProduceMore(.failure(error)) } } - #else - @inlinable - public func send( - contentsOf sequence: S, - onProduceMore: @escaping @Sendable (Result) -> Void - ) where Element == S.Element, S: Sequence & Sendable, Element: Sendable { - do { - let sendResult = try self.send(contentsOf: sequence) - - switch sendResult { - case .produceMore: - onProduceMore(Result.success(())) - - case .enqueueCallback(let callbackToken): - self.enqueueCallback(callbackToken: callbackToken, onProduceMore: onProduceMore) - } - } catch { - onProduceMore(.failure(error)) - } - } - #endif /// Sends the element to the channel. /// @@ -370,23 +340,13 @@ extension MultiProducerSingleConsumerChannel { /// - element: The element to send to the channel. /// - onProduceMore: The callback which gets invoked once more elements should be produced. This callback might be /// invoked during the call to ``send(_:onProduceMore:)``. - #if compiler(>=6.0) @inlinable - public func send( + public mutating func send( _ element: sending Element, onProduceMore: @escaping @Sendable (Result) -> Void ) { self.send(contentsOf: CollectionOfOne(element), onProduceMore: onProduceMore) } - #else - @inlinable - public func send( - _ element: Element, - onProduceMore: @escaping @Sendable (Result) -> Void - ) where Element: Sendable { - self.send(contentsOf: CollectionOfOne(element), onProduceMore: onProduceMore) - } - #endif /// Send new elements to the channel. /// @@ -398,65 +358,29 @@ extension MultiProducerSingleConsumerChannel { /// /// - Parameters: /// - sequence: The elements to send to the channel. - #if compiler(>=6.0) - @inlinable - public func send(contentsOf sequence: sending S) async throws where Element == S.Element, S: Sequence { - let sendResult = try { try self.send(contentsOf: sequence) }() - - switch sendResult { - case .produceMore: - return - - case .enqueueCallback(let callbackToken): - try await withTaskCancellationHandler { - try await withUnsafeThrowingContinuation { 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) - } - } - } - #else @inlinable - public func send(contentsOf sequence: S) async throws where Element == S.Element, S: Sequence & Sendable, Element: Sendable { + public mutating func send(contentsOf sequence: sending S) async throws where Element == S.Element, S: Sequence { let sendResult = try { try self.send(contentsOf: sequence) }() - switch sendResult { + switch consume sendResult { case .produceMore: - return + return () case .enqueueCallback(let callbackToken): + let id = callbackToken._id + let storage = self._storage try await withTaskCancellationHandler { try await withUnsafeThrowingContinuation { continuation in - self.enqueueCallback( - callbackToken: callbackToken, - onProduceMore: { result in - switch result { - case .success(): - continuation.resume(returning: ()) - case .failure(let error): - continuation.resume(throwing: error) - } - } + self._storage.enqueueProducer( + callbackToken: id, + continuation: continuation ) } } onCancel: { - self.cancelCallback(callbackToken: callbackToken) + storage.cancelProducer(callbackToken: id) } } } - #endif /// Send new element to the channel. /// @@ -468,17 +392,10 @@ extension MultiProducerSingleConsumerChannel { /// /// - Parameters: /// - element: The element to send to the channel. - #if compiler(>=6.0) @inlinable - public func send(_ element: sending Element) async throws { + public mutating func send(_ element: sending Element) async throws { try await self.send(contentsOf: CollectionOfOne(element)) } - #else - @inlinable - public func send(_ element: Element) async throws where Element: Sendable { - try await self.send(contentsOf: CollectionOfOne(element)) - } - #endif /// Send the elements of the asynchronous sequence to the channel. /// @@ -488,21 +405,12 @@ extension MultiProducerSingleConsumerChannel { /// /// - Parameters: /// - sequence: The elements to send to the channel. - #if compiler(>=6.0) - @inlinable - public func send(contentsOf sequence: sending S) async throws where Element == S.Element, S: AsyncSequence { - for try await element in sequence { - try await self.send(contentsOf: CollectionOfOne(element)) - } - } - #else @inlinable - public func send(contentsOf sequence: S) async throws where Element == S.Element, S: AsyncSequence & Sendable, Element: Sendable{ + public mutating func send(contentsOf sequence: sending S) async throws where Element == S.Element, S: AsyncSequence { for try await element in sequence { try await self.send(contentsOf: CollectionOfOne(element)) } } - #endif /// Indicates that the production terminated. /// @@ -514,8 +422,8 @@ extension MultiProducerSingleConsumerChannel { /// - Parameters: /// - error: The error to throw, or `nil`, to finish normally. @inlinable - public func finish(throwing error: Failure? = nil) { - self._backing.storage.finish(error) + public consuming func finish(throwing error: Failure? = nil) { + self._storage.finish(error) } } } @@ -550,37 +458,12 @@ extension MultiProducerSingleConsumerChannel { self._backing = .init(storage: storage) } - /// The next value from the channel. - /// - /// When `next()` returns `nil`, this signifies the end of the channel. - /// - /// 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 channel terminates. In this case, - /// `next()` may return `nil` immediately, or else return `nil` on - /// subsequent calls. @_disfavoredOverload @inlinable public mutating func next() async throws -> Element? { - try await self._backing.storage.next() + try await self._backing.storage.next(isolation: nil) } - #if compiler(>=6.0) - /// The next value from the channel. - /// - /// When `next()` returns `nil`, this signifies the end of the channel. - /// - /// 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 channel terminates. In this case, - /// `next()` may return `nil` immediately, or else return `nil` on - /// subsequent calls. @available(macOS 15.0, iOS 18.0, watchOS 11.0, tvOS 18.0, visionOS 2.0, *) @inlinable public mutating func next( @@ -592,7 +475,6 @@ extension MultiProducerSingleConsumerChannel { throw error as! Failure } } - #endif } /// Creates the asynchronous iterator that produces elements of this @@ -606,3 +488,4 @@ extension MultiProducerSingleConsumerChannel: Sendable where Element: Sendable { @available(*, unavailable) extension MultiProducerSingleConsumerChannel.Iterator: Sendable {} +#endif From bb5586a5f20fafd6a36e6800ac6105ba46701518 Mon Sep 17 00:00:00 2001 From: Franz Busch Date: Thu, 13 Jun 2024 14:38:42 +0200 Subject: [PATCH 4/7] Use ~Copyable for the state and state machine --- ...oducerSingleConsumerChannel+Internal.swift | 2409 +++++++++-------- .../MultiProducerSingleConsumerChannel.swift | 21 +- ...tiProducerSingleConsumerChannelTests.swift | 1093 ++++---- 3 files changed, 1773 insertions(+), 1750 deletions(-) diff --git a/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel+Internal.swift b/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel+Internal.swift index ba717e60..5cea923a 100644 --- a/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel+Internal.swift +++ b/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel+Internal.swift @@ -12,1471 +12,1505 @@ #if compiler(>=6.0) import DequeModule -@usableFromInline -struct _MultiProducerSingleConsumerChannelWatermarkBackpressureStrategy: Sendable, CustomStringConvertible { - /// 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 level. - @usableFromInline - var _currentWatermark: Int = 0 - - /// A closure that can be used to calculate the watermark impact of a single element +extension MultiProducerSingleConsumerChannel { @usableFromInline - let _waterLevelForElement: (@Sendable (Element) -> Int)? - - @usableFromInline - var description: String { - "watermark(\(self._currentWatermark))" - } + enum _InternalBackpressureStrategy: Sendable, CustomStringConvertible { + @usableFromInline + struct _Watermark: Sendable, CustomStringConvertible { + /// 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 level. + @usableFromInline + var _currentWatermark: Int = 0 + + /// A closure that can be used to calculate the watermark impact of a single element + @usableFromInline + let _waterLevelForElement: (@Sendable (Element) -> Int)? + + @usableFromInline + var description: String { + "watermark(\(self._currentWatermark))" + } - init(low: Int, high: Int, waterLevelForElement: (@Sendable (Element) -> Int)?) { - precondition(low <= high) - self._low = low - self._high = high - self._waterLevelForElement = waterLevelForElement - } + init(low: Int, high: Int, waterLevelForElement: (@Sendable (Element) -> Int)?) { + precondition(low <= high) + self._low = low + self._high = high + self._waterLevelForElement = waterLevelForElement + } - @inlinable - mutating func didSend(elements: Deque.SubSequence) -> Bool { - if let waterLevelForElement = self._waterLevelForElement { - self._currentWatermark += elements.reduce(0) { $0 + waterLevelForElement($1) } - } else { - self._currentWatermark += elements.count - } - precondition(self._currentWatermark >= 0) - // We are demanding more until we reach the high watermark - return self._currentWatermark < self._high - } + @inlinable + mutating func didSend(elements: Deque.SubSequence) -> Bool { + if let waterLevelForElement = self._waterLevelForElement { + self._currentWatermark += elements.reduce(0) { $0 + waterLevelForElement($1) } + } else { + self._currentWatermark += elements.count + } + precondition(self._currentWatermark >= 0) + // We are demanding more until we reach the high watermark + return self._currentWatermark < self._high + } - @inlinable - mutating func didConsume(element: Element) -> Bool { - if let waterLevelForElement = self._waterLevelForElement { - self._currentWatermark -= waterLevelForElement(element) - } else { - self._currentWatermark -= 1 + @inlinable + mutating func didConsume(element: Element) -> Bool { + if let waterLevelForElement = self._waterLevelForElement { + self._currentWatermark -= waterLevelForElement(element) + } else { + self._currentWatermark -= 1 + } + precondition(self._currentWatermark >= 0) + // We start demanding again once we are below the low watermark + return self._currentWatermark < self._low + } } - precondition(self._currentWatermark >= 0) - // We start demanding again once we are below the low watermark - return self._currentWatermark < self._low - } -} -@usableFromInline -struct _MultiProducerSingleConsumerChannelUnboundedBackpressureStrategy: Sendable, CustomStringConvertible { - @usableFromInline - var description: String { - return "unbounded" - } + @usableFromInline + struct _Unbounded: Sendable, CustomStringConvertible { + @usableFromInline + var description: String { + return "unbounded" + } - init() { } + init() { } - @inlinable - mutating func didSend(elements: Deque.SubSequence) -> Bool { - return true - } + @inlinable + mutating func didSend(elements: Deque.SubSequence) -> Bool { + return true + } - @inlinable - mutating func didConsume(element: Element) -> Bool { - return true - } -} + @inlinable + mutating func didConsume(element: Element) -> Bool { + return true + } + } -@usableFromInline -enum _MultiProducerSingleConsumerChannelInternalBackpressureStrategy: Sendable, CustomStringConvertible { - /// A watermark based strategy. - case watermark(_MultiProducerSingleConsumerChannelWatermarkBackpressureStrategy) - /// An unbounded based strategy. - case unbounded(_MultiProducerSingleConsumerChannelUnboundedBackpressureStrategy) + /// A watermark based strategy. + case watermark(_Watermark) + /// An unbounded based strategy. + case unbounded(_Unbounded) - @usableFromInline - var description: String { - switch self { - case .watermark(let strategy): - return strategy.description - case .unbounded(let unbounded): - return unbounded.description + @usableFromInline + var description: String { + switch consume self { + case .watermark(let strategy): + return strategy.description + case .unbounded(let unbounded): + return unbounded.description + } } - } - @inlinable - mutating func didSend(elements: Deque.SubSequence) -> Bool { - switch self { - case .watermark(var strategy): - let result = strategy.didSend(elements: elements) - self = .watermark(strategy) - return result - case .unbounded(var strategy): - let result = strategy.didSend(elements: elements) - self = .unbounded(strategy) - return result + @inlinable + mutating func didSend(elements: Deque.SubSequence) -> Bool { + switch consume self { + case .watermark(var strategy): + let result = strategy.didSend(elements: elements) + self = .watermark(strategy) + return result + case .unbounded(var strategy): + let result = strategy.didSend(elements: elements) + self = .unbounded(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 - case .unbounded(var strategy): - let result = strategy.didConsume(element: element) - self = .unbounded(strategy) - return result + @inlinable + mutating func didConsume(element: Element) -> Bool { + switch consume self { + case .watermark(var strategy): + let result = strategy.didConsume(element: element) + self = .watermark(strategy) + return result + case .unbounded(var strategy): + let result = strategy.didConsume(element: element) + self = .unbounded(strategy) + return result + } } } } -@usableFromInline -final class _MultiProducerSingleConsumerChannelBackpressuredStorage { - /// The state machine +extension MultiProducerSingleConsumerChannel { @usableFromInline - var _stateMachine: ManagedCriticalState<_MultiProducerSingleConsumerStateMachine> + final class _Storage { + @usableFromInline + let _lock = Lock.allocate() + /// The state machine + @usableFromInline + var _stateMachine: _StateMachine - var onTermination: (@Sendable () -> Void)? { - set { - self._stateMachine.withCriticalRegion { - $0._onTermination = newValue + var onTermination: (@Sendable () -> Void)? { + set { + self._lock.withLockVoid { + self._stateMachine._onTermination = newValue + } } - } - get { - self._stateMachine.withCriticalRegion { - $0._onTermination + get { + self._lock.withLock { + self._stateMachine._onTermination + } } } - } - init( - backpressureStrategy: _MultiProducerSingleConsumerChannelInternalBackpressureStrategy - ) { - self._stateMachine = .init(.init(backpressureStrategy: backpressureStrategy)) - } - - func sequenceDeinitialized() { - let action = self._stateMachine.withCriticalRegion { - $0.sequenceDeinitialized() + init( + backpressureStrategy: _InternalBackpressureStrategy + ) { + self._stateMachine = .init(backpressureStrategy: backpressureStrategy) } - switch action { - case .callOnTermination(let onTermination): - onTermination?() + func sequenceDeinitialized() { + let action = self._lock.withLock { + self._stateMachine.sequenceDeinitialized() + } - case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): - for producerContinuation in producerContinuations { - switch producerContinuation { - case .closure(let onProduceMore): - onProduceMore(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) - case .continuation(let continuation): - continuation.resume(throwing: MultiProducerSingleConsumerChannelAlreadyFinishedError()) + switch action { + case .callOnTermination(let onTermination): + onTermination?() + + case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): + for producerContinuation in producerContinuations { + switch producerContinuation { + case .closure(let onProduceMore): + onProduceMore(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) + case .continuation(let continuation): + continuation.resume(throwing: MultiProducerSingleConsumerChannelAlreadyFinishedError()) + } } - } - onTermination?() + onTermination?() - case .none: - break + case .none: + break + } } - } - func iteratorInitialized() { - self._stateMachine.withCriticalRegion { - $0.iteratorInitialized() + func iteratorInitialized() { + self._lock.withLockVoid { + self._stateMachine.iteratorInitialized() + } } - } - func iteratorDeinitialized() { - let action = self._stateMachine.withCriticalRegion { - $0.iteratorDeinitialized() - } + func iteratorDeinitialized() { + let action = self._lock.withLock { + self._stateMachine.iteratorDeinitialized() + } - switch action { - case .callOnTermination(let onTermination): - onTermination?() + switch action { + case .callOnTermination(let onTermination): + onTermination?() - case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): - for producerContinuation in producerContinuations { - switch producerContinuation { - case .closure(let onProduceMore): - onProduceMore(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) - case .continuation(let continuation): - continuation.resume(throwing: MultiProducerSingleConsumerChannelAlreadyFinishedError()) + case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): + for producerContinuation in producerContinuations { + switch producerContinuation { + case .closure(let onProduceMore): + onProduceMore(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) + case .continuation(let continuation): + continuation.resume(throwing: MultiProducerSingleConsumerChannelAlreadyFinishedError()) + } } - } - onTermination?() - - case .none: - break - } - } + onTermination?() - func sourceDeinitialized() { - let action = self._stateMachine.withCriticalRegion { - $0.sourceDeinitialized() + case .none: + break + } } - switch action { - case .callOnTermination(let onTermination): - onTermination?() - - case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): - for producerContinuation in producerContinuations { - producerContinuation(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) + func sourceDeinitialized() { + let action = self._lock.withLock { + self._stateMachine.sourceDeinitialized() } - onTermination?() - case .failProducers(let producerContinuations): - for producerContinuation in producerContinuations { - producerContinuation(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) - } + switch action { + case .callOnTermination(let onTermination): + onTermination?() - case .none: - break - } - } + case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): + for producerContinuation in producerContinuations { + producerContinuation(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) + } + onTermination?() - @inlinable - func send( - contentsOf sequence: some Sequence - ) throws -> MultiProducerSingleConsumerChannel.Source.SendResult { - let action = self._stateMachine.withCriticalRegion { - return $0.send(sequence) + case .failProducers(let producerContinuations): + for producerContinuation in producerContinuations { + producerContinuation(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) + } + + case .none: + break + } } - switch action { - case .returnProduceMore: - return .produceMore + @inlinable + func send( + contentsOf sequence: some Sequence + ) throws -> MultiProducerSingleConsumerChannel.Source.SendResult { + let action = self._lock.withLock { + return self._stateMachine.send(sequence) + } - case .returnEnqueue(let callbackToken): - return .enqueueCallback(.init(id: callbackToken)) + switch action { + case .returnProduceMore: + return .produceMore - case .resumeConsumerAndReturnProduceMore(let continuation, let element): - continuation.resume(returning: element) - return .produceMore + case .returnEnqueue(let callbackToken): + return .enqueueCallback(.init(id: callbackToken)) - case .resumeConsumerAndReturnEnqueue(let continuation, let element, let callbackToken): - continuation.resume(returning: element) - return .enqueueCallback(.init(id: callbackToken)) + case .resumeConsumerAndReturnProduceMore(let continuation, let element): + continuation.resume(returning: element) + return .produceMore - case .throwFinishedError: - throw MultiProducerSingleConsumerChannelAlreadyFinishedError() - } - } + case .resumeConsumerAndReturnEnqueue(let continuation, let element, let callbackToken): + continuation.resume(returning: element) + return .enqueueCallback(.init(id: callbackToken)) - @inlinable - func enqueueProducer( - callbackToken: UInt64, - continuation: UnsafeContinuation - ) { - let action = self._stateMachine.withCriticalRegion { - $0.enqueueContinuation(callbackToken: callbackToken, continuation: continuation) + case .throwFinishedError: + throw MultiProducerSingleConsumerChannelAlreadyFinishedError() + } } - switch action { - case .resumeProducer(let continuation): - continuation.resume() + @inlinable + func enqueueProducer( + callbackToken: UInt64, + continuation: UnsafeContinuation + ) { + let action = self._lock.withLock { + self._stateMachine.enqueueContinuation(callbackToken: callbackToken, continuation: continuation) + } - case .resumeProducerWithError(let continuation, let error): - continuation.resume(throwing: error) + switch action { + case .resumeProducer(let continuation): + continuation.resume() - case .none: - break - } - } + case .resumeProducerWithError(let continuation, let error): + continuation.resume(throwing: error) - @inlinable - func enqueueProducer( - callbackToken: UInt64, - onProduceMore: sending @escaping (Result) -> Void - ) { - let action = self._stateMachine.withCriticalRegion { - $0.enqueueProducer(callbackToken: callbackToken, onProduceMore: onProduceMore) + case .none: + break + } } - switch action { - case .resumeProducer(let onProduceMore): - onProduceMore(Result.success(())) + @inlinable + func enqueueProducer( + callbackToken: UInt64, + onProduceMore: sending @escaping (Result) -> Void + ) { + let action = self._lock.withLock { + self._stateMachine.enqueueProducer(callbackToken: callbackToken, onProduceMore: onProduceMore) + } - case .resumeProducerWithError(let onProduceMore, let error): - onProduceMore(Result.failure(error)) + switch action { + case .resumeProducer(let onProduceMore): + onProduceMore(Result.success(())) - case .none: - break - } - } + case .resumeProducerWithError(let onProduceMore, let error): + onProduceMore(Result.failure(error)) - @inlinable - func cancelProducer( - callbackToken: UInt64 - ) { - let action = self._stateMachine.withCriticalRegion { - $0.cancelProducer(callbackToken: callbackToken) + case .none: + break + } } - switch action { - case .resumeProducerWithCancellationError(let onProduceMore): - switch onProduceMore { - case .closure(let onProduceMore): - onProduceMore(.failure(CancellationError())) - case .continuation(let continuation): - continuation.resume(throwing: CancellationError()) + @inlinable + func cancelProducer( + callbackToken: UInt64 + ) { + let action = self._lock.withLock { + self._stateMachine.cancelProducer(callbackToken: callbackToken) } - case .none: - break - } - } + switch action { + case .resumeProducerWithCancellationError(let onProduceMore): + switch onProduceMore { + case .closure(let onProduceMore): + onProduceMore(.failure(CancellationError())) + case .continuation(let continuation): + continuation.resume(throwing: CancellationError()) + } - @inlinable - func finish(_ failure: Failure?) { - let action = self._stateMachine.withCriticalRegion { - $0.finish(failure) + case .none: + break + } } - 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) + @inlinable + func finish(_ failure: Failure?) { + let action = self._lock.withLock { + self._stateMachine.finish(failure) } - onTermination?() + switch action { + case .callOnTermination(let onTermination): + onTermination?() - case .resumeProducers(let producerContinuations): - for producerContinuation in producerContinuations { - switch producerContinuation { - case .closure(let onProduceMore): - onProduceMore(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) - case .continuation(let continuation): - continuation.resume(throwing: MultiProducerSingleConsumerChannelAlreadyFinishedError()) + case .resumeConsumerAndCallOnTermination(let consumerContinuation, let failure, let onTermination): + switch failure { + case .some(let error): + consumerContinuation.resume(throwing: error) + case .none: + consumerContinuation.resume(returning: nil) } - } - case .none: - break - } - } + onTermination?() - @inlinable - func next(isolation actor: isolated (any Actor)?) async throws -> Element? { - let action = self._stateMachine.withCriticalRegion { - $0.next() + case .resumeProducers(let producerContinuations): + for producerContinuation in producerContinuations { + switch producerContinuation { + case .closure(let onProduceMore): + onProduceMore(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) + case .continuation(let continuation): + continuation.resume(throwing: MultiProducerSingleConsumerChannelAlreadyFinishedError()) + } + } + + case .none: + break + } } - switch action { - case .returnElement(let element): - return element + @inlinable + func next(isolation actor: isolated (any Actor)?) async throws -> Element? { + let action = self._lock.withLock { + self._stateMachine.next() + } - case .returnElementAndResumeProducers(let element, let producerContinuations): - for producerContinuation in producerContinuations { - switch producerContinuation { - case .closure(let onProduceMore): - onProduceMore(.success(())) - case .continuation(let continuation): - continuation.resume() + switch action { + case .returnElement(let element): + return element + + case .returnElementAndResumeProducers(let element, let producerContinuations): + for producerContinuation in producerContinuations { + switch producerContinuation { + case .closure(let onProduceMore): + onProduceMore(.success(())) + case .continuation(let continuation): + continuation.resume() + } } - } - return element + return element - case .returnFailureAndCallOnTermination(let failure, let onTermination): - onTermination?() - switch failure { - case .some(let error): - throw error + case .returnFailureAndCallOnTermination(let failure, let onTermination): + onTermination?() + switch failure { + case .some(let error): + throw error - case .none: + case .none: + return nil + } + + case .returnNil: return nil + + case .suspendTask: + return try await self.suspendNext(isolation: actor) } + } - case .returnNil: - return nil + @inlinable + func suspendNext(isolation actor: isolated (any Actor)?) async throws -> Element? { + return try await withTaskCancellationHandler { + return try await withUnsafeThrowingContinuation { continuation in + let action = self._lock.withLock { + self._stateMachine.suspendNext(continuation: continuation) + } - case .suspendTask: - return try await self.suspendNext(isolation: actor) - } - } + 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 { + switch producerContinuation { + case .closure(let onProduceMore): + onProduceMore(.failure(CancellationError())) + case .continuation(let continuation): + continuation.resume() + } + } - @inlinable - func suspendNext(isolation actor: isolated (any Actor)?) async throws -> Element? { - return try await withTaskCancellationHandler { - return try await withUnsafeThrowingContinuation { continuation in - let action = self._stateMachine.withCriticalRegion { - $0.suspendNext(continuation: continuation) + case .resumeConsumerWithFailureAndCallOnTermination(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._lock.withLock { + self._stateMachine.cancelNext() } switch action { - case .resumeConsumerWithElement(let continuation, let element): - continuation.resume(returning: element) + case .resumeConsumerWithNilAndCallOnTermination(let continuation, let onTermination): + continuation.resume(returning: nil) + onTermination?() - case .resumeConsumerWithElementAndProducers(let continuation, let element, let producerContinuations): - continuation.resume(returning: element) + case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): for producerContinuation in producerContinuations { switch producerContinuation { case .closure(let onProduceMore): - onProduceMore(.failure(CancellationError())) + onProduceMore(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) case .continuation(let continuation): - continuation.resume() + continuation.resume(throwing: MultiProducerSingleConsumerChannelAlreadyFinishedError()) } } - - case .resumeConsumerWithFailureAndCallOnTermination(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 .resumeConsumerWithNilAndCallOnTermination(let continuation, let onTermination): - continuation.resume(returning: nil) - onTermination?() - - case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): - for producerContinuation in producerContinuations { - switch producerContinuation { - case .closure(let onProduceMore): - onProduceMore(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) - case .continuation(let continuation): - continuation.resume(throwing: MultiProducerSingleConsumerChannelAlreadyFinishedError()) - } - } - onTermination?() - - case .none: - break - } } } } -/// The state machine of the channel. -@usableFromInline -struct _MultiProducerSingleConsumerStateMachine { - /// The state machine's current state. +extension MultiProducerSingleConsumerChannel._Storage { + /// The state machine of the channel. @usableFromInline - var _state: _MultiProducerSingleConsumerState + struct _StateMachine: ~Copyable { + /// The state machine's current state. + @usableFromInline + var _state: _State? - // The ID used for the next CallbackToken. - var _nextCallbackTokenID: UInt64 = 0 + // The ID used for the next CallbackToken. + @usableFromInline + var _nextCallbackTokenID: UInt64 = 0 - var _onTermination: (@Sendable () -> Void)? { - set { - switch self._state { - case .initial(var initial): - initial.onTermination = newValue - self._state = .initial(initial) + @usableFromInline + var _onTermination: (@Sendable () -> Void)? { + set { + switch self._state.take()! { + case .initial(var initial): + initial.onTermination = newValue - case .channeling(var channeling): - channeling.onTermination = newValue - self._state = .channeling(channeling) + self._state = .initial(initial) - case .sourceFinished(var sourceFinished): - sourceFinished.onTermination = newValue - self._state = .sourceFinished(sourceFinished) + case .channeling(var channeling): + channeling.onTermination = newValue + self._state = .channeling(channeling) - case .finished: - break + case .sourceFinished(var sourceFinished): + sourceFinished.onTermination = newValue + self._state = .sourceFinished(sourceFinished) - case .modify: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + case .finished: + break + } + } + get { + switch self._state { + case .initial(let initial): + return initial.onTermination + + case .channeling(let channeling): + return channeling.onTermination + + case .sourceFinished(let sourceFinished): + return sourceFinished.onTermination + + case .finished: + return nil + + case .none: + fatalError() + } } } - get { - switch self._state { + + /// Initializes a new `StateMachine`. + /// + /// We are passing and holding the backpressure strategy here because + /// it is a customizable extension of the state machine. + /// + /// - Parameter backpressureStrategy: The backpressure strategy. + init( + backpressureStrategy: MultiProducerSingleConsumerChannel._InternalBackpressureStrategy + ) { + self._state = .initial( + .init( + backpressureStrategy: backpressureStrategy, + iteratorInitialized: false, + onTermination: nil + ) + ) + } + + /// Generates the next callback token. + @usableFromInline + mutating func nextCallbackToken() -> UInt64 { + let id = self._nextCallbackTokenID + self._nextCallbackTokenID += 1 + return 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( + _TinyArray<_MultiProducerSingleConsumerSuspendedProducer>, + (@Sendable () -> Void)? + ) + } + + @inlinable + mutating func sequenceDeinitialized() -> SequenceDeinitializedAction? { + switch self._state.take()! { case .initial(let initial): - return initial.onTermination + guard initial.iteratorInitialized else { + // No iterator was created so we can transition to finished right away. + self._state = .finished(.init(iteratorInitialized: false, sourceFinished: false)) + + return .callOnTermination(initial.onTermination) + } + // An iterator was created and we deinited the sequence. + // This is an expected pattern and we just continue on normal. + self._state = .initial(initial) + + return .none case .channeling(let channeling): - return channeling.onTermination + guard channeling.iteratorInitialized else { + // No iterator was created so we can transition to finished right away. + self._state = .finished(.init(iteratorInitialized: false, sourceFinished: false)) + + return .failProducersAndCallOnTermination( + .init(channeling.suspendedProducers.lazy.map { $0.1 }), + channeling.onTermination + ) + } + // An iterator was created and we deinited the sequence. + // This is an expected pattern and we just continue on normal. + self._state = .channeling(channeling) + + return .none case .sourceFinished(let sourceFinished): - return sourceFinished.onTermination + guard sourceFinished.iteratorInitialized else { + // No iterator was created so we can transition to finished right away. + self._state = .finished(.init(iteratorInitialized: false, sourceFinished: true)) - case .finished: - return nil + return .callOnTermination(sourceFinished.onTermination) + } + // An iterator was created and we deinited the sequence. + // This is an expected pattern and we just continue on normal. + self._state = .sourceFinished(sourceFinished) - case .modify: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + return .none + + case .finished(let finished): + // We are already finished so there is nothing left to clean up. + // This is just the references dropping afterwards. + self._state = .finished(finished) + + return .none } } - } - - /// Initializes a new `StateMachine`. - /// - /// We are passing and holding the backpressure strategy here because - /// it is a customizable extension of the state machine. - /// - /// - Parameter backpressureStrategy: The backpressure strategy. - init( - backpressureStrategy: _MultiProducerSingleConsumerChannelInternalBackpressureStrategy - ) { - self._state = .initial( - .init( - backpressureStrategy: backpressureStrategy, - iteratorInitialized: false, - onTermination: nil - ) - ) - } - /// Generates the next callback token. - @usableFromInline - mutating func nextCallbackToken() -> UInt64 { - let id = self._nextCallbackTokenID - self._nextCallbackTokenID += 1 - return id - } + @inlinable + mutating func iteratorInitialized() { + switch self._state.take()! { + 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) + } - /// 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( - _TinyArray<_MultiProducerSingleConsumerSuspendedProducer>, - (@Sendable () -> Void)? - ) - } + case .channeling(var channeling): + if channeling.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. + channeling.iteratorInitialized = true + self._state = .channeling(channeling) + } - @inlinable - mutating func sequenceDeinitialized() -> SequenceDeinitializedAction? { - switch self._state { - case .initial(let initial): - guard initial.iteratorInitialized else { - // No iterator was created so we can transition to finished right away. - self._state = .finished(.init(iteratorInitialized: false, sourceFinished: false)) + 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) + } - return .callOnTermination(initial.onTermination) + case .finished(let finished): + if finished.iteratorInitialized { + // Our sequence is a unicast sequence and does not support multiple AsyncIterator's + fatalError("Only a single AsyncIterator can be created") + } else { + self._state = .finished(.init(iteratorInitialized: true, sourceFinished: finished.sourceFinished)) + } } - // An iterator was created and we deinited the sequence. - // This is an expected pattern and we just continue on normal. - return .none + } - case .channeling(let channeling): - guard channeling.iteratorInitialized else { - // No iterator was created so we can transition to finished right away. - self._state = .finished(.init(iteratorInitialized: false, sourceFinished: false)) + /// 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( + _TinyArray<_MultiProducerSingleConsumerSuspendedProducer>, + (@Sendable () -> Void)? + ) + } - return .failProducersAndCallOnTermination( - .init(channeling.suspendedProducers.lazy.map { $0.1 }), - channeling.onTermination - ) - } - // An iterator was created and we deinited the sequence. - // This is an expected pattern and we just continue on normal. - return .none + @inlinable + mutating func iteratorDeinitialized() -> IteratorDeinitializedAction? { + switch self._state.take()! { + 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(.init(iteratorInitialized: true, sourceFinished: false)) + + return .callOnTermination(initial.onTermination) + } else { + // An iterator needs to be initialized before it can be deinitialized. + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + } - case .sourceFinished(let sourceFinished): - guard sourceFinished.iteratorInitialized else { - // No iterator was created so we can transition to finished right away. - self._state = .finished(.init(iteratorInitialized: false, sourceFinished: true)) + case .channeling(let channeling): + if channeling.iteratorInitialized { + // An iterator was created and deinited. Since we only support + // a single iterator we can now transition to finish. + self._state = .finished(.init(iteratorInitialized: true, sourceFinished: false)) + + return .failProducersAndCallOnTermination( + .init(channeling.suspendedProducers.lazy.map { $0.1 }), + channeling.onTermination + ) + } else { + // An iterator needs to be initialized before it can be deinitialized. + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + } - return .callOnTermination(sourceFinished.onTermination) - } - // An iterator was created and we deinited the sequence. - // This is an expected pattern and we just continue on normal. - return .none + 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(.init(iteratorInitialized: true, sourceFinished: true)) + + return .callOnTermination(sourceFinished.onTermination) + } else { + // An iterator needs to be initialized before it can be deinitialized. + fatalError("MultiProducerSingleConsumerChannel 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 .finished(let finished): + // We are already finished so there is nothing left to clean up. + // This is just the references dropping afterwards. + self._state = .finished(finished) - case .modify: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + return .none + } } - } - @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) - } + /// 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( + _TinyArray<(Result) -> Void>, + (@Sendable () -> Void)? + ) + /// Indicates that all producers should be failed. + case failProducers(_TinyArray<(Result) -> Void>) + } - case .channeling(var channeling): - if channeling.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. - channeling.iteratorInitialized = true + @inlinable + mutating func sourceDeinitialized() -> SourceDeinitializedAction? { + switch self._state.take()! { + case .initial: + fatalError("The channel's source hasn't been finished but deinited") + + case .channeling(let channeling): self._state = .channeling(channeling) - } - 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) - } + return nil - case .finished(let finished): - if finished.iteratorInitialized { - // Our sequence is a unicast sequence and does not support multiple AsyncIterator's - fatalError("Only a single AsyncIterator can be created") - } else { - self._state = .finished(.init(iteratorInitialized: true, sourceFinished: finished.sourceFinished)) - } + case .sourceFinished(let sourceFinished): + // This is the expected case where finish was called and then the source deinited + self._state = .sourceFinished(sourceFinished) - case .modify: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - } - } + return .none - /// 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( - _TinyArray<_MultiProducerSingleConsumerSuspendedProducer>, - (@Sendable () -> Void)? - ) - } + case .finished(let finished): + if finished.sourceFinished { + // The source already got finished so this is fine. + self._state = .finished(finished) - @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(.init(iteratorInitialized: true, sourceFinished: false)) - return .callOnTermination(initial.onTermination) - } else { - // An iterator needs to be initialized before it can be deinitialized. - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + return .none + } else { + fatalError("The channel's source hasn't been finished but deinited") + } } + } - case .channeling(let channeling): - if channeling.iteratorInitialized { - // An iterator was created and deinited. Since we only support - // a single iterator we can now transition to finish. - self._state = .finished(.init(iteratorInitialized: true, sourceFinished: false)) + /// Actions returned by `send()`. + @usableFromInline + enum SendAction { + /// 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: UInt64 + ) + /// Indicates that the consumer should be resumed and the producer should be notified to produce more. + case resumeConsumerAndReturnProduceMore( + continuation: UnsafeContinuation, + element: Element + ) + /// Indicates that the consumer should be resumed and the producer should be suspended. + case resumeConsumerAndReturnEnqueue( + continuation: UnsafeContinuation, + element: Element, + callbackToken: UInt64 + ) + /// Indicates that the producer has been finished. + case throwFinishedError + + @inlinable + init( + callbackToken: UInt64?, + continuationAndElement: (UnsafeContinuation, 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 + ) - return .failProducersAndCallOnTermination( - .init(channeling.suspendedProducers.lazy.map { $0.1 }), - channeling.onTermination - ) - } else { - // An iterator needs to be initialized before it can be deinitialized. - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + case (.some(let callbackToken), .some((let continuation, let element))): + self = .resumeConsumerAndReturnEnqueue( + continuation: continuation, + element: element, + callbackToken: callbackToken + ) + } } + } - 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(.init(iteratorInitialized: true, sourceFinished: true)) - return .callOnTermination(sourceFinished.onTermination) - } else { - // An iterator needs to be initialized before it can be deinitialized. - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - } + @inlinable + mutating func send(_ sequence: some Sequence) -> SendAction { + switch self._state.take()! { + case .initial(var initial): + var buffer = Deque() + buffer.append(contentsOf: sequence) - case .finished: - // We are already finished so there is nothing left to clean up. - // This is just the references dropping afterwards. - return .none + let shouldProduceMore = initial.backpressureStrategy.didSend(elements: buffer[...]) + let callbackToken = shouldProduceMore ? nil : self.nextCallbackToken() - case .modify: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - } - } + self._state = .channeling( + .init( + backpressureStrategy: initial.backpressureStrategy, + iteratorInitialized: initial.iteratorInitialized, + onTermination: initial.onTermination, + buffer: buffer, + consumerContinuation: nil, + producerContinuations: .init(), + cancelledAsyncProducers: .init(), + hasOutstandingDemand: shouldProduceMore + ) + ) - /// 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( - _TinyArray<(Result) -> Void>, - (@Sendable () -> Void)? - ) - /// Indicates that all producers should be failed. - case failProducers(_TinyArray<(Result) -> Void>) - } + return .init(callbackToken: callbackToken) - @inlinable - mutating func sourceDeinitialized() -> SourceDeinitializedAction? { - switch self._state { - case .initial, .channeling: - fatalError("The channel's source hasn't been finished but deinited") - - case .sourceFinished: - // This is the expected case where finish was called and then the source deinited - return .none - - case .finished(let finished): - if finished.sourceFinished { - // The source already got finished so this is fine. - return .none - } else { - fatalError("The channel's source hasn't been finished but deinited") - } + case .channeling(var channeling): + // We have an element and can resume the continuation + let bufferEndIndexBeforeAppend = channeling.buffer.endIndex + channeling.buffer.append(contentsOf: sequence) + var shouldProduceMore = channeling.backpressureStrategy.didSend( + elements: channeling.buffer[bufferEndIndexBeforeAppend...] + ) + channeling.hasOutstandingDemand = shouldProduceMore - case .modify: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - } - } + guard let consumerContinuation = channeling.consumerContinuation else { + // We don't have a suspended consumer so we just buffer the elements + self._state = .channeling(channeling) - /// Actions returned by `send()`. - @usableFromInline - enum SendAction { - /// 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: UInt64 - ) - /// Indicates that the consumer should be resumed and the producer should be notified to produce more. - case resumeConsumerAndReturnProduceMore( - continuation: UnsafeContinuation, - element: Element - ) - /// Indicates that the consumer should be resumed and the producer should be suspended. - case resumeConsumerAndReturnEnqueue( - continuation: UnsafeContinuation, - element: Element, - callbackToken: UInt64 - ) - /// Indicates that the producer has been finished. - case throwFinishedError + return .init( + callbackToken: shouldProduceMore ? nil : self.nextCallbackToken() + ) + } + guard let element = channeling.buffer.popFirst() else { + // We got a send of an empty sequence. We just tolerate this. + self._state = .channeling(channeling) - @inlinable - init( - callbackToken: UInt64?, - continuationAndElement: (UnsafeContinuation, Element)? = nil - ) { - switch (callbackToken, continuationAndElement) { - case (.none, .none): - self = .returnProduceMore + return .init(callbackToken: shouldProduceMore ? nil : self.nextCallbackToken()) + } + // We need to tell the back pressure strategy that we consumed + shouldProduceMore = channeling.backpressureStrategy.didConsume(element: element) + channeling.hasOutstandingDemand = shouldProduceMore - case (.some(let callbackToken), .none): - self = .returnEnqueue(callbackToken: callbackToken) + // We got a consumer continuation and an element. We can resume the consumer now + channeling.consumerContinuation = nil + self._state = .channeling(channeling) - case (.none, .some((let continuation, let element))): - self = .resumeConsumerAndReturnProduceMore( - continuation: continuation, - element: element + return .init( + callbackToken: shouldProduceMore ? nil : self.nextCallbackToken(), + continuationAndElement: (consumerContinuation, element) ) - case (.some(let callbackToken), .some((let continuation, let element))): - self = .resumeConsumerAndReturnEnqueue( - continuation: continuation, - element: element, - callbackToken: callbackToken - ) + case .sourceFinished(let sourceFinished): + // If the source has finished we are dropping the elements. + self._state = .sourceFinished(sourceFinished) + + return .throwFinishedError + + case .finished(let finished): + // If the source has finished we are dropping the elements. + self._state = .finished(finished) + + return .throwFinishedError } } - } - @inlinable - mutating func send(_ sequence: some Sequence) -> SendAction { - switch self._state { - case .initial(var initial): - var buffer = Deque() - buffer.append(contentsOf: sequence) + /// 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, Error) + } - let shouldProduceMore = initial.backpressureStrategy.didSend(elements: buffer[...]) - let callbackToken = shouldProduceMore ? nil : self.nextCallbackToken() + @inlinable + mutating func enqueueProducer( + callbackToken: UInt64, + onProduceMore: sending @escaping (Result) -> Void + ) -> EnqueueProducerAction? { + switch self._state.take()! { + case .initial: + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - self._state = .channeling( - .init( - backpressureStrategy: initial.backpressureStrategy, - iteratorInitialized: initial.iteratorInitialized, - onTermination: initial.onTermination, - buffer: buffer, - consumerContinuation: nil, - producerContinuations: .init(), - cancelledAsyncProducers: .init(), - hasOutstandingDemand: shouldProduceMore - ) - ) + case .channeling(var channeling): + if let index = channeling.cancelledAsyncProducers.firstIndex(of: callbackToken) { + // Our producer got marked as cancelled. + channeling.cancelledAsyncProducers.remove(at: index) + self._state = .channeling(channeling) + + return .resumeProducerWithError(onProduceMore, CancellationError()) + } else if channeling.hasOutstandingDemand { + // We hit an edge case here where we wrote but the consuming thread got interleaved + self._state = .channeling(channeling) + + return .resumeProducer(onProduceMore) + } else { + channeling.suspendedProducers.append((callbackToken, .closure(onProduceMore))) + self._state = .channeling(channeling) + + return .none + } - return .init(callbackToken: callbackToken) + case .sourceFinished(let sourceFinished): + // Since we are unlocking between sending elements and suspending the send + // It can happen that the source got finished or the consumption fully finishes. + self._state = .sourceFinished(sourceFinished) - case .channeling(var channeling): - self._state = .modify + return .resumeProducerWithError(onProduceMore, MultiProducerSingleConsumerChannelAlreadyFinishedError()) - // We have an element and can resume the continuation - let bufferEndIndexBeforeAppend = channeling.buffer.endIndex - channeling.buffer.append(contentsOf: sequence) - var shouldProduceMore = channeling.backpressureStrategy.didSend( - elements: channeling.buffer[bufferEndIndexBeforeAppend...] - ) - channeling.hasOutstandingDemand = shouldProduceMore + case .finished(let finished): + // Since we are unlocking between sending elements and suspending the send + // It can happen that the source got finished or the consumption fully finishes. + self._state = .finished(finished) - guard let consumerContinuation = channeling.consumerContinuation else { - // We don't have a suspended consumer so we just buffer the elements - self._state = .channeling(channeling) - return .init( - callbackToken: shouldProduceMore ? nil : self.nextCallbackToken() - ) + return .resumeProducerWithError(onProduceMore, MultiProducerSingleConsumerChannelAlreadyFinishedError()) } - guard let element = channeling.buffer.popFirst() else { - // We got a send of an empty sequence. We just tolerate this. - self._state = .channeling(channeling) - return .init(callbackToken: shouldProduceMore ? nil : self.nextCallbackToken()) - } - // We need to tell the back pressure strategy that we consumed - shouldProduceMore = channeling.backpressureStrategy.didConsume(element: element) - channeling.hasOutstandingDemand = shouldProduceMore - - // We got a consumer continuation and an element. We can resume the consumer now - channeling.consumerContinuation = nil - self._state = .channeling(channeling) - return .init( - callbackToken: shouldProduceMore ? nil : self.nextCallbackToken(), - continuationAndElement: (consumerContinuation, element) - ) - - case .sourceFinished, .finished: - // If the source has finished we are dropping the elements. - return .throwFinishedError + } - case .modify: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + /// Actions returned by `enqueueContinuation()`. + @usableFromInline + enum EnqueueContinuationAction { + /// Indicates that the producer should be notified to produce more. + case resumeProducer(UnsafeContinuation) + /// Indicates that the producer should be notified about an error. + case resumeProducerWithError(UnsafeContinuation, Error) } - } - /// 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, Error) - } + @inlinable + mutating func enqueueContinuation( + callbackToken: UInt64, + continuation: UnsafeContinuation + ) -> EnqueueContinuationAction? { + switch self._state.take()! { + case .initial: + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - @inlinable - mutating func enqueueProducer( - callbackToken: UInt64, - onProduceMore: sending @escaping (Result) -> Void - ) -> EnqueueProducerAction? { - switch self._state { - case .initial: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - - case .channeling(var channeling): - if let index = channeling.cancelledAsyncProducers.firstIndex(of: callbackToken) { - // Our producer got marked as cancelled. - self._state = .modify - channeling.cancelledAsyncProducers.remove(at: index) - self._state = .channeling(channeling) + case .channeling(var channeling): + if let index = channeling.cancelledAsyncProducers.firstIndex(of: callbackToken) { + // Our producer got marked as cancelled. + channeling.cancelledAsyncProducers.remove(at: index) + self._state = .channeling(channeling) + + return .resumeProducerWithError(continuation, CancellationError()) + } else if channeling.hasOutstandingDemand { + // We hit an edge case here where we wrote but the consuming thread got interleaved + self._state = .channeling(channeling) + + return .resumeProducer(continuation) + } else { + channeling.suspendedProducers.append((callbackToken, .continuation(continuation))) + self._state = .channeling(channeling) + + return .none + } - return .resumeProducerWithError(onProduceMore, CancellationError()) - } else if channeling.hasOutstandingDemand { - // We hit an edge case here where we wrote but the consuming thread got interleaved - return .resumeProducer(onProduceMore) - } else { - self._state = .modify - channeling.suspendedProducers.append((callbackToken, .closure(onProduceMore))) + case .sourceFinished(let sourceFinished): + // Since we are unlocking between sending elements and suspending the send + // It can happen that the source got finished or the consumption fully finishes. + self._state = .sourceFinished(sourceFinished) - self._state = .channeling(channeling) - return .none - } + return .resumeProducerWithError(continuation, MultiProducerSingleConsumerChannelAlreadyFinishedError()) - case .sourceFinished, .finished: - // Since we are unlocking between sending elements and suspending the send - // It can happen that the source got finished or the consumption fully finishes. - return .resumeProducerWithError(onProduceMore, MultiProducerSingleConsumerChannelAlreadyFinishedError()) + case .finished(let finished): + // Since we are unlocking between sending elements and suspending the send + // It can happen that the source got finished or the consumption fully finishes. + self._state = .finished(finished) - case .modify: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + return .resumeProducerWithError(continuation, MultiProducerSingleConsumerChannelAlreadyFinishedError()) + } } - } - /// Actions returned by `enqueueContinuation()`. - @usableFromInline - enum EnqueueContinuationAction { - /// Indicates that the producer should be notified to produce more. - case resumeProducer(UnsafeContinuation) - /// Indicates that the producer should be notified about an error. - case resumeProducerWithError(UnsafeContinuation, Error) - } + /// Actions returned by `cancelProducer()`. + @usableFromInline + enum CancelProducerAction { + /// Indicates that the producer should be notified about cancellation. + case resumeProducerWithCancellationError(_MultiProducerSingleConsumerSuspendedProducer) + } - @inlinable - mutating func enqueueContinuation( - callbackToken: UInt64, - continuation: UnsafeContinuation - ) -> EnqueueContinuationAction? { - switch self._state { - case .initial: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - - case .channeling(var channeling): - if let index = channeling.cancelledAsyncProducers.firstIndex(of: callbackToken) { - // Our producer got marked as cancelled. - self._state = .modify - channeling.cancelledAsyncProducers.remove(at: index) - self._state = .channeling(channeling) + @inlinable + mutating func cancelProducer( + callbackToken: UInt64 + ) -> CancelProducerAction? { + //print(#function, self._state.description) + switch self._state.take()! { + case .initial: + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - return .resumeProducerWithError(continuation, CancellationError()) - } else if channeling.hasOutstandingDemand { - // We hit an edge case here where we wrote but the consuming thread got interleaved - return .resumeProducer(continuation) - } else { - self._state = .modify - channeling.suspendedProducers.append((callbackToken, .continuation(continuation))) + case .channeling(var channeling): + guard let index = channeling.suspendedProducers.firstIndex(where: { $0.0 == callbackToken }) else { + // The task that sends was cancelled before sending elements so the cancellation handler + // got invoked right away + channeling.cancelledAsyncProducers.append(callbackToken) + self._state = .channeling(channeling) + return .none + } + // We have an enqueued producer that we need to resume now + let continuation = channeling.suspendedProducers.remove(at: index).1 self._state = .channeling(channeling) - return .none - } - case .sourceFinished, .finished: - // Since we are unlocking between sending elements and suspending the send - // It can happen that the source got finished or the consumption fully finishes. - return .resumeProducerWithError(continuation, MultiProducerSingleConsumerChannelAlreadyFinishedError()) + return .resumeProducerWithCancellationError(continuation) - case .modify: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - } - } + case .sourceFinished(let sourceFinished): + // Since we are unlocking between sending elements and suspending the send + // It can happen that the source got finished or the consumption fully finishes. + self._state = .sourceFinished(sourceFinished) - /// Actions returned by `cancelProducer()`. - @usableFromInline - enum CancelProducerAction { - /// Indicates that the producer should be notified about cancellation. - case resumeProducerWithCancellationError(_MultiProducerSingleConsumerSuspendedProducer) - } + return .none - @inlinable - mutating func cancelProducer( - callbackToken: UInt64 - ) -> CancelProducerAction? { - //print(#function, self._state.description) - switch self._state { - case .initial: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - - case .channeling(var channeling): - guard let index = channeling.suspendedProducers.firstIndex(where: { $0.0 == callbackToken }) else { - // The task that sends was cancelled before sending elements so the cancellation handler - // got invoked right away - self._state = .modify - channeling.cancelledAsyncProducers.append(callbackToken) - self._state = .channeling(channeling) + case .finished(let finished): + // Since we are unlocking between sending elements and suspending the send + // It can happen that the source got finished or the consumption fully finishes. + self._state = .finished(finished) return .none } - // We have an enqueued producer that we need to resume now - self._state = .modify - let continuation = channeling.suspendedProducers.remove(at: index).1 - self._state = .channeling(channeling) - - return .resumeProducerWithCancellationError(continuation) - - case .sourceFinished, .finished: - // Since we are unlocking between sending elements and suspending the send - // It can happen that the source got finished or the consumption fully finishes. - return .none - - case .modify: - fatalError("MultiProducerSingleConsumerChannel 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: UnsafeContinuation, - failure: Failure?, - onTermination: (() -> Void)? - ) - /// Indicates that the producers should be resumed with an error. - case resumeProducers( - producerContinuations: _TinyArray<_MultiProducerSingleConsumerSuspendedProducer> - ) - } - - @inlinable - mutating func finish(_ failure: Failure?) -> FinishAction? { - switch self._state { - case .initial(let initial): - // Nothing was sent 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 - ) + /// 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: UnsafeContinuation, + failure: Failure?, + onTermination: (() -> Void)? ) + /// Indicates that the producers should be resumed with an error. + case resumeProducers( + producerContinuations: _TinyArray<_MultiProducerSingleConsumerSuspendedProducer> + ) + } - return .callOnTermination(initial.onTermination) - - case .channeling(let channeling): - guard let consumerContinuation = channeling.consumerContinuation else { - // We don't have a suspended consumer so we are just going to mark - // the source as finished and terminate the current suspended producers. + @inlinable + mutating func finish(_ failure: Failure?) -> FinishAction? { + switch self._state.take()! { + case .initial(let initial): + // Nothing was sent nor did anybody call next + // This means we can transition to sourceFinished and store the failure self._state = .sourceFinished( .init( - iteratorInitialized: channeling.iteratorInitialized, - buffer: channeling.buffer, + iteratorInitialized: initial.iteratorInitialized, + buffer: .init(), failure: failure, - onTermination: channeling.onTermination + onTermination: initial.onTermination ) ) - return .resumeProducers(producerContinuations: .init(channeling.suspendedProducers.lazy.map { $0.1 })) - } - // 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(channeling.buffer.isEmpty, "Expected an empty buffer") - - self._state = .finished(.init(iteratorInitialized: channeling.iteratorInitialized, sourceFinished: true)) + return .callOnTermination(initial.onTermination) - return .resumeConsumerAndCallOnTermination( - consumerContinuation: consumerContinuation, - failure: failure, - onTermination: channeling.onTermination - ) + case .channeling(let channeling): + guard let consumerContinuation = channeling.consumerContinuation else { + // We don't have a suspended consumer so we are just going to mark + // the source as finished and terminate the current suspended producers. + self._state = .sourceFinished( + .init( + iteratorInitialized: channeling.iteratorInitialized, + buffer: channeling.buffer, + failure: failure, + onTermination: channeling.onTermination + ) + ) - case .sourceFinished: - // If the source has finished, finishing again has no effect. - return .none + return .resumeProducers(producerContinuations: .init(channeling.suspendedProducers.lazy.map { $0.1 })) + } + // 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(channeling.buffer.isEmpty, "Expected an empty buffer") - case .finished(var finished): - finished.sourceFinished = true - self._state = .finished(finished) - return .none + self._state = .finished(.init(iteratorInitialized: channeling.iteratorInitialized, sourceFinished: true)) - case .modify: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - } - } + return .resumeConsumerAndCallOnTermination( + consumerContinuation: consumerContinuation, + failure: failure, + onTermination: channeling.onTermination + ) - /// 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, _TinyArray<_MultiProducerSingleConsumerSuspendedProducer>) - /// Indicates that the `Failure` should be returned to the caller and that `onTermination` should be called. - case returnFailureAndCallOnTermination(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 - } + case .sourceFinished(let sourceFinished): + // If the source has finished, finishing again has no effect. + self._state = .sourceFinished(sourceFinished) - @inlinable - mutating func next() -> NextAction { - switch self._state { - case .initial(let initial): - // We are not interacting with the backpressure strategy here because - // we are doing this inside `suspendNext` - self._state = .channeling( - .init( - backpressureStrategy: initial.backpressureStrategy, - iteratorInitialized: initial.iteratorInitialized, - onTermination: initial.onTermination, - buffer: Deque(), - consumerContinuation: nil, - producerContinuations: .init(), - cancelledAsyncProducers: .init(), - hasOutstandingDemand: false - ) - ) + return .none - return .suspendTask - case .channeling(var channeling): - guard channeling.consumerContinuation == nil else { - // We have multiple AsyncIterators iterating the sequence - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + case .finished(var finished): + finished.sourceFinished = true + self._state = .finished(finished) + return .none } + } - self._state = .modify + /// 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, _TinyArray<_MultiProducerSingleConsumerSuspendedProducer>) + /// Indicates that the `Failure` should be returned to the caller and that `onTermination` should be called. + case returnFailureAndCallOnTermination(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 + } - guard let element = channeling.buffer.popFirst() else { - // There is nothing in the buffer to fulfil the demand so we need to suspend. + @inlinable + mutating func next() -> NextAction { + switch self._state.take()! { + case .initial(let initial): // We are not interacting with the backpressure strategy here because // we are doing this inside `suspendNext` - self._state = .channeling(channeling) + self._state = .channeling( + .init( + backpressureStrategy: initial.backpressureStrategy, + iteratorInitialized: initial.iteratorInitialized, + onTermination: initial.onTermination, + buffer: Deque(), + consumerContinuation: nil, + producerContinuations: .init(), + cancelledAsyncProducers: .init(), + hasOutstandingDemand: false + ) + ) return .suspendTask - } - // We have an element to fulfil the demand right away. - let shouldProduceMore = channeling.backpressureStrategy.didConsume(element: element) - channeling.hasOutstandingDemand = shouldProduceMore + case .channeling(var channeling): + guard channeling.consumerContinuation == nil else { + // We have multiple AsyncIterators iterating the sequence + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + } + + guard let element = channeling.buffer.popFirst() else { + // There is nothing in the buffer to fulfil the demand so we need to suspend. + // We are not interacting with the backpressure strategy here because + // we are doing this inside `suspendNext` + self._state = .channeling(channeling) - guard shouldProduceMore else { - // We don't have any new demand, so we can just return the element. + return .suspendTask + } + // We have an element to fulfil the demand right away. + let shouldProduceMore = channeling.backpressureStrategy.didConsume(element: element) + channeling.hasOutstandingDemand = shouldProduceMore + + guard shouldProduceMore else { + // We don't have any new demand, so we can just return the element. + self._state = .channeling(channeling) + + return .returnElement(element) + } + // There is demand and we have to resume our producers + let producers = _TinyArray(channeling.suspendedProducers.lazy.map { $0.1 }) + channeling.suspendedProducers.removeAll(keepingCapacity: true) self._state = .channeling(channeling) - return .returnElement(element) - } - // There is demand and we have to resume our producers - let producers = _TinyArray(channeling.suspendedProducers.lazy.map { $0.1 }) - channeling.suspendedProducers.removeAll(keepingCapacity: true) - self._state = .channeling(channeling) - return .returnElementAndResumeProducers(element, producers) - case .sourceFinished(var sourceFinished): - // Check if we have an element left in the buffer and return it - self._state = .modify + return .returnElementAndResumeProducers(element, producers) - guard let element = sourceFinished.buffer.popFirst() else { - // We are returning the queued failure now and can transition to finished - self._state = .finished(.init(iteratorInitialized: sourceFinished.iteratorInitialized, sourceFinished: true)) + case .sourceFinished(var sourceFinished): + // Check if we have an element left in the buffer and return it + guard let element = sourceFinished.buffer.popFirst() else { + // We are returning the queued failure now and can transition to finished + self._state = .finished(.init(iteratorInitialized: sourceFinished.iteratorInitialized, sourceFinished: true)) - return .returnFailureAndCallOnTermination(sourceFinished.failure, sourceFinished.onTermination) - } - self._state = .sourceFinished(sourceFinished) + return .returnFailureAndCallOnTermination(sourceFinished.failure, sourceFinished.onTermination) + } + self._state = .sourceFinished(sourceFinished) - return .returnElement(element) + return .returnElement(element) - case .finished: - return .returnNil + case .finished(let finished): + self._state = .finished(finished) - case .modify: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + return .returnNil + } } - } - /// Actions returned by `suspendNext()`. - @usableFromInline - enum SuspendNextAction { - /// Indicates that the consumer should be resumed. - case resumeConsumerWithElement(UnsafeContinuation, Element) - /// Indicates that the consumer and all producers should be resumed. - case resumeConsumerWithElementAndProducers( - UnsafeContinuation, - Element, - _TinyArray<_MultiProducerSingleConsumerSuspendedProducer> - ) - /// Indicates that the consumer should be resumed with the failure and that `onTermination` should be called. - case resumeConsumerWithFailureAndCallOnTermination( - UnsafeContinuation, - Failure?, - (() -> Void)? - ) - /// Indicates that the consumer should be resumed with `nil`. - case resumeConsumerWithNil(UnsafeContinuation) - } + /// Actions returned by `suspendNext()`. + @usableFromInline + enum SuspendNextAction { + /// Indicates that the consumer should be resumed. + case resumeConsumerWithElement(UnsafeContinuation, Element) + /// Indicates that the consumer and all producers should be resumed. + case resumeConsumerWithElementAndProducers( + UnsafeContinuation, + Element, + _TinyArray<_MultiProducerSingleConsumerSuspendedProducer> + ) + /// Indicates that the consumer should be resumed with the failure and that `onTermination` should be called. + case resumeConsumerWithFailureAndCallOnTermination( + UnsafeContinuation, + Failure?, + (() -> Void)? + ) + /// Indicates that the consumer should be resumed with `nil`. + case resumeConsumerWithNil(UnsafeContinuation) + } - @inlinable - mutating func suspendNext(continuation: UnsafeContinuation) -> SuspendNextAction? { - switch self._state { - case .initial: - preconditionFailure("MultiProducerSingleConsumerChannel internal inconsistency") - - case .channeling(var channeling): - guard channeling.consumerContinuation == nil else { - // We have multiple AsyncIterators iterating the sequence - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - } + @inlinable + mutating func suspendNext(continuation: UnsafeContinuation) -> SuspendNextAction? { + switch self._state.take()! { + case .initial: + preconditionFailure("MultiProducerSingleConsumerChannel internal inconsistency") - self._state = .modify + case .channeling(var channeling): + guard channeling.consumerContinuation == nil else { + // We have multiple AsyncIterators iterating the sequence + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + } - // We have to check here again since we might have a producer interleave next and suspendNext - guard let element = channeling.buffer.popFirst() else { - // There is nothing in the buffer to fulfil the demand so we to store the continuation. - channeling.consumerContinuation = continuation - self._state = .channeling(channeling) + // We have to check here again since we might have a producer interleave next and suspendNext + guard let element = channeling.buffer.popFirst() else { + // There is nothing in the buffer to fulfil the demand so we to store the continuation. + channeling.consumerContinuation = continuation + self._state = .channeling(channeling) - return .none - } - // We have an element to fulfil the demand right away. + return .none + } + // We have an element to fulfil the demand right away. + + let shouldProduceMore = channeling.backpressureStrategy.didConsume(element: element) + channeling.hasOutstandingDemand = shouldProduceMore - let shouldProduceMore = channeling.backpressureStrategy.didConsume(element: element) - channeling.hasOutstandingDemand = shouldProduceMore + guard shouldProduceMore else { + // We don't have any new demand, so we can just return the element. + self._state = .channeling(channeling) - guard shouldProduceMore else { - // We don't have any new demand, so we can just return the element. + return .resumeConsumerWithElement(continuation, element) + } + // There is demand and we have to resume our producers + let producers = _TinyArray(channeling.suspendedProducers.lazy.map { $0.1 }) + channeling.suspendedProducers.removeAll(keepingCapacity: true) self._state = .channeling(channeling) - return .resumeConsumerWithElement(continuation, element) - } - // There is demand and we have to resume our producers - let producers = _TinyArray(channeling.suspendedProducers.lazy.map { $0.1 }) - channeling.suspendedProducers.removeAll(keepingCapacity: true) - self._state = .channeling(channeling) - return .resumeConsumerWithElementAndProducers(continuation, element, producers) - - case .sourceFinished(var sourceFinished): - // Check if we have an element left in the buffer and return it - self._state = .modify - - guard let element = sourceFinished.buffer.popFirst() else { - // We are returning the queued failure now and can transition to finished - self._state = .finished(.init(iteratorInitialized: sourceFinished.iteratorInitialized, sourceFinished: true)) - - return .resumeConsumerWithFailureAndCallOnTermination( - continuation, - sourceFinished.failure, - sourceFinished.onTermination - ) - } - self._state = .sourceFinished(sourceFinished) - return .resumeConsumerWithElement(continuation, element) + return .resumeConsumerWithElementAndProducers(continuation, element, producers) + + case .sourceFinished(var sourceFinished): + // Check if we have an element left in the buffer and return it + guard let element = sourceFinished.buffer.popFirst() else { + // We are returning the queued failure now and can transition to finished + self._state = .finished(.init(iteratorInitialized: sourceFinished.iteratorInitialized, sourceFinished: true)) + + return .resumeConsumerWithFailureAndCallOnTermination( + continuation, + sourceFinished.failure, + sourceFinished.onTermination + ) + } + self._state = .sourceFinished(sourceFinished) - case .finished: - return .resumeConsumerWithNil(continuation) + return .resumeConsumerWithElement(continuation, element) + + case .finished(let finished): + self._state = .finished(finished) - case .modify: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + return .resumeConsumerWithNil(continuation) + } } - } - /// Actions returned by `cancelNext()`. - @usableFromInline - enum CancelNextAction { - /// Indicates that the continuation should be resumed with nil, the producers should be finished and call onTermination. - case resumeConsumerWithNilAndCallOnTermination(UnsafeContinuation, (() -> Void)?) - /// Indicates that the producers should be finished and call onTermination. - case failProducersAndCallOnTermination(_TinyArray<_MultiProducerSingleConsumerSuspendedProducer>, (() -> Void)?) - } + /// Actions returned by `cancelNext()`. + @usableFromInline + enum CancelNextAction { + /// Indicates that the continuation should be resumed with nil, the producers should be finished and call onTermination. + case resumeConsumerWithNilAndCallOnTermination(UnsafeContinuation, (() -> Void)?) + /// Indicates that the producers should be finished and call onTermination. + case failProducersAndCallOnTermination(_TinyArray<_MultiProducerSingleConsumerSuspendedProducer>, (() -> Void)?) + } - @inlinable - mutating func cancelNext() -> CancelNextAction? { - switch self._state { - case .initial: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + @inlinable + mutating func cancelNext() -> CancelNextAction? { + switch self._state.take()! { + case .initial: + fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - case .channeling(let channeling): - self._state = .finished(.init(iteratorInitialized: channeling.iteratorInitialized, sourceFinished: false)) + case .channeling(let channeling): + self._state = .finished(.init(iteratorInitialized: channeling.iteratorInitialized, sourceFinished: false)) - guard let consumerContinuation = channeling.consumerContinuation else { - return .failProducersAndCallOnTermination( - .init(channeling.suspendedProducers.lazy.map { $0.1 }), + guard let consumerContinuation = channeling.consumerContinuation else { + return .failProducersAndCallOnTermination( + .init(channeling.suspendedProducers.lazy.map { $0.1 }), + channeling.onTermination + ) + } + precondition( + channeling.suspendedProducers.isEmpty, + "Internal inconsistency. Unexpected producer continuations." + ) + return .resumeConsumerWithNilAndCallOnTermination( + consumerContinuation, channeling.onTermination ) - } - precondition( - channeling.suspendedProducers.isEmpty, - "Internal inconsistency. Unexpected producer continuations." - ) - return .resumeConsumerWithNilAndCallOnTermination( - consumerContinuation, - channeling.onTermination - ) - case .sourceFinished, .finished: - return .none + case .sourceFinished(let sourceFinished): + self._state = .sourceFinished(sourceFinished) + + return .none + + case .finished(let finished): + self._state = .finished(finished) - case .modify: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") + return .none + } } } } -@usableFromInline -enum _MultiProducerSingleConsumerState: CustomStringConvertible { +extension MultiProducerSingleConsumerChannel._Storage._StateMachine { @usableFromInline - struct Initial: CustomStringConvertible { - /// The backpressure strategy. - @usableFromInline - var backpressureStrategy: _MultiProducerSingleConsumerChannelInternalBackpressureStrategy - - /// Indicates if the iterator was initialized. + enum _State: ~Copyable { @usableFromInline - var iteratorInitialized: Bool - - /// The onTermination callback. - @usableFromInline - var onTermination: (@Sendable () -> Void)? - - @usableFromInline - var description: String { - "backpressure:\(self.backpressureStrategy.description) iteratorInitialized:\(self.iteratorInitialized)" + struct Initial: ~Copyable { + /// The backpressure strategy. + @usableFromInline + var backpressureStrategy: MultiProducerSingleConsumerChannel._InternalBackpressureStrategy + + /// Indicates if the iterator was initialized. + @usableFromInline + var iteratorInitialized: Bool + + /// The onTermination callback. + @usableFromInline + var onTermination: (@Sendable () -> Void)? + + @usableFromInline + var description: String { + "backpressure:\(self.backpressureStrategy.description) iteratorInitialized:\(self.iteratorInitialized)" + } } - } - @usableFromInline - struct Channeling { - /// The backpressure strategy. @usableFromInline - var backpressureStrategy: _MultiProducerSingleConsumerChannelInternalBackpressureStrategy + struct Channeling: ~Copyable { + /// The backpressure strategy. + @usableFromInline + var backpressureStrategy: MultiProducerSingleConsumerChannel._InternalBackpressureStrategy - /// Indicates if the iterator was initialized. - @usableFromInline - var iteratorInitialized: Bool + /// Indicates if the iterator was initialized. + @usableFromInline + var iteratorInitialized: Bool - /// The onTermination callback. - @usableFromInline - var onTermination: (@Sendable () -> Void)? + /// The onTermination callback. + @usableFromInline + var onTermination: (@Sendable () -> Void)? - /// The buffer of elements. - @usableFromInline - var buffer: Deque + /// The buffer of elements. + @usableFromInline + var buffer: Deque - /// The optional consumer continuation. - @usableFromInline - var consumerContinuation: UnsafeContinuation? + /// The optional consumer continuation. + @usableFromInline + var consumerContinuation: UnsafeContinuation? - /// The producer continuations. - @usableFromInline - var suspendedProducers: Deque<(UInt64, _MultiProducerSingleConsumerSuspendedProducer)> + /// The producer continuations. + @usableFromInline + var suspendedProducers: Deque<(UInt64, _MultiProducerSingleConsumerSuspendedProducer)> - /// The producers that have been cancelled. - @usableFromInline - var cancelledAsyncProducers: Deque + /// The producers that have been cancelled. + @usableFromInline + var cancelledAsyncProducers: Deque - /// Indicates if we currently have outstanding demand. - @usableFromInline - var hasOutstandingDemand: Bool + /// Indicates if we currently have outstanding demand. + @usableFromInline + var hasOutstandingDemand: Bool - var description: String { - "backpressure:\(self.backpressureStrategy.description) iteratorInitialized:\(self.iteratorInitialized) buffer:\(self.buffer.count) consumerContinuation:\(self.consumerContinuation == nil) producerContinuations:\(self.suspendedProducers.count) cancelledProducers:\(self.cancelledAsyncProducers.count) hasOutstandingDemand:\(self.hasOutstandingDemand)" - } + var description: String { + "backpressure:\(self.backpressureStrategy.description) iteratorInitialized:\(self.iteratorInitialized) buffer:\(self.buffer.count) consumerContinuation:\(self.consumerContinuation == nil) producerContinuations:\(self.suspendedProducers.count) cancelledProducers:\(self.cancelledAsyncProducers.count) hasOutstandingDemand:\(self.hasOutstandingDemand)" + } - @usableFromInline - init( - backpressureStrategy: _MultiProducerSingleConsumerChannelInternalBackpressureStrategy, iteratorInitialized: Bool, - onTermination: (@Sendable () -> Void)? = nil, - buffer: Deque, - consumerContinuation: UnsafeContinuation? = nil, - producerContinuations: Deque<(UInt64, _MultiProducerSingleConsumerSuspendedProducer)>, - cancelledAsyncProducers: Deque, - hasOutstandingDemand: Bool) { - self.backpressureStrategy = backpressureStrategy - self.iteratorInitialized = iteratorInitialized - self.onTermination = onTermination - self.buffer = buffer - self.consumerContinuation = consumerContinuation - self.suspendedProducers = producerContinuations - self.cancelledAsyncProducers = cancelledAsyncProducers - self.hasOutstandingDemand = hasOutstandingDemand + @usableFromInline + init( + backpressureStrategy: MultiProducerSingleConsumerChannel._InternalBackpressureStrategy, + iteratorInitialized: Bool, + onTermination: (@Sendable () -> Void)? = nil, + buffer: Deque, + consumerContinuation: UnsafeContinuation? = nil, + producerContinuations: Deque<(UInt64, _MultiProducerSingleConsumerSuspendedProducer)>, + cancelledAsyncProducers: Deque, + hasOutstandingDemand: Bool) { + self.backpressureStrategy = backpressureStrategy + self.iteratorInitialized = iteratorInitialized + self.onTermination = onTermination + self.buffer = buffer + self.consumerContinuation = consumerContinuation + self.suspendedProducers = producerContinuations + self.cancelledAsyncProducers = cancelledAsyncProducers + self.hasOutstandingDemand = hasOutstandingDemand + } } - } - @usableFromInline - struct SourceFinished { - /// Indicates if the iterator was initialized. @usableFromInline - var iteratorInitialized: Bool + struct SourceFinished: ~Copyable { + /// Indicates if the iterator was initialized. + @usableFromInline + var iteratorInitialized: Bool - /// The buffer of elements. - @usableFromInline - var buffer: Deque + /// The buffer of elements. + @usableFromInline + var buffer: Deque - /// The failure that should be thrown after the last element has been consumed. - @usableFromInline - var failure: Failure? + /// The failure that should be thrown after the last element has been consumed. + @usableFromInline + var failure: Failure? - /// The onTermination callback. - @usableFromInline - var onTermination: (@Sendable () -> Void)? + /// The onTermination callback. + @usableFromInline + var onTermination: (@Sendable () -> Void)? - var description: String { - "iteratorInitialized:\(self.iteratorInitialized) buffer:\(self.buffer.count) failure:\(self.failure == nil)" - } + var description: String { + "iteratorInitialized:\(self.iteratorInitialized) buffer:\(self.buffer.count) failure:\(self.failure == nil)" + } - @usableFromInline - init( - iteratorInitialized: Bool, - buffer: Deque, - failure: Failure? = nil, - onTermination: (@Sendable () -> Void)? = nil - ) { - self.iteratorInitialized = iteratorInitialized - self.buffer = buffer - self.failure = failure - self.onTermination = onTermination + @usableFromInline + init( + iteratorInitialized: Bool, + buffer: Deque, + failure: Failure? = nil, + onTermination: (@Sendable () -> Void)? = nil + ) { + self.iteratorInitialized = iteratorInitialized + self.buffer = buffer + self.failure = failure + self.onTermination = onTermination + } } - } - @usableFromInline - struct Finished { - /// Indicates if the iterator was initialized. @usableFromInline - var iteratorInitialized: Bool + struct Finished: ~Copyable { + /// Indicates if the iterator was initialized. + @usableFromInline + var iteratorInitialized: Bool - /// Indicates if the source was finished. - @usableFromInline - var sourceFinished: Bool + /// Indicates if the source was finished. + @usableFromInline + var sourceFinished: Bool - var description: String { - "iteratorInitialized:\(self.iteratorInitialized) sourceFinished:\(self.sourceFinished)" - } + var description: String { + "iteratorInitialized:\(self.iteratorInitialized) sourceFinished:\(self.sourceFinished)" + } - @inlinable - init( - iteratorInitialized: Bool, - sourceFinished: Bool - ) { - self.iteratorInitialized = iteratorInitialized - self.sourceFinished = sourceFinished + @inlinable + init( + iteratorInitialized: Bool, + sourceFinished: Bool + ) { + self.iteratorInitialized = iteratorInitialized + self.sourceFinished = sourceFinished + } } - } - /// The initial state. - case initial(Initial) + /// The initial state. + case initial(Initial) - /// The state once either any element was sent or `next()` was called. - case channeling(Channeling) + /// The state once either any element was sent or `next()` was called. + case channeling(Channeling) - /// The state once the underlying source signalled that it is finished. - case sourceFinished(SourceFinished) + /// 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(Finished) + /// 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(Finished) - /// An intermediate state to avoid CoWs. - case modify - - @usableFromInline - var description: String { - switch self { - case .initial(let initial): - return "initial \(initial.description)" - case .channeling(let channeling): - return "channeling \(channeling.description)" - case .sourceFinished(let sourceFinished): - return "sourceFinished \(sourceFinished.description)" - case .finished(let finished): - return "finished \(finished.description)" - case .modify: - fatalError() + @usableFromInline + var description: String { + switch self { + case .initial(let initial): + return "initial \(initial.description)" + case .channeling(let channeling): + return "channeling \(channeling.description)" + case .sourceFinished(let sourceFinished): + return "sourceFinished \(sourceFinished.description)" + case .finished(let finished): + return "finished \(finished.description)" + } } } } @@ -1486,4 +1520,13 @@ enum _MultiProducerSingleConsumerSuspendedProducer { case closure((Result) -> Void) case continuation(UnsafeContinuation) } + +extension Optional where Wrapped: ~Copyable { + @inlinable + mutating func take() -> Self { + let result = consume self + self = nil + return result + } +} #endif diff --git a/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel.swift b/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel.swift index e4fb7782..39fc8fd9 100644 --- a/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel.swift +++ b/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel.swift @@ -72,9 +72,9 @@ public struct MultiProducerSingleConsumerChannel: Async /// can tell the producer when any potential consumer went away. private final class _Backing: Sendable { /// The underlying storage. - fileprivate let storage: _MultiProducerSingleConsumerChannelBackpressuredStorage + fileprivate let storage: _Storage - init(storage: _MultiProducerSingleConsumerChannelBackpressuredStorage) { + init(storage: _Storage) { self.storage = storage } @@ -113,7 +113,7 @@ public struct MultiProducerSingleConsumerChannel: Async throwing failureType: Failure.Type = Never.self, backpressureStrategy: Source.BackpressureStrategy ) -> NewChannel { - let storage = _MultiProducerSingleConsumerChannelBackpressuredStorage( + let storage = _Storage( backpressureStrategy: backpressureStrategy.internalBackpressureStrategy ) let source = Source(storage: storage) @@ -121,7 +121,7 @@ public struct MultiProducerSingleConsumerChannel: Async return .init(channel: .init(storage: storage), source: source) } - init(storage: _MultiProducerSingleConsumerChannelBackpressuredStorage) { + init(storage: _Storage) { self.backing = .init(storage: storage) } } @@ -135,7 +135,7 @@ extension MultiProducerSingleConsumerChannel { public struct Source: ~Copyable, Sendable { /// A strategy that handles the backpressure of the channel. public struct BackpressureStrategy: Sendable { - var internalBackpressureStrategy: _MultiProducerSingleConsumerChannelInternalBackpressureStrategy + var internalBackpressureStrategy: _InternalBackpressureStrategy /// A backpressure strategy using a high and low watermark to suspend and resume production respectively. /// @@ -223,9 +223,9 @@ extension MultiProducerSingleConsumerChannel { } @usableFromInline - let _storage: _MultiProducerSingleConsumerChannelBackpressuredStorage + let _storage: _Storage - internal init(storage: _MultiProducerSingleConsumerChannelBackpressuredStorage) { + internal init(storage: _Storage) { self._storage = storage } @@ -439,9 +439,9 @@ extension MultiProducerSingleConsumerChannel { @usableFromInline final class _Backing { @usableFromInline - let storage: _MultiProducerSingleConsumerChannelBackpressuredStorage + let storage: _Storage - init(storage: _MultiProducerSingleConsumerChannelBackpressuredStorage) { + init(storage: _Storage) { self.storage = storage self.storage.iteratorInitialized() } @@ -454,7 +454,7 @@ extension MultiProducerSingleConsumerChannel { @usableFromInline let _backing: _Backing - init(storage: _MultiProducerSingleConsumerChannelBackpressuredStorage) { + init(storage: _Storage) { self._backing = .init(storage: storage) } @@ -464,7 +464,6 @@ extension MultiProducerSingleConsumerChannel { try await self._backing.storage.next(isolation: nil) } - @available(macOS 15.0, iOS 18.0, watchOS 11.0, tvOS 18.0, visionOS 2.0, *) @inlinable public mutating func next( isolation actor: isolated (any Actor)? = #isolation diff --git a/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift b/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift index 93e5c3e8..155a063c 100644 --- a/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift +++ b/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift @@ -17,14 +17,13 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { // MARK: - sequenceDeinitialized func testSequenceDeinitialized_whenNoIterator() async throws { - var (channel, source): (MultiProducerSingleConsumerChannel?, MultiProducerSingleConsumerChannel.Source) = - MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 5, high: 10) - ) - defer { - source.finish() - } + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + backpressureStrategy: .watermark(low: 5, high: 10) + ) + var channel: MultiProducerSingleConsumerChannel? = channelAndStream.channel + var source = consume channelAndStream.source + let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() source.onTermination = { @@ -60,14 +59,12 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { } func testSequenceDeinitialized_whenIterator() async throws { - var (channel, source): (MultiProducerSingleConsumerChannel?, MultiProducerSingleConsumerChannel.Source) = - MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 5, high: 10) - ) - defer { - source.finish() - } + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + backpressureStrategy: .watermark(low: 5, high: 10) + ) + var channel: MultiProducerSingleConsumerChannel? = channelAndStream.channel + var source = consume channelAndStream.source var iterator = channel?.makeAsyncIterator() @@ -111,11 +108,12 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { } func testSequenceDeinitialized_whenFinished() async throws { - var (channel, source): (MultiProducerSingleConsumerChannel?, MultiProducerSingleConsumerChannel.Source) = - MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 5, high: 10) - ) + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + backpressureStrategy: .watermark(low: 5, high: 10) + ) + var channel: MultiProducerSingleConsumerChannel? = channelAndStream.channel + var source = consume channelAndStream.source let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() source.onTermination = { @@ -133,35 +131,23 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { var onTerminationIterator = onTerminationStream.makeAsyncIterator() _ = await onTerminationIterator.next() - withExtendedLifetime(channel) { - source.finish(throwing: nil) - } - channel = nil let terminationResult: Void? = await onTerminationIterator.next() XCTAssertNil(terminationResult) - - do { - _ = try { try source.send(1) }() - XCTFail("Expected an error to be thrown") - } catch { - XCTAssertTrue(error is MultiProducerSingleConsumerChannelAlreadyFinishedError) - } + XCTAssertNil(channel) group.cancelAll() } } func testSequenceDeinitialized_whenChanneling_andSuspendedProducer() async throws { - var (channel, source): (MultiProducerSingleConsumerChannel?, MultiProducerSingleConsumerChannel.Source) = - MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 1, high: 2) - ) - defer { - source.finish() - } + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + backpressureStrategy: .watermark(low: 1, high: 2) + ) + var channel: MultiProducerSingleConsumerChannel? = channelAndStream.channel + var source = consume channelAndStream.source _ = try { try source.send(1) }() @@ -182,25 +168,23 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { // MARK: - iteratorInitialized func testIteratorInitialized_whenInitial() async throws { - let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 5, high: 10) ) - defer { - source.finish() - } + let channel = channelAndStream.channel + let source = consume channelAndStream.source _ = channel.makeAsyncIterator() } func testIteratorInitialized_whenChanneling() async throws { - let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 5, high: 10) ) - defer { - source.finish() - } + let channel = channelAndStream.channel + var source = consume channelAndStream.source try await source.send(1) @@ -210,10 +194,12 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { } func testIteratorInitialized_whenSourceFinished() async throws { - let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 5, high: 10) ) + let channel = channelAndStream.channel + var source = consume channelAndStream.source try await source.send(1) source.finish(throwing: nil) @@ -226,10 +212,12 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { } func testIteratorInitialized_whenFinished() async throws { - let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 5, high: 10) ) + let channel = channelAndStream.channel + var source = consume channelAndStream.source source.finish(throwing: nil) @@ -241,13 +229,12 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { // MARK: - iteratorDeinitialized func testIteratorDeinitialized_whenInitial() async throws { - var (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 5, high: 10) ) - defer { - source.finish() - } + let channel = channelAndStream.channel + var source = consume channelAndStream.source let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() source.onTermination = { @@ -277,13 +264,12 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { } func testIteratorDeinitialized_whenChanneling() async throws { - var (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 5, high: 10) ) - defer { - source.finish() - } + let channel = channelAndStream.channel + var source = consume channelAndStream.source let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() source.onTermination = { @@ -315,10 +301,12 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { } func testIteratorDeinitialized_whenSourceFinished() async throws { - var (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 5, high: 10) ) + let channel = channelAndStream.channel + var source = consume channelAndStream.source let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() source.onTermination = { @@ -351,11 +339,13 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { } func testIteratorDeinitialized_whenFinished() async throws { - var (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, throwing: Error.self, backpressureStrategy: .watermark(low: 5, high: 10) ) + let channel = channelAndStream.channel + var source = consume channelAndStream.source let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() source.onTermination = { @@ -387,15 +377,13 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { } func testIteratorDeinitialized_whenChanneling_andSuspendedProducer() async throws { - var (channel, source): (MultiProducerSingleConsumerChannel?, MultiProducerSingleConsumerChannel.Source) = - MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - throwing: Error.self, - backpressureStrategy: .watermark(low: 1, high: 2) - ) - defer { - source.finish() - } + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + throwing: Error.self, + backpressureStrategy: .watermark(low: 5, high: 10) + ) + var channel: MultiProducerSingleConsumerChannel? = channelAndStream.channel + var source = consume channelAndStream.source var iterator: MultiProducerSingleConsumerChannel.AsyncIterator? = channel?.makeAsyncIterator() channel = nil @@ -420,12 +408,13 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { // MARK: - sourceDeinitialized func testSourceDeinitialized_whenSourceFinished() async throws { - var (channel, source): (MultiProducerSingleConsumerChannel, MultiProducerSingleConsumerChannel.Source?) = - MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - throwing: Error.self, - backpressureStrategy: .watermark(low: 5, high: 10) - ) + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + throwing: Error.self, + backpressureStrategy: .watermark(low: 5, high: 10) + ) + let channel = channelAndStream.channel + var source: MultiProducerSingleConsumerChannel.Source? = consume channelAndStream.source let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() source?.onTermination = { @@ -450,8 +439,6 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { var iterator: MultiProducerSingleConsumerChannel.AsyncIterator? = channel.makeAsyncIterator() _ = try await iterator?.next() - source = nil - _ = await onTerminationIterator.next() _ = try await iterator?.next() @@ -465,11 +452,13 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { } func testSourceDeinitialized_whenFinished() async throws { - var (channel, source): (MultiProducerSingleConsumerChannel, MultiProducerSingleConsumerChannel.Source?) = - MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 5, high: 10) - ) + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + throwing: Error.self, + backpressureStrategy: .watermark(low: 5, high: 10) + ) + let channel = channelAndStream.channel + var source: MultiProducerSingleConsumerChannel.Source? = consume channelAndStream.source let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() source?.onTermination = { @@ -491,8 +480,6 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { _ = channel.makeAsyncIterator() - source = nil - _ = await onTerminationIterator.next() let terminationResult: Void? = await onTerminationIterator.next() @@ -505,13 +492,12 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { // MARK: - write func testWrite_whenInitial() async throws { - let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 2, high: 5) ) - defer { - source.finish() - } + let channel = channelAndStream.channel + var source = consume channelAndStream.source try await source.send(1) @@ -521,18 +507,17 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { } func testWrite_whenChanneling_andNoConsumer() async throws { - let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 2, high: 5) ) - defer { - source.finish() - } + let channel = channelAndStream.channel + var source = consume channelAndStream.source try await source.send(1) try await source.send(2) - var iterator = stream.makeAsyncIterator() + var iterator = channel.makeAsyncIterator() let element1 = await iterator.next() XCTAssertEqual(element1, 1) let element2 = await iterator.next() @@ -540,17 +525,16 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { } func testWrite_whenChanneling_andSuspendedConsumer() async throws { - let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 2, high: 5) ) - defer { - source.finish() - } + let channel = channelAndStream.channel + var source = consume channelAndStream.source try await withThrowingTaskGroup(of: Int?.self) { group in group.addTask { - return await stream.first { _ in true } + return await channel.first { _ in true } } // This is always going to be a bit racy since we need the call to next() suspend @@ -563,17 +547,16 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { } func testWrite_whenChanneling_andSuspendedConsumer_andEmptySequence() async throws { - let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 2, high: 5) ) - defer { - source.finish() - } + let channel = channelAndStream.channel + var source = consume channelAndStream.source try await withThrowingTaskGroup(of: Int?.self) { group in group.addTask { - return await stream.first { _ in true } + return await channel.first { _ in true } } // This is always going to be a bit racy since we need the call to next() suspend @@ -589,13 +572,12 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { // MARK: - enqueueProducer func testEnqueueProducer_whenChanneling_andAndCancelled() async throws { - let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 1, high: 2) ) - defer { - source.finish() - } + let channel = channelAndStream.channel + var source = consume channelAndStream.source let (producerStream, producerSource) = AsyncThrowingStream.makeStream() @@ -603,15 +585,15 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { let writeResult = try { try source.send(2) }() - switch writeResult { + switch consume writeResult { case .produceMore: preconditionFailure() case .enqueueCallback(let callbackToken): source.cancelCallback(callbackToken: callbackToken) - source.enqueueCallback(callbackToken: callbackToken) { result in - producerSource.yield(with: result) - } +// source.enqueueCallback(callbackToken: callbackToken) { result in +// producerSource.yield(with: result) +// } } do { @@ -621,18 +603,17 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { XCTAssertTrue(error is CancellationError) } - let element = await stream.first { _ in true } + let element = await channel.first { _ in true } XCTAssertEqual(element, 1) } func testEnqueueProducer_whenChanneling_andAndCancelled_andAsync() async throws { - let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( of: Int.self, backpressureStrategy: .watermark(low: 1, high: 2) ) - defer { - source.finish() - } + let channel = channelAndStream.channel + var source = consume channelAndStream.source try await source.send(1) @@ -650,446 +631,446 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { } } - let element = await stream.first { _ in true } + let element = await channel.first { _ in true } XCTAssertEqual(element, 1) } - - func testEnqueueProducer_whenChanneling_andInterleaving() async throws { - let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 1, high: 1) - ) - defer { - source.finish() - } - var iterator = stream.makeAsyncIterator() - - let (producerStream, producerSource) = AsyncThrowingStream.makeStream() - - let writeResult = try { try source.send(1) }() - - switch writeResult { - case .produceMore: - preconditionFailure() - case .enqueueCallback(let callbackToken): - let element = 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_whenChanneling_andSuspending() async throws { - let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 1, high: 1) - ) - defer { - source.finish() - } - var iterator = stream.makeAsyncIterator() - - let (producerStream, producerSource) = AsyncThrowingStream.makeStream() - - let writeResult = try { try source.send(1) }() - - switch writeResult { - case .produceMore: - preconditionFailure() - case .enqueueCallback(let callbackToken): - source.enqueueCallback(callbackToken: callbackToken) { result in - producerSource.yield(with: result) - } - } - - let element = 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) = MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 1, high: 1) - ) - defer { - source.finish() - } - var iterator = stream.makeAsyncIterator() - - let (producerStream, producerSource) = AsyncThrowingStream.makeStream() - - let writeResult = try { try source.send(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 = await iterator.next() - XCTAssertEqual(element, 1) - - do { - _ = try await producerStream.first { _ in true } - XCTFail("Expected an error to be thrown") - } catch { - XCTAssertTrue(error is MultiProducerSingleConsumerChannelAlreadyFinishedError) - } - } - - // MARK: - cancelProducer - - func testCancelProducer_whenChanneling() async throws { - let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 1, high: 2) - ) - defer { - source.finish() - } - - let (producerStream, producerSource) = AsyncThrowingStream.makeStream() - - try await source.send(1) - - let writeResult = try { try source.send(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 = await stream.first { _ in true } - XCTAssertEqual(element, 1) - } - - func testCancelProducer_whenSourceFinished() async throws { - let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 1, high: 2) - ) - - let (producerStream, producerSource) = AsyncThrowingStream.makeStream() - - try await source.send(1) - - let writeResult = try { try source.send(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 MultiProducerSingleConsumerChannelAlreadyFinishedError) - } - - let element = await stream.first { _ in true } - XCTAssertEqual(element, 1) - } - - // MARK: - finish - - func testFinish_whenChanneling_andConsumerSuspended() async throws { - let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 1, high: 1) - ) - - try await withThrowingTaskGroup(of: Int?.self) { group in - group.addTask { - return 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(for: .seconds(0.5)) - - source.finish(throwing: nil) - let element = try await group.next() - XCTAssertEqual(element, .some(nil)) - } - } - - func testFinish_whenInitial() async throws { - let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - throwing: Error.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 (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 2, high: 4) - ) - defer { - source.finish() - } - - let (backpressureEventStream, backpressureEventContinuation) = AsyncStream.makeStream(of: Void.self) - - await withThrowingTaskGroup(of: Void.self) { group in - group.addTask { - while true { - backpressureEventContinuation.yield(()) - try await source.send(contentsOf: [1]) - } - } - - var backpressureEventIterator = backpressureEventStream.makeAsyncIterator() - var iterator = channel.makeAsyncIterator() - - await backpressureEventIterator.next() - await backpressureEventIterator.next() - await backpressureEventIterator.next() - await backpressureEventIterator.next() - - _ = await iterator.next() - _ = await iterator.next() - _ = await iterator.next() - - await backpressureEventIterator.next() - await backpressureEventIterator.next() - await backpressureEventIterator.next() - - group.cancelAll() - } - } - - func testBackpressureSync() async throws { - let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 2, high: 4) - ) - defer { - source.finish() - } - - let (backpressureEventStream, backpressureEventContinuation) = AsyncStream.makeStream(of: Void.self) - - await withThrowingTaskGroup(of: Void.self) { group in - group.addTask { - @Sendable func yield() { - backpressureEventContinuation.yield(()) - source.send(contentsOf: [1]) { result in - switch result { - case .success: - yield() - - case .failure: - break - } - } - } - - yield() - } - - var backpressureEventIterator = backpressureEventStream.makeAsyncIterator() - var iterator = channel.makeAsyncIterator() - - await backpressureEventIterator.next() - await backpressureEventIterator.next() - await backpressureEventIterator.next() - await backpressureEventIterator.next() - - _ = await iterator.next() - _ = await iterator.next() - _ = await iterator.next() - - await backpressureEventIterator.next() - await backpressureEventIterator.next() - await backpressureEventIterator.next() - - group.cancelAll() - } - } - - func testWatermarkWithCustomCoount() async throws { - let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( - of: [Int].self, - backpressureStrategy: .watermark(low: 2, high: 4, waterLevelForElement: { $0.count }) - ) - defer { - source.finish() - } - var iterator = channel.makeAsyncIterator() - - try await source.send([1, 1, 1]) - - _ = await iterator.next() - - try await source.send([1, 1, 1]) - - _ = await iterator.next() - } - - func testWatermarWithLotsOfElements() async throws { - // This test should in the future use a custom task executor to schedule to avoid sending - // 1000 elements. - let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 2, high: 4) - ) - await withThrowingTaskGroup(of: Void.self) { group in - group.addTask { - for i in 0...10000 { - try await source.send(i) - } - source.finish() - } - - group.addTask { - var sum = 0 - for try await element in channel { - sum += element - } - } - } - } - - func testThrowsError() async throws { - let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - throwing: Error.self, - backpressureStrategy: .watermark(low: 2, high: 4) - ) - - try await source.send(1) - try await source.send(2) - source.finish(throwing: CancellationError()) - - var elements = [Int]() - var iterator = channel.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 (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 2, high: 4) - ) - - continuation.yield(1) - continuation.yield(2) - continuation.finish() - - try await source.send(contentsOf: stream) - source.finish(throwing: nil) - - let elements = await channel.collect() - XCTAssertEqual(elements, [1, 2]) - } - - // MARK: NonThrowing - - func testNonThrowing() async throws { - let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 2, high: 4) - ) - defer { - source.finish() - } - - let (backpressureEventStream, backpressureEventContinuation) = AsyncStream.makeStream(of: Void.self) - - await withThrowingTaskGroup(of: Void.self) { group in - group.addTask { - while true { - backpressureEventContinuation.yield(()) - try await source.send(contentsOf: [1]) - } - } - - var backpressureEventIterator = backpressureEventStream.makeAsyncIterator() - var iterator = channel.makeAsyncIterator() - - await backpressureEventIterator.next() - await backpressureEventIterator.next() - await backpressureEventIterator.next() - await backpressureEventIterator.next() - - _ = await iterator.next() - _ = await iterator.next() - _ = await iterator.next() - - await backpressureEventIterator.next() - await backpressureEventIterator.next() - await backpressureEventIterator.next() - - group.cancelAll() - } - } +// +// func testEnqueueProducer_whenChanneling_andInterleaving() async throws { +// let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( +// of: Int.self, +// backpressureStrategy: .watermark(low: 1, high: 1) +// ) +// defer { +// source.finish() +// } +// var iterator = stream.makeAsyncIterator() +// +// let (producerStream, producerSource) = AsyncThrowingStream.makeStream() +// +// let writeResult = try { try source.send(1) }() +// +// switch writeResult { +// case .produceMore: +// preconditionFailure() +// case .enqueueCallback(let callbackToken): +// let element = 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_whenChanneling_andSuspending() async throws { +// let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( +// of: Int.self, +// backpressureStrategy: .watermark(low: 1, high: 1) +// ) +// defer { +// source.finish() +// } +// var iterator = stream.makeAsyncIterator() +// +// let (producerStream, producerSource) = AsyncThrowingStream.makeStream() +// +// let writeResult = try { try source.send(1) }() +// +// switch writeResult { +// case .produceMore: +// preconditionFailure() +// case .enqueueCallback(let callbackToken): +// source.enqueueCallback(callbackToken: callbackToken) { result in +// producerSource.yield(with: result) +// } +// } +// +// let element = 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) = MultiProducerSingleConsumerChannel.makeChannel( +// of: Int.self, +// backpressureStrategy: .watermark(low: 1, high: 1) +// ) +// defer { +// source.finish() +// } +// var iterator = stream.makeAsyncIterator() +// +// let (producerStream, producerSource) = AsyncThrowingStream.makeStream() +// +// let writeResult = try { try source.send(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 = await iterator.next() +// XCTAssertEqual(element, 1) +// +// do { +// _ = try await producerStream.first { _ in true } +// XCTFail("Expected an error to be thrown") +// } catch { +// XCTAssertTrue(error is MultiProducerSingleConsumerChannelAlreadyFinishedError) +// } +// } +// +// // MARK: - cancelProducer +// +// func testCancelProducer_whenChanneling() async throws { +// let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( +// of: Int.self, +// backpressureStrategy: .watermark(low: 1, high: 2) +// ) +// defer { +// source.finish() +// } +// +// let (producerStream, producerSource) = AsyncThrowingStream.makeStream() +// +// try await source.send(1) +// +// let writeResult = try { try source.send(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 = await stream.first { _ in true } +// XCTAssertEqual(element, 1) +// } +// +// func testCancelProducer_whenSourceFinished() async throws { +// let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( +// of: Int.self, +// backpressureStrategy: .watermark(low: 1, high: 2) +// ) +// +// let (producerStream, producerSource) = AsyncThrowingStream.makeStream() +// +// try await source.send(1) +// +// let writeResult = try { try source.send(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 MultiProducerSingleConsumerChannelAlreadyFinishedError) +// } +// +// let element = await stream.first { _ in true } +// XCTAssertEqual(element, 1) +// } +// +// // MARK: - finish +// +// func testFinish_whenChanneling_andConsumerSuspended() async throws { +// let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( +// of: Int.self, +// backpressureStrategy: .watermark(low: 1, high: 1) +// ) +// +// try await withThrowingTaskGroup(of: Int?.self) { group in +// group.addTask { +// return 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(for: .seconds(0.5)) +// +// source.finish(throwing: nil) +// let element = try await group.next() +// XCTAssertEqual(element, .some(nil)) +// } +// } +// +// func testFinish_whenInitial() async throws { +// let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( +// of: Int.self, +// throwing: Error.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 (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( +// of: Int.self, +// backpressureStrategy: .watermark(low: 2, high: 4) +// ) +// defer { +// source.finish() +// } +// +// let (backpressureEventStream, backpressureEventContinuation) = AsyncStream.makeStream(of: Void.self) +// +// await withThrowingTaskGroup(of: Void.self) { group in +// group.addTask { +// while true { +// backpressureEventContinuation.yield(()) +// try await source.send(contentsOf: [1]) +// } +// } +// +// var backpressureEventIterator = backpressureEventStream.makeAsyncIterator() +// var iterator = channel.makeAsyncIterator() +// +// await backpressureEventIterator.next() +// await backpressureEventIterator.next() +// await backpressureEventIterator.next() +// await backpressureEventIterator.next() +// +// _ = await iterator.next() +// _ = await iterator.next() +// _ = await iterator.next() +// +// await backpressureEventIterator.next() +// await backpressureEventIterator.next() +// await backpressureEventIterator.next() +// +// group.cancelAll() +// } +// } +// +// func testBackpressureSync() async throws { +// let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( +// of: Int.self, +// backpressureStrategy: .watermark(low: 2, high: 4) +// ) +// defer { +// source.finish() +// } +// +// let (backpressureEventStream, backpressureEventContinuation) = AsyncStream.makeStream(of: Void.self) +// +// await withThrowingTaskGroup(of: Void.self) { group in +// group.addTask { +// @Sendable func yield() { +// backpressureEventContinuation.yield(()) +// source.send(contentsOf: [1]) { result in +// switch result { +// case .success: +// yield() +// +// case .failure: +// break +// } +// } +// } +// +// yield() +// } +// +// var backpressureEventIterator = backpressureEventStream.makeAsyncIterator() +// var iterator = channel.makeAsyncIterator() +// +// await backpressureEventIterator.next() +// await backpressureEventIterator.next() +// await backpressureEventIterator.next() +// await backpressureEventIterator.next() +// +// _ = await iterator.next() +// _ = await iterator.next() +// _ = await iterator.next() +// +// await backpressureEventIterator.next() +// await backpressureEventIterator.next() +// await backpressureEventIterator.next() +// +// group.cancelAll() +// } +// } +// +// func testWatermarkWithCustomCoount() async throws { +// let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( +// of: [Int].self, +// backpressureStrategy: .watermark(low: 2, high: 4, waterLevelForElement: { $0.count }) +// ) +// defer { +// source.finish() +// } +// var iterator = channel.makeAsyncIterator() +// +// try await source.send([1, 1, 1]) +// +// _ = await iterator.next() +// +// try await source.send([1, 1, 1]) +// +// _ = await iterator.next() +// } +// +// func testWatermarWithLotsOfElements() async throws { +// // This test should in the future use a custom task executor to schedule to avoid sending +// // 1000 elements. +// let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( +// of: Int.self, +// backpressureStrategy: .watermark(low: 2, high: 4) +// ) +// await withThrowingTaskGroup(of: Void.self) { group in +// group.addTask { +// for i in 0...10000 { +// try await source.send(i) +// } +// source.finish() +// } +// +// group.addTask { +// var sum = 0 +// for try await element in channel { +// sum += element +// } +// } +// } +// } +// +// func testThrowsError() async throws { +// let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( +// of: Int.self, +// throwing: Error.self, +// backpressureStrategy: .watermark(low: 2, high: 4) +// ) +// +// try await source.send(1) +// try await source.send(2) +// source.finish(throwing: CancellationError()) +// +// var elements = [Int]() +// var iterator = channel.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 (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( +// of: Int.self, +// backpressureStrategy: .watermark(low: 2, high: 4) +// ) +// +// continuation.yield(1) +// continuation.yield(2) +// continuation.finish() +// +// try await source.send(contentsOf: stream) +// source.finish(throwing: nil) +// +// let elements = await channel.collect() +// XCTAssertEqual(elements, [1, 2]) +// } +// +// // MARK: NonThrowing +// +// func testNonThrowing() async throws { +// let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( +// of: Int.self, +// backpressureStrategy: .watermark(low: 2, high: 4) +// ) +// defer { +// source.finish() +// } +// +// let (backpressureEventStream, backpressureEventContinuation) = AsyncStream.makeStream(of: Void.self) +// +// await withThrowingTaskGroup(of: Void.self) { group in +// group.addTask { +// while true { +// backpressureEventContinuation.yield(()) +// try await source.send(contentsOf: [1]) +// } +// } +// +// var backpressureEventIterator = backpressureEventStream.makeAsyncIterator() +// var iterator = channel.makeAsyncIterator() +// +// await backpressureEventIterator.next() +// await backpressureEventIterator.next() +// await backpressureEventIterator.next() +// await backpressureEventIterator.next() +// +// _ = await iterator.next() +// _ = await iterator.next() +// _ = await iterator.next() +// +// await backpressureEventIterator.next() +// await backpressureEventIterator.next() +// await backpressureEventIterator.next() +// +// group.cancelAll() +// } +// } } extension AsyncSequence { @@ -1105,7 +1086,7 @@ extension MultiProducerSingleConsumerChannel.Source.SendResult { func assertIsProducerMore() { switch self { case .produceMore: - return + return () case .enqueueCallback: XCTFail("Expected produceMore") @@ -1118,7 +1099,7 @@ extension MultiProducerSingleConsumerChannel.Source.SendResult { XCTFail("Expected enqueueCallback") case .enqueueCallback: - return + return () } } } From 662693eceab65b0d2f63dedc14eb318e6be60537 Mon Sep 17 00:00:00 2001 From: Franz Busch Date: Thu, 13 Jun 2024 21:40:20 +0200 Subject: [PATCH 5/7] More refactorings --- Package.swift | 5 +- Sources/AsyncAlgorithms/Locking.swift | 12 +- ...oducerSingleConsumerChannel+Internal.swift | 473 +++---- .../MultiProducerSingleConsumerChannel.swift | 9 +- ...tiProducerSingleConsumerChannelTests.swift | 1112 ++++++++--------- 5 files changed, 717 insertions(+), 894 deletions(-) diff --git a/Package.swift b/Package.swift index 2932e199..c8b857fa 100644 --- a/Package.swift +++ b/Package.swift @@ -20,7 +20,10 @@ let package = Package( targets: [ .target( name: "AsyncAlgorithms", - dependencies: [.product(name: "Collections", package: "swift-collections")], + dependencies: [ + .product(name: "DequeModule", package: "swift-collections"), + .product(name: "OrderedCollections", package: "swift-collections"), + ], swiftSettings: [ .enableExperimentalFeature("StrictConcurrency=complete"), ] diff --git a/Sources/AsyncAlgorithms/Locking.swift b/Sources/AsyncAlgorithms/Locking.swift index ac8d34a3..669c9b8c 100644 --- a/Sources/AsyncAlgorithms/Locking.swift +++ b/Sources/AsyncAlgorithms/Locking.swift @@ -18,7 +18,7 @@ import WinSDK #endif @usableFromInline -internal struct Lock { +internal class Lock { #if canImport(Darwin) @usableFromInline typealias Primitive = os_unfair_lock @@ -35,12 +35,17 @@ internal struct Lock { @usableFromInline typealias PlatformLock = UnsafeMutablePointer + @usableFromInline let platformLock: PlatformLock private init(_ platformLock: PlatformLock) { self.platformLock = platformLock } - + + deinit { + self.deinitialize() + } + fileprivate static func initialize(_ platformLock: PlatformLock) { #if canImport(Darwin) platformLock.initialize(to: os_unfair_lock()) @@ -93,10 +98,12 @@ internal struct Lock { Lock.deinitialize(platformLock) } + @inlinable func lock() { Lock.lock(platformLock) } + @inlinable func unlock() { Lock.unlock(platformLock) } @@ -109,6 +116,7 @@ internal struct Lock { /// /// - 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 { diff --git a/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel+Internal.swift b/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel+Internal.swift index 5cea923a..58b41ae1 100644 --- a/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel+Internal.swift +++ b/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel+Internal.swift @@ -230,15 +230,15 @@ extension MultiProducerSingleConsumerChannel { case .failProducersAndCallOnTermination(let producerContinuations, let onTermination): for producerContinuation in producerContinuations { - producerContinuation(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) + switch producerContinuation { + case .closure(let onProduceMore): + onProduceMore(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) + case .continuation(let continuation): + continuation.resume(throwing: MultiProducerSingleConsumerChannelAlreadyFinishedError()) + } } onTermination?() - case .failProducers(let producerContinuations): - for producerContinuation in producerContinuations { - producerContinuation(.failure(MultiProducerSingleConsumerChannelAlreadyFinishedError())) - } - case .none: break } @@ -486,38 +486,26 @@ extension MultiProducerSingleConsumerChannel._Storage { struct _StateMachine: ~Copyable { /// The state machine's current state. @usableFromInline - var _state: _State? - - // The ID used for the next CallbackToken. - @usableFromInline - var _nextCallbackTokenID: UInt64 = 0 + var _state: _State - @usableFromInline + @inlinable var _onTermination: (@Sendable () -> Void)? { set { - switch self._state.take()! { - case .initial(var initial): - initial.onTermination = newValue - - self._state = .initial(initial) - + switch consume self._state { case .channeling(var channeling): channeling.onTermination = newValue - self._state = .channeling(channeling) + self = .init(state: .channeling(channeling)) case .sourceFinished(var sourceFinished): sourceFinished.onTermination = newValue - self._state = .sourceFinished(sourceFinished) + self = .init(state: .sourceFinished(sourceFinished)) - case .finished: - break + case .finished(let finished): + self = .init(state: .finished(finished)) } } get { switch self._state { - case .initial(let initial): - return initial.onTermination - case .channeling(let channeling): return channeling.onTermination @@ -526,42 +514,35 @@ extension MultiProducerSingleConsumerChannel._Storage { case .finished: return nil - - case .none: - fatalError() } } } - /// Initializes a new `StateMachine`. - /// - /// We are passing and holding the backpressure strategy here because - /// it is a customizable extension of the state machine. - /// - /// - Parameter backpressureStrategy: The backpressure strategy. init( backpressureStrategy: MultiProducerSingleConsumerChannel._InternalBackpressureStrategy ) { - self._state = .initial( + self._state = .channeling( .init( backpressureStrategy: backpressureStrategy, iteratorInitialized: false, - onTermination: nil + buffer: .init(), + producerContinuations: .init(), + cancelledAsyncProducers: .init(), + hasOutstandingDemand: true, + activeProducers: 1, + nextCallbackTokenID: 0 ) ) } - /// Generates the next callback token. - @usableFromInline - mutating func nextCallbackToken() -> UInt64 { - let id = self._nextCallbackTokenID - self._nextCallbackTokenID += 1 - return id + @inlinable + init(state: consuming _State) { + self._state = state } - /// Actions returned by `sequenceDeinitialized()`. + /// Actions returned by `sourceDeinitialized()`. @usableFromInline - enum SequenceDeinitializedAction { + enum SourceDeinitializedAction { /// Indicates that `onTermination` should be called. case callOnTermination((@Sendable () -> Void)?) /// Indicates that all producers should be failed and `onTermination` should be called. @@ -572,25 +553,65 @@ extension MultiProducerSingleConsumerChannel._Storage { } @inlinable - mutating func sequenceDeinitialized() -> SequenceDeinitializedAction? { - switch self._state.take()! { - case .initial(let initial): - guard initial.iteratorInitialized else { - // No iterator was created so we can transition to finished right away. - self._state = .finished(.init(iteratorInitialized: false, sourceFinished: false)) + mutating func sourceDeinitialized() -> SourceDeinitializedAction? { + switch consume self._state { + case .channeling(var channeling): + channeling.activeProducers -= 1 + + if channeling.activeProducers == 0 { + // This was the last producer so we can transition to source finished now + + self = .init(state: .sourceFinished(.init( + iteratorInitialized: channeling.iteratorInitialized, + buffer: channeling.buffer + ))) + + if channeling.suspendedProducers.isEmpty { + return .callOnTermination(channeling.onTermination) + } else { + return .failProducersAndCallOnTermination( + .init(channeling.suspendedProducers.lazy.map { $0.1 }), + channeling.onTermination + ) + } + } else { + // We still have more producers + self = .init(state: .channeling(channeling)) - return .callOnTermination(initial.onTermination) + return nil } - // An iterator was created and we deinited the sequence. - // This is an expected pattern and we just continue on normal. - self._state = .initial(initial) + case .sourceFinished(let sourceFinished): + // This can happen if one producer calls finish and another deinits afterwards + self = .init(state: .sourceFinished(sourceFinished)) - return .none + return nil + case .finished(let finished): + // This can happen if the consumer finishes and the producers deinit + self = .init(state: .finished(finished)) + return nil + } + } + + /// 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( + _TinyArray<_MultiProducerSingleConsumerSuspendedProducer>, + (@Sendable () -> Void)? + ) + } + + @inlinable + mutating func sequenceDeinitialized() -> SequenceDeinitializedAction? { + switch consume self._state { case .channeling(let channeling): guard channeling.iteratorInitialized else { // No iterator was created so we can transition to finished right away. - self._state = .finished(.init(iteratorInitialized: false, sourceFinished: false)) + self = .init(state: .finished(.init(iteratorInitialized: false, sourceFinished: false))) return .failProducersAndCallOnTermination( .init(channeling.suspendedProducers.lazy.map { $0.1 }), @@ -599,27 +620,27 @@ extension MultiProducerSingleConsumerChannel._Storage { } // An iterator was created and we deinited the sequence. // This is an expected pattern and we just continue on normal. - self._state = .channeling(channeling) + self = .init(state: .channeling(channeling)) return .none case .sourceFinished(let sourceFinished): guard sourceFinished.iteratorInitialized else { // No iterator was created so we can transition to finished right away. - self._state = .finished(.init(iteratorInitialized: false, sourceFinished: true)) + self = .init(state: .finished(.init(iteratorInitialized: false, sourceFinished: true))) return .callOnTermination(sourceFinished.onTermination) } // An iterator was created and we deinited the sequence. // This is an expected pattern and we just continue on normal. - self._state = .sourceFinished(sourceFinished) + self = .init(state: .sourceFinished(sourceFinished)) return .none case .finished(let finished): // We are already finished so there is nothing left to clean up. // This is just the references dropping afterwards. - self._state = .finished(finished) + self = .init(state: .finished(finished)) return .none } @@ -627,17 +648,7 @@ extension MultiProducerSingleConsumerChannel._Storage { @inlinable mutating func iteratorInitialized() { - switch self._state.take()! { - 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) - } - + switch consume self._state { case .channeling(var channeling): if channeling.iteratorInitialized { // Our sequence is a unicast sequence and does not support multiple AsyncIterator's @@ -645,7 +656,7 @@ extension MultiProducerSingleConsumerChannel._Storage { } else { // The first and only iterator was initialized. channeling.iteratorInitialized = true - self._state = .channeling(channeling) + self = .init(state: .channeling(channeling)) } case .sourceFinished(var sourceFinished): @@ -655,7 +666,7 @@ extension MultiProducerSingleConsumerChannel._Storage { } else { // The first and only iterator was initialized. sourceFinished.iteratorInitialized = true - self._state = .sourceFinished(sourceFinished) + self = .init(state: .sourceFinished(sourceFinished)) } case .finished(let finished): @@ -663,7 +674,7 @@ extension MultiProducerSingleConsumerChannel._Storage { // Our sequence is a unicast sequence and does not support multiple AsyncIterator's fatalError("Only a single AsyncIterator can be created") } else { - self._state = .finished(.init(iteratorInitialized: true, sourceFinished: finished.sourceFinished)) + self = .init(state: .finished(.init(iteratorInitialized: true, sourceFinished: finished.sourceFinished))) } } } @@ -682,24 +693,12 @@ extension MultiProducerSingleConsumerChannel._Storage { @inlinable mutating func iteratorDeinitialized() -> IteratorDeinitializedAction? { - switch self._state.take()! { - 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(.init(iteratorInitialized: true, sourceFinished: false)) - - return .callOnTermination(initial.onTermination) - } else { - // An iterator needs to be initialized before it can be deinitialized. - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - } - + switch consume self._state { case .channeling(let channeling): if channeling.iteratorInitialized { // An iterator was created and deinited. Since we only support // a single iterator we can now transition to finish. - self._state = .finished(.init(iteratorInitialized: true, sourceFinished: false)) + self = .init(state: .finished(.init(iteratorInitialized: true, sourceFinished: false))) return .failProducersAndCallOnTermination( .init(channeling.suspendedProducers.lazy.map { $0.1 }), @@ -714,7 +713,7 @@ extension MultiProducerSingleConsumerChannel._Storage { 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(.init(iteratorInitialized: true, sourceFinished: true)) + self = .init(state: .finished(.init(iteratorInitialized: true, sourceFinished: true))) return .callOnTermination(sourceFinished.onTermination) } else { @@ -725,55 +724,12 @@ extension MultiProducerSingleConsumerChannel._Storage { case .finished(let finished): // We are already finished so there is nothing left to clean up. // This is just the references dropping afterwards. - self._state = .finished(finished) + self = .init(state: .finished(finished)) return .none } } - /// 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( - _TinyArray<(Result) -> Void>, - (@Sendable () -> Void)? - ) - /// Indicates that all producers should be failed. - case failProducers(_TinyArray<(Result) -> Void>) - } - - @inlinable - mutating func sourceDeinitialized() -> SourceDeinitializedAction? { - switch self._state.take()! { - case .initial: - fatalError("The channel's source hasn't been finished but deinited") - - case .channeling(let channeling): - self._state = .channeling(channeling) - - return nil - - case .sourceFinished(let sourceFinished): - // This is the expected case where finish was called and then the source deinited - self._state = .sourceFinished(sourceFinished) - - return .none - - case .finished(let finished): - if finished.sourceFinished { - // The source already got finished so this is fine. - self._state = .finished(finished) - - return .none - } else { - fatalError("The channel's source hasn't been finished but deinited") - } - } - } - /// Actions returned by `send()`. @usableFromInline enum SendAction { @@ -827,29 +783,7 @@ extension MultiProducerSingleConsumerChannel._Storage { @inlinable mutating func send(_ sequence: some Sequence) -> SendAction { - switch self._state.take()! { - case .initial(var initial): - var buffer = Deque() - buffer.append(contentsOf: sequence) - - let shouldProduceMore = initial.backpressureStrategy.didSend(elements: buffer[...]) - let callbackToken = shouldProduceMore ? nil : self.nextCallbackToken() - - self._state = .channeling( - .init( - backpressureStrategy: initial.backpressureStrategy, - iteratorInitialized: initial.iteratorInitialized, - onTermination: initial.onTermination, - buffer: buffer, - consumerContinuation: nil, - producerContinuations: .init(), - cancelledAsyncProducers: .init(), - hasOutstandingDemand: shouldProduceMore - ) - ) - - return .init(callbackToken: callbackToken) - + switch consume self._state { case .channeling(var channeling): // We have an element and can resume the continuation let bufferEndIndexBeforeAppend = channeling.buffer.endIndex @@ -861,17 +795,19 @@ extension MultiProducerSingleConsumerChannel._Storage { guard let consumerContinuation = channeling.consumerContinuation else { // We don't have a suspended consumer so we just buffer the elements - self._state = .channeling(channeling) + let callbackToken = shouldProduceMore ? nil : channeling.nextCallbackToken() + self = .init(state: .channeling(channeling)) return .init( - callbackToken: shouldProduceMore ? nil : self.nextCallbackToken() + callbackToken: callbackToken ) } guard let element = channeling.buffer.popFirst() else { // We got a send of an empty sequence. We just tolerate this. - self._state = .channeling(channeling) + let callbackToken = shouldProduceMore ? nil : channeling.nextCallbackToken() + self = .init(state: .channeling(channeling)) - return .init(callbackToken: shouldProduceMore ? nil : self.nextCallbackToken()) + return .init(callbackToken: callbackToken) } // We need to tell the back pressure strategy that we consumed shouldProduceMore = channeling.backpressureStrategy.didConsume(element: element) @@ -879,22 +815,23 @@ extension MultiProducerSingleConsumerChannel._Storage { // We got a consumer continuation and an element. We can resume the consumer now channeling.consumerContinuation = nil - self._state = .channeling(channeling) + let callbackToken = shouldProduceMore ? nil : channeling.nextCallbackToken() + self = .init(state: .channeling(channeling)) return .init( - callbackToken: shouldProduceMore ? nil : self.nextCallbackToken(), + callbackToken: callbackToken, continuationAndElement: (consumerContinuation, element) ) case .sourceFinished(let sourceFinished): // If the source has finished we are dropping the elements. - self._state = .sourceFinished(sourceFinished) + self = .init(state: .sourceFinished(sourceFinished)) return .throwFinishedError case .finished(let finished): // If the source has finished we are dropping the elements. - self._state = .finished(finished) + self = .init(state: .finished(finished)) return .throwFinishedError } @@ -914,25 +851,22 @@ extension MultiProducerSingleConsumerChannel._Storage { callbackToken: UInt64, onProduceMore: sending @escaping (Result) -> Void ) -> EnqueueProducerAction? { - switch self._state.take()! { - case .initial: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - + switch consume self._state { case .channeling(var channeling): if let index = channeling.cancelledAsyncProducers.firstIndex(of: callbackToken) { // Our producer got marked as cancelled. channeling.cancelledAsyncProducers.remove(at: index) - self._state = .channeling(channeling) + self = .init(state: .channeling(channeling)) return .resumeProducerWithError(onProduceMore, CancellationError()) } else if channeling.hasOutstandingDemand { // We hit an edge case here where we wrote but the consuming thread got interleaved - self._state = .channeling(channeling) + self = .init(state: .channeling(channeling)) return .resumeProducer(onProduceMore) } else { channeling.suspendedProducers.append((callbackToken, .closure(onProduceMore))) - self._state = .channeling(channeling) + self = .init(state: .channeling(channeling)) return .none } @@ -940,14 +874,14 @@ extension MultiProducerSingleConsumerChannel._Storage { case .sourceFinished(let sourceFinished): // Since we are unlocking between sending elements and suspending the send // It can happen that the source got finished or the consumption fully finishes. - self._state = .sourceFinished(sourceFinished) + self = .init(state: .sourceFinished(sourceFinished)) return .resumeProducerWithError(onProduceMore, MultiProducerSingleConsumerChannelAlreadyFinishedError()) case .finished(let finished): // Since we are unlocking between sending elements and suspending the send // It can happen that the source got finished or the consumption fully finishes. - self._state = .finished(finished) + self = .init(state: .finished(finished)) return .resumeProducerWithError(onProduceMore, MultiProducerSingleConsumerChannelAlreadyFinishedError()) } @@ -967,25 +901,22 @@ extension MultiProducerSingleConsumerChannel._Storage { callbackToken: UInt64, continuation: UnsafeContinuation ) -> EnqueueContinuationAction? { - switch self._state.take()! { - case .initial: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - + switch consume self._state { case .channeling(var channeling): if let index = channeling.cancelledAsyncProducers.firstIndex(of: callbackToken) { // Our producer got marked as cancelled. channeling.cancelledAsyncProducers.remove(at: index) - self._state = .channeling(channeling) + self = .init(state: .channeling(channeling)) return .resumeProducerWithError(continuation, CancellationError()) } else if channeling.hasOutstandingDemand { // We hit an edge case here where we wrote but the consuming thread got interleaved - self._state = .channeling(channeling) + self = .init(state: .channeling(channeling)) return .resumeProducer(continuation) } else { channeling.suspendedProducers.append((callbackToken, .continuation(continuation))) - self._state = .channeling(channeling) + self = .init(state: .channeling(channeling)) return .none } @@ -993,14 +924,14 @@ extension MultiProducerSingleConsumerChannel._Storage { case .sourceFinished(let sourceFinished): // Since we are unlocking between sending elements and suspending the send // It can happen that the source got finished or the consumption fully finishes. - self._state = .sourceFinished(sourceFinished) + self = .init(state: .sourceFinished(sourceFinished)) return .resumeProducerWithError(continuation, MultiProducerSingleConsumerChannelAlreadyFinishedError()) case .finished(let finished): // Since we are unlocking between sending elements and suspending the send // It can happen that the source got finished or the consumption fully finishes. - self._state = .finished(finished) + self = .init(state: .finished(finished)) return .resumeProducerWithError(continuation, MultiProducerSingleConsumerChannelAlreadyFinishedError()) } @@ -1017,37 +948,33 @@ extension MultiProducerSingleConsumerChannel._Storage { mutating func cancelProducer( callbackToken: UInt64 ) -> CancelProducerAction? { - //print(#function, self._state.description) - switch self._state.take()! { - case .initial: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - + switch consume self._state { case .channeling(var channeling): guard let index = channeling.suspendedProducers.firstIndex(where: { $0.0 == callbackToken }) else { // The task that sends was cancelled before sending elements so the cancellation handler // got invoked right away channeling.cancelledAsyncProducers.append(callbackToken) - self._state = .channeling(channeling) + self = .init(state: .channeling(channeling)) return .none } // We have an enqueued producer that we need to resume now let continuation = channeling.suspendedProducers.remove(at: index).1 - self._state = .channeling(channeling) + self = .init(state: .channeling(channeling)) return .resumeProducerWithCancellationError(continuation) case .sourceFinished(let sourceFinished): // Since we are unlocking between sending elements and suspending the send // It can happen that the source got finished or the consumption fully finishes. - self._state = .sourceFinished(sourceFinished) + self = .init(state: .sourceFinished(sourceFinished)) return .none case .finished(let finished): // Since we are unlocking between sending elements and suspending the send // It can happen that the source got finished or the consumption fully finishes. - self._state = .finished(finished) + self = .init(state: .finished(finished)) return .none } @@ -1073,32 +1000,18 @@ extension MultiProducerSingleConsumerChannel._Storage { @inlinable mutating func finish(_ failure: Failure?) -> FinishAction? { - switch self._state.take()! { - case .initial(let initial): - // Nothing was sent 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) - + switch consume self._state { case .channeling(let channeling): guard let consumerContinuation = channeling.consumerContinuation else { // We don't have a suspended consumer so we are just going to mark // the source as finished and terminate the current suspended producers. - self._state = .sourceFinished( + self = .init(state: .sourceFinished( .init( iteratorInitialized: channeling.iteratorInitialized, buffer: channeling.buffer, failure: failure, onTermination: channeling.onTermination - ) + )) ) return .resumeProducers(producerContinuations: .init(channeling.suspendedProducers.lazy.map { $0.1 })) @@ -1108,7 +1021,7 @@ extension MultiProducerSingleConsumerChannel._Storage { // and resume the continuation with the failure precondition(channeling.buffer.isEmpty, "Expected an empty buffer") - self._state = .finished(.init(iteratorInitialized: channeling.iteratorInitialized, sourceFinished: true)) + self = .init(state: .finished(.init(iteratorInitialized: channeling.iteratorInitialized, sourceFinished: true))) return .resumeConsumerAndCallOnTermination( consumerContinuation: consumerContinuation, @@ -1118,13 +1031,13 @@ extension MultiProducerSingleConsumerChannel._Storage { case .sourceFinished(let sourceFinished): // If the source has finished, finishing again has no effect. - self._state = .sourceFinished(sourceFinished) + self = .init(state: .sourceFinished(sourceFinished)) return .none case .finished(var finished): finished.sourceFinished = true - self._state = .finished(finished) + self = .init(state: .finished(finished)) return .none } } @@ -1146,24 +1059,7 @@ extension MultiProducerSingleConsumerChannel._Storage { @inlinable mutating func next() -> NextAction { - switch self._state.take()! { - case .initial(let initial): - // We are not interacting with the backpressure strategy here because - // we are doing this inside `suspendNext` - self._state = .channeling( - .init( - backpressureStrategy: initial.backpressureStrategy, - iteratorInitialized: initial.iteratorInitialized, - onTermination: initial.onTermination, - buffer: Deque(), - consumerContinuation: nil, - producerContinuations: .init(), - cancelledAsyncProducers: .init(), - hasOutstandingDemand: false - ) - ) - - return .suspendTask + switch consume self._state { case .channeling(var channeling): guard channeling.consumerContinuation == nil else { // We have multiple AsyncIterators iterating the sequence @@ -1174,7 +1070,7 @@ extension MultiProducerSingleConsumerChannel._Storage { // There is nothing in the buffer to fulfil the demand so we need to suspend. // We are not interacting with the backpressure strategy here because // we are doing this inside `suspendNext` - self._state = .channeling(channeling) + self = .init(state: .channeling(channeling)) return .suspendTask } @@ -1184,14 +1080,14 @@ extension MultiProducerSingleConsumerChannel._Storage { guard shouldProduceMore else { // We don't have any new demand, so we can just return the element. - self._state = .channeling(channeling) + self = .init(state: .channeling(channeling)) return .returnElement(element) } // There is demand and we have to resume our producers let producers = _TinyArray(channeling.suspendedProducers.lazy.map { $0.1 }) channeling.suspendedProducers.removeAll(keepingCapacity: true) - self._state = .channeling(channeling) + self = .init(state: .channeling(channeling)) return .returnElementAndResumeProducers(element, producers) @@ -1199,16 +1095,16 @@ extension MultiProducerSingleConsumerChannel._Storage { // Check if we have an element left in the buffer and return it guard let element = sourceFinished.buffer.popFirst() else { // We are returning the queued failure now and can transition to finished - self._state = .finished(.init(iteratorInitialized: sourceFinished.iteratorInitialized, sourceFinished: true)) + self = .init(state: .finished(.init(iteratorInitialized: sourceFinished.iteratorInitialized, sourceFinished: true))) return .returnFailureAndCallOnTermination(sourceFinished.failure, sourceFinished.onTermination) } - self._state = .sourceFinished(sourceFinished) + self = .init(state: .sourceFinished(sourceFinished)) return .returnElement(element) case .finished(let finished): - self._state = .finished(finished) + self = .init(state: .finished(finished)) return .returnNil } @@ -1237,10 +1133,7 @@ extension MultiProducerSingleConsumerChannel._Storage { @inlinable mutating func suspendNext(continuation: UnsafeContinuation) -> SuspendNextAction? { - switch self._state.take()! { - case .initial: - preconditionFailure("MultiProducerSingleConsumerChannel internal inconsistency") - + switch consume self._state { case .channeling(var channeling): guard channeling.consumerContinuation == nil else { // We have multiple AsyncIterators iterating the sequence @@ -1251,7 +1144,7 @@ extension MultiProducerSingleConsumerChannel._Storage { guard let element = channeling.buffer.popFirst() else { // There is nothing in the buffer to fulfil the demand so we to store the continuation. channeling.consumerContinuation = continuation - self._state = .channeling(channeling) + self = .init(state: .channeling(channeling)) return .none } @@ -1262,14 +1155,14 @@ extension MultiProducerSingleConsumerChannel._Storage { guard shouldProduceMore else { // We don't have any new demand, so we can just return the element. - self._state = .channeling(channeling) + self = .init(state: .channeling(channeling)) return .resumeConsumerWithElement(continuation, element) } // There is demand and we have to resume our producers let producers = _TinyArray(channeling.suspendedProducers.lazy.map { $0.1 }) channeling.suspendedProducers.removeAll(keepingCapacity: true) - self._state = .channeling(channeling) + self = .init(state: .channeling(channeling)) return .resumeConsumerWithElementAndProducers(continuation, element, producers) @@ -1277,7 +1170,7 @@ extension MultiProducerSingleConsumerChannel._Storage { // Check if we have an element left in the buffer and return it guard let element = sourceFinished.buffer.popFirst() else { // We are returning the queued failure now and can transition to finished - self._state = .finished(.init(iteratorInitialized: sourceFinished.iteratorInitialized, sourceFinished: true)) + self = .init(state: .finished(.init(iteratorInitialized: sourceFinished.iteratorInitialized, sourceFinished: true))) return .resumeConsumerWithFailureAndCallOnTermination( continuation, @@ -1285,12 +1178,12 @@ extension MultiProducerSingleConsumerChannel._Storage { sourceFinished.onTermination ) } - self._state = .sourceFinished(sourceFinished) + self = .init(state: .sourceFinished(sourceFinished)) return .resumeConsumerWithElement(continuation, element) case .finished(let finished): - self._state = .finished(finished) + self = .init(state: .finished(finished)) return .resumeConsumerWithNil(continuation) } @@ -1307,12 +1200,9 @@ extension MultiProducerSingleConsumerChannel._Storage { @inlinable mutating func cancelNext() -> CancelNextAction? { - switch self._state.take()! { - case .initial: - fatalError("MultiProducerSingleConsumerChannel internal inconsistency") - + switch consume self._state { case .channeling(let channeling): - self._state = .finished(.init(iteratorInitialized: channeling.iteratorInitialized, sourceFinished: false)) + self = .init(state: .finished(.init(iteratorInitialized: channeling.iteratorInitialized, sourceFinished: false))) guard let consumerContinuation = channeling.consumerContinuation else { return .failProducersAndCallOnTermination( @@ -1330,12 +1220,12 @@ extension MultiProducerSingleConsumerChannel._Storage { ) case .sourceFinished(let sourceFinished): - self._state = .sourceFinished(sourceFinished) + self = .init(state: .sourceFinished(sourceFinished)) return .none case .finished(let finished): - self._state = .finished(finished) + self = .init(state: .finished(finished)) return .none } @@ -1346,26 +1236,6 @@ extension MultiProducerSingleConsumerChannel._Storage { extension MultiProducerSingleConsumerChannel._Storage._StateMachine { @usableFromInline enum _State: ~Copyable { - @usableFromInline - struct Initial: ~Copyable { - /// The backpressure strategy. - @usableFromInline - var backpressureStrategy: MultiProducerSingleConsumerChannel._InternalBackpressureStrategy - - /// Indicates if the iterator was initialized. - @usableFromInline - var iteratorInitialized: Bool - - /// The onTermination callback. - @usableFromInline - var onTermination: (@Sendable () -> Void)? - - @usableFromInline - var description: String { - "backpressure:\(self.backpressureStrategy.description) iteratorInitialized:\(self.iteratorInitialized)" - } - } - @usableFromInline struct Channeling: ~Copyable { /// The backpressure strategy. @@ -1400,11 +1270,19 @@ extension MultiProducerSingleConsumerChannel._Storage._StateMachine { @usableFromInline var hasOutstandingDemand: Bool + /// The number of active producers. + @usableFromInline + var activeProducers: UInt64 + + /// The next callback token. + @usableFromInline + var nextCallbackTokenID: UInt64 + var description: String { "backpressure:\(self.backpressureStrategy.description) iteratorInitialized:\(self.iteratorInitialized) buffer:\(self.buffer.count) consumerContinuation:\(self.consumerContinuation == nil) producerContinuations:\(self.suspendedProducers.count) cancelledProducers:\(self.cancelledAsyncProducers.count) hasOutstandingDemand:\(self.hasOutstandingDemand)" } - @usableFromInline + @inlinable init( backpressureStrategy: MultiProducerSingleConsumerChannel._InternalBackpressureStrategy, iteratorInitialized: Bool, @@ -1413,16 +1291,29 @@ extension MultiProducerSingleConsumerChannel._Storage._StateMachine { consumerContinuation: UnsafeContinuation? = nil, producerContinuations: Deque<(UInt64, _MultiProducerSingleConsumerSuspendedProducer)>, cancelledAsyncProducers: Deque, - hasOutstandingDemand: Bool) { - self.backpressureStrategy = backpressureStrategy - self.iteratorInitialized = iteratorInitialized - self.onTermination = onTermination - self.buffer = buffer - self.consumerContinuation = consumerContinuation - self.suspendedProducers = producerContinuations - self.cancelledAsyncProducers = cancelledAsyncProducers - self.hasOutstandingDemand = hasOutstandingDemand - } + hasOutstandingDemand: Bool, + activeProducers: UInt64, + nextCallbackTokenID: UInt64 + ) { + self.backpressureStrategy = backpressureStrategy + self.iteratorInitialized = iteratorInitialized + self.onTermination = onTermination + self.buffer = buffer + self.consumerContinuation = consumerContinuation + self.suspendedProducers = producerContinuations + self.cancelledAsyncProducers = cancelledAsyncProducers + self.hasOutstandingDemand = hasOutstandingDemand + self.activeProducers = activeProducers + self.nextCallbackTokenID = nextCallbackTokenID + } + + /// Generates the next callback token. + @inlinable + mutating func nextCallbackToken() -> UInt64 { + let id = self.nextCallbackTokenID + self.nextCallbackTokenID += 1 + return id + } } @usableFromInline @@ -1447,7 +1338,7 @@ extension MultiProducerSingleConsumerChannel._Storage._StateMachine { "iteratorInitialized:\(self.iteratorInitialized) buffer:\(self.buffer.count) failure:\(self.failure == nil)" } - @usableFromInline + @inlinable init( iteratorInitialized: Bool, buffer: Deque, @@ -1485,9 +1376,6 @@ extension MultiProducerSingleConsumerChannel._Storage._StateMachine { } } - /// The initial state. - case initial(Initial) - /// The state once either any element was sent or `next()` was called. case channeling(Channeling) @@ -1502,8 +1390,6 @@ extension MultiProducerSingleConsumerChannel._Storage._StateMachine { @usableFromInline var description: String { switch self { - case .initial(let initial): - return "initial \(initial.description)" case .channeling(let channeling): return "channeling \(channeling.description)" case .sourceFinished(let sourceFinished): @@ -1520,13 +1406,4 @@ enum _MultiProducerSingleConsumerSuspendedProducer { case closure((Result) -> Void) case continuation(UnsafeContinuation) } - -extension Optional where Wrapped: ~Copyable { - @inlinable - mutating func take() -> Self { - let result = consume self - self = nil - return result - } -} #endif diff --git a/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel.swift b/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel.swift index 39fc8fd9..5e860a89 100644 --- a/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel.swift +++ b/Sources/AsyncAlgorithms/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannel.swift @@ -87,7 +87,7 @@ public struct MultiProducerSingleConsumerChannel: Async private let backing: _Backing @frozen - public struct NewChannel: ~Copyable { + public struct ChannelAndStream: ~Copyable { public var channel: MultiProducerSingleConsumerChannel public var source: Source @@ -112,7 +112,7 @@ public struct MultiProducerSingleConsumerChannel: Async of elementType: Element.Type = Element.self, throwing failureType: Failure.Type = Never.self, backpressureStrategy: Source.BackpressureStrategy - ) -> NewChannel { + ) -> ChannelAndStream { let storage = _Storage( backpressureStrategy: backpressureStrategy.internalBackpressureStrategy ) @@ -187,7 +187,7 @@ extension MultiProducerSingleConsumerChannel { public enum SendResult: ~Copyable, Sendable { /// A token that is returned when the channel's backpressure strategy indicated that production should /// be suspended. Use this token to enqueue a callback by calling the ``enqueueCallback(_:)`` method. - public struct CallbackToken: ~Copyable, Sendable { + public struct CallbackToken: Sendable { @usableFromInline let _id: UInt64 @@ -230,8 +230,7 @@ extension MultiProducerSingleConsumerChannel { } deinit { - // TODO: We can't finish here. -// self.finish() + self._storage.sourceDeinitialized() } diff --git a/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift b/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift index 155a063c..e15fccf7 100644 --- a/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift +++ b/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift @@ -16,154 +16,154 @@ import XCTest final class MultiProducerSingleConsumerChannelTests: XCTestCase { // MARK: - sequenceDeinitialized - func testSequenceDeinitialized_whenNoIterator() async throws { - let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 5, high: 10) - ) - var channel: MultiProducerSingleConsumerChannel? = channelAndStream.channel - var source = consume channelAndStream.source - - - 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(for: .seconds(0.2)) - } - } - - var onTerminationIterator = onTerminationStream.makeAsyncIterator() - _ = await onTerminationIterator.next() - - withExtendedLifetime(channel) {} - channel = nil - - let terminationResult: Void? = await onTerminationIterator.next() - XCTAssertNil(terminationResult) - - do { - _ = try { try source.send(2) }() - XCTFail("Expected an error to be thrown") - } catch { - XCTAssertTrue(error is MultiProducerSingleConsumerChannelAlreadyFinishedError) - } - - group.cancelAll() - } - } - - func testSequenceDeinitialized_whenIterator() async throws { - let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 5, high: 10) - ) - var channel: MultiProducerSingleConsumerChannel? = channelAndStream.channel - var source = consume channelAndStream.source - - var iterator = channel?.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(for: .seconds(0.2)) - } - } - - var onTerminationIterator = onTerminationStream.makeAsyncIterator() - _ = await onTerminationIterator.next() - - try withExtendedLifetime(channel) { - let writeResult = try source.send(1) - writeResult.assertIsProducerMore() - } - - channel = nil - - do { - let writeResult = try { try source.send(2) }() - writeResult.assertIsProducerMore() - } catch { - XCTFail("Expected no error to be thrown") - } - - let element1 = await iterator?.next() - XCTAssertEqual(element1, 1) - let element2 = await iterator?.next() - XCTAssertEqual(element2, 2) - - group.cancelAll() - } - } - - func testSequenceDeinitialized_whenFinished() async throws { - let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 5, high: 10) - ) - var channel: MultiProducerSingleConsumerChannel? = channelAndStream.channel - var source = consume channelAndStream.source - - 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(for: .seconds(0.2)) - } - } - - var onTerminationIterator = onTerminationStream.makeAsyncIterator() - _ = await onTerminationIterator.next() - - channel = nil + // Following tests are disabled since the channel is not getting deinited due to a known bug - let terminationResult: Void? = await onTerminationIterator.next() - XCTAssertNil(terminationResult) - XCTAssertNil(channel) - - group.cancelAll() - } - } - - func testSequenceDeinitialized_whenChanneling_andSuspendedProducer() async throws { - let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( - of: Int.self, - backpressureStrategy: .watermark(low: 1, high: 2) - ) - var channel: MultiProducerSingleConsumerChannel? = channelAndStream.channel - var source = consume channelAndStream.source - - _ = try { try source.send(1) }() - - do { - try await withCheckedThrowingContinuation { continuation in - source.send(1) { result in - continuation.resume(with: result) - } - - channel = nil - _ = channel?.makeAsyncIterator() - } - } catch { - XCTAssertTrue(error is MultiProducerSingleConsumerChannelAlreadyFinishedError) - } - } +// func testSequenceDeinitialized_whenNoIterator() async throws { +// var channelAndStream: MultiProducerSingleConsumerChannel.ChannelAndStream! = MultiProducerSingleConsumerChannel.makeChannel( +// of: Int.self, +// backpressureStrategy: .watermark(low: 5, high: 10) +// ) +// var channel: MultiProducerSingleConsumerChannel? = channelAndStream.channel +// var source = channelAndStream.source +// channelAndStream = nil +// +// let (onTerminationStream, onTerminationContinuation) = AsyncStream.makeStream() +// source.onTermination = { +// onTerminationContinuation.finish() +// } +// +// await withThrowingTaskGroup(of: Void.self) { group in +// group.addTask { +// onTerminationContinuation.yield() +// try await Task.sleep(for: .seconds(10)) +// } +// +// var onTerminationIterator = onTerminationStream.makeAsyncIterator() +// _ = await onTerminationIterator.next() +// +// withExtendedLifetime(channel) {} +// channel = nil +// +// let terminationResult: Void? = await onTerminationIterator.next() +// XCTAssertNil(terminationResult) +// +// do { +// _ = try { try source.send(2) }() +// XCTFail("Expected an error to be thrown") +// } catch { +// XCTAssertTrue(error is MultiProducerSingleConsumerChannelAlreadyFinishedError) +// } +// +// group.cancelAll() +// } +// } +// +// func testSequenceDeinitialized_whenIterator() async throws { +// let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( +// of: Int.self, +// backpressureStrategy: .watermark(low: 5, high: 10) +// ) +// var channel: MultiProducerSingleConsumerChannel? = channelAndStream.channel +// var source = consume channelAndStream.source +// +// var iterator = channel?.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(for: .seconds(0.2)) +// } +// } +// +// var onTerminationIterator = onTerminationStream.makeAsyncIterator() +// _ = await onTerminationIterator.next() +// +// try withExtendedLifetime(channel) { +// let writeResult = try source.send(1) +// writeResult.assertIsProducerMore() +// } +// +// channel = nil +// +// do { +// let writeResult = try { try source.send(2) }() +// writeResult.assertIsProducerMore() +// } catch { +// XCTFail("Expected no error to be thrown") +// } +// +// let element1 = await iterator?.next() +// XCTAssertEqual(element1, 1) +// let element2 = await iterator?.next() +// XCTAssertEqual(element2, 2) +// +// group.cancelAll() +// } +// } +// +// func testSequenceDeinitialized_whenFinished() async throws { +// let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( +// of: Int.self, +// backpressureStrategy: .watermark(low: 5, high: 10) +// ) +// var channel: MultiProducerSingleConsumerChannel? = channelAndStream.channel +// var source = consume channelAndStream.source +// +// 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(for: .seconds(0.2)) +// } +// } +// +// var onTerminationIterator = onTerminationStream.makeAsyncIterator() +// _ = await onTerminationIterator.next() +// +// channel = nil +// +// let terminationResult: Void? = await onTerminationIterator.next() +// XCTAssertNil(terminationResult) +// XCTAssertNil(channel) +// +// group.cancelAll() +// } +// } +// +// func testSequenceDeinitialized_whenChanneling_andSuspendedProducer() async throws { +// let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( +// of: Int.self, +// backpressureStrategy: .watermark(low: 1, high: 2) +// ) +// var channel: MultiProducerSingleConsumerChannel? = channelAndStream.channel +// var source = consume channelAndStream.source +// +// _ = try { try source.send(1) }() +// +// do { +// try await withCheckedThrowingContinuation { continuation in +// source.send(1) { result in +// continuation.resume(with: result) +// } +// +// channel = nil +// _ = channel?.makeAsyncIterator() +// } +// } catch { +// XCTAssertTrue(error is MultiProducerSingleConsumerChannelAlreadyFinishedError) +// } +// } // MARK: - iteratorInitialized @@ -217,7 +217,7 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { backpressureStrategy: .watermark(low: 5, high: 10) ) let channel = channelAndStream.channel - var source = consume channelAndStream.source + let source = consume channelAndStream.source source.finish(throwing: nil) @@ -591,9 +591,9 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { case .enqueueCallback(let callbackToken): source.cancelCallback(callbackToken: callbackToken) -// source.enqueueCallback(callbackToken: callbackToken) { result in -// producerSource.yield(with: result) -// } + source.enqueueCallback(callbackToken: callbackToken) { result in + producerSource.yield(with: result) + } } do { @@ -634,443 +634,371 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { let element = await channel.first { _ in true } XCTAssertEqual(element, 1) } -// -// func testEnqueueProducer_whenChanneling_andInterleaving() async throws { -// let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( -// of: Int.self, -// backpressureStrategy: .watermark(low: 1, high: 1) -// ) -// defer { -// source.finish() -// } -// var iterator = stream.makeAsyncIterator() -// -// let (producerStream, producerSource) = AsyncThrowingStream.makeStream() -// -// let writeResult = try { try source.send(1) }() -// -// switch writeResult { -// case .produceMore: -// preconditionFailure() -// case .enqueueCallback(let callbackToken): -// let element = 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_whenChanneling_andSuspending() async throws { -// let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( -// of: Int.self, -// backpressureStrategy: .watermark(low: 1, high: 1) -// ) -// defer { -// source.finish() -// } -// var iterator = stream.makeAsyncIterator() -// -// let (producerStream, producerSource) = AsyncThrowingStream.makeStream() -// -// let writeResult = try { try source.send(1) }() -// -// switch writeResult { -// case .produceMore: -// preconditionFailure() -// case .enqueueCallback(let callbackToken): -// source.enqueueCallback(callbackToken: callbackToken) { result in -// producerSource.yield(with: result) -// } -// } -// -// let element = 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) = MultiProducerSingleConsumerChannel.makeChannel( -// of: Int.self, -// backpressureStrategy: .watermark(low: 1, high: 1) -// ) -// defer { -// source.finish() -// } -// var iterator = stream.makeAsyncIterator() -// -// let (producerStream, producerSource) = AsyncThrowingStream.makeStream() -// -// let writeResult = try { try source.send(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 = await iterator.next() -// XCTAssertEqual(element, 1) -// -// do { -// _ = try await producerStream.first { _ in true } -// XCTFail("Expected an error to be thrown") -// } catch { -// XCTAssertTrue(error is MultiProducerSingleConsumerChannelAlreadyFinishedError) -// } -// } -// -// // MARK: - cancelProducer -// -// func testCancelProducer_whenChanneling() async throws { -// let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( -// of: Int.self, -// backpressureStrategy: .watermark(low: 1, high: 2) -// ) -// defer { -// source.finish() -// } -// -// let (producerStream, producerSource) = AsyncThrowingStream.makeStream() -// -// try await source.send(1) -// -// let writeResult = try { try source.send(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 = await stream.first { _ in true } -// XCTAssertEqual(element, 1) -// } -// -// func testCancelProducer_whenSourceFinished() async throws { -// let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( -// of: Int.self, -// backpressureStrategy: .watermark(low: 1, high: 2) -// ) -// -// let (producerStream, producerSource) = AsyncThrowingStream.makeStream() -// -// try await source.send(1) -// -// let writeResult = try { try source.send(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 MultiProducerSingleConsumerChannelAlreadyFinishedError) -// } -// -// let element = await stream.first { _ in true } -// XCTAssertEqual(element, 1) -// } -// -// // MARK: - finish -// -// func testFinish_whenChanneling_andConsumerSuspended() async throws { -// let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( -// of: Int.self, -// backpressureStrategy: .watermark(low: 1, high: 1) -// ) -// -// try await withThrowingTaskGroup(of: Int?.self) { group in -// group.addTask { -// return 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(for: .seconds(0.5)) -// -// source.finish(throwing: nil) -// let element = try await group.next() -// XCTAssertEqual(element, .some(nil)) -// } -// } -// -// func testFinish_whenInitial() async throws { -// let (stream, source) = MultiProducerSingleConsumerChannel.makeChannel( -// of: Int.self, -// throwing: Error.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 (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( -// of: Int.self, -// backpressureStrategy: .watermark(low: 2, high: 4) -// ) -// defer { -// source.finish() -// } -// -// let (backpressureEventStream, backpressureEventContinuation) = AsyncStream.makeStream(of: Void.self) -// -// await withThrowingTaskGroup(of: Void.self) { group in -// group.addTask { -// while true { -// backpressureEventContinuation.yield(()) -// try await source.send(contentsOf: [1]) -// } -// } -// -// var backpressureEventIterator = backpressureEventStream.makeAsyncIterator() -// var iterator = channel.makeAsyncIterator() -// -// await backpressureEventIterator.next() -// await backpressureEventIterator.next() -// await backpressureEventIterator.next() -// await backpressureEventIterator.next() -// -// _ = await iterator.next() -// _ = await iterator.next() -// _ = await iterator.next() -// -// await backpressureEventIterator.next() -// await backpressureEventIterator.next() -// await backpressureEventIterator.next() -// -// group.cancelAll() -// } -// } -// -// func testBackpressureSync() async throws { -// let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( -// of: Int.self, -// backpressureStrategy: .watermark(low: 2, high: 4) -// ) -// defer { -// source.finish() -// } -// -// let (backpressureEventStream, backpressureEventContinuation) = AsyncStream.makeStream(of: Void.self) -// -// await withThrowingTaskGroup(of: Void.self) { group in -// group.addTask { -// @Sendable func yield() { -// backpressureEventContinuation.yield(()) -// source.send(contentsOf: [1]) { result in -// switch result { -// case .success: -// yield() -// -// case .failure: -// break -// } -// } -// } -// -// yield() -// } -// -// var backpressureEventIterator = backpressureEventStream.makeAsyncIterator() -// var iterator = channel.makeAsyncIterator() -// -// await backpressureEventIterator.next() -// await backpressureEventIterator.next() -// await backpressureEventIterator.next() -// await backpressureEventIterator.next() -// -// _ = await iterator.next() -// _ = await iterator.next() -// _ = await iterator.next() -// -// await backpressureEventIterator.next() -// await backpressureEventIterator.next() -// await backpressureEventIterator.next() -// -// group.cancelAll() -// } -// } -// -// func testWatermarkWithCustomCoount() async throws { -// let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( -// of: [Int].self, -// backpressureStrategy: .watermark(low: 2, high: 4, waterLevelForElement: { $0.count }) -// ) -// defer { -// source.finish() -// } -// var iterator = channel.makeAsyncIterator() -// -// try await source.send([1, 1, 1]) -// -// _ = await iterator.next() -// -// try await source.send([1, 1, 1]) -// -// _ = await iterator.next() -// } -// -// func testWatermarWithLotsOfElements() async throws { -// // This test should in the future use a custom task executor to schedule to avoid sending -// // 1000 elements. -// let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( -// of: Int.self, -// backpressureStrategy: .watermark(low: 2, high: 4) -// ) -// await withThrowingTaskGroup(of: Void.self) { group in -// group.addTask { -// for i in 0...10000 { -// try await source.send(i) -// } -// source.finish() -// } -// -// group.addTask { -// var sum = 0 -// for try await element in channel { -// sum += element -// } -// } -// } -// } -// -// func testThrowsError() async throws { -// let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( -// of: Int.self, -// throwing: Error.self, -// backpressureStrategy: .watermark(low: 2, high: 4) -// ) -// -// try await source.send(1) -// try await source.send(2) -// source.finish(throwing: CancellationError()) -// -// var elements = [Int]() -// var iterator = channel.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 (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( -// of: Int.self, -// backpressureStrategy: .watermark(low: 2, high: 4) -// ) -// -// continuation.yield(1) -// continuation.yield(2) -// continuation.finish() -// -// try await source.send(contentsOf: stream) -// source.finish(throwing: nil) -// -// let elements = await channel.collect() -// XCTAssertEqual(elements, [1, 2]) -// } -// -// // MARK: NonThrowing -// -// func testNonThrowing() async throws { -// let (channel, source) = MultiProducerSingleConsumerChannel.makeChannel( -// of: Int.self, -// backpressureStrategy: .watermark(low: 2, high: 4) -// ) -// defer { -// source.finish() -// } -// -// let (backpressureEventStream, backpressureEventContinuation) = AsyncStream.makeStream(of: Void.self) -// -// await withThrowingTaskGroup(of: Void.self) { group in -// group.addTask { -// while true { -// backpressureEventContinuation.yield(()) -// try await source.send(contentsOf: [1]) -// } -// } -// -// var backpressureEventIterator = backpressureEventStream.makeAsyncIterator() -// var iterator = channel.makeAsyncIterator() -// -// await backpressureEventIterator.next() -// await backpressureEventIterator.next() -// await backpressureEventIterator.next() -// await backpressureEventIterator.next() -// -// _ = await iterator.next() -// _ = await iterator.next() -// _ = await iterator.next() -// -// await backpressureEventIterator.next() -// await backpressureEventIterator.next() -// await backpressureEventIterator.next() -// -// group.cancelAll() -// } -// } + + func testEnqueueProducer_whenChanneling_andInterleaving() async throws { + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + backpressureStrategy: .watermark(low: 1, high: 1) + ) + let channel = channelAndStream.channel + var source = consume channelAndStream.source + var iterator = channel.makeAsyncIterator() + + let (producerStream, producerSource) = AsyncThrowingStream.makeStream() + + let writeResult = try { try source.send(1) }() + + switch writeResult { + case .produceMore: + preconditionFailure() + case .enqueueCallback(let callbackToken): + let element = 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_whenChanneling_andSuspending() async throws { + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + backpressureStrategy: .watermark(low: 1, high: 1) + ) + let channel = channelAndStream.channel + var source = consume channelAndStream.source + var iterator = channel.makeAsyncIterator() + + let (producerStream, producerSource) = AsyncThrowingStream.makeStream() + + let writeResult = try { try source.send(1) }() + + switch writeResult { + case .produceMore: + preconditionFailure() + case .enqueueCallback(let callbackToken): + source.enqueueCallback(callbackToken: callbackToken) { result in + producerSource.yield(with: result) + } + } + + let element = await iterator.next() + XCTAssertEqual(element, 1) + + do { + _ = try await producerStream.first { _ in true } + } catch { + XCTFail("Expected no error to be thrown") + } + } + + // MARK: - cancelProducer + + func testCancelProducer_whenChanneling() async throws { + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + backpressureStrategy: .watermark(low: 1, high: 2) + ) + let channel = channelAndStream.channel + var source = consume channelAndStream.source + + let (producerStream, producerSource) = AsyncThrowingStream.makeStream() + + try await source.send(1) + + let writeResult = try { try source.send(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 = await channel.first { _ in true } + XCTAssertEqual(element, 1) + } + + // MARK: - finish + + func testFinish_whenChanneling_andConsumerSuspended() async throws { + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + backpressureStrategy: .watermark(low: 1, high: 1) + ) + let channel = channelAndStream.channel + var source: MultiProducerSingleConsumerChannel.Source? = consume channelAndStream.source + + try await withThrowingTaskGroup(of: Int?.self) { group in + group.addTask { + return await channel.first { $0 == 2 } + } + + // This is always going to be a bit racy since we need the call to next() suspend + try await Task.sleep(for: .seconds(0.5)) + + source?.finish(throwing: nil) + source = nil + + let element = try await group.next() + XCTAssertEqual(element, .some(nil)) + } + } + + func testFinish_whenInitial() async throws { + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + throwing: Error.self, + backpressureStrategy: .watermark(low: 1, high: 1) + ) + let channel = channelAndStream.channel + var source = consume channelAndStream.source + + source.finish(throwing: CancellationError()) + + do { + for try await _ in channel {} + XCTFail("Expected an error to be thrown") + } catch { + XCTAssertTrue(error is CancellationError) + } + + } + + // MARK: - Backpressure + + func testBackpressure() async throws { + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + backpressureStrategy: .watermark(low: 2, high: 4) + ) + let channel = channelAndStream.channel + var source = consume channelAndStream.source + + let (backpressureEventStream, backpressureEventContinuation) = AsyncStream.makeStream(of: Void.self) + + await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + while true { + backpressureEventContinuation.yield(()) + try await source.send(contentsOf: [1]) + } + } + + var backpressureEventIterator = backpressureEventStream.makeAsyncIterator() + var iterator = channel.makeAsyncIterator() + + await backpressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() + + _ = await iterator.next() + _ = await iterator.next() + _ = await iterator.next() + + await backpressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() + + group.cancelAll() + } + } + + func testBackpressureSync() async throws { + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + backpressureStrategy: .watermark(low: 2, high: 4) + ) + let channel = channelAndStream.channel + var source = consume channelAndStream.source + + let (backpressureEventStream, backpressureEventContinuation) = AsyncStream.makeStream(of: Void.self) + + await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + while true { + backpressureEventContinuation.yield(()) + try await withCheckedThrowingContinuation { continuation in + source.send(contentsOf: [1]) { result in + continuation.resume(with: result) + } + } + } + } + + var backpressureEventIterator = backpressureEventStream.makeAsyncIterator() + var iterator = channel.makeAsyncIterator() + + await backpressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() + + _ = await iterator.next() + _ = await iterator.next() + _ = await iterator.next() + + await backpressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() + + group.cancelAll() + } + } + + func testWatermarkWithCustomCoount() async throws { + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: [Int].self, + backpressureStrategy: .watermark(low: 2, high: 4, waterLevelForElement: { $0.count }) + ) + let channel = channelAndStream.channel + var source = consume channelAndStream.source + var iterator = channel.makeAsyncIterator() + + try await source.send([1, 1, 1]) + + _ = await iterator.next() + + try await source.send([1, 1, 1]) + + _ = await iterator.next() + } + + func testWatermarWithLotsOfElements() async throws { + // This test should in the future use a custom task executor to schedule to avoid sending + // 1000 elements. + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + backpressureStrategy: .watermark(low: 2, high: 4) + ) + let channel = channelAndStream.channel + var source: MultiProducerSingleConsumerChannel.Source! = consume channelAndStream.source + await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + var source = source.take()! + for i in 0...10000 { + try await source.send(i) + } + source.finish() + } + + group.addTask { + var sum = 0 + for try await element in channel { + sum += element + } + } + } + } + + func testThrowsError() async throws { + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + throwing: Error.self, + backpressureStrategy: .watermark(low: 2, high: 4) + ) + let channel = channelAndStream.channel + var source = consume channelAndStream.source + + try await source.send(1) + try await source.send(2) + source.finish(throwing: CancellationError()) + + var elements = [Int]() + var iterator = channel.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 channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + backpressureStrategy: .watermark(low: 2, high: 4) + ) + let channel = channelAndStream.channel + var source = consume channelAndStream.source + + continuation.yield(1) + continuation.yield(2) + continuation.finish() + + try await source.send(contentsOf: stream) + source.finish(throwing: nil) + + let elements = await channel.collect() + XCTAssertEqual(elements, [1, 2]) + } + + // MARK: NonThrowing + + func testNonThrowing() async throws { + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + backpressureStrategy: .watermark(low: 2, high: 4) + ) + let channel = channelAndStream.channel + var source = consume channelAndStream.source + + let (backpressureEventStream, backpressureEventContinuation) = AsyncStream.makeStream(of: Void.self) + + await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + while true { + backpressureEventContinuation.yield(()) + try await source.send(contentsOf: [1]) + } + } + + var backpressureEventIterator = backpressureEventStream.makeAsyncIterator() + var iterator = channel.makeAsyncIterator() + + await backpressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() + + _ = await iterator.next() + _ = await iterator.next() + _ = await iterator.next() + + await backpressureEventIterator.next() + await backpressureEventIterator.next() + await backpressureEventIterator.next() + + group.cancelAll() + } + } } extension AsyncSequence { @@ -1103,3 +1031,11 @@ extension MultiProducerSingleConsumerChannel.Source.SendResult { } } } + +extension Optional where Wrapped: ~Copyable { + fileprivate mutating func take() -> Self { + let result = consume self + self = nil + return result + } +} From e64fef726ce166078dd3866c5b3db53997ab1a48 Mon Sep 17 00:00:00 2001 From: ser-0xff <122270051+ser-0xff@users.noreply.github.com> Date: Thu, 31 Oct 2024 10:52:12 +0300 Subject: [PATCH 6/7] Add regression test for case when one of two channel sources deinitialized. --- ...tiProducerSingleConsumerChannelTests.swift | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift b/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift index e15fccf7..1596280e 100644 --- a/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift +++ b/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift @@ -489,6 +489,41 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { } } + func testOneOfTwoSourcesDeinitialized() async throws { + let channelAndStream = MultiProducerSingleConsumerChannel.makeChannel( + of: Int.self, + throwing: Never.self, + backpressureStrategy: .watermark(low: 5, high: 10) + ) + + let channel = channelAndStream.channel + let consumerTask = Task { + var count = 0 + for await _ in channel { + count += 1 + if count == 2 { + break + } + } + return count + } + + var source = consume channelAndStream.source + _ = try await { + var source = source.copy() + _ = try await source.send(1) + }() + + do { + _ = try await source.send(2) + } catch { + XCTFail("source.send() unexpectedly failed \(error)") + } + + let consumedEvents = await consumerTask.value + XCTAssertEqual(consumedEvents, 2) + } + // MARK: - write func testWrite_whenInitial() async throws { From 698b51b1ccf4de736f38710276ce1b486c4187a7 Mon Sep 17 00:00:00 2001 From: ser-0xff <122270051+ser-0xff@users.noreply.github.com> Date: Thu, 31 Oct 2024 11:25:58 +0300 Subject: [PATCH 7/7] Cosmetic fix. --- .../MultiProducerSingleConsumerChannelTests.swift | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift b/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift index 1596280e..cca9e25c 100644 --- a/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift +++ b/Tests/AsyncAlgorithmsTests/MultiProducerSingleConsumerChannel/MultiProducerSingleConsumerChannelTests.swift @@ -510,8 +510,8 @@ final class MultiProducerSingleConsumerChannelTests: XCTestCase { var source = consume channelAndStream.source _ = try await { - var source = source.copy() - _ = try await source.send(1) + var sourceCopy = source.copy() + _ = try await sourceCopy.send(1) }() do {