Call didSendRequestPart at the right time

FranzBusch · FranzBusch · commit 5780c100e49e · 2022-03-17T10:08:57.000+01:00
# Motivation
Right now, we call `didSendRequestPart` after passing the write to the executor. However, this does not mean that the write hit the socket. To implement proper backpressure using the delegate, we should only call this method once the write was successful.

# Modification
Pass a promise to the actual channel write and only call the delegate once that promise succeeds.

# Result
The delegate method `didSendRequestPart` is only called after the write was successful.
diff --git a/Sources/AsyncHTTPClient/AsyncAwait/Transaction.swift b/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:
diff --git a/Sources/AsyncHTTPClient/ConnectionPool/HTTP1/HTTP1ClientChannelHandler.swift b/Sources/AsyncHTTPClient/ConnectionPool/HTTP1/HTTP1ClientChannelHandler.swift
@@ -201,8 +201,8 @@ final class HTTP1ClientChannelHandler: ChannelDuplexHandler {
                 }
             }
 
-        case .sendBodyPart(let part):
-            context.writeAndFlush(self.wrapOutboundOut(.body(part)), promise: nil)
+        case .sendBodyPart(let part, let promise):
+            context.writeAndFlush(self.wrapOutboundOut(.body(part)), promise: promise)
 
         case .sendRequestEnd:
             context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil)
@@ -330,7 +330,7 @@ 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
@@ -340,7 +340,7 @@ final class HTTP1ClientChannelHandler: ChannelDuplexHandler {
             return
         }
 
-        let action = self.state.requestStreamPartReceived(data)
+        let action = self.state.requestStreamPartReceived(data, promise: promise)
         self.run(action, context: context)
     }
 
@@ -380,12 +380,12 @@ 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)
             }
         }
     }
diff --git a/Sources/AsyncHTTPClient/ConnectionPool/HTTP1/HTTP1ConnectionStateMachine.swift b/Sources/AsyncHTTPClient/ConnectionPool/HTTP1/HTTP1ConnectionStateMachine.swift
@@ -41,7 +41,7 @@ struct HTTP1ConnectionStateMachine {
         }
 
         case sendRequestHead(HTTPRequestHead, startBody: Bool)
-        case sendBodyPart(IOData)
+        case sendBodyPart(IOData, EventLoopPromise<Void>?)
         case sendRequestEnd
 
         case pauseRequestBodyStream
@@ -189,13 +189,13 @@ 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)
         }
@@ -377,8 +377,8 @@ extension HTTP1ConnectionStateMachine.State {
             return .pauseRequestBodyStream
         case .resumeRequestBodyStream:
             return .resumeRequestBodyStream
-        case .sendBodyPart(let part):
-            return .sendBodyPart(part)
+        case .sendBodyPart(let part, let promise):
+            return .sendBodyPart(part, promise)
         case .sendRequestEnd:
             return .sendRequestEnd
         case .forwardResponseHead(let head, let pauseRequestBodyStream):
diff --git a/Sources/AsyncHTTPClient/ConnectionPool/HTTP2/HTTP2ClientRequestHandler.swift b/Sources/AsyncHTTPClient/ConnectionPool/HTTP2/HTTP2ClientRequestHandler.swift
@@ -148,8 +148,8 @@ final class HTTP2ClientRequestHandler: ChannelDuplexHandler {
             // that the request is neither failed nor finished yet
             self.request!.pauseRequestBodyStream()
 
-        case .sendBodyPart(let data):
-            context.writeAndFlush(self.wrapOutboundOut(.body(data)), promise: nil)
+        case .sendBodyPart(let data, let promise):
+            context.writeAndFlush(self.wrapOutboundOut(.body(data)), promise: promise)
 
         case .sendRequestEnd:
             context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil)
@@ -281,7 +281,7 @@ final class HTTP2ClientRequestHandler: 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
@@ -291,7 +291,7 @@ final class HTTP2ClientRequestHandler: ChannelDuplexHandler {
             return
         }
 
-        let action = self.state.requestStreamPartReceived(data)
+        let action = self.state.requestStreamPartReceived(data, promise: promise)
         self.run(action, context: context)
     }
 
