Fix HTTP2StreamChannel leak

Author: fabianfett

Sources/AsyncHTTPClient/ConnectionPool/HTTP2/HTTP2ClientRequestHandler.swift

@@ -239,21 +239,29 @@ final class HTTP2ClientRequestHandler: ChannelDuplexHandler {
 
     private func runSuccessfulFinalAction(_ action: HTTPRequestStateMachine.Action.FinalSuccessfulRequestAction, context: ChannelHandlerContext) {
         switch action {
-        case .close:
+        case .close, .none:
+            // We must close the http2 stream after the request has finished. This breaks a reference
+            // cycle in HTTP2Connection.
             context.close(promise: nil)
 
         case .sendRequestEnd(let writePromise):
-            context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: writePromise)
-
-        case .none:
-            break
+            let promise = writePromise ?? context.eventLoop.makePromise(of: Void.self)
+            context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: promise)
+            promise.futureResult.whenComplete { _ in
+                // Once we have written the last message we must close the http2 stream, since this
+                // break a reference cycle in HTTP2Connection.
+                context.close(promise: nil)
+            }
         }
     }
 
     private func runFailedFinalAction(_ action: HTTPRequestStateMachine.Action.FinalFailedRequestAction, context: ChannelHandlerContext, error: Error) {
+        // We must close the http2 stream after the request has finished. This breaks a reference
+        // cycle in HTTP2Connection.
+        context.close(promise: nil)
+
         switch action {
         case .close(let writePromise):
-            context.close(promise: nil)
             writePromise?.fail(error)
 
         case .none:

Sources/AsyncHTTPClient/ConnectionPool/HTTP2/HTTP2Connection.swift

@@ -77,7 +77,7 @@ final class HTTP2Connection {
 
     /// We use this channel set to remember, which open streams we need to inform that
     /// we want to close the connection. The channels shall than cancel their currently running
-    /// request.
+    /// request. This property must only be accessed from the connections `EventLoop`.
     private var openStreams = Set<ChannelBox>()
     let id: HTTPConnectionPool.Connection.ID
     let decompression: HTTPClient.Decompression
@@ -241,7 +241,7 @@ final class HTTP2Connection {
                     // before.
                     let box = ChannelBox(channel)
                     self.openStreams.insert(box)
-                    self.channel.closeFuture.whenComplete { _ in
+                    channel.closeFuture.whenComplete { _ in
                         self.openStreams.remove(box)
                     }
 

Tests/AsyncHTTPClientTests/HTTP2ClientRequestHandlerTests+XCTest.swift

@@ -30,6 +30,8 @@ extension HTTP2ClientRequestHandlerTests {
             ("testIdleReadTimeout", testIdleReadTimeout),
             ("testIdleReadTimeoutIsCanceledIfRequestIsCanceled", testIdleReadTimeoutIsCanceledIfRequestIsCanceled),
             ("testWriteHTTPHeadFails", testWriteHTTPHeadFails),
+            ("testChannelIsClosedAfterRequest", testChannelIsClosedAfterRequest),
+            ("testRequestFinishesAfterResponseWasReceived", testRequestFinishesAfterResponseWasReceived),
         ]
     }
 }

Tests/AsyncHTTPClientTests/HTTP2ClientRequestHandlerTests.swift

@@ -275,8 +275,10 @@ class HTTP2ClientRequestHandlerTests: XCTestCase {
         XCTAssertEqual(readEventHandler.readHitCounter, 1)
         XCTAssertNoThrow(try embedded.writeInbound(HTTPClientResponsePart.head(responseHead)))
 
+        XCTAssertTrue(embedded.isActive)
         // canceling the request
         requestBag.cancel()
+        XCTAssertFalse(embedded.isActive, "The stream is closed")
         XCTAssertThrowsError(try requestBag.task.futureResult.wait()) {
             XCTAssertEqual($0 as? HTTPClientError, .cancelled)
         }
@@ -335,14 +337,142 @@ class HTTP2ClientRequestHandlerTests: XCTestCase {
 
             // the handler only writes once the channel is writable
             XCTAssertEqual(try embedded.readOutbound(as: HTTPClientRequestPart.self), .none)
+            XCTAssertTrue(embedded.isActive)
             embedded.isWritable = true
             embedded.pipeline.fireChannelWritabilityChanged()
 
             XCTAssertThrowsError(try requestBag.task.futureResult.wait()) {
                 XCTAssertEqual($0 as? WriteError, WriteError())
             }
 
-            XCTAssertEqual(embedded.isActive, false)
+            XCTAssertFalse(embedded.isActive)
         }
     }
+
+    func testChannelIsClosedAfterRequest() {
+        let embedded = EmbeddedChannel()
+        let requestHandler = HTTP2ClientRequestHandler(eventLoop: embedded.eventLoop)
+        let logger = Logger(label: "test")
+
+        XCTAssertNoThrow(try embedded.pipeline.syncOperations.addHandlers([requestHandler]))
+
+        var maybeRequest: HTTPClient.Request?
+        XCTAssertNoThrow(maybeRequest = try HTTPClient.Request(url: "http://localhost/"))
+        guard let request = maybeRequest else { return XCTFail("Expected to be able to create a request") }
+
+        let delegate = ResponseAccumulator(request: request)
+        var maybeRequestBag: RequestBag<ResponseAccumulator>?
+        XCTAssertNoThrow(maybeRequestBag = try RequestBag(
+            request: request,
+            eventLoopPreference: .delegate(on: embedded.eventLoop),
+            task: .init(eventLoop: embedded.eventLoop, logger: logger),
+            redirectHandler: nil,
+            connectionDeadline: .now() + .seconds(30),
+            requestOptions: .forTests(),
+            delegate: delegate
+        ))
+        guard let requestBag = maybeRequestBag else { return XCTFail("Expected to be able to create a request bag") }
+
+        embedded.write(requestBag, promise: nil)
+        XCTAssertNoThrow(try embedded.connect(to: .makeAddressResolvingHost("localhost", port: 0)).wait())
+
+        XCTAssertNoThrow(try embedded.receiveHeadAndVerify {
+            XCTAssertEqual($0.method, .GET)
+            XCTAssertEqual($0.uri, "/")
+            XCTAssertEqual($0.headers.first(name: "host"), "localhost")
+        })
+        XCTAssertEqual(try embedded.readOutbound(as: HTTPClientRequestPart.self), .end(nil))
+
+        let responseHead = HTTPResponseHead(version: .http1_1, status: .ok, headers: HTTPHeaders([("content-length", "4")]))
+
+        XCTAssertNoThrow(try embedded.writeInbound(HTTPClientResponsePart.head(responseHead)))
+        XCTAssertNoThrow(try embedded.writeInbound(HTTPClientResponsePart.body(ByteBuffer(bytes: 0...3))))
+        XCTAssertTrue(embedded.isActive)
+        XCTAssertNoThrow(try embedded.writeInbound(HTTPClientResponsePart.end(nil)))
+        XCTAssertFalse(embedded.isActive)
+
+        XCTAssertNoThrow(try requestBag.task.futureResult.wait())
+    }
+
+    func testRequestFinishesAfterResponseWasReceived() {
+        let embedded = EmbeddedChannel()
+        let requestHandler = HTTP2ClientRequestHandler(eventLoop: embedded.eventLoop)
+        XCTAssertNoThrow(try embedded.pipeline.syncOperations.addHandler(requestHandler))
+        let logger = Logger(label: "test")
+
+        let testWriter = TestBackpressureWriter(eventLoop: embedded.eventLoop, parts: 50)
+
+        var maybeRequest: HTTPClient.Request?
+        XCTAssertNoThrow(maybeRequest = try HTTPClient.Request(url: "http://localhost/", method: .POST, body: .stream(length: 100) { writer in
+            testWriter.start(writer: writer)
+        }))
+        guard let request = maybeRequest else { return XCTFail("Expected to be able to create a request") }
+
+        let delegate = ResponseAccumulator(request: request)
+        var maybeRequestBag: RequestBag<ResponseAccumulator>?
+        XCTAssertNoThrow(maybeRequestBag = try RequestBag(
+            request: request,
+            eventLoopPreference: .delegate(on: embedded.eventLoop),
+            task: .init(eventLoop: embedded.eventLoop, logger: logger),
+            redirectHandler: nil,
+            connectionDeadline: .now() + .seconds(30),
+            requestOptions: .forTests(idleReadTimeout: .milliseconds(200)),
+            delegate: delegate
+        ))
+        guard let requestBag = maybeRequestBag else { return XCTFail("Expected to be able to create a request bag") }
+
+        embedded.isWritable = false
+        XCTAssertNoThrow(try embedded.connect(to: .makeAddressResolvingHost("localhost", port: 0)).wait())
+        embedded.write(requestBag, promise: nil)
+
+        // the handler only writes once the channel is writable
+        XCTAssertEqual(try embedded.readOutbound(as: HTTPClientRequestPart.self), .none)
+        embedded.isWritable = true
+        testWriter.writabilityChanged(true)
+        embedded.pipeline.fireChannelWritabilityChanged()
+
+        // receive the head
+        XCTAssertNoThrow(try embedded.receiveHeadAndVerify {
+            XCTAssertEqual($0.method, .POST)
+            XCTAssertEqual($0.uri, "/")
+            XCTAssertEqual($0.headers.first(name: "host"), "localhost")
+            XCTAssertEqual($0.headers.first(name: "content-length"), "100")
+        })
+
+        // send full response right away
+        let responseHead = HTTPResponseHead(version: .http1_1, status: .ok)
+        XCTAssertNoThrow(try embedded.writeInbound(HTTPClientResponsePart.head(responseHead)))
+        embedded.read()
+        XCTAssertNoThrow(try embedded.writeInbound(HTTPClientResponsePart.end(nil)))
+        XCTAssertTrue(embedded.isActive)
+
+        // the next body write will be executed once we tick the el. before we make the channel
+        // unwritable
+
+        for index in 0..<50 {
+            embedded.isWritable = false
+            testWriter.writabilityChanged(false)
+            embedded.pipeline.fireChannelWritabilityChanged()
+
+            XCTAssertEqual(testWriter.written, index)
+
+            embedded.embeddedEventLoop.run()
+
+            XCTAssertNoThrow(try embedded.receiveBodyAndVerify {
+                XCTAssertEqual($0.readableBytes, 2)
+            })
+
+            XCTAssertEqual(testWriter.written, index + 1)
+
+            embedded.isWritable = true
+            testWriter.writabilityChanged(true)
+            embedded.pipeline.fireChannelWritabilityChanged()
+        }
+
+        embedded.embeddedEventLoop.run()
+        XCTAssertNoThrow(try embedded.receiveEnd())
+        XCTAssertFalse(embedded.isActive, "Expected the channel to be closed after last write")
+
+        XCTAssertNoThrow(try requestBag.task.futureResult.wait())
+    }
 }