fix: Ensure all data is written before closing FoundationStreamBridge (#713)

Author: jbelkins

Sources/ClientRuntime/Networking/Http/URLSession/FoundationStreamBridge.swift

@@ -19,20 +19,26 @@ import class Foundation.Timer
 import struct Foundation.TimeInterval
 import protocol Foundation.StreamDelegate
 
-/// Reads data from a smithy-swift native `ReadableStream` and streams the data to a Foundation `InputStream`.
+/// Reads data from a smithy-swift native `ReadableStream` and streams the data through to a Foundation `InputStream`.
+///
+/// A pair of Foundation "bound streams" is created.  Data from the `ReadableStream` is transferred into the Foundation
+/// `OutputStream` until the `ReadableStream` is closed and all data has been read from it.  The Foundation
+/// `InputStream` is exposed as a property, and may be used to stream the data to other components.
 ///
 /// Used to permit SDK streaming request bodies to be used with `URLSession`-based HTTP requests.
 class FoundationStreamBridge: NSObject, StreamDelegate {
 
-    /// The max number of bytes to buffer internally (and transfer) at any given time.
-    let bufferSize: Int
+    /// The max number of bytes to buffer between the `ReadableStream` and the Foundation `OutputStream`
+    /// at any given time.
+    let bridgeBufferSize: Int
 
-    /// A buffer to hold data that has been read from the ReadableStream but not yet written to the OutputStream.
+    /// A buffer to hold data that has been read from the `ReadableStream` but not yet written to the
+    /// Foundation `OutputStream`.  At most, it will contain `bridgeBufferSize` bytes.
     private var buffer: Data
 
     /// The `ReadableStream` that will serve as the input to this bridge.
     /// The bridge will read bytes from this stream and dump them to the Foundation stream
-    /// pair  as they become available.
+    /// pair as they become available.
     let readableStream: ReadableStream
 
     /// A Foundation stream that will carry the bytes read from the readableStream as they become available.
@@ -44,80 +50,100 @@ class FoundationStreamBridge: NSObject, StreamDelegate {
     /// A Logger for logging events.
     private let logger: LogAgent
 
-    /// Actor used to ensure writes are performed in series.
-    actor WriteCoordinator {
+    /// Actor used to ensure writes are performed in series, one at a time.
+    private actor WriteCoordinator {
         var task: Task<Void, Error>?
 
-        /// `true` if the readable stream has been found to be empty, `false` otherwise.  Will flip to `true` if the readable stream is read,
-        /// and `nil` is returned.
-        var readableStreamIsEmpty = false
-
-        /// Sets stream status to indicate the stream is empty.
-        func setReadableStreamIsEmpty() async {
-            readableStreamIsEmpty = true
-        }
-
         /// Creates a new concurrent Task that executes the passed block, ensuring that the previous Task
         /// finishes before this task starts.
         ///
         /// Acts as a sort of "serial queue" of Swift concurrency tasks.
         /// - Parameter block: The code to be performed in this task.
-        func perform(_ block: @escaping @Sendable (WriteCoordinator) async throws -> Void) {
-            self.task = Task { [task] in
+        func perform(_ block: @escaping @Sendable () async throws -> Void) async throws {
+            let task = Task { [task] in
                 _ = await task?.result
-                try await block(self)
+                try await block()
             }
+            self.task = task
+            _ = try await task.value
         }
     }
 
     /// Actor used to enforce the order of multiple concurrent stream writes.
     private let writeCoordinator = WriteCoordinator()
 
-    /// A shared serial DispatchQueue to run the stream operations.
-    /// Performing operations on an async queue allows Swift concurrency tasks to not block.
+    /// A serial `DispatchQueue` to run the stream operations for the Foundation `OutputStream`.
+    ///
+    /// Operations performed on the queue include:
+    /// - Opening the stream
+    /// - Closing the stream
+    /// - Writing to the stream
+    /// - Receiving `StreamDelegate` callbacks
+    ///
+    /// Queue operations are run in the order they are placed on the queue, and only one operation
+    /// runs at a time (i.e. this is a "serial queue".)
     private let queue = DispatchQueue(label: "AWSFoundationStreamBridge")
 
+    /// `true` if the readable stream has been closed, `false` otherwise.  Will be flipped to `true` once the readable stream is read,
+    /// and `nil` is returned.
+    ///
+    /// Access this variable only during a write operation to ensure exclusive access.
+    private var readableStreamIsClosed = false
+
     // MARK: - init & deinit
 
-    /// Creates a stream bridge taking the passed `ReadableStream` as its input
+    /// Creates a stream bridge taking the passed `ReadableStream` as its input., and exposing a Foundation `InputStream`
+    /// that may be used for streaming data on to Foundation components.
     ///
     /// Data will be buffered in an internal, in-memory buffer.  The Foundation `InputStream` that exposes `readableStream`
     /// is exposed by the `inputStream` property after creation.
     /// - Parameters:
     ///   - readableStream: The `ReadableStream` that serves as the input to the bridge.
-    ///   - bufferSize: The number of bytes in the in-memory buffer.  The buffer is allocated for this size no matter if in use or not.
-    ///   Defaults to 65536 bytes.
-    init(readableStream: ReadableStream, bufferSize: Int = 65_536, logger: LogAgent) {
+    ///   - bridgeBufferSize: The number of bytes in the in-memory buffer.  The buffer is allocated for this size no matter if in use or not.
+    ///   Defaults to 65536 bytes (64 kb).
+    ///   - boundStreamBufferSize: The number of bytes in the buffer between the bound Foundation streams.  If `nil`, uses the
+    ///   same size as `bridgeBufferSize`.  Defaults to `nil`.  Primary use of this parameter is for testing.
+    ///   - logger: A logger that can be used to log stream events.
+    init(
+        readableStream: ReadableStream,
+        bridgeBufferSize: Int = 65_536,
+        boundStreamBufferSize: Int? = nil,
+        logger: LogAgent
+    ) {
         var inputStream: InputStream?
         var outputStream: OutputStream?
 
         // Create a "bound stream pair" of Foundation streams.
         // Data written into the output stream will automatically flow to the inputStream for reading.
         // The bound streams have a buffer between them of size equal to the buffer held by this bridge.
         Foundation.Stream.getBoundStreams(
-            withBufferSize: bufferSize, inputStream: &inputStream, outputStream: &outputStream
+            withBufferSize: boundStreamBufferSize ?? bridgeBufferSize,
+            inputStream: &inputStream,
+            outputStream: &outputStream
         )
         guard let inputStream, let outputStream else {
             // Fail with fatalError since this is not a failure that would happen in normal operation.
             fatalError("Get pair of bound streams failed.  Please file a bug with AWS SDK for Swift.")
         }
-        self.bufferSize = bufferSize
-        self.buffer = Data(capacity: bufferSize)
+        self.bridgeBufferSize = bridgeBufferSize
+        self.buffer = Data(capacity: bridgeBufferSize)
         self.readableStream = readableStream
         self.inputStream = inputStream
         self.outputStream = outputStream
         self.logger = logger
 
-        // The output stream is configured to deliver its callbacks on the dispatch queue.
+        // The Foundation `OutputStream` is configured to deliver its callbacks on the dispatch queue.
         // This precludes the need for a Thread with RunLoop.
         // For safety, all interactions with the output stream will be performed on this queue.
         CFWriteStreamSetDispatchQueue(outputStream, queue)
     }
 
     // MARK: - Opening & closing
 
-    /// Schedule the output stream on the queue for stream callbacks.
-    /// Do not wait to complete opening before returning.
+    /// Open the output stream and schedule this bridge to receive stream delegate callbacks.
+    ///
+    /// Stream operations are performed on the stream's queue.
+    /// Stream opening is completed before asynchronous return to the caller.
     func open() async {
         await withCheckedContinuation { continuation in
             queue.async {
@@ -128,8 +154,10 @@ class FoundationStreamBridge: NSObject, StreamDelegate {
         }
     }
 
-    /// Unschedule the output stream on the special stream callback thread.
-    /// Do not wait to complete closing before returning.
+    /// Close the output stream and unschedule this bridge from receiving stream delegate callbacks.
+    ///
+    /// Stream operations are performed on the stream's queue.
+    /// Stream closing is completed before asynchronous return to the caller.
     func close() async {
         await withCheckedContinuation { continuation in
             queue.async {
@@ -142,50 +170,97 @@ class FoundationStreamBridge: NSObject, StreamDelegate {
 
     // MARK: - Writing to bridge
 
-    /// Tries to read from the readable stream if possible, then transfer the data to the output stream.
+    /// Writes buffered data to the output stream.
+    /// If the buffer is empty, the `ReadableStream` will be read first to replenish the buffer.
+    ///
+    /// If the buffer is empty and the readable stream is closed, there is no more data to bridge, and the output stream is closed.
     private func writeToOutput() async throws {
-        await writeCoordinator.perform { [self] writeCoordinator in
-            var data = Data()
-            if await !writeCoordinator.readableStreamIsEmpty {
-                if let newData = try await readableStream.readAsync(upToCount: bufferSize) {
-                    data = newData
+
+        // Perform the write on the `WriteCoordinator` to ensure that writes happen in-order
+        // and one at a time.
+        //
+        // Note that it is safe to access `buffer` and `readableStreamIsClosed` instance vars
+        // from inside the block passed to `perform()` because this is the only place
+        // these instance vars are accessed, and the code in the `perform()` block runs
+        // in series with any other calls to `perform()`.
+        try await writeCoordinator.perform { [self] in
+
+            // If there is no data in the buffer and the `ReadableStream` is still open,
+            // attempt to read the stream.  Otherwise, skip reading the `ReadableStream` and
+            // write what's in the buffer immediately.
+            if !readableStreamIsClosed && buffer.isEmpty {
+                if let newData = try await readableStream.readAsync(upToCount: bridgeBufferSize - buffer.count) {
+                    buffer.append(newData)
                 } else {
-                    await writeCoordinator.setReadableStreamIsEmpty()
-                    await close()
+                    readableStreamIsClosed = true
                 }
             }
-            try await writeToOutputStream(data: data)
+
+            // Write the previously buffered data and/or newly read data, if any, to the Foundation `OutputStream`.
+            // Capture the error from the stream write, if any.
+            var streamError: Error?
+            if !buffer.isEmpty {
+                streamError = await writeToOutputStream()
+            }
+
+            // If the readable stream has closed and there is no data in the buffer,
+            // there is nothing left to forward to the output stream, so close it.
+            if readableStreamIsClosed && buffer.isEmpty {
+                await close()
+            }
+
+            // If the output stream write produced an error, throw it now, else just return.
+            if let streamError { throw streamError }
         }
     }
 
-    /// Write the passed data to the output stream, using the reserved thread.
-    private func writeToOutputStream(data: Data) async throws {
-        try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
+    /// Using the output stream's callback queue, write the buffered data to the Foundation `OutputStream`.
+    ///
+    /// After writing, remove the written data from the buffer.
+    /// - Returns: The error resulting from the write to the Foundation `OutputStream`, or `nil` if no error occurred.
+    private func writeToOutputStream() async -> Error? {
+
+        // Suspend the caller while the write is performed on the Foundation `OutputStream`'s queue.
+        await withCheckedContinuation { continuation in
+
+            // Perform the write to the Foundation `OutputStream` on its queue.
             queue.async { [self] in
-                guard !buffer.isEmpty || !data.isEmpty else { continuation.resume(); return }
-                buffer.append(data)
+
+                // Write to the output stream.  It may not accept all data, so get the number of bytes
+                // it accepted in `writeCount`.
                 var writeCount = 0
                 buffer.withUnsafeBytes { bufferPtr in
                     guard let bytePtr = bufferPtr.bindMemory(to: UInt8.self).baseAddress else { return }
                     writeCount = outputStream.write(bytePtr, maxLength: buffer.count)
                 }
+
+                // `writeCount` will be a positive number if bytes were written.
+                // Remove the written bytes from the front of the buffer.
                 if writeCount > 0 {
                     logger.info("FoundationStreamBridge: wrote \(writeCount) bytes to request body")
                     buffer.removeFirst(writeCount)
                 }
-                if let error = outputStream.streamError {
-                    continuation.resume(throwing: error)
-                } else {
-                    continuation.resume()
-                }
+
+                // Resume the caller now that the write is complete, returning the stream error, if any.
+                continuation.resume(returning: outputStream.streamError)
             }
         }
     }
 
     // MARK: - StreamDelegate protocol
 
-    /// The stream places this callback when appropriate.  Call will be delivered on the GCD queue for stream callbacks.
-    /// `.hasSpaceAvailable` prompts this type to query the readable stream for more data.
+    /// The stream places this callback when an event happens.
+    ///
+    /// The `FoundationStreamBridge` sets itself as the delegate of the Foundation `OutputStream` whenever the
+    /// `OutputStream` is open.  Stream callbacks will be delivered on the GCD serial queue.
+    ///
+    /// `.hasSpaceAvailable` is the only event where the `FoundationStreamBridge` takes action; in response to
+    /// this event, the `FoundationStreamBridge` will write data to the `OutputStream`.
+    ///
+    /// This method is implemented for the Foundation `StreamDelegate` protocol.
+    /// - Parameters:
+    ///   - aStream: The stream which experienced the event.
+    ///   - eventCode: A code describing the type of event that happened.
     @objc func stream(_ aStream: Foundation.Stream, handle eventCode: Foundation.Stream.Event) {
         switch eventCode {
         case .openCompleted:
@@ -199,7 +274,7 @@ class FoundationStreamBridge: NSObject, StreamDelegate {
             Task { try await writeToOutput() }
         case .errorOccurred:
             logger.info("FoundationStreamBridge: .errorOccurred event")
-            logger.info("FoundationStreamBridge: Stream error: \(String(describing: aStream.streamError))")
+            logger.info("FoundationStreamBridge: Stream error: \(aStream.streamError.debugDescription)")
         case .endEncountered:
             break
         default:

Sources/ClientRuntime/Networking/Http/URLSession/URLSessionHTTPClient.swift

@@ -410,7 +410,7 @@ public final class URLSessionHTTPClient: HTTPClient {
             // that URLSession can stream its request body from.
             // Allow 16kb of in-memory buffer for request body streaming
             let streamBridge = requestStream.map {
-                FoundationStreamBridge(readableStream: $0, bufferSize: 16_384, logger: logger)
+                FoundationStreamBridge(readableStream: $0, bridgeBufferSize: 16_384, logger: logger)
             }
 
             // Create the request (with a streaming body when needed.)