Call didSendRequestPart at the right time

Sources/AsyncHTTPClient/AsyncAwait/Transaction.swift

@@ -63,7 +63,7 @@ final class Transaction: @unchecked Sendable {
 
         switch writeAction {
         case .writeAndWait(let executor), .writeAndContinue(let executor):
-            executor.writeRequestBodyPart(.byteBuffer(byteBuffer), request: self)
+            executor.writeRequestBodyPart(.byteBuffer(byteBuffer), request: self, promise: nil)
 
         case .fail:
             // an error/cancellation has happened. we don't need to continue here
@@ -105,14 +105,14 @@ final class Transaction: @unchecked Sendable {
         switch self.state.writeNextRequestPart() {
         case .writeAndContinue(let executor):
             self.stateLock.unlock()
-            executor.writeRequestBodyPart(.byteBuffer(part), request: self)
+            executor.writeRequestBodyPart(.byteBuffer(part), request: self, promise: nil)
 
         case .writeAndWait(let executor):
             try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
                 self.state.waitForRequestBodyDemand(continuation: continuation)
                 self.stateLock.unlock()
 
-                executor.writeRequestBodyPart(.byteBuffer(part), request: self)
+                executor.writeRequestBodyPart(.byteBuffer(part), request: self, promise: nil)
             }
 
         case .fail:
@@ -132,7 +132,7 @@ final class Transaction: @unchecked Sendable {
             break
 
         case .forwardStreamFinished(let executor, let succeedContinuation):
-            executor.finishRequestBodyStream(self)
+            executor.finishRequestBodyStream(self, promise: nil)
             succeedContinuation?.resume(returning: nil)
         }
         return

Sources/AsyncHTTPClient/ConnectionPool/HTTP1/HTTP1ClientChannelHandler.swift

@@ -185,11 +185,11 @@ final class HTTP1ClientChannelHandler: ChannelDuplexHandler {
         case .sendRequestHead(let head, startBody: let startBody):
             self.sendRequestHead(head, startBody: startBody, context: context)
 
-        case .sendBodyPart(let part):
-            context.writeAndFlush(self.wrapOutboundOut(.body(part)), promise: nil)
+        case .sendBodyPart(let part, let writePromise):
+            context.writeAndFlush(self.wrapOutboundOut(.body(part)), promise: writePromise)
 
-        case .sendRequestEnd:
-            context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil)
+        case .sendRequestEnd(let writePromise):
+            context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: writePromise)
 
             if let timeoutAction = self.idleReadTimeoutStateMachine?.requestEndSent() {
                 self.runTimeoutAction(timeoutAction, context: context)
@@ -260,34 +260,52 @@ final class HTTP1ClientChannelHandler: ChannelDuplexHandler {
             switch finalAction {
             case .close:
                 context.close(promise: nil)
-            case .sendRequestEnd:
-                context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil)
+                oldRequest.succeedRequest(buffer)
+            case .sendRequestEnd(let writePromise):
+                let writePromise = writePromise ?? context.eventLoop.makePromise(of: Void.self)
+                // We need to defer succeeding the old request to avoid ordering issues
+                writePromise.futureResult.whenComplete { result in
+                    switch result {
+                    case .success:
+                        oldRequest.succeedRequest(buffer)
+                    case .failure(let error):
+                        oldRequest.fail(error)
+                    }
+                }
+
+                context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: writePromise)
             case .informConnectionIsIdle:
                 self.connection.taskCompleted()
-            case .none:
-                break
+                oldRequest.succeedRequest(buffer)
             }
 
-            oldRequest.succeedRequest(buffer)
-
         case .failRequest(let error, let finalAction):
             // see comment in the `succeedRequest` case.
             let oldRequest = self.request!
             self.request = nil
             self.runTimeoutAction(.clearIdleReadTimeoutTimer, context: context)
 
             switch finalAction {
-            case .close:
+            case .close(let writePromise):
                 context.close(promise: nil)
-            case .sendRequestEnd:
-                context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil)
-            case .informConnectionIsIdle:
+                writePromise?.fail(error)
+                oldRequest.fail(error)
+
+            case .informConnectionIsIdle(let writePromise):
                 self.connection.taskCompleted()
+                writePromise?.fail(error)
+                oldRequest.fail(error)
+
+            case .failWritePromise(let writePromise):
+                writePromise?.fail(error)
+                oldRequest.fail(error)
+
             case .none:
-                break
+                oldRequest.fail(error)
             }
 
-            oldRequest.fail(error)
+        case .failSendBodyPart(let error, let writePromise), .failSendStreamFinished(let error, let writePromise):
+            writePromise?.fail(error)
         }
     }
 