@@ -327,12 +327,12 @@ final class HTTP2ClientRequestHandler: ChannelDuplexHandler {
 }
 
 extension HTTP2ClientRequestHandler: 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)
             }
         }
     }
diff --git a/Sources/AsyncHTTPClient/ConnectionPool/HTTPExecutableRequest.swift b/Sources/AsyncHTTPClient/ConnectionPool/HTTPExecutableRequest.swift
@@ -180,7 +180,7 @@ protocol HTTPRequestExecutor {
     /// Writes a body part into the channel pipeline
     ///
     /// This method may be **called on any thread**. The executor needs to ensure thread safety.
-    func writeRequestBodyPart(_: IOData, request: HTTPExecutableRequest)
+    func writeRequestBodyPart(_: IOData, request: HTTPExecutableRequest, promise: EventLoopPromise<Void>?)
 
     /// Signals that the request body stream has finished
     ///
diff --git a/Sources/AsyncHTTPClient/ConnectionPool/HTTPRequestStateMachine.swift b/Sources/AsyncHTTPClient/ConnectionPool/HTTPRequestStateMachine.swift
@@ -83,7 +83,7 @@ struct HTTPRequestStateMachine {
         }
 
         case sendRequestHead(HTTPRequestHead, startBody: Bool)
-        case sendBodyPart(IOData)
+        case sendBodyPart(IOData, EventLoopPromise<Void>?)
         case sendRequestEnd
 
         case pauseRequestBodyStream
@@ -261,7 +261,7 @@ struct HTTPRequestStateMachine {
         }
     }
 
-    mutating func requestStreamPartReceived(_ part: IOData) -> Action {
+    mutating func requestStreamPartReceived(_ part: IOData, promise: EventLoopPromise<Void>?) -> Action {
         switch self.state {
         case .initialized,
              .waitForChannelToBecomeWritable,
@@ -303,7 +303,7 @@ struct HTTPRequestStateMachine {
 
             self.state = .running(requestState, responseState)
 
-            return .sendBodyPart(part)
+            return .sendBodyPart(part, promise)
 
         case .failed:
             return .wait
diff --git a/Sources/AsyncHTTPClient/RequestBag.swift b/Sources/AsyncHTTPClient/RequestBag.swift
@@ -154,8 +154,11 @@ final class RequestBag<Delegate: HTTPClientResponseDelegate> {
             return self.task.eventLoop.makeFailedFuture(error)
 
         case .write(let part, let writer, let future):
-            writer.writeRequestBodyPart(part, request: self)
-            self.delegate.didSendRequestPart(task: self.task, part)
+            let promise = self.task.eventLoop.makePromise(of: Void.self)
+            promise.futureResult.whenSuccess {
+                self.delegate.didSendRequestPart(task: self.task, part)
+            }
+            writer.writeRequestBodyPart(part, request: self, promise: promise)
             return future
         }
     }
diff --git a/Tests/AsyncHTTPClientTests/HTTP1ConnectionStateMachineTests.swift b/Tests/AsyncHTTPClientTests/HTTP1ConnectionStateMachineTests.swift
@@ -32,21 +32,21 @@ class HTTP1ConnectionStateMachineTests: XCTestCase {
         let part1 = IOData.byteBuffer(ByteBuffer(bytes: [1]))
         let part2 = IOData.byteBuffer(ByteBuffer(bytes: [2]))
         let part3 = IOData.byteBuffer(ByteBuffer(bytes: [3]))
-        XCTAssertEqual(state.requestStreamPartReceived(part0), .sendBodyPart(part0))
-        XCTAssertEqual(state.requestStreamPartReceived(part1), .sendBodyPart(part1))
+        XCTAssertEqual(state.requestStreamPartReceived(part0, promise: nil), .sendBodyPart(part0, nil))
+        XCTAssertEqual(state.requestStreamPartReceived(part1, promise: nil), .sendBodyPart(part1, nil))
 
         // oh the channel reports... we should slow down producing...
         XCTAssertEqual(state.writabilityChanged(writable: false), .pauseRequestBodyStream)
 
         // but we issued a .produceMoreRequestBodyData before... Thus, we must accept more produced
         // data
-        XCTAssertEqual(state.requestStreamPartReceived(part2), .sendBodyPart(part2))
+        XCTAssertEqual(state.requestStreamPartReceived(part2, promise: nil), .sendBodyPart(part2, nil))
         // however when we have put the data on the channel, we should not issue further
         // .produceMoreRequestBodyData events
 
         // once we receive a writable event again, we can allow the producer to produce more data
         XCTAssertEqual(state.writabilityChanged(writable: true), .resumeRequestBodyStream)
-        XCTAssertEqual(state.requestStreamPartReceived(part3), .sendBodyPart(part3))
+        XCTAssertEqual(state.requestStreamPartReceived(part3, promise: nil), .sendBodyPart(part3, nil))
         XCTAssertEqual(state.requestStreamFinished(), .sendRequestEnd)
 
         let responseHead = HTTPResponseHead(version: .http1_1, status: .ok)
@@ -186,8 +186,8 @@ class HTTP1ConnectionStateMachineTests: XCTestCase {
 
         let part0 = IOData.byteBuffer(ByteBuffer(bytes: [0]))
         let part1 = IOData.byteBuffer(ByteBuffer(bytes: [1]))
-        XCTAssertEqual(state.requestStreamPartReceived(part0), .sendBodyPart(part0))
-        XCTAssertEqual(state.requestStreamPartReceived(part1), .sendBodyPart(part1))
+        XCTAssertEqual(state.requestStreamPartReceived(part0, promise: nil), .sendBodyPart(part0, nil))
+        XCTAssertEqual(state.requestStreamPartReceived(part1, promise: nil), .sendBodyPart(part1, nil))
         XCTAssertEqual(state.requestCancelled(closeConnection: false), .failRequest(HTTPClientError.cancelled, .close))
     }
 
@@ -295,7 +295,7 @@ extension HTTP1ConnectionStateMachine.Action: Equatable {
         case (.sendRequestHead(let lhsHead, let lhsStartBody), .sendRequestHead(let rhsHead, let rhsStartBody)):
             return lhsHead == rhsHead && lhsStartBody == rhsStartBody
 
-        case (.sendBodyPart(let lhsData), .sendBodyPart(let rhsData)):
+        case (.sendBodyPart(let lhsData, _), .sendBodyPart(let rhsData, _)):
             return lhsData == rhsData
 
         case (.sendRequestEnd, .sendRequestEnd):
diff --git a/Tests/AsyncHTTPClientTests/HTTPRequestStateMachineTests.swift b/Tests/AsyncHTTPClientTests/HTTPRequestStateMachineTests.swift
diff --git a/Tests/AsyncHTTPClientTests/Mocks/MockRequestExecutor.swift b/Tests/AsyncHTTPClientTests/Mocks/MockRequestExecutor.swift
diff --git a/Tests/AsyncHTTPClientTests/RequestBagTests.swift b/Tests/AsyncHTTPClientTests/RequestBagTests.swift

Original file line number	Diff line number	Diff line change
`@@ -201,8 +201,8 @@ final class HTTP1ClientChannelHandler: ChannelDuplexHandler {`
`201`	`201`	`}`
`202`	`202`	`}`
`203`	`203`
`204`		`- case .sendBodyPart(let part):`
`205`		`- context.writeAndFlush(self.wrapOutboundOut(.body(part)), promise: nil)`
	`204`	`+ case .sendBodyPart(let part, let promise):`
	`205`	`+ context.writeAndFlush(self.wrapOutboundOut(.body(part)), promise: promise)`
`206`	`206`
`207`	`207`	`case .sendRequestEnd:`
`208`	`208`	`context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil)`
`@@ -330,7 +330,7 @@ final class HTTP1ClientChannelHandler: ChannelDuplexHandler {`
`330`	`330`
`331`	`331`	`// MARK: Private HTTPRequestExecutor`
`332`	`332`
`333`		`- private func writeRequestBodyPart0(_ data: IOData, request: HTTPExecutableRequest) {`
	`333`	`+ private func writeRequestBodyPart0(_ data: IOData, request: HTTPExecutableRequest, promise: EventLoopPromise<Void>?) {`
`334`	`334`	`guard self.request === request, let context = self.channelContext else {`
`335`	`335`	`// Because the HTTPExecutableRequest may run in a different thread to our eventLoop,`
`336`	`336`	`// calls from the HTTPExecutableRequest to our ChannelHandler may arrive here after`
`@@ -340,7 +340,7 @@ final class HTTP1ClientChannelHandler: ChannelDuplexHandler {`
`340`	`340`	`return`
`341`	`341`	`}`
`342`	`342`
`343`		`- let action = self.state.requestStreamPartReceived(data)`
	`343`	`+ let action = self.state.requestStreamPartReceived(data, promise: promise)`
`344`	`344`	`self.run(action, context: context)`
`345`	`345`	`}`
`346`	`346`
`@@ -380,12 +380,12 @@ final class HTTP1ClientChannelHandler: ChannelDuplexHandler {`
`380`	`380`	`}`
`381`	`381`
`382`	`382`	`extension HTTP1ClientChannelHandler: HTTPRequestExecutor {`
`383`		`- func writeRequestBodyPart(_ data: IOData, request: HTTPExecutableRequest) {`
	`383`	`+ func writeRequestBodyPart(_ data: IOData, request: HTTPExecutableRequest, promise: EventLoopPromise<Void>?) {`
`384`	`384`	`if self.eventLoop.inEventLoop {`
`385`		`- self.writeRequestBodyPart0(data, request: request)`
	`385`	`+ self.writeRequestBodyPart0(data, request: request, promise: promise)`
`386`	`386`	`} else {`
`387`	`387`	`self.eventLoop.execute {`
`388`		`- self.writeRequestBodyPart0(data, request: request)`
	`388`	`+ self.writeRequestBodyPart0(data, request: request, promise: promise)`
`389`	`389`	`}`
`390`	`390`	`}`
`391`	`391`	`}`
Original file line number	Diff line number	Diff line change
`@@ -83,7 +83,7 @@ struct HTTPRequestStateMachine {`
`83`	`83`	`}`
`84`	`84`
`85`	`85`	`case sendRequestHead(HTTPRequestHead, startBody: Bool)`
`86`		`- case sendBodyPart(IOData)`
	`86`	`+ case sendBodyPart(IOData, EventLoopPromise<Void>?)`
`87`	`87`	`case sendRequestEnd`
`88`	`88`
`89`	`89`	`case pauseRequestBodyStream`
`@@ -261,7 +261,7 @@ struct HTTPRequestStateMachine {`
`261`	`261`	`}`
`262`	`262`	`}`
`263`	`263`
`264`		`- mutating func requestStreamPartReceived(_ part: IOData) -> Action {`
	`264`	`+ mutating func requestStreamPartReceived(_ part: IOData, promise: EventLoopPromise<Void>?) -> Action {`
`265`	`265`	`switch self.state {`
`266`	`266`	`case .initialized,`
`267`	`267`	`.waitForChannelToBecomeWritable,`
`@@ -303,7 +303,7 @@ struct HTTPRequestStateMachine {`
`303`	`303`
`304`	`304`	`self.state = .running(requestState, responseState)`
`305`	`305`
`306`		`- return .sendBodyPart(part)`
	`306`	`+ return .sendBodyPart(part, promise)`
`307`	`307`
`308`	`308`	`case .failed:`
`309`	`309`	`return .wait`
Original file line number	Diff line number	Diff line change
`@@ -154,8 +154,11 @@ final class RequestBag<Delegate: HTTPClientResponseDelegate> {`
`154`	`154`	`return self.task.eventLoop.makeFailedFuture(error)`
`155`	`155`
`156`	`156`	`case .write(let part, let writer, let future):`
`157`		`- writer.writeRequestBodyPart(part, request: self)`
`158`		`- self.delegate.didSendRequestPart(task: self.task, part)`
	`157`	`+ let promise = self.task.eventLoop.makePromise(of: Void.self)`
	`158`	`+ promise.futureResult.whenSuccess {`
	`159`	`+ self.delegate.didSendRequestPart(task: self.task, part)`
	`160`	`+ }`
	`161`	`+ writer.writeRequestBodyPart(part, request: self, promise: promise)`
`159`	`162`	`return future`
`160`	`163`	`}`
`161`	`164`	`}`