diff --git a/Sources/FeatherStorage/StorageSequence.swift b/Sources/FeatherStorage/StorageSequence.swift index fa85794..e97041f 100644 --- a/Sources/FeatherStorage/StorageSequence.swift +++ b/Sources/FeatherStorage/StorageSequence.swift @@ -8,13 +8,13 @@ import NIOCore /// A type-erased async sequence of storage byte buffers. public struct StorageSequence: Sendable, AsyncSequence { - typealias BaseAsyncSequence = AsyncStream> - /// An async iterator over a `StorageSequence`. public struct AsyncIterator: AsyncIteratorProtocol { - private var base: BaseAsyncSequence.AsyncIterator + private var base: any AsyncIteratorProtocol - init(base: BaseAsyncSequence.AsyncIterator) { + init( + base: any AsyncIteratorProtocol + ) { self.base = base } @@ -25,7 +25,7 @@ public struct StorageSequence: Sendable, AsyncSequence { /// - Throws: Any error emitted by the underlying async sequence. @concurrent public mutating func next() async throws -> ByteBuffer? { - try await self.base.next(isolation: nil)?.get() + try await base.next(isolation: nil) } #else /// Returns the next available byte buffer from the sequence. @@ -33,7 +33,7 @@ public struct StorageSequence: Sendable, AsyncSequence { /// - Returns: The next `ByteBuffer`, or `nil` when the sequence is finished. /// - Throws: Any error emitted by the underlying async sequence. public mutating func next() async throws -> ByteBuffer? { - try await self.base.next()?.get() + try await base.next() } #endif @@ -45,11 +45,35 @@ public struct StorageSequence: Sendable, AsyncSequence { public mutating func next( isolation actor: isolated (any Actor)? ) async throws -> Element? { - try await self.base.next(isolation: actor)?.get() + try await base.next(isolation: actor) } } - private let makeIteratorCallback: @Sendable () -> BaseAsyncSequence + private struct FailureErasingAsyncSequence: + AsyncSequence, + Sendable + where Base.Element == ByteBuffer { + typealias Element = ByteBuffer + typealias Failure = any Error + + struct AsyncIterator: AsyncIteratorProtocol { + var base: Base.AsyncIterator + + mutating func next( + isolation actor: isolated (any Actor)? + ) async throws(any Error) -> ByteBuffer? { + try await base.next(isolation: actor) + } + } + + let base: Base + + func makeAsyncIterator() -> AsyncIterator { + .init(base: base.makeAsyncIterator()) + } + } + + private let makeIterator: @Sendable () -> AsyncIterator /// Optional known byte length of the sequence. public let length: UInt64? @@ -64,23 +88,11 @@ public struct StorageSequence: Sendable, AsyncSequence { length: UInt64? = nil ) where S.Element == ByteBuffer { self.length = length - self.makeIteratorCallback = { - BaseAsyncSequence { continuation in - let task = Task { - do { - for try await element in asyncSequence { - continuation.yield(.success(element)) - } - } - catch { - continuation.yield(.failure(error)) - } - continuation.finish() - } - continuation.onTermination = { _ in - task.cancel() - } - } + self.makeIterator = { + AsyncIterator( + base: FailureErasingAsyncSequence(base: asyncSequence) + .makeAsyncIterator() + ) } } @@ -106,6 +118,6 @@ public struct StorageSequence: Sendable, AsyncSequence { /// /// - Returns: A new `AsyncIterator` instance. public func makeAsyncIterator() -> AsyncIterator { - AsyncIterator(base: makeIteratorCallback().makeAsyncIterator()) + makeIterator() } } diff --git a/Tests/FeatherStorageTests/StorageSequenceTestSuite.swift b/Tests/FeatherStorageTests/StorageSequenceTestSuite.swift index d5b0c6f..7a05d97 100644 --- a/Tests/FeatherStorageTests/StorageSequenceTestSuite.swift +++ b/Tests/FeatherStorageTests/StorageSequenceTestSuite.swift @@ -16,6 +16,41 @@ struct StorageSequenceTestSuite { case failed } + private actor DemandProbe { + private(set) var requestCount = 0 + + func recordRequest() { + requestCount += 1 + } + } + + private struct DemandTrackingSequence: AsyncSequence, Sendable { + typealias Element = ByteBuffer + + struct AsyncIterator: AsyncIteratorProtocol { + let probe: DemandProbe + var remainingCount: Int + + mutating func next( + isolation actor: isolated (any Actor)? + ) async -> ByteBuffer? { + guard remainingCount > 0 else { + return nil + } + remainingCount -= 1 + await probe.recordRequest() + return ByteBuffer(bytes: [UInt8(remainingCount)]) + } + } + + let probe: DemandProbe + let count: Int + + func makeAsyncIterator() -> AsyncIterator { + .init(probe: probe, remainingCount: count) + } + } + @Test func initFromAsyncSequencePreservesElementsAndLength() async throws { let allocator = ByteBufferAllocator() @@ -93,6 +128,29 @@ struct StorageSequenceTestSuite { } } + @Test + func requestsUpstreamElementsOnlyWhenConsumerAdvances() async throws { + let probe = DemandProbe() + let sequence = StorageSequence( + asyncSequence: DemandTrackingSequence(probe: probe, count: 3) + ) + var iterator = sequence.makeAsyncIterator() + + #expect(await probe.requestCount == 0) + + _ = try await iterator.next() + for _ in 0..<10 { + await Task.yield() + } + #expect(await probe.requestCount == 1) + + _ = try await iterator.next() + for _ in 0..<10 { + await Task.yield() + } + #expect(await probe.requestCount == 2) + } + private static func makeBuffer( _ bytes: [UInt8], allocator: ByteBufferAllocator