@@ -355,27 +373,28 @@ final class HTTP1ClientChannelHandler: ChannelDuplexHandler {
 
     // MARK: Private HTTPRequestExecutor
 
-    private func writeRequestBodyPart0(_ data: IOData, request: HTTPExecutableRequest) {
+    private func writeRequestBodyPart0(_ data: IOData, request: HTTPExecutableRequest, promise: EventLoopPromise<Void>?) {
         guard self.request === request, let context = self.channelContext else {
             // Because the HTTPExecutableRequest may run in a different thread to our eventLoop,
             // calls from the HTTPExecutableRequest to our ChannelHandler may arrive here after
             // the request has been popped by the state machine or the ChannelHandler has been
             // removed from the Channel pipeline. This is a normal threading issue, noone has
             // screwed up.
+            promise?.fail(HTTPClientError.requestStreamCancelled)
             return
         }
 
-        let action = self.state.requestStreamPartReceived(data)
+        let action = self.state.requestStreamPartReceived(data, promise: promise)
         self.run(action, context: context)
     }
 
-    private func finishRequestBodyStream0(_ request: HTTPExecutableRequest) {
+    private func finishRequestBodyStream0(_ request: HTTPExecutableRequest, promise: EventLoopPromise<Void>?) {
         guard self.request === request, let context = self.channelContext else {
             // See code comment in `writeRequestBodyPart0`
             return
         }
 
-        let action = self.state.requestStreamFinished()
+        let action = self.state.requestStreamFinished(promise: promise)
         self.run(action, context: context)
     }
 
@@ -405,22 +424,22 @@ final class HTTP1ClientChannelHandler: ChannelDuplexHandler {
 }
 
 extension HTTP1ClientChannelHandler: HTTPRequestExecutor {
-    func writeRequestBodyPart(_ data: IOData, request: HTTPExecutableRequest) {
+    func writeRequestBodyPart(_ data: IOData, request: HTTPExecutableRequest, promise: EventLoopPromise<Void>?) {
         if self.eventLoop.inEventLoop {
-            self.writeRequestBodyPart0(data, request: request)
+            self.writeRequestBodyPart0(data, request: request, promise: promise)
         } else {
             self.eventLoop.execute {
-                self.writeRequestBodyPart0(data, request: request)
+                self.writeRequestBodyPart0(data, request: request, promise: promise)
             }
         }
     }
 
-    func finishRequestBodyStream(_ request: HTTPExecutableRequest) {
+    func finishRequestBodyStream(_ request: HTTPExecutableRequest, promise: EventLoopPromise<Void>?) {
         if self.eventLoop.inEventLoop {
-            self.finishRequestBodyStream0(request)
+            self.finishRequestBodyStream0(request, promise: promise)
         } else {
             self.eventLoop.execute {
-                self.finishRequestBodyStream0(request)
+                self.finishRequestBodyStream0(request, promise: promise)
             }
         }
     }

Sources/AsyncHTTPClient/ConnectionPool/HTTP1/HTTP1ConnectionStateMachine.swift

@@ -28,30 +28,48 @@ struct HTTP1ConnectionStateMachine {
 
     enum Action {
         /// A action to execute, when we consider a request "done".
-        enum FinalStreamAction {
+        enum FinalSuccessfulStreamAction {
             /// Close the connection
             case close
             /// If the server has replied, with a status of 200...300 before all data was sent, a request is considered succeeded,
             /// as soon as we wrote the request end onto the wire.
-            case sendRequestEnd
+            ///
+            /// The promise is an optional write promise.
+            case sendRequestEnd(EventLoopPromise<Void>?)
             /// Inform an observer that the connection has become idle
             case informConnectionIsIdle
+        }
+
+        /// A action to execute, when we consider a request "done".
+        enum FinalFailedStreamAction {
+            /// Close the connection
+            ///
+            /// The promise is an optional write promise.
+            case close(EventLoopPromise<Void>?)
+            /// Inform an observer that the connection has become idle
+            ///
+            /// The promise is an optional write promise.
+            case informConnectionIsIdle(EventLoopPromise<Void>?)
+            /// Fail the write promise
+            case failWritePromise(EventLoopPromise<Void>?)
             /// Do nothing.
             case none
         }
 
         case sendRequestHead(HTTPRequestHead, startBody: Bool)
-        case sendBodyPart(IOData)
-        case sendRequestEnd
+        case sendBodyPart(IOData, EventLoopPromise<Void>?)
+        case sendRequestEnd(EventLoopPromise<Void>?)
+        case failSendBodyPart(Error, EventLoopPromise<Void>?)
+        case failSendStreamFinished(Error, EventLoopPromise<Void>?)
 
         case pauseRequestBodyStream
         case resumeRequestBodyStream
 
         case forwardResponseHead(HTTPResponseHead, pauseRequestBodyStream: Bool)
         case forwardResponseBodyParts(CircularBuffer<ByteBuffer>)
 
-        case failRequest(Error, FinalStreamAction)
-        case succeedRequest(FinalStreamAction, CircularBuffer<ByteBuffer>)
+        case failRequest(Error, FinalFailedStreamAction)
+        case succeedRequest(FinalSuccessfulStreamAction, CircularBuffer<ByteBuffer>)
 
         case read
         case close
@@ -173,7 +191,7 @@ struct HTTP1ConnectionStateMachine {
             // as closed.
             //
             // TODO: AHC should support a fast rescheduling mechanism here.
-            return .failRequest(HTTPClientError.remoteConnectionClosed, .none)
+            return .failRequest(HTTPClientError.remoteConnectionClosed, .failWritePromise(nil))
 
         case .idle:
             var requestStateMachine = HTTPRequestStateMachine(isChannelWritable: self.isChannelWritable)
@@ -189,25 +207,25 @@ struct HTTP1ConnectionStateMachine {
         }
     }
 
-    mutating func requestStreamPartReceived(_ part: IOData) -> Action {
+    mutating func requestStreamPartReceived(_ part: IOData, promise: EventLoopPromise<Void>?) -> Action {
         guard case .inRequest(var requestStateMachine, let close) = self.state else {
             preconditionFailure("Invalid state: \(self.state)")
         }
 
         return self.avoidingStateMachineCoW { state -> Action in
-            let action = requestStateMachine.requestStreamPartReceived(part)
+            let action = requestStateMachine.requestStreamPartReceived(part, promise: promise)
             state = .inRequest(requestStateMachine, close: close)
             return state.modify(with: action)
         }
     }
 
-    mutating func requestStreamFinished() -> Action {
+    mutating func requestStreamFinished(promise: EventLoopPromise<Void>?) -> Action {
         guard case .inRequest(var requestStateMachine, let close) = self.state else {
             preconditionFailure("Invalid state: \(self.state)")
         }
 
         return self.avoidingStateMachineCoW { state -> Action in
-            let action = requestStateMachine.requestStreamFinished()
+            let action = requestStateMachine.requestStreamFinished(promise: promise)
             state = .inRequest(requestStateMachine, close: close)
             return state.modify(with: action)
         }
@@ -377,10 +395,10 @@ extension HTTP1ConnectionStateMachine.State {
             return .pauseRequestBodyStream
         case .resumeRequestBodyStream:
             return .resumeRequestBodyStream
-        case .sendBodyPart(let part):
-            return .sendBodyPart(part)
-        case .sendRequestEnd:
-            return .sendRequestEnd
+        case .sendBodyPart(let part, let writePromise):
+            return .sendBodyPart(part, writePromise)
+        case .sendRequestEnd(let writePromise):
+            return .sendRequestEnd(writePromise)
         case .forwardResponseHead(let head, let pauseRequestBodyStream):
             return .forwardResponseHead(head, pauseRequestBodyStream: pauseRequestBodyStream)
         case .forwardResponseBodyParts(let parts):
@@ -390,41 +408,57 @@ extension HTTP1ConnectionStateMachine.State {
                 preconditionFailure("Invalid state: \(self)")
             }
 
-            let newFinalAction: HTTP1ConnectionStateMachine.Action.FinalStreamAction
+            let newFinalAction: HTTP1ConnectionStateMachine.Action.FinalSuccessfulStreamAction
             switch finalAction {
             case .close:
                 self = .closing
                 newFinalAction = .close
-            case .sendRequestEnd:
-                newFinalAction = .sendRequestEnd
+            case .sendRequestEnd(let writePromise):
+                newFinalAction = .sendRequestEnd(writePromise)
             case .none:
                 self = .idle
                 newFinalAction = close ? .close : .informConnectionIsIdle
             }
             return .succeedRequest(newFinalAction, finalParts)
 
         case .failRequest(let error, let finalAction):
-            switch self {
-            case .initialized:
+            switch (self, finalAction) {
+            case (.initialized, _):
                 preconditionFailure("Invalid state: \(self)")
-            case .idle:
+
+            case (.idle, _):
                 preconditionFailure("How can we fail a task, if we are idle")
-            case .inRequest(_, close: let close):
-                if close || finalAction == .close {
-                    self = .closing
-                    return .failRequest(error, .close)
-                } else {
-                    self = .idle
-                    return .failRequest(error, .informConnectionIsIdle)
-                }
 
-            case .closing:
+            // If we are either in .inRequest(_, close: true) or the final action is .close
+            // we have to fail the request with .close()
+            case (.inRequest(_, let close), .none) where close:
+                self = .closing
+                return .failRequest(error, .close(nil))
+
+            case (.inRequest(_, _), .close(let writePromise)):
+                self = .closing
+                return .failRequest(error, .close(writePromise))
+
+            // otherwise we fail with .informConnectionIsIdle
+            case (.inRequest(_, _), .none):
+                self = .idle
+                return .failRequest(error, .informConnectionIsIdle(nil))
+
+            case (.closing, .close(let writePromise)):
+                return .failRequest(error, .failWritePromise(writePromise))
+
+            case (.closing, .none):
                 return .failRequest(error, .none)
-            case .closed:
+
+            case (.closed, .close(let writePromise)):
+                // this state can be reached, if the connection was unexpectedly closed by remote
+                return .failRequest(error, .failWritePromise(writePromise))
+
+            case (.closed, .none):
                 // this state can be reached, if the connection was unexpectedly closed by remote
                 return .failRequest(error, .none)
 
-            case .modifying:
+            case (.modifying, _):
                 preconditionFailure("Invalid state: \(self)")
             }
 
@@ -433,6 +467,12 @@ extension HTTP1ConnectionStateMachine.State {
 
         case .wait:
             return .wait
+
+        case .failSendBodyPart(let error, let writePromise):
+            return .failSendBodyPart(error, writePromise)
+
+        case .failSendStreamFinished(let error, let writePromise):
+            return .failSendStreamFinished(error, writePromise)
         }
     }
 }