Fixes bi-directional streaming

Motivation:
When we stream request body, current implementation expects that body
will finish streaming _before_ we start to receive response body parts.
This is not correct, reponse body parts can start to arrive before we
finish sending the request.

Modifications:
 - Simplifies state machine, we only case about request being fully sent
   to prevent sending body parts after .end, but response state machine
   is mostly ignored and correct flow will be handled by NIOHTTP1
   pipeline
 - Adds HTTPEchoHandler, that replies to each response body part
 - Adds bi-directional streaming test

Result:
Closes #327

Author: artemredkin

Sources/AsyncHTTPClient/HTTPHandler.swift

@@ -664,12 +664,12 @@ internal struct TaskCancelEvent {}
 
 internal class TaskHandler<Delegate: HTTPClientResponseDelegate>: RemovableChannelHandler {
     enum State {
-        case idle
-        case bodySent
-        case sent
-        case head
+        case inactive
+        case sending_waiting
+        case sending_head
+        case sent_waiting
+        case sent_head
         case redirected(HTTPResponseHead, URL)
-        case body
         case endOrError
     }
 
@@ -679,7 +679,7 @@ internal class TaskHandler<Delegate: HTTPClientResponseDelegate>: RemovableChann
     let ignoreUncleanSSLShutdown: Bool
     let logger: Logger // We are okay to store the logger here because a TaskHandler is just for one request.
 
-    var state: State = .idle
+    var state: State = .inactive
     var expectedBodyLength: Int?
     var actualBodyLength: Int = 0
     var pendingRead = false
@@ -794,7 +794,8 @@ extension TaskHandler: ChannelDuplexHandler {
     typealias OutboundOut = HTTPClientRequestPart
 
     func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
-        self.state = .idle
+        self.state = .sending_waiting
+
         let request = self.unwrapOutboundIn(data)
 
         var head = HTTPRequestHead(version: HTTPVersion(major: 1, minor: 1),
@@ -840,23 +841,27 @@ extension TaskHandler: ChannelDuplexHandler {
             self.writeBody(request: request, context: context)
         }.flatMap {
             context.eventLoop.assertInEventLoop()
+
             if case .endOrError = self.state {
                 return context.eventLoop.makeSucceededFuture(())
+            } else if case .sending_waiting = self.state {
+                self.state = .sent_waiting
+            } else if case .sending_head = self.state {
+                self.state = .sent_head
             }
 
-            self.state = .bodySent
             if let expectedBodyLength = self.expectedBodyLength, expectedBodyLength != self.actualBodyLength {
                 let error = HTTPClientError.bodyLengthMismatch
                 return context.eventLoop.makeFailedFuture(error)
             }
             return context.writeAndFlush(self.wrapOutboundOut(.end(nil)))
         }.map {
             context.eventLoop.assertInEventLoop()
+
             if case .endOrError = self.state {
                 return
             }
 
-            self.state = .sent
             self.callOutToDelegateFireAndForget(self.delegate.didSendRequest)
         }.flatMapErrorThrowing { error in
             context.eventLoop.assertInEventLoop()
@@ -902,7 +907,7 @@ extension TaskHandler: ChannelDuplexHandler {
 
     private func writeBodyPart(context: ChannelHandlerContext, part: IOData, promise: EventLoopPromise<Void>) {
         switch self.state {
-        case .idle:
+        case .sending_waiting, .sending_head:
             if let limit = self.expectedBodyLength, self.actualBodyLength + part.readableBytes > limit {
                 let error = HTTPClientError.bodyLengthMismatch
                 self.errorCaught(context: context, error: error)
@@ -933,6 +938,10 @@ extension TaskHandler: ChannelDuplexHandler {
         case .head(let head):
             if case .endOrError = self.state {
                 return
+            } else if case .sending_waiting = self.state {
+                self.state = .sending_head
+            } else if case .sent_waiting = self.state {
+                self.state = .sent_head
             }
 
             if !head.isKeepAlive {
@@ -942,7 +951,6 @@ extension TaskHandler: ChannelDuplexHandler {
             if let redirectURL = self.redirectHandler?.redirectTarget(status: head.status, headers: head.headers) {
                 self.state = .redirected(head, redirectURL)
             } else {
-                self.state = .head
                 self.mayRead = false
                 self.callOutToDelegate(value: head, channelEventLoop: context.eventLoop, self.delegate.didReceiveHead)
                     .whenComplete { result in
@@ -954,7 +962,6 @@ extension TaskHandler: ChannelDuplexHandler {
             case .redirected, .endOrError:
                 break
             default:
-                self.state = .body
                 self.mayRead = false
                 self.callOutToDelegate(value: body, channelEventLoop: context.eventLoop, self.delegate.didReceiveBodyPart)
                     .whenComplete { result in
@@ -1009,10 +1016,10 @@ extension TaskHandler: ChannelDuplexHandler {
 
     func channelInactive(context: ChannelHandlerContext) {
         switch self.state {
+        case .inactive, .sending_waiting, .sending_head, .sent_waiting, .sent_head, .redirected:
+            self.errorCaught(context: context, error: HTTPClientError.remoteConnectionClosed)
         case .endOrError:
             break
-        case .body, .head, .idle, .redirected, .sent, .bodySent:
-            self.errorCaught(context: context, error: HTTPClientError.remoteConnectionClosed)
         }
         context.fireChannelInactive()
     }
@@ -1025,8 +1032,7 @@ extension TaskHandler: ChannelDuplexHandler {
                 /// Some HTTP Servers can 'forget' to respond with CloseNotify when client is closing connection,
                 /// this could lead to incomplete SSL shutdown. But since request is already processed, we can ignore this error.
                 break
-            case .head where self.ignoreUncleanSSLShutdown,
-                 .body where self.ignoreUncleanSSLShutdown:
+            case .sent_head where self.ignoreUncleanSSLShutdown:
                 /// We can also ignore this error like `.end`.
                 break
             default:
@@ -1035,7 +1041,7 @@ extension TaskHandler: ChannelDuplexHandler {
             }
         default:
             switch self.state {
-            case .idle, .bodySent, .sent, .head, .redirected, .body:
+            case .inactive, .sending_waiting, .sending_head, .sent_waiting, .sent_head, .redirected:
                 self.state = .endOrError
                 self.failTaskAndNotifyDelegate(error: error, self.delegate.didReceiveError)
             case .endOrError:

Tests/AsyncHTTPClientTests/HTTPClientInternalTests.swift

@@ -145,7 +145,7 @@ class HTTPClientInternalTests: XCTestCase {
 
         try channel.pipeline.addHandler(handler).wait()
 
-        handler.state = .sent
+        handler.state = .sent_waiting
         var body = channel.allocator.buffer(capacity: 4)
         body.writeStaticString("1234")
 

Tests/AsyncHTTPClientTests/HTTPClientTestUtils.swift

@@ -896,6 +896,30 @@ struct CollectEverythingLogHandler: LogHandler {
     }
 }
 
+class HTTPEchoHandler: ChannelInboundHandler {
+    typealias InboundIn = HTTPServerRequestPart
+    typealias OutboundOut = HTTPServerResponsePart
+
+    var promises: CircularBuffer<EventLoopPromise<Void>> = CircularBuffer()
+
+    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
+        let request = self.unwrapInboundIn(data)
+        switch request {
+        case .head:
+            context.writeAndFlush(self.wrapOutboundOut(.head(.init(version: .init(major: 1, minor: 1), status: .ok))), promise: nil)
+        case .body(let bytes):
+            context.writeAndFlush(self.wrapOutboundOut(.body(.byteBuffer(bytes)))).whenSuccess {
+                if let promise = self.promises.popFirst() {
+                    promise.succeed(())
+                }
+            }
+        case .end:
+            context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil)
+            context.close(promise: nil)
+        }
+    }
+}
+
 private let cert = """
 -----BEGIN CERTIFICATE-----
 MIICmDCCAYACCQCPC8JDqMh1zzANBgkqhkiG9w0BAQsFADANMQswCQYDVQQGEwJ1

Tests/AsyncHTTPClientTests/HTTPClientTests+XCTest.swift

@@ -128,6 +128,7 @@ extension HTTPClientTests {
             ("testSSLHandshakeErrorPropagation", testSSLHandshakeErrorPropagation),
             ("testSSLHandshakeErrorPropagationDelayedClose", testSSLHandshakeErrorPropagationDelayedClose),
             ("testWeCloseConnectionsWhenConnectionCloseSetByServer", testWeCloseConnectionsWhenConnectionCloseSetByServer),
+            ("testBiDirectionalStreaming", testBiDirectionalStreaming),
         ]
     }
 }

Tests/AsyncHTTPClientTests/HTTPClientTests.swift

@@ -2708,15 +2708,12 @@ class HTTPClientTests: XCTestCase {
         let task = client.execute(request: request, delegate: TestHTTPDelegate())
 
         XCTAssertThrowsError(try task.wait()) { error in
-            #if os(Linux)
+            if isTestingNIOTS() {
+                XCTAssertEqual(error as? ChannelError, .connectTimeout(.milliseconds(100)))
+            } else {
+                //XCTAssertEqual((error as? IOError).map { $0.errnoCode }, ECONNRESET)
                 XCTAssertEqual(error as? NIOSSLError, NIOSSLError.uncleanShutdown)
-            #else
-                if isTestingNIOTS() {
-                    XCTAssertEqual(error as? ChannelError, .connectTimeout(.milliseconds(100)))
-                } else {
-                    XCTAssertEqual((error as? IOError).map { $0.errnoCode }, ECONNRESET)
-                }
-            #endif
+            }
         }
     }
 
@@ -2756,15 +2753,12 @@ class HTTPClientTests: XCTestCase {
         let task = client.execute(request: request, delegate: TestHTTPDelegate())
 
         XCTAssertThrowsError(try task.wait()) { error in
-            #if os(Linux)
+            if isTestingNIOTS() {
+                XCTAssertEqual(error as? ChannelError, .connectTimeout(.milliseconds(200)))
+            } else {
+                //XCTAssertEqual((error as? IOError).map { $0.errnoCode }, ECONNRESET)
                 XCTAssertEqual(error as? NIOSSLError, NIOSSLError.uncleanShutdown)
-            #else
-                if isTestingNIOTS() {
-                    XCTAssertEqual(error as? ChannelError, .connectTimeout(.milliseconds(200)))
-                } else {
-                    XCTAssertEqual((error as? IOError).map { $0.errnoCode }, ECONNRESET)
-                }
-            #endif
+            }
         }
     }
 
@@ -2793,4 +2787,40 @@ class HTTPClientTests: XCTestCase {
         let result = group.wait(timeout: DispatchTime.now() + DispatchTimeInterval.milliseconds(100))
         XCTAssertEqual(result, .success, "we never closed the connection!")
     }
+
+    func testBiDirectionalStreaming() throws {
+        let handler = HTTPEchoHandler()
+
+        let server = try ServerBootstrap(group: self.serverGroup)
+            .serverChannelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1)
+            .childChannelInitializer { channel in
+                channel.pipeline.configureHTTPServerPipeline().flatMap {
+                    channel.pipeline.addHandler(handler)
+                }
+            }
+            .bind(host: "localhost", port: 0)
+            .wait()
+
+        defer {
+            server.close(promise: nil)
+        }
+
+        let body: HTTPClient.Body = .stream { writer in
+            let promise = self.clientGroup.next().makePromise(of: Void.self)
+            handler.promises.append(promise)
+            return writer.write(.byteBuffer(ByteBuffer(string: "hello"))).flatMap {
+                promise.futureResult
+            }.flatMap {
+                let promise = self.clientGroup.next().makePromise(of: Void.self)
+                handler.promises.append(promise)
+                return writer.write(.byteBuffer(ByteBuffer(string: "hello2"))).flatMap {
+                    promise.futureResult
+                }
+            }
+        }
+
+        let future = self.defaultClient.execute(url: "http://localhost:\(server.localAddress!.port!)", body: body)
+
+        XCTAssertNoThrow(try future.wait())
+    }
 }