operators: implement multicast

Author: Thibault Wittemberg

CHANGELOG.md

@@ -1,6 +1,7 @@
-**v0.1.1 **
+**v0.2.0 **
 
-- AsyncStreams.CurrentValue `element` made public
+- AsyncStreams.CurrentValue `element` made public and available with get/set
+- new Multicast operator
 
 **v0.1.0 - Hydrogen:**
 

README.md

@@ -42,6 +42,7 @@ AsyncSequences
 * [FlatMapLatest](#FlatMapLatest)
 * [HandleEvents](#HandleEvents)
 * [Assign](#Assign)
+* [Multicast] (#Multicast)
 * [EraseToAnyAsyncSequence](#EraseToAnyAsyncSequence)
 
 More operators and extensions are to come. Pull requests are of course welcome.
@@ -237,13 +238,13 @@ The current value is replayed for any new async loop.
 let currentValue = AsyncStreams.CurrentValue<Int>(1)
 
 Task {
-    for try await element in passthrough {
+    for try await element in currentValue {
         print(element) // will print 1 2
     }
 }
 
 Task {
-    for try await element in passthrough {
+    for try await element in currentValue {
         print(element) // will print 1 2
     }
 }
@@ -391,6 +392,42 @@ let fromSequence = AsyncSequences.From(["1", "2", "3"])
 try await fromSequence.assign(to: \.property, on: root) // will set the property value to "1", "2", "3"
 ```
 
+### Multicast
+
+`multicast(_:)` is useful when you have multiple client loops, but you want the upstream async sequence to only produce a single `AsyncIterator`.
+
+```swift
+let stream = AsyncStreams.Passthrough<(String, Int)>()
+let multicastedAsyncSequence = ["First", "Second", "Third"]
+    .asyncElements
+    .map { ($0, Int.random(in: 0...100)) }
+    .handleEvents(onElement: { print("AsyncSequence produces: ($0)") })
+    .multicast(stream)
+
+Task {
+    try await multicastedAsyncSequence
+        .collect { print ("Task 1 received: \($0)") }
+}
+
+Task {
+    try await multicastedAsyncSequence
+        .collect { print ("Task 2 received: \($0)") }
+}
+
+multicastedAsyncSequence.connect()
+
+// will print:
+// AsyncSequence produces: ("First", 78)
+// Stream 2 received: ("First", 78)
+// Stream 1 received: ("First", 78)
+// AsyncSequence produces: ("Second", 98)
+// Stream 2 received: ("Second", 98)
+// Stream 1 received: ("Second", 98)
+// AsyncSequence produces: ("Third", 61)
+// Stream 2 received: ("Third", 61)
+// Stream 1 received: ("Third", 61)
+```
+
 ### EraseToAnyAsyncSequence
 
 `eraseToAnyAsyncSequence()` type-erases the async sequence into an AnyAsyncSequence.

Sources/AsyncSequences/AsyncSequences+From.swift

@@ -50,7 +50,7 @@ public struct AsyncFromSequence<BaseSequence: Sequence>: AsyncSequence {
     }
 
     public func makeAsyncIterator() -> AsyncIterator {
-        Iterator(baseIterator: self.baseSequence.makeIterator(), interval: self.interval)
+        return Iterator(baseIterator: self.baseSequence.makeIterator(), interval: self.interval)
     }
 
     public struct Iterator: AsyncIteratorProtocol {

Sources/AsyncSequences/AsyncSequences+Merge.swift

@@ -4,6 +4,7 @@
 //
 //  Created by Thibault Wittemberg on 31/12/2021.
 //
+import Foundation
 
 public extension AsyncSequences {
     /// `Merge` is an AsyncSequence that merges several async sequences respecting
@@ -63,63 +64,89 @@ public struct AsyncMergeSequence<UpstreamAsyncSequence: AsyncSequence>: AsyncSeq
     }
 
     public func makeAsyncIterator() -> Iterator {
-        return Iterator(upstreamIterators: self.upstreamAsyncSequences.map { $0.makeAsyncIterator() })
+        return Iterator(upstreamAsyncSequences: self.upstreamAsyncSequences)
     }
 
     enum UpstreamElement {
         case element(Element)
         case finished
     }
 
-    actor ElementCounter {
-        var counter = 0
-
-        func increaseCounter() {
-            self.counter += 1
-        }
-
-        func decreaseCounter() {
-            guard self.counter > 0 else { return }
-            self.counter -= 1
-        }
-
-        func hasElement() -> Bool {
-            self.counter > 0
-        }
-    }
+    //    actor ElementCounter {
+    //        var counter = 0
+    //
+    //        func increaseCounter() {
+    //            self.counter += 1
+    //        }
+    //
+    //        func decreaseCounter() {
+    //            guard self.counter > 0 else { return }
+    //            self.counter -= 1
+    //        }
+    //
+    //        func hasElement() -> Bool {
+    //            self.counter > 0
+    //        }
+    //    }
 
     public struct Iterator: AsyncIteratorProtocol {
-        let sink = AsyncStreams.Passthrough<UpstreamElement>()
-        var sinkIterator: AsyncStreams.Passthrough<UpstreamElement>.AsyncIterator
-        let upstreamIterators: [SharedAsyncIterator<UpstreamAsyncSequence.AsyncIterator>]
-        let elementCounter = ElementCounter()
+
+        var downstreamIterator: AsyncStreams.Passthrough<UpstreamElement>.AsyncIterator
+        let upstreamAsyncSequenceRegulators: [ConcurrentAccessRegulator<UpstreamAsyncSequence>]
+        // let elementCounter = ElementCounter()
         var numberOfFinished = 0
 
-        public init(upstreamIterators: [UpstreamAsyncSequence.AsyncIterator]) {
-            self.upstreamIterators = upstreamIterators.map { SharedAsyncIterator(iterator: $0) }
-            self.sinkIterator = self.sink.makeAsyncIterator()
+        public init(upstreamAsyncSequences: [UpstreamAsyncSequence]) {
+            let downstreamStream = AsyncStreams.Passthrough<UpstreamElement>()
+            self.downstreamIterator = downstreamStream.makeAsyncIterator()
+            self.upstreamAsyncSequenceRegulators = upstreamAsyncSequences.map { upstreamAsyncSequence in
+                var isAlreadyFinished = false
+                return ConcurrentAccessRegulator(
+                    upstreamAsyncSequence,
+                    onNext: { element in
+                        if let nonNilElement = element {
+                            // await localElementCounter.increaseCounter()
+                            downstreamStream.send(.element(nonNilElement))
+                            return
+                        }
+
+                        guard !isAlreadyFinished else { return }
+
+                        isAlreadyFinished = true
+                        downstreamStream.send(.finished)
+                    },
+                    onError: { error in
+                        downstreamStream.send(termination: .failure(error))
+                    },
+                    onCancel: {
+                        guard !isAlreadyFinished else { return }
+                        isAlreadyFinished = true
+                        downstreamStream.send(.finished)
+                    })
+            }
         }
 
-        mutating func nextElementFromSink() async throws -> Element? {
+        mutating func nextElementFromDownstreamStream() async throws -> Element? {
             var noValue = true
             var value: Element?
 
             // we now have to eliminate the intermediate ".finished" values until the next
             // true value is found.
             // if every upstream iterator has finished, then the zipped async sequence is also finished
             while noValue {
-                guard let nextChildElement = try await self.sinkIterator.next() else {
-                    // the sink stream is finished, so is the zipped async sequence
+                guard let nextChildElement = try await self.downstreamIterator.next() else {
+                    // the downstream stream is finished, so is the zipped async sequence
                     noValue = false
                     value = nil
                     break
                 }
 
                 switch nextChildElement {
                 case .finished:
+
                     self.numberOfFinished += 1
 
-                    if self.numberOfFinished == self.upstreamIterators.count {
+                    if self.numberOfFinished == self.upstreamAsyncSequenceRegulators.count {
                         noValue = false
                         value = nil
                         break
@@ -128,7 +155,7 @@ public struct AsyncMergeSequence<UpstreamAsyncSequence: AsyncSequence>: AsyncSeq
                     // nominal case: a next element is available
                     noValue = false
                     value = element
-                    await self.elementCounter.decreaseCounter()
+                    // await self.elementCounter.decreaseCounter()
                 }
             }
 
@@ -138,57 +165,27 @@ public struct AsyncMergeSequence<UpstreamAsyncSequence: AsyncSequence>: AsyncSeq
         public mutating func next() async throws -> Element? {
             guard !Task.isCancelled else { return nil }
 
-            // before requesting elements from the upstream iterators, we should reauest the next element from the sink iterator
+            // before requesting elements from the upstream iterators, we should request the next element from the sink iterator
             // if it has some stacked values
-
             // for now we leave it commented as I'm not sure it is not counterproductive.
             // This "early" drain might prevent from requesting the next available upstream iterators as soon as possible
             // since the sink iterator might deliver a value and the next will return right away
 
-//            guard await !self.elementCounter.hasElement() else {
-//                return try await self.nextElementFromSink()
-//            }
-
-            let localSink = self.sink
-            let localElementCounter = self.elementCounter
+            //            guard await !self.elementCounter.hasElement() else {
+            //                return try await self.nextElementFromSink()
+            //            }
 
             // iterating over the upstream iterators to ask for their next element. Only
-            // the available iterators are requested (not already being computing the next
+            // the available iterators will respond (not already being computing the next
             // element from the previous iteration and not already finished)
-            for upstreamIterator in self.upstreamIterators {
-                guard !Task.isCancelled else { break }
-
-                let localUpstreamIterator = upstreamIterator
-                guard await !localUpstreamIterator.isFinished() else { continue }
+            for upstreamAsyncSequenceRegulator in self.upstreamAsyncSequenceRegulators {
                 Task {
-                    do {
-                        let nextSharedElement = try await localUpstreamIterator.next()
-
-                        // if the next element is nil, it means one of the upstream iterator
-                        // is finished ... its does not mean the zipped async sequence is finished (all upstream iterators have to be finished)
-                        guard let nextNonNilSharedElement = nextSharedElement else {
-                            await localSink.send(.finished)
-                            return
-                        }
-
-                        guard case let .value(nextElement) = nextNonNilSharedElement else {
-                            // the upstream iterator was not available ... see you at the next iteration
-                            return
-                        }
-
-                        // we have a next element from an upstream iterator, pushing it in the sink stream
-                        await localSink.send(.element(nextElement))
-                        await localElementCounter.increaseCounter()
-                    } catch is CancellationError {
-                        await localSink.send(.finished)
-                    } catch {
-                        await localSink.send(termination: .failure(error))
-                    }
+                    await upstreamAsyncSequenceRegulator.requestNextIfAvailable()
                 }
             }
 
             // we wait for the sink iterator to deliver the next element
-            return try await self.nextElementFromSink()
+            return try await self.nextElementFromDownstreamStream()
         }
     }
 }

Sources/AsyncStreams/AsyncStreams+CurrentValue.swift

@@ -36,8 +36,8 @@ public extension AsyncStreams {
     typealias CurrentValue<Element> = AsyncCurrentValueStream<Element>
 }
 
-public final class AsyncCurrentValueStream<Element>: Stream, Sendable {
-    actor Storage {
+public final class AsyncCurrentValueStream<Element>: Stream, @unchecked Sendable {
+    final class Storage {
         var element: Element
 
         init(_ element: Element) {
@@ -54,12 +54,27 @@ public final class AsyncCurrentValueStream<Element>: Stream, Sendable {
     }
 
     public typealias AsyncIterator = AsyncStreams.Iterator<Element>
+
+    // we must make sure the inner continuations and storage can be used in a concurrent context since there can be multiple
+    // operations happening at the same time (concurrent registrations and sendings).
+    // we could use an Actor to enforce that BUT there is a drawback. If we use an Actor to handle Continuations,
+    // when registering a new continuation, the register function would have to be called within a Task
+    // because of its async nature. Doing so, it means that we could call `send` while the registration is not done and we
+    // would loose the value.
+    let serialQueue = DispatchQueue(label: UUID().uuidString)
+
     let continuations = AsyncStreams.Continuations<Element>()
     let storage: Storage
 
     public var element: Element {
-        get async {
-            await self.storage.retrieve()
+        get {
+            self.serialQueue.sync {
+                self.storage.retrieve()
+            }
+        }
+
+        set {
+            self.send(newValue)
         }
     }
 
@@ -69,25 +84,27 @@ public final class AsyncCurrentValueStream<Element>: Stream, Sendable {
 
     /// Sends a value to all underlying async sequences
     /// - Parameter element: the value to send
-    public func send(_ element: Element) async {
-        await self.storage.update(element)
-        await self.continuations.send(element)
+    public func send(_ element: Element) {
+        self.serialQueue.async { [weak self] in
+            self?.storage.update(element)
+            self?.continuations.send(element)
+        }
     }
 
     /// Finishes the async sequences with either a normal ending or an error.
     /// - Parameter termination: The termination to finish the async sequence.
-    public func send(termination: Termination) async {
-        await self.continuations.send(termination)
+    public func send(termination: Termination) {
+        self.serialQueue.async { [weak self] in
+            self?.continuations.send(termination)
+        }
     }
 
     func makeStream(forClientId clientId: UUID) -> AsyncThrowingStream<Element, Error> {
-        return AsyncThrowingStream<Element, Error>(Element.self, bufferingPolicy: .unbounded) { [continuations, storage] continuation in
-            Task {
-                // registration is async because the continuations are managed by an actor (to avoid race conditions on its internal storage).
-                // registering a continuation is possible only when the actor is available.
-                let element = await storage.retrieve()
+        return AsyncThrowingStream<Element, Error>(Element.self, bufferingPolicy: .unbounded) { [weak self] continuation in
+            self?.serialQueue.async { [weak self] in
+                self?.continuations.register(continuation: continuation, forId: clientId)
+                guard let element = self?.storage.retrieve() else { return }
                 continuation.yield(element)
-                await continuations.register(continuation: continuation, forId: clientId)
             }
         }
     }
@@ -96,9 +113,12 @@ public final class AsyncCurrentValueStream<Element>: Stream, Sendable {
         let clientId = UUID()
         let stream = self.makeStream(forClientId: clientId)
         return AsyncStreams.Iterator<Element>(
-            clientId: clientId,
             baseIterator: stream.makeAsyncIterator(),
-            continuations: self.continuations
+            onCancelOrFinish: { [weak self] in
+                self?.serialQueue.async { [weak self] in
+                    self?.continuations.unregister(id: clientId)
+                }
+            }
         )
     }
 }