Wait on end being sent before closing connection (#68)

hummingbird-project · Apr 18, 2024 · 2a16f07 · 2a16f07
1 parent ce1c3f1
commit 2a16f07
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 6 deletions.
diff --git a/Sources/HummingbirdCore/Server/HTTPServerHandler.swift b/Sources/HummingbirdCore/Server/HTTPServerHandler.swift
@@ -217,24 +217,25 @@ final class HBHTTPServerHandler: ChannelDuplexHandler, RemovableChannelHandler {
         }
         context.write(self.wrapOutboundOut(.head(head)), promise: nil)
         assert(!isInvalidHeaderError(self.propagatedError), "Invalid header")
+        let promise = context.eventLoop.makePromise(of: Void.self)
         switch response.body {
         case .byteBuffer(let buffer):
             context.write(self.wrapOutboundOut(.body(.byteBuffer(buffer))), promise: nil)
-            context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil)
+            context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: promise)
             // don't use error from writeAndFlush so return static version instead of allocating
             // a new EventLoopFuture.
-            return context.eventLoop.makeSucceededVoidFuture()
+            return promise.futureResult
         case .stream(let streamer):
             return streamer.write(on: context.eventLoop) { buffer in
                 context.writeAndFlush(self.wrapOutboundOut(.body(.byteBuffer(buffer))), promise: nil)
             }
             .flatAlways { _ in
-                context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil)
-                return context.eventLoop.makeSucceededVoidFuture()
+                context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: promise)
+                return promise.futureResult
             }
         case .empty:
-            context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil)
-            return context.eventLoop.makeSucceededVoidFuture()
+            context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: promise)
+            return promise.futureResult
         }
     }
 

diff --git a/Tests/HummingbirdCoreTests/CoreTests.swift b/Tests/HummingbirdCoreTests/CoreTests.swift
@@ -608,6 +608,28 @@ class HummingBirdCoreTests: XCTestCase {
         XCTAssertNoThrow(try timeoutPromise.futureResult.wait())
     }
 
+    /// test server sends full payload if "connection" header is set to "close"
+    func testConnectionCloseLargePayload() throws {
+        struct HelloResponder: HBHTTPResponder {
+            func respond(to request: HBHTTPRequest, context: ChannelHandlerContext, onComplete: @escaping (Result<HBHTTPResponse, Error>) -> Void) {
+                let responseHead = HTTPResponseHead(version: .init(major: 1, minor: 1), status: .ok)
+                let responseBody = context.channel.allocator.buffer(repeating: 1, count: 2_000_000)
+                let response = HBHTTPResponse(head: responseHead, body: .byteBuffer(responseBody))
+                onComplete(.success(response))
+            }
+        }
+        let server = HBHTTPServer(group: Self.eventLoopGroup, configuration: .init(address: .hostname(port: 0)))
+        XCTAssertNoThrow(try server.start(responder: HelloResponder()).wait())
+        defer { XCTAssertNoThrow(try server.stop().wait()) }
+
+        let client = HBXCTClient(host: "localhost", port: server.port!, eventLoopGroupProvider: .createNew)
+        client.connect()
+        defer { XCTAssertNoThrow(try client.syncShutdown()) }
+
+        let response = try client.get("/", headers: ["connection": "close"]).wait()
+        XCTAssertEqual(response.body?.readableBytes, 2_000_000)
+    }
+
     /// Test we can run with an embedded channel. HummingbirdXCT uses this quite a lot
     func testEmbeddedChannel() {
         enum HTTPError: Error {