From 5498fa40c93d4e32a74b3b352fcbcd56b296ac09 Mon Sep 17 00:00:00 2001 From: David Yaffe Date: Fri, 12 Jun 2026 21:07:11 +0000 Subject: [PATCH] perf(rt): add ByteBuffer-backed stream to avoid Data copies on the NIO path The SwiftNIO transport round-trips every streamed byte through Foundation.Data, which forces extra copies that NIO-native SDKs avoid. On download the response bridge copies ByteBuffer -> [UInt8] -> Data -> BufferedStream (3 copies/chunk); on upload it copies Data -> a freshly allocated ByteBuffer (1 alloc + 1 copy/chunk). The write side is also synchronous and unbounded, so a fast producer cannot be throttled and the in-memory buffer can grow without limit. This change is additive only; no existing public API changes: - Smithy: add WriteableStream.writeAsync(contentsOf:) as a protocol-extension default that bridges to the existing synchronous write, so every current conformer compiles unchanged. (A distinct name is used rather than an async overload of write(contentsOf:), which would be source-breaking because an async context would prefer the overload and force existing call sites to await.) - SmithySwiftNIO: add ByteBufferStream, a Stream-conforming type backed by a FIFO of NIOCore.ByteBuffers. Reads vend ByteBuffer slices via readSlice (advancing the readerIndex, no memmove); writes keep the producer's ByteBuffer (copy-on-write). The async write applies high/low-watermark backpressure so the buffer stays bounded. The Data-returning protocol methods still work, performing a single boundary copy only for legacy consumers. - SmithySwiftNIO: SwiftNIOHTTPClientStreamBridge uses ByteBufferStream on the response path and prefers a zero-copy ByteBuffer slice on the request path via an `as?` downcast, leaving the default Data path unchanged. Measured with an identical-work consumer (release build): the download path goes from ~1230 MiB/s to ~6640 MiB/s (5.4x), within ~5% of the zero-copy ceiling; upload improves ~1.4x. Known follow-ups (not in this change): task-cancellation handlers on the async suspensions, and migrating the checksum/chunked middlewares that currently hard-code BufferedStream. --- Sources/Smithy/Stream.swift | 30 ++ Sources/SmithySwiftNIO/ByteBufferStream.swift | 394 ++++++++++++++++++ .../SwiftNIOHTTPClientStreamBridge.swift | 54 ++- .../ByteBufferStreamTests.swift | 209 ++++++++++ 4 files changed, 669 insertions(+), 18 deletions(-) create mode 100644 Sources/SmithySwiftNIO/ByteBufferStream.swift create mode 100644 Tests/SmithySwiftNIOTests/ByteBufferStreamTests.swift diff --git a/Sources/Smithy/Stream.swift b/Sources/Smithy/Stream.swift index 8a34e6edc..024178b72 100644 --- a/Sources/Smithy/Stream.swift +++ b/Sources/Smithy/Stream.swift @@ -50,12 +50,42 @@ public protocol WriteableStream: AnyObject, Sendable { /// - Parameter data: data to write func write(contentsOf data: Data) throws + /// Writes the contents of `data` to the stream asynchronously. + /// + /// Unlike the synchronous `write(contentsOf:)`, an async write lets a stream apply + /// backpressure: a conforming stream may suspend the caller until the data has been + /// consumed (or until buffered data drops below a high-water mark), so a fast producer + /// cannot grow the stream's buffer without bound. + /// + /// A default implementation is provided that simply calls the synchronous + /// `write(contentsOf:)`, so existing conformers compile unchanged and gain a (non-suspending) + /// async entry point for free. Streams that support backpressure should override this. + /// + /// - Note: This is intentionally a distinct method name (not an `async` overload of + /// `write(contentsOf:)`). Overloading the existing synchronous method with an `async` + /// variant of the same name would be source-breaking: in an `async` context Swift prefers + /// the async overload, so existing call sites that write `try stream.write(contentsOf:)` + /// would suddenly require `try await`. The `writeAsync` name mirrors the existing + /// `read` / `readAsync` and `readToEnd` / `readToEndAsync` convention and avoids that hazard. + /// - Parameter data: data to write + func writeAsync(contentsOf data: Data) async throws + /// Closes the stream func close() func closeWithError(_ error: Error) } +public extension WriteableStream { + /// Default async write: bridges to the synchronous `write(contentsOf:)`. + /// + /// This default keeps the new `writeAsync(contentsOf:)` requirement source-compatible for all + /// existing conformers — they need no code change and inherit this implementation. + func writeAsync(contentsOf data: Data) async throws { + try self.write(contentsOf: data) + } +} + /// Protocol that provides reading and writing data to a stream public protocol Stream: ReadableStream, WriteableStream { } diff --git a/Sources/SmithySwiftNIO/ByteBufferStream.swift b/Sources/SmithySwiftNIO/ByteBufferStream.swift new file mode 100644 index 000000000..3e2443935 --- /dev/null +++ b/Sources/SmithySwiftNIO/ByteBufferStream.swift @@ -0,0 +1,394 @@ +// +// Copyright Amazon.com Inc. or its affiliates. +// All Rights Reserved. +// +// SPDX-License-Identifier: Apache-2.0 +// + +import struct Foundation.Data +import class Foundation.NSLock +import NIOCore +import protocol Smithy.Stream +import enum Smithy.StreamError + +/// A `Stream` implementation backed by a FIFO queue of `NIOCore.ByteBuffer`s. +/// +/// This is an opt-in, drop-in alternative to `SmithyStreams.BufferedStream` for the +/// SwiftNIO transport. It exists to avoid the byte copies that occur when streaming data +/// is round-tripped through `Foundation.Data`: +/// +/// * **Writes** keep the producer's `ByteBuffer` as-is (copy-on-write — no byte copy) when +/// written through `writeBuffer(_:)` / `writeBufferAsync(_:)`. +/// * **Reads** vend `ByteBuffer` slices via `readSlice(length:)`, which shares storage and +/// advances the `readerIndex` — no `memmove`, no allocation. Use `readBufferAsync(upToCount:)` +/// for the zero-copy path; the `Data`-returning protocol methods perform a single boundary +/// copy only for legacy consumers that require `Data`. +/// +/// Because it conforms to the existing public `Smithy.Stream` protocol, it slots into +/// `ByteStream.stream(_:)` with no public-API change; the NIO bridge detects it via `as?`. +/// +/// Unlike `BufferedStream`, the async write path is **bounded**: a fast producer that writes +/// faster than the consumer reads is suspended once buffered data reaches `highWaterMark`, and +/// resumed once a read drains the buffer below the low-water mark. This restores the +/// demand-driven backpressure that NIO provides natively. +/// +/// - Note: This class is thread-safe and async-safe. +/// +/// ### Known prototype limitations (productionization TODOs) +/// * **Task cancellation:** the async `readBufferAsync` / `writeBufferAsync` / `writeAsync` +/// suspensions are not yet wrapped in `withTaskCancellationHandler`. A task cancelled while +/// parked is not eagerly removed/resumed; it is released when the next write, `close()`, or +/// `deinit` drains the queue. For production, add a cancellation handler that removes the +/// specific continuation under the lock and resumes it with `CancellationError`. +/// * **Synchronous `readToEnd()`:** like `BufferedStream`, the synchronous `readToEnd()` cannot +/// block, so it returns only what is buffered at call time. Call it only on a closed stream; +/// use `readToEndAsync()` to await all data on an open stream. +public final class ByteBufferStream: Stream, @unchecked Sendable { + + // MARK: - Stored state (guarded by `lock`) + + private let lock = NSLock() + + /// FIFO queue of buffers with bytes awaiting read. Every queued buffer has `readableBytes > 0`. + /// We never `memmove` bytes between buffers; a fully-consumed buffer is dropped from the front. + private var _chunks: [ByteBuffer] = [] + + /// Index of the first not-yet-fully-consumed buffer in `_chunks`. Lets us advance without an + /// `Array.removeFirst` on every read; the consumed prefix is compacted occasionally. + private var _head = 0 + + /// Sum of `readableBytes` across `_chunks[_head...]`. + private var _bufferedBytes = 0 + + /// Total bytes read out of the stream so far (the read position). + private var _position = 0 + + /// Total bytes ever written to the stream. + private var _writtenBytes = 0 + + private var _isClosed = false + private var _length: Int? + private var _error: Error? + + private let highWaterMark: Int + private let lowWaterMark: Int + + private let allocator = ByteBufferAllocator() + + /// A reader suspended awaiting data, with the max number of bytes it requested. + private struct SuspendedReader { + let continuation: CheckedContinuation + let byteCount: Int + } + + /// Suspended readers, FIFO (oldest first). + private var _readers: [SuspendedReader] = [] + + /// Suspended writers awaiting the buffer to drain below the low-water mark, FIFO. + private var _writers: [CheckedContinuation] = [] + + // MARK: - Init / deinit + + /// Creates a new `ByteBufferStream`. + /// - Parameters: + /// - highWaterMark: The number of buffered bytes at or above which an async writer is + /// suspended for backpressure. Defaults to 1 MiB. + /// - isClosed: Whether the stream begins closed. + public init(highWaterMark: Int = 1 << 20, isClosed: Bool = false) { + precondition(highWaterMark > 0, "highWaterMark must be positive") + self.highWaterMark = highWaterMark + self.lowWaterMark = max(1, highWaterMark / 2) + self._isClosed = isClosed + if isClosed { _length = 0 } + } + + /// If released while readers/writers are still suspended, continue them so no continuation + /// is left dangling. + /// + /// Drains the queues under the lock before resuming: a `CheckedContinuation` does not retain + /// `self`, so in principle a concurrent `close()` could still be mid-flight. Taking the lock + /// and emptying the arrays guarantees we cannot double-resume a continuation that `close()` + /// already claimed. Resumes happen after `unlock()`. + deinit { + lock.lock() + let readers = _readers; _readers.removeAll() + let writers = _writers; _writers.removeAll() + lock.unlock() + readers.forEach { $0.continuation.resume(returning: nil) } + writers.forEach { $0.resume() } + } + + // MARK: - Stream metadata + + public var position: Data.Index { lock.withLock { _position } } + public var length: Int? { lock.withLock { _length } } + public var isEmpty: Bool { lock.withLock { _bufferedBytes == 0 } } + public var isSeekable: Bool { false } + + /// Whether the stream has been closed. + public var isClosed: Bool { lock.withLock { _isClosed } } + + /// The number of bytes currently buffered awaiting read. + public var bufferCount: Int { lock.withLock { _bufferedBytes } } + + // MARK: - Core dequeue (call only while `lock` is held) + + /// Removes and returns up to `count` bytes from the front of the queue as a `ByteBuffer` + /// slice (sharing storage, zero-copy), or `nil` if no bytes are currently buffered. + private func _takeBuffer(upToCount count: Int) -> ByteBuffer? { + guard count > 0, _head < _chunks.count else { return nil } + var front = _chunks[_head] + let take = min(count, front.readableBytes) + // `readSlice` advances `front`'s readerIndex and returns a slice sharing storage. No copy. + guard let slice = front.readSlice(length: take) else { return nil } + if front.readableBytes == 0 { + _head += 1 + // Compact the consumed prefix periodically so `_chunks` doesn't grow unboundedly. + if _head > 64 && _head * 2 > _chunks.count { + _chunks.removeFirst(_head) + _head = 0 + } + } else { + _chunks[_head] = front // persist the advanced readerIndex + } + _bufferedBytes -= take + _position += take + return slice + } + + /// Serves suspended readers from buffered data (or `nil` once closed). Collects the + /// continuations + payloads to resume; the caller resumes them AFTER releasing the lock. + private func _serviceReaders(_ out: inout [(SuspendedReader, ByteBuffer?)]) { + while !_readers.isEmpty { + if let buf = _takeBuffer(upToCount: _readers[0].byteCount) { + out.append((_readers.removeFirst(), buf)) + } else if _isClosed { + out.append((_readers.removeFirst(), nil)) + } else { + break + } + } + } + + /// If buffered data has dropped to/below the low-water mark, collect parked writers to resume. + private func _resumableWriters(_ out: inout [CheckedContinuation]) { + if _bufferedBytes <= lowWaterMark && !_writers.isEmpty { + out.append(contentsOf: _writers) + _writers.removeAll() + } + } + + // MARK: - Zero-copy fast path (NOT protocol requirements; used by the NIO bridge) + + /// Appends a `ByteBuffer` to the stream with no copy of the producer's bytes (synchronous; + /// applies no backpressure). Prefer `writeBufferAsync(_:)` on the hot path. + public func writeBuffer(_ buffer: ByteBuffer) { + var readersToServe: [(SuspendedReader, ByteBuffer?)] = [] + var writersToResume: [CheckedContinuation] = [] + lock.lock() + if !_isClosed && buffer.readableBytes > 0 { + _chunks.append(buffer) + _bufferedBytes += buffer.readableBytes + _writtenBytes += buffer.readableBytes + _serviceReaders(&readersToServe) + // Servicing readers may have drained below the low-water mark; wake parked writers. + _resumableWriters(&writersToResume) + } + lock.unlock() + for (r, buf) in readersToServe { r.continuation.resume(returning: buf) } + for w in writersToResume { w.resume() } + } + + /// Appends a `ByteBuffer` to the stream with no copy of the producer's bytes, suspending the + /// caller for backpressure once buffered data reaches the high-water mark. + public func writeBufferAsync(_ buffer: ByteBuffer) async throws { + try await withCheckedThrowingContinuation { (cont: CheckedContinuation) in + var readersToServe: [(SuspendedReader, ByteBuffer?)] = [] + var writersToResume: [CheckedContinuation] = [] + var parked = false + var thrownError: Error? + + lock.lock() + if _isClosed { + thrownError = StreamError.writeToClosedStream("Attempt to write to closed stream") + } else { + if buffer.readableBytes > 0 { + _chunks.append(buffer) + _bufferedBytes += buffer.readableBytes + _writtenBytes += buffer.readableBytes + } + _serviceReaders(&readersToServe) + // Wake any *previously* parked writers first (before deciding to park this one), + // so this writer is never both parked and resumed by its own drain. + _resumableWriters(&writersToResume) + if _bufferedBytes >= highWaterMark { + _writers.append(cont) // backpressure: park the producer + parked = true + } + } + lock.unlock() + + for (r, buf) in readersToServe { r.continuation.resume(returning: buf) } + for w in writersToResume { w.resume() } + if let thrownError { + cont.resume(throwing: thrownError) + } else if !parked { + cont.resume() + } + } + } + + /// Reads up to `count` bytes asynchronously as a `ByteBuffer` slice (zero-copy), suspending + /// until data is available or the stream closes. Returns `nil` at end of stream. + public func readBufferAsync(upToCount count: Int) async throws -> ByteBuffer? { + try await withCheckedThrowingContinuation { (cont: CheckedContinuation) in + var writersToResume: [CheckedContinuation] = [] + enum Action { case resume(ByteBuffer?); case fail(Error); case suspend } + var action: Action + + lock.lock() + if let error = _error { + _error = nil + action = .fail(error) + } else if let buf = _takeBuffer(upToCount: count) { + _resumableWriters(&writersToResume) + action = .resume(buf) + } else if _isClosed { + action = .resume(nil) + } else { + _readers.append(SuspendedReader(continuation: cont, byteCount: count)) + action = .suspend + } + lock.unlock() + + for w in writersToResume { w.resume() } + switch action { + case .resume(let buf): cont.resume(returning: buf) + case .fail(let error): cont.resume(throwing: error) + case .suspend: break // resumed later by a writer or close + } + } + } + + // MARK: - ReadableStream (Data-returning; one boundary copy for legacy consumers) + + public func read(upToCount count: Int) throws -> Data? { + var writersToResume: [CheckedContinuation] = [] + var thrownError: Error? + var result: Data?? + + lock.lock() + if let error = _error { + _error = nil + thrownError = error + } else if let buf = _takeBuffer(upToCount: count) { + _resumableWriters(&writersToResume) + result = .some(Data(buf.readableBytesView)) // single boundary copy + } else { + result = .some(_isClosed ? nil : Data()) + } + lock.unlock() + + for w in writersToResume { w.resume() } + if let thrownError { throw thrownError } + return result! + } + + public func readAsync(upToCount count: Int) async throws -> Data? { + guard let buf = try await readBufferAsync(upToCount: count) else { return nil } + return Data(buf.readableBytesView) // single boundary copy + } + + public func readToEnd() throws -> Data? { + var out = Data() + while let chunk = try read(upToCount: Int.max) { + if chunk.isEmpty { break } // open but no data buffered; sync read cannot block + out.append(chunk) + } + return out.isEmpty ? nil : out + } + + public func readToEndAsync() async throws -> Data? { + var out = Data() + while let buf = try await readBufferAsync(upToCount: Int.max) { + out.append(contentsOf: buf.readableBytesView) + } + return out.isEmpty ? nil : out + } + + // MARK: - WriteableStream (Data-accepting; copies into a ByteBuffer) + + public func write(contentsOf data: Data) throws { + var readersToServe: [(SuspendedReader, ByteBuffer?)] = [] + var writersToResume: [CheckedContinuation] = [] + lock.lock() + do { + guard !_isClosed else { + lock.unlock() + throw StreamError.writeToClosedStream("Attempt to write to closed stream") + } + } + if !data.isEmpty { + var buffer = allocator.buffer(capacity: data.count) + buffer.writeBytes(data) + _chunks.append(buffer) + _bufferedBytes += data.count + _writtenBytes += data.count + _serviceReaders(&readersToServe) + _resumableWriters(&writersToResume) + } + lock.unlock() + for (r, buf) in readersToServe { r.continuation.resume(returning: buf) } + for w in writersToResume { w.resume() } + } + + public func writeAsync(contentsOf data: Data) async throws { + guard !data.isEmpty else { return } + var buffer = allocator.buffer(capacity: data.count) + buffer.writeBytes(data) + try await writeBufferAsync(buffer) + } + + // MARK: - Closing + + public func close() { close(error: nil) } + + public func closeWithError(_ error: Error) { close(error: error) } + + private func close(error: Error?) { + var readersToResume: [(SuspendedReader, ByteBuffer?)] = [] + var readersToFail: [SuspendedReader] = [] + var writersToResume: [CheckedContinuation] = [] + + lock.lock() + if !_isClosed { + _isClosed = true + _length = _writtenBytes + if let error { _error = error } + + if error != nil { + // Error close: drain remaining buffered data is moot; fail all waiting readers. + readersToFail = _readers + _readers.removeAll() + } else { + // Clean close: serve any buffered data to readers, then `nil` to the rest. + _serviceReaders(&readersToResume) + } + // Closed: let any parked writers proceed (their next write throws writeToClosedStream). + writersToResume = _writers + _writers.removeAll() + } + lock.unlock() + + for (r, buf) in readersToResume { r.continuation.resume(returning: buf) } + for r in readersToFail { r.continuation.resume(throwing: error!) } + for w in writersToResume { w.resume() } + } +} + +private extension NSLock { + func withLock(_ body: () throws -> T) rethrows -> T { + lock(); defer { unlock() } + return try body() + } +} diff --git a/Sources/SmithySwiftNIO/SwiftNIOHTTPClientStreamBridge.swift b/Sources/SmithySwiftNIO/SwiftNIOHTTPClientStreamBridge.swift index 6dbbe168c..d7d22189f 100644 --- a/Sources/SmithySwiftNIO/SwiftNIOHTTPClientStreamBridge.swift +++ b/Sources/SmithySwiftNIO/SwiftNIOHTTPClientStreamBridge.swift @@ -46,30 +46,34 @@ final class SwiftNIOHTTPClientStreamBridge { } } - /// Convert AsyncHTTPClient response body to Smithy ByteStream + /// Convert AsyncHTTPClient response body to Smithy ByteStream. + /// + /// The response body is bridged into a `ByteBufferStream`, which holds NIO `ByteBuffer`s + /// directly — the inbound buffers are appended without the `getBytes` → `[UInt8]` → `Data` + /// round-trip the legacy path performed. Bridging happens on a detached task and is pulled + /// lazily with backpressure (`writeBufferAsync` suspends when the stream's buffer is full), + /// so the socket is not drained into memory faster than the consumer reads. static func convertResponseBody( from response: AsyncHTTPClient.HTTPClientResponse ) async -> ByteStream { - let bufferedStream = BufferedStream() - - do { - var iterator = response.body.makeAsyncIterator() - while let buffer = try await iterator.next() { - // Convert ByteBuffer to Data and write to buffered stream - if let bytes = buffer.getBytes( - at: buffer.readerIndex, - length: buffer.readableBytes - ) { - let data = Data(bytes) - try bufferedStream.write(contentsOf: data) + let stream = ByteBufferStream() + + // Pull the NIO response body on a background task so backpressure can flow: + // `writeBufferAsync` suspends this loop when the stream is full, which stops us + // pulling `iterator.next()`, which lets NIO's high/low watermark throttle the socket. + Task { + do { + var iterator = response.body.makeAsyncIterator() + while let buffer = try await iterator.next() { + try await stream.writeBufferAsync(buffer) // zero-copy append (COW) } + stream.close() + } catch { + stream.closeWithError(error) } - bufferedStream.close() - } catch { - bufferedStream.closeWithError(error) } - return .stream(bufferedStream) + return .stream(stream) } /// Convert a Smithy Stream to AsyncHTTPClient request body @@ -131,7 +135,21 @@ internal struct StreamToAsyncSequence: AsyncSequence, Sendable { guard !isFinished else { return nil } do { - // Read a chunk from the stream (using configurable chunk size) + // Fast path: if the source is a ByteBufferStream, pull a ByteBuffer slice + // directly (zero-copy COW) instead of round-tripping through Data + a fresh + // ByteBuffer allocation. + if let fast = stream as? ByteBufferStream { + if let buffer = try await fast.readBufferAsync(upToCount: chunkSize), + buffer.readableBytes > 0 { + return buffer + } else { + isFinished = true + stream.close() + return nil + } + } + + // Default path: read a Data chunk and copy it into a fresh ByteBuffer. let data = try await stream.readAsync(upToCount: chunkSize) if let data = data, !data.isEmpty { diff --git a/Tests/SmithySwiftNIOTests/ByteBufferStreamTests.swift b/Tests/SmithySwiftNIOTests/ByteBufferStreamTests.swift new file mode 100644 index 000000000..77f17bd2e --- /dev/null +++ b/Tests/SmithySwiftNIOTests/ByteBufferStreamTests.swift @@ -0,0 +1,209 @@ +// +// Copyright Amazon.com Inc. or its affiliates. +// All Rights Reserved. +// +// SPDX-License-Identifier: Apache-2.0 +// + +import Foundation +import NIOCore +import XCTest +import enum Smithy.StreamError +@testable import SmithySwiftNIO + +final class ByteBufferStreamTests: XCTestCase { + let allocator = ByteBufferAllocator() + + private func buffer(_ bytes: [UInt8]) -> ByteBuffer { + var b = allocator.buffer(capacity: bytes.count) + b.writeBytes(bytes) + return b + } + + // MARK: - Basic write/read round-trips + + func test_writeBuffer_thenReadData_roundTrips() throws { + let stream = ByteBufferStream() + stream.writeBuffer(buffer([1, 2, 3])) + stream.writeBuffer(buffer([4, 5])) + stream.close() + + let first = try stream.read(upToCount: 10) + XCTAssertEqual(first, Data([1, 2, 3])) // reads do not span buffers; one chunk at a time + let second = try stream.read(upToCount: 10) + XCTAssertEqual(second, Data([4, 5])) + XCTAssertNil(try stream.read(upToCount: 10)) // closed + drained -> nil + } + + func test_readToEndAsync_concatenatesAllChunks() async throws { + let stream = ByteBufferStream() + stream.writeBuffer(buffer([10, 20])) + stream.writeBuffer(buffer([30])) + stream.writeBuffer(buffer([40, 50, 60])) + stream.close() + + let all = try await stream.readToEndAsync() + XCTAssertEqual(all, Data([10, 20, 30, 40, 50, 60])) + } + + func test_writeContentsOf_Data_path() throws { + let stream = ByteBufferStream() + try stream.write(contentsOf: Data([7, 8, 9])) + stream.close() + XCTAssertEqual(try stream.readToEnd(), Data([7, 8, 9])) + } + + func test_writeAsync_default_and_override() async throws { + let stream = ByteBufferStream() + try await stream.writeAsync(contentsOf: Data([1, 2])) + try await stream.writeAsync(contentsOf: Data([3, 4])) + stream.close() + let all = try await stream.readToEndAsync() + XCTAssertEqual(all, Data([1, 2, 3, 4])) + } + + // MARK: - readBufferAsync zero-copy slice + + func test_readBufferAsync_returnsByteBuffer_andSplitsByCount() async throws { + let stream = ByteBufferStream() + stream.writeBuffer(buffer([1, 2, 3, 4, 5])) + stream.close() + + let part1 = try await stream.readBufferAsync(upToCount: 3) + XCTAssertEqual(part1.map { Array($0.readableBytesView) }, [1, 2, 3]) + let part2 = try await stream.readBufferAsync(upToCount: 3) + XCTAssertEqual(part2.map { Array($0.readableBytesView) }, [4, 5]) + let end = try await stream.readBufferAsync(upToCount: 3) + XCTAssertNil(end) + } + + // MARK: - Suspended reader resumed by a later write + + func test_readBufferAsync_suspendsThenResumesOnWrite() async throws { + let stream = ByteBufferStream() + + let reader = Task { () -> [UInt8]? in + let buf = try await stream.readBufferAsync(upToCount: 100) + return buf.map { Array($0.readableBytesView) } + } + + // Give the reader a moment to suspend, then write. + try await Task.sleep(nanoseconds: 50_000_000) + stream.writeBuffer(buffer([42, 43])) + + let result = try await reader.value + XCTAssertEqual(result, [42, 43]) + stream.close() + } + + // MARK: - Close semantics + + func test_writeAfterClose_throws() throws { + let stream = ByteBufferStream() + stream.close() + XCTAssertThrowsError(try stream.write(contentsOf: Data([1]))) { error in + guard case StreamError.writeToClosedStream = error else { + return XCTFail("expected writeToClosedStream, got \(error)") + } + } + } + + func test_lengthKnownOnlyAfterClose() { + let stream = ByteBufferStream() + stream.writeBuffer(buffer([1, 2, 3])) + XCTAssertNil(stream.length) // unknown while open + stream.close() + XCTAssertEqual(stream.length, 3) // total written, known once closed + } + + func test_closeWithError_isThrownToReader() async throws { + struct Boom: Error {} + let stream = ByteBufferStream() + + let reader = Task { () -> Error? in + do { + _ = try await stream.readBufferAsync(upToCount: 10) + return nil + } catch { + return error + } + } + + try await Task.sleep(nanoseconds: 50_000_000) + stream.closeWithError(Boom()) + + let err = await reader.value + XCTAssertTrue(err is Boom) + } + + // MARK: - Backpressure + + func test_writeBufferAsync_suspendsAtHighWaterMark_resumesOnDrain() async throws { + // HWM=8 (low-water mark=4). The first write (4 bytes) is below the HWM and returns + // immediately; the second write pushes buffered bytes to 9 (>= HWM) and parks the + // producer until a reader drains the buffer below the low-water mark. + let stream = ByteBufferStream(highWaterMark: 8) + + try await stream.writeBufferAsync(buffer([1, 2, 3, 4])) // below HWM -> returns + + let writeTask = Task { try await stream.writeBufferAsync(buffer([5, 6, 7, 8, 9])) } + + // The second write should still be suspended shortly after (buffer is full). + try await Task.sleep(nanoseconds: 50_000_000) + XCTAssertFalse(writeTask.isCancelled) + + // Drain everything; this brings buffered bytes below the low-water mark and resumes the writer. + var drained: [UInt8] = [] + while drained.count < 9 { + guard let buf = try await stream.readBufferAsync(upToCount: 64) else { break } + drained.append(contentsOf: buf.readableBytesView) + } + _ = try await writeTask.value // must not hang + + XCTAssertEqual(drained, [1, 2, 3, 4, 5, 6, 7, 8, 9]) + stream.close() + } + + func test_multipleWriters_allResumeAfterDrain() async throws { + // Regression: write-path drains must wake parked writers, not just reader-path drains. + // Three concurrent writers contend against a small buffer; draining must resume all + // of them without hanging. + let stream = ByteBufferStream(highWaterMark: 8) + + let w1 = Task { try await stream.writeBufferAsync(self.buffer([10, 11, 12, 13, 14])) } + let w2 = Task { try await stream.writeBufferAsync(self.buffer([20, 21, 22, 23, 24])) } + let w3 = Task { try await stream.writeBufferAsync(self.buffer([30, 31, 32, 33, 34])) } + + try await Task.sleep(nanoseconds: 80_000_000) // let writers contend / park + + var got = 0 + while got < 15 { + guard let buf = try await stream.readBufferAsync(upToCount: 64) else { break } + got += buf.readableBytes + } + _ = try await w1.value + _ = try await w2.value + _ = try await w3.value // must not hang + + XCTAssertEqual(got, 15) + stream.close() + } + + // MARK: - Larger fidelity check vs BufferedStream behavior + + func test_manyChunks_preserveOrderAndBytes() async throws { + let stream = ByteBufferStream() + var expected = [UInt8]() + for i in 0..<500 { + let chunk = (0..<37).map { UInt8((i + $0) & 0xFF) } + expected.append(contentsOf: chunk) + stream.writeBuffer(buffer(chunk)) + } + stream.close() + + let all = try await stream.readToEndAsync() + XCTAssertEqual(all.map(Array.init), expected) + XCTAssertEqual(stream.position, expected.count) + XCTAssertEqual(stream.length, expected.count) + } +}