Fixes bi-directional streaming

artemredkin · artemredkin · commit 06a4098f1db6 · 2021-03-03T12:40:31.000Z
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
diff --git a/Sources/AsyncHTTPClient/HTTPHandler.swift b/Sources/AsyncHTTPClient/HTTPHandler.swift
@@ -665,11 +665,11 @@ internal struct TaskCancelEvent {}
 internal class TaskHandler<Delegate: HTTPClientResponseDelegate>: RemovableChannelHandler {
     enum State {
         case idle
-        case bodySent
-        case sent
-        case head
+        case sendingBodyWaitingResponseHead
+        case sendingBodyResponseHeadReceived
+        case bodySentWaitingResponseHead
+        case bodySentResponseHeadReceived
         case redirected(HTTPResponseHead, URL)
-        case body
         case endOrError
     }
 
@@ -794,7 +794,8 @@ extension TaskHandler: ChannelDuplexHandler {
     typealias OutboundOut = HTTPClientRequestPart
 
     func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
-        self.state = .idle
+        self.state = .sendingBodyWaitingResponseHead
+
         let request = self.unwrapOutboundIn(data)
 
         var head = HTTPRequestHead(version: HTTPVersion(major: 1, minor: 1),
@@ -840,23 +841,32 @@ extension TaskHandler: ChannelDuplexHandler {
             self.writeBody(request: request, context: context)
         }.flatMap {
             context.eventLoop.assertInEventLoop()
-            if case .endOrError = self.state {
+
+            switch self.state {
+            case .idle:
+                preconditionFailure("should not happen")
+            case .sendingBodyWaitingResponseHead:
+                self.state = .bodySentWaitingResponseHead
+            case .sendingBodyResponseHeadReceived:
+                self.state = .bodySentResponseHeadReceived
+            case .bodySentWaitingResponseHead, .bodySentResponseHeadReceived:
+                preconditionFailure("should not happen")
+            case .endOrError, .redirected:
                 return context.eventLoop.makeSucceededFuture(())
             }
 
-            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()
@@ -903,6 +913,8 @@ extension TaskHandler: ChannelDuplexHandler {
     private func writeBodyPart(context: ChannelHandlerContext, part: IOData, promise: EventLoopPromise<Void>) {
         switch self.state {
         case .idle:
+            preconditionFailure("should not happen")
+        case .sendingBodyWaitingResponseHead, .sendingBodyResponseHeadReceived:
             if let limit = self.expectedBodyLength, self.actualBodyLength + part.readableBytes > limit {
                 let error = HTTPClientError.bodyLengthMismatch
                 self.errorCaught(context: context, error: error)
@@ -911,7 +923,7 @@ extension TaskHandler: ChannelDuplexHandler {
             }
             self.actualBodyLength += part.readableBytes
             context.writeAndFlush(self.wrapOutboundOut(.body(part)), promise: promise)
-        default:
+        case .bodySentWaitingResponseHead, .bodySentResponseHeadReceived, .redirected, .endOrError:
             let error = HTTPClientError.writeAfterRequestSent
             self.errorCaught(context: context, error: error)
             promise.fail(error)
@@ -931,7 +943,16 @@ extension TaskHandler: ChannelDuplexHandler {
         let response = self.unwrapInboundIn(data)
         switch response {
         case .head(let head):
-            if case .endOrError = self.state {
+            switch self.state {
+            case .idle:
+                preconditionFailure("should not happen")
+            case .sendingBodyWaitingResponseHead:
+                self.state = .sendingBodyResponseHeadReceived
+            case .bodySentWaitingResponseHead:
+                self.state = .bodySentResponseHeadReceived
+            case .sendingBodyResponseHeadReceived, .bodySentResponseHeadReceived, .redirected:
+                preconditionFailure("should not happen")
+            case .endOrError:
                 return
             }
 
@@ -942,7 +963,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 +974,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 +1028,10 @@ extension TaskHandler: ChannelDuplexHandler {
 
     func channelInactive(context: ChannelHandlerContext) {
         switch self.state {
+        case .idle, .sendingBodyWaitingResponseHead, .sendingBodyResponseHeadReceived, .bodySentWaitingResponseHead, .bodySentResponseHeadReceived, .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 +1044,8 @@ 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 .sendingBodyResponseHeadReceived where self.ignoreUncleanSSLShutdown,
+                 .bodySentResponseHeadReceived where self.ignoreUncleanSSLShutdown:
                 /// We can also ignore this error like `.end`.
                 break
             default:
@@ -1035,7 +1054,7 @@ extension TaskHandler: ChannelDuplexHandler {
             }
         default:
             switch self.state {
-            case .idle, .bodySent, .sent, .head, .redirected, .body:
+            case .idle, .sendingBodyWaitingResponseHead, .sendingBodyResponseHeadReceived, .bodySentWaitingResponseHead, .bodySentResponseHeadReceived, .redirected:
                 self.state = .endOrError
                 self.failTaskAndNotifyDelegate(error: error, self.delegate.didReceiveError)
             case .endOrError:
diff --git a/Tests/AsyncHTTPClientTests/HTTPClientInternalTests.swift b/Tests/AsyncHTTPClientTests/HTTPClientInternalTests.swift
@@ -145,7 +145,7 @@ class HTTPClientInternalTests: XCTestCase {
 
         try channel.pipeline.addHandler(handler).wait()
 
-        handler.state = .sent
+        handler.state = .bodySentWaitingResponseHead
         var body = channel.allocator.buffer(capacity: 4)
         body.writeStaticString("1234")
 
diff --git a/Tests/AsyncHTTPClientTests/HTTPClientTestUtils.swift b/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
diff --git a/Tests/AsyncHTTPClientTests/HTTPClientTests+XCTest.swift b/Tests/AsyncHTTPClientTests/HTTPClientTests+XCTest.swift
@@ -128,6 +128,7 @@ extension HTTPClientTests {
             ("testSSLHandshakeErrorPropagation", testSSLHandshakeErrorPropagation),
             ("testSSLHandshakeErrorPropagationDelayedClose", testSSLHandshakeErrorPropagationDelayedClose),
             ("testWeCloseConnectionsWhenConnectionCloseSetByServer", testWeCloseConnectionsWhenConnectionCloseSetByServer),
+            ("testBiDirectionalStreaming", testBiDirectionalStreaming),
         ]
     }
 }
diff --git a/Tests/AsyncHTTPClientTests/HTTPClientTests.swift b/Tests/AsyncHTTPClientTests/HTTPClientTests.swift
@@ -2708,15 +2708,11 @@ 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? 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 +2752,11 @@ 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? 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 +2785,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())
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -128,6 +128,7 @@ extension HTTPClientTests {`
`128`	`128`	`("testSSLHandshakeErrorPropagation", testSSLHandshakeErrorPropagation),`
`129`	`129`	`("testSSLHandshakeErrorPropagationDelayedClose", testSSLHandshakeErrorPropagationDelayedClose),`
`130`	`130`	`("testWeCloseConnectionsWhenConnectionCloseSetByServer", testWeCloseConnectionsWhenConnectionCloseSetByServer),`
	`131`	`+ ("testBiDirectionalStreaming", testBiDirectionalStreaming),`
`131`	`132`	`]`
`132`	`133`	`}`
`133`	`134`	`}`