Stop the RequestBag body stream after an early >= 300 response

Sources/AsyncHTTPClient/RequestBag+StateMachine.swift

@@ -68,6 +68,7 @@ extension RequestBag {
 
         private var state: State
         private let requestFramingMetadata: RequestFramingMetadata
+        private var responseStatusIsAtLeast300 = false
 
         init(redirectHandler: RedirectHandler<Delegate.Response>?, requestFramingMetadata: RequestFramingMetadata) {
             self.state = .initialized(redirectHandler)
@@ -359,6 +360,7 @@ extension RequestBag.StateMachine {
                     return .redirect(executor, redirectHandler, head, redirectURL)
                 }
             } else {
+                self.responseStatusIsAtLeast300 = head.status.code >= 300
                 self.state = .executing(executor, requestState, .buffering(.init(), next: .askExecutorForMore))
                 return .forwardResponseHead(head)
             }
@@ -435,7 +437,7 @@ extension RequestBag.StateMachine {
     enum ReceiveResponseEndAction {
         case consume(ByteBuffer)
         case redirect(RedirectHandler<Delegate.Response>, HTTPResponseHead, URL)
-        case succeedRequest
+        case succeedRequest(failBackpressurePromise: EventLoopPromise<Void>?)
         case none
     }
 
@@ -465,13 +467,17 @@ extension RequestBag.StateMachine {
         case .executing(let executor, let requestState, .waitingForRemote):
             guard var newChunks = newChunks, !newChunks.isEmpty else {
                 switch requestState {
+                case .paused(let promise) where self.responseStatusIsAtLeast300:
+                    self.state = .finished(error: nil)
+                    return .succeedRequest(failBackpressurePromise: promise)
+
                 case .initialized, .paused, .producing:
                     self.state = .executing(executor, requestState, .endReceived)
                     return .none
 
                 case .finished:
                     self.state = .finished(error: nil)
-                    return .succeedRequest
+                    return .succeedRequest(failBackpressurePromise: nil)
                 }
             }
 
@@ -496,7 +502,7 @@ extension RequestBag.StateMachine {
     enum ConsumeAction {
         case requestMoreFromExecutor(HTTPRequestExecutor)
         case consume(ByteBuffer)
-        case finishStream
+        case finishStream(failBackpressurePromise: EventLoopPromise<Void>?)
         case failTask(Error, executorToCancel: HTTPRequestExecutor?)
         case doNothing
     }
@@ -584,7 +590,12 @@ extension RequestBag.StateMachine {
             }
 
             self.state = .finished(error: nil)
-            return .finishStream
+            switch requestState {
+            case .paused(let promise) where self.responseStatusIsAtLeast300:
+                return .finishStream(failBackpressurePromise: promise)
+            case .initialized, .producing, .paused, .finished:
+                return .finishStream(failBackpressurePromise: nil)
+            }
 
         case .executing(_, _, .buffering(_, next: .error(let error))):
             self.state = .finished(error: error)

Sources/AsyncHTTPClient/RequestBag.swift

@@ -349,13 +349,14 @@ final class RequestBag<Delegate: HTTPClientResponseDelegate & Sendable>: Sendabl
                     self.consumeMoreBodyData0(resultOfPreviousConsume: $0)
                 }
 
-        case .succeedRequest:
+        case .succeedRequest(let backpressurePromise):
             do {
                 let response = try self.delegate.didFinishRequest(task: self.task)
                 self.task.promise.succeed(response)
             } catch {
                 self.task.promise.fail(error)
             }
+            backpressurePromise?.fail(HTTPClientError.requestStreamCancelled)
 
         case .redirect(let handler, let head, let newURL):
             self.loopBoundState.value.redirectTask = handler.redirect(
@@ -405,13 +406,14 @@ final class RequestBag<Delegate: HTTPClientResponseDelegate & Sendable>: Sendabl
 
         case .doNothing:
             break
-        case .finishStream:
+        case .finishStream(let backpressurePromise):
             do {
                 let response = try self.delegate.didFinishRequest(task: self.task)
                 self.task.promise.assumeIsolated().succeed(response)
             } catch {
                 self.task.promise.fail(error)
             }
+            backpressurePromise?.fail(HTTPClientError.requestStreamCancelled)
 
         case .failTask(let error, let executor):
             executor?.cancelRequest(self)

Tests/AsyncHTTPClientTests/HTTP1ClientChannelHandlerTests.swift

@@ -1062,6 +1062,58 @@ class HTTP1ClientChannelHandlerTests: XCTestCase {
         )
     }
 
+    func testUploadIsCancelledIfResponseAtLeast300EndsWhileUploadIsPaused() throws {
+        let embedded = EmbeddedChannel()
+        let testUtils = try embedded.setupHTTP1Connection()
+
+        let writeSecondPartPromise = embedded.eventLoop.makePromise(of: Void.self)
+        let uploadResult = NIOLockedValueBox<Result<Void, Error>?>(nil)
+        let request = try HTTPClient.Request(
+            url: "http://localhost/",
+            method: .POST,
+            body: .stream(contentLength: 8) { writer in
+                writer.write(.byteBuffer(.init(bytes: 0...3))).flatMap {
+                    writeSecondPartPromise.futureResult
+                }.flatMap {
+                    writer.write(.byteBuffer(.init(bytes: 4...7)))
+                }.always { result in
+                    uploadResult.withLockedValue { $0 = result }
+                }
+            }
+        )
+
+        let requestBag = try RequestBag(
+            request: request,
+            eventLoopPreference: .delegate(on: embedded.eventLoop),
+            task: .init(eventLoop: embedded.eventLoop, logger: testUtils.logger),
+            redirectHandler: nil,
+            connectionDeadline: .now() + .seconds(30),
+            requestOptions: .forTests(),
+            delegate: ResponseAccumulator(request: request)
+        )
+
+        testUtils.connection.executeRequest(requestBag)
+        embedded.embeddedEventLoop.run()
+        XCTAssertNoThrow(try embedded.receiveHeadAndVerify { XCTAssertEqual($0.method, .POST) })
+        XCTAssertNoThrow(try embedded.receiveBodyAndVerify { XCTAssertEqual($0, .byteBuffer(.init(bytes: 0...3))) })
+
+        let responseHead = HTTPResponseHead(version: .http1_1, status: .payloadTooLarge)
+        XCTAssertNoThrow(try embedded.writeInbound(HTTPClientResponsePart.head(responseHead)))
+        writeSecondPartPromise.succeed(())
+        embedded.embeddedEventLoop.run()
+        XCTAssertNil(try embedded.readOutbound(as: HTTPClientRequestPart.self))
+        XCTAssertNil(uploadResult.withLockedValue { $0 })
+
+        XCTAssertNoThrow(try embedded.writeInbound(HTTPClientResponsePart.end(nil)))
+        embedded.embeddedEventLoop.run()
+        XCTAssertEqual(testUtils.connectionDelegate.hitConnectionClosed, 1)
+        let upload = try XCTUnwrap(uploadResult.withLockedValue { $0 })
+        XCTAssertThrowsError(try upload.get()) {
+            XCTAssertEqual($0 as? HTTPClientError, .requestStreamCancelled)
+        }
+        XCTAssertEqual(try requestBag.task.futureResult.wait().status, .payloadTooLarge)
+    }
+
     func testDefaultMaxBufferSize() {
         if MemoryLayout<Int>.size == 8 {
             XCTAssertEqual(ResponseAccumulator.maxByteBufferSize, Int(UInt32.max))

Tests/AsyncHTTPClientTests/RequestBagTests.swift

@@ -654,6 +654,227 @@ final class RequestBagTests: XCTestCase {
         XCTAssertNoThrow(try bag.task.futureResult.wait())
     }
 
+    func testUploadIsCancelledIfResponseAtLeast300EndsWhileUploadIsPaused() throws {
+        let embeddedEventLoop = EmbeddedEventLoop()
+        defer { XCTAssertNoThrow(try embeddedEventLoop.syncShutdownGracefully()) }
+        let logger = Logger(label: "test")
+
+        let writeSecondPartPromise = embeddedEventLoop.makePromise(of: Void.self)
+        let uploadResult = NIOLockedValueBox<Result<Void, Error>?>(nil)
+        let request = try HTTPClient.Request(
+            url: "https://swift.org",
+            method: .POST,
+            body: .stream(contentLength: 8) { writer -> EventLoopFuture<Void> in
+                writer.write(.byteBuffer(.init(bytes: 0...3))).flatMap {
+                    writeSecondPartPromise.futureResult
+                }.flatMap {
+                    writer.write(.byteBuffer(.init(bytes: 4...7)))
+                }.always { result in
+                    uploadResult.withLockedValue { $0 = result }
+                }
+            }
+        )
+
+        let delegate = UploadCountingDelegate(eventLoop: embeddedEventLoop)
+        let bag = try RequestBag(
+            request: request,
+            eventLoopPreference: .delegate(on: embeddedEventLoop),
+            task: .init(eventLoop: embeddedEventLoop, logger: logger),
+            redirectHandler: nil,
+            connectionDeadline: .now() + .seconds(30),
+            requestOptions: .forTests(),
+            delegate: delegate
+        )
+
+        let executor = MockRequestExecutor(eventLoop: embeddedEventLoop)
+        executor.runRequest(bag)
+        bag.resumeRequestBodyStream()
+        XCTAssertNoThrow(try executor.receiveRequestBody { XCTAssertEqual($0, ByteBuffer(bytes: 0...3)) })
+
+        bag.receiveResponseHead(.init(version: .http1_1, status: .payloadTooLarge))
+        bag.pauseRequestBodyStream()
+        XCTAssertNoThrow(try XCTUnwrap(delegate.backpressurePromise).succeed(()))
+        XCTAssertTrue(executor.signalledDemandForResponseBody)
+
+        writeSecondPartPromise.succeed(())
+        XCTAssertNil(uploadResult.withLockedValue { $0 })
+
+        bag.receiveResponseEnd(nil, trailers: nil)
+        XCTAssertEqual(delegate.hitDidReceiveResponse, 1)
+        let upload = try XCTUnwrap(uploadResult.withLockedValue { $0 })
+        XCTAssertThrowsError(try upload.get()) {
+            XCTAssertEqual($0 as? HTTPClientError, .requestStreamCancelled)
+        }
+        XCTAssertNoThrow(try bag.task.futureResult.wait())
+    }
+
+    func testUploadIsCancelledIfBufferedResponseAtLeast300EndsWhileUploadIsPaused() throws {
+        let embeddedEventLoop = EmbeddedEventLoop()
+        defer { XCTAssertNoThrow(try embeddedEventLoop.syncShutdownGracefully()) }
+        let logger = Logger(label: "test")
+
+        let writeSecondPartPromise = embeddedEventLoop.makePromise(of: Void.self)
+        let uploadResult = NIOLockedValueBox<Result<Void, Error>?>(nil)
+        let request = try HTTPClient.Request(
+            url: "https://swift.org",
+            method: .POST,
+            body: .stream(contentLength: 8) { writer -> EventLoopFuture<Void> in
+                writer.write(.byteBuffer(.init(bytes: 0...3))).flatMap {
+                    writeSecondPartPromise.futureResult
+                }.flatMap {
+                    writer.write(.byteBuffer(.init(bytes: 4...7)))
+                }.always { result in
+                    uploadResult.withLockedValue { $0 = result }
+                }
+            }
+        )
+
+        let delegate = UploadCountingDelegate(eventLoop: embeddedEventLoop)
+        let bag = try RequestBag(
+            request: request,
+            eventLoopPreference: .delegate(on: embeddedEventLoop),
+            task: .init(eventLoop: embeddedEventLoop, logger: logger),
+            redirectHandler: nil,
+            connectionDeadline: .now() + .seconds(30),
+            requestOptions: .forTests(),
+            delegate: delegate
+        )
+
+        let executor = MockRequestExecutor(eventLoop: embeddedEventLoop)
+        executor.runRequest(bag)
+        bag.resumeRequestBodyStream()
+        XCTAssertNoThrow(try executor.receiveRequestBody { XCTAssertEqual($0, ByteBuffer(bytes: 0...3)) })
+
+        bag.receiveResponseHead(.init(version: .http1_1, status: .payloadTooLarge))
+        bag.pauseRequestBodyStream()
+        writeSecondPartPromise.succeed(())
+        bag.receiveResponseEnd([ByteBuffer(string: "foo")], trailers: nil)
+        XCTAssertNil(uploadResult.withLockedValue { $0 })
+
+        XCTAssertNoThrow(try XCTUnwrap(delegate.backpressurePromise).succeed(()))
+        XCTAssertEqual(delegate.hitDidReceiveBodyPart, 1)
+        XCTAssertNoThrow(try XCTUnwrap(delegate.backpressurePromise).succeed(()))
+        XCTAssertEqual(delegate.hitDidReceiveResponse, 1)
+        let upload = try XCTUnwrap(uploadResult.withLockedValue { $0 })
+        XCTAssertThrowsError(try upload.get()) {
+            XCTAssertEqual($0 as? HTTPClientError, .requestStreamCancelled)
+        }
+        XCTAssertNoThrow(try bag.task.futureResult.wait())
+    }
+
+    func testUploadIsCancelledIfResponseAtLeast300EndsWhileWaitingForNextBodyPart() throws {
+        let embeddedEventLoop = EmbeddedEventLoop()
+        defer { XCTAssertNoThrow(try embeddedEventLoop.syncShutdownGracefully()) }
+        let logger = Logger(label: "test")
+
+        let writeSecondPartPromise = embeddedEventLoop.makePromise(of: Void.self)
+        let uploadResult = NIOLockedValueBox<Result<Void, Error>?>(nil)
+        let request = try HTTPClient.Request(
+            url: "https://swift.org",
+            method: .POST,
+            body: .stream(contentLength: 8) { writer -> EventLoopFuture<Void> in
+                writer.write(.byteBuffer(.init(bytes: 0...3))).flatMap {
+                    writeSecondPartPromise.futureResult
+                }.flatMap {
+                    writer.write(.byteBuffer(.init(bytes: 4...7)))
+                }.always { result in
+                    uploadResult.withLockedValue { $0 = result }
+                }
+            }
+        )
+
+        let delegate = UploadCountingDelegate(eventLoop: embeddedEventLoop)
+        let bag = try RequestBag(
+            request: request,
+            eventLoopPreference: .delegate(on: embeddedEventLoop),
+            task: .init(eventLoop: embeddedEventLoop, logger: logger),
+            redirectHandler: nil,
+            connectionDeadline: .now() + .seconds(30),
+            requestOptions: .forTests(),
+            delegate: delegate
+        )
+
+        let executor = MockRequestExecutor(eventLoop: embeddedEventLoop)
+        executor.runRequest(bag)
+        bag.resumeRequestBodyStream()
+        XCTAssertNoThrow(try executor.receiveRequestBody { XCTAssertEqual($0, ByteBuffer(bytes: 0...3)) })
+
+        bag.receiveResponseHead(.init(version: .http1_1, status: .payloadTooLarge))
+        bag.pauseRequestBodyStream()
+        XCTAssertNoThrow(try XCTUnwrap(delegate.backpressurePromise).succeed(()))
+        XCTAssertTrue(executor.signalledDemandForResponseBody)
+
+        let taskResult = NIOLockedValueBox<Result<Void, Error>?>(nil)
+        bag.task.futureResult.whenComplete { result in taskResult.withLockedValue { $0 = result } }
+        bag.receiveResponseEnd(nil, trailers: nil)
+        XCTAssertEqual(delegate.hitDidReceiveResponse, 1)
+        XCTAssertEqual(delegate.receivedHead?.status, .payloadTooLarge)
+        XCTAssertNoThrow(try XCTUnwrap(taskResult.withLockedValue { $0 }).get())
+
+        writeSecondPartPromise.succeed(())
+        XCTAssertEqual(executor.requestBodyPartsCount, 0)
+        let upload = try XCTUnwrap(uploadResult.withLockedValue { $0 })
+        XCTAssertThrowsError(try upload.get()) {
+            XCTAssertEqual($0 as? HTTPClientError, .requestStreamCancelled)
+        }
+        XCTAssertEqual(delegate.hitDidReceiveError, 0)
+    }
+
+    func testUploadResumesIfResponseBelow300EndsWhileUploadIsPaused() throws {
+        let embeddedEventLoop = EmbeddedEventLoop()
+        defer { XCTAssertNoThrow(try embeddedEventLoop.syncShutdownGracefully()) }
+        let logger = Logger(label: "test")
+
+        let writeSecondPartPromise = embeddedEventLoop.makePromise(of: Void.self)
+        let uploadResult = NIOLockedValueBox<Result<Void, Error>?>(nil)
+        let request = try HTTPClient.Request(
+            url: "https://swift.org",
+            method: .POST,
+            body: .stream(contentLength: 8) { writer -> EventLoopFuture<Void> in
+                writer.write(.byteBuffer(.init(bytes: 0...3))).flatMap {
+                    writeSecondPartPromise.futureResult
+                }.flatMap {
+                    writer.write(.byteBuffer(.init(bytes: 4...7)))
+                }.always { result in
+                    uploadResult.withLockedValue { $0 = result }
+                }
+            }
+        )
+
+        let delegate = UploadCountingDelegate(eventLoop: embeddedEventLoop)
+        let bag = try RequestBag(
+            request: request,
+            eventLoopPreference: .delegate(on: embeddedEventLoop),
+            task: .init(eventLoop: embeddedEventLoop, logger: logger),
+            redirectHandler: nil,
+            connectionDeadline: .now() + .seconds(30),
+            requestOptions: .forTests(),
+            delegate: delegate
+        )
+
+        let executor = MockRequestExecutor(eventLoop: embeddedEventLoop)
+        executor.runRequest(bag)
+        bag.resumeRequestBodyStream()
+        XCTAssertNoThrow(try executor.receiveRequestBody { XCTAssertEqual($0, ByteBuffer(bytes: 0...3)) })
+
+        bag.receiveResponseHead(.init(version: .http1_1, status: .ok))
+        bag.pauseRequestBodyStream()
+        XCTAssertNoThrow(try XCTUnwrap(delegate.backpressurePromise).succeed(()))
+        writeSecondPartPromise.succeed(())
+        XCTAssertNoThrow(try executor.receiveRequestBody { XCTAssertEqual($0, ByteBuffer(bytes: 4...7)) })
+
+        bag.receiveResponseEnd(nil, trailers: nil)
+        XCTAssertEqual(delegate.hitDidReceiveResponse, 0)
+        XCTAssertNil(uploadResult.withLockedValue { $0 })
+
+        bag.resumeRequestBodyStream()
+        XCTAssertNoThrow(try XCTUnwrap(uploadResult.withLockedValue { $0 }).get())
+        XCTAssertNoThrow(try executor.receiveEndOfStream())
+        XCTAssertEqual(delegate.hitDidSendRequest, 1)
+        XCTAssertEqual(delegate.hitDidReceiveResponse, 1)
+        XCTAssertNoThrow(try bag.task.futureResult.wait())
+    }
+
     func testRaceBetweenConnectionCloseAndDemandMoreData() {
         let embeddedEventLoop = EmbeddedEventLoop()
         defer { XCTAssertNoThrow(try embeddedEventLoop.syncShutdownGracefully()) }