From 0032344e7a1cc02c5a6413415f7eedc4723d8b74 Mon Sep 17 00:00:00 2001
From: Ilya Teterin <iteterin@apple.com>
Date: Fri, 25 Sep 2020 17:43:59 +0100
Subject: [PATCH 1/4] Introduce helper methods for test of ConnectionsState

Motivation:

Issue #234 highlights that we directly manipulate ConnectionsState and
this commit prepares tests to be refactored to manipulate the state
by exposed APIs instead.

Modifications:

* introduce helper methods in ConnectionPoolTestsSupport.swift

Result:

* no observable changes
---
 .../ConnectionPoolTests.swift                 |   2 -
 .../ConnectionPoolTestsSupport.swift          | 119 ++++++++++++++++++
 2 files changed, 119 insertions(+), 2 deletions(-)
 create mode 100644 Tests/AsyncHTTPClientTests/ConnectionPoolTestsSupport.swift

diff --git a/Tests/AsyncHTTPClientTests/ConnectionPoolTests.swift b/Tests/AsyncHTTPClientTests/ConnectionPoolTests.swift
index aee407b28..e46e9d841 100644
--- a/Tests/AsyncHTTPClientTests/ConnectionPoolTests.swift
+++ b/Tests/AsyncHTTPClientTests/ConnectionPoolTests.swift
@@ -27,8 +27,6 @@ class ConnectionPoolTests: XCTestCase {
     var eventLoop: EmbeddedEventLoop!
     var http1ConnectionProvider: HTTP1ConnectionProvider!
 
-    struct TempError: Error {}
-
     func testPending() {
         var state = HTTP1ConnectionProvider.ConnectionsState(eventLoop: self.eventLoop)
 
diff --git a/Tests/AsyncHTTPClientTests/ConnectionPoolTestsSupport.swift b/Tests/AsyncHTTPClientTests/ConnectionPoolTestsSupport.swift
new file mode 100644
index 000000000..dc5281e60
--- /dev/null
+++ b/Tests/AsyncHTTPClientTests/ConnectionPoolTestsSupport.swift
@@ -0,0 +1,119 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the AsyncHTTPClient open source project
+//
+// Copyright (c) 2018-2019 Apple Inc. and the AsyncHTTPClient project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of AsyncHTTPClient project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+@testable import AsyncHTTPClient
+import NIO
+import XCTest
+
+extension ConnectionPoolTests {
+    func buildState(count: Int, release: Bool = true, eventLoop: EventLoop? = nil) -> (HTTP1ConnectionProvider.ConnectionsState, [Connection]) {
+        let eventLoop = eventLoop ?? self.eventLoop!
+
+        var state = HTTP1ConnectionProvider.ConnectionsState(eventLoop: eventLoop)
+        var items: [Connection] = []
+
+        if count == 0 {
+            return (state, items)
+        }
+
+        let channel = ActiveChannel(eventLoop: self.eventLoop)
+
+        for _ in 1...count {
+            // Set up connection pool to have one available connection
+            do {
+                let connection = Connection(channel: channel, provider: self.http1ConnectionProvider)
+                items.append(connection)
+                // First, we ask the empty pool for a connection, triggering connection creation
+                XCTAssertTrue(state.enqueue())
+                let action = state.acquire(waiter: .init(promise: eventLoop.makePromise(), setupComplete: eventLoop.makeSucceededFuture(()), preference: .indifferent))
+
+                switch action {
+                case .create(let waiter):
+                    waiter.promise.succeed(connection)
+                default:
+                    XCTFail("Unexpected action: \(action)")
+                }
+
+                // We offer the connection to the pool so that it can be tracked
+                _ = state.offer(connection: connection)
+            }
+        }
+
+        if release {
+            for item in items {
+                // No we release the connection, making it available for the next caller
+                _ = state.release(connection: item, closing: false)
+            }
+        }
+        return (state, items)
+    }
+}
+
+func XCTAssertState(_ state: HTTP1ConnectionProvider.ConnectionsState, available: Int, leased: Int, waiters: Int, pending: Int, opened: Int) {
+    let snapshot = state.testsOnly_getInternalState()
+    XCTAssertEqual(available, snapshot.availableConnections.count)
+    XCTAssertEqual(leased, snapshot.leasedConnections.count)
+    XCTAssertEqual(waiters, snapshot.waiters.count)
+    XCTAssertEqual(pending, snapshot.pending)
+    XCTAssertEqual(opened, snapshot.openedConnectionsCount)
+}
+
+func XCTAssertState(_ state: HTTP1ConnectionProvider.ConnectionsState, available: Int, leased: Int, waiters: Int, pending: Int, opened: Int, isLeased connection: Connection) {
+    let snapshot = state.testsOnly_getInternalState()
+    XCTAssertEqual(available, snapshot.availableConnections.count)
+    XCTAssertEqual(leased, snapshot.leasedConnections.count)
+    XCTAssertEqual(waiters, snapshot.waiters.count)
+    XCTAssertEqual(pending, snapshot.pending)
+    XCTAssertEqual(opened, snapshot.openedConnectionsCount)
+    XCTAssertTrue(snapshot.leasedConnections.contains(ConnectionKey(connection)))
+}
+
+func XCTAssertState(_ state: HTTP1ConnectionProvider.ConnectionsState, available: Int, leased: Int, waiters: Int, pending: Int, opened: Int, isNotLeased connection: Connection) {
+    let snapshot = state.testsOnly_getInternalState()
+    XCTAssertEqual(available, snapshot.availableConnections.count)
+    XCTAssertEqual(leased, snapshot.leasedConnections.count)
+    XCTAssertEqual(waiters, snapshot.waiters.count)
+    XCTAssertEqual(pending, snapshot.pending)
+    XCTAssertEqual(opened, snapshot.openedConnectionsCount)
+    XCTAssertFalse(snapshot.leasedConnections.contains(ConnectionKey(connection)))
+}
+
+struct XCTEmptyError: Error {}
+
+func XCTUnwrap<T>(_ value: T?) throws -> T {
+    if let unwrapped = value {
+        return unwrapped
+    }
+    throw XCTEmptyError()
+}
+
+struct TempError: Error {}
+
+func XCTAssertStateClose(_ state: HTTP1ConnectionProvider.ConnectionsState, available: Int, leased: Int, waiters: Int, clean: Bool) throws {
+    var state = state
+
+    let (foundWaiters, foundAvailable, foundLeased, foundClean) = try XCTUnwrap(state.close())
+    XCTAssertEqual(waiters, foundWaiters.count)
+    XCTAssertEqual(available, foundAvailable.count)
+    XCTAssertEqual(leased, foundLeased.count)
+    XCTAssertEqual(clean, foundClean)
+
+    for waiter in foundWaiters {
+        waiter.promise.fail(TempError())
+    }
+
+    for lease in foundLeased {
+        try lease.cancel().wait()
+    }
+}

From e04891716d733d31e12d7b2f33aee85ca80f7535 Mon Sep 17 00:00:00 2001
From: Ilya Teterin <iteterin@apple.com>
Date: Fri, 25 Sep 2020 18:04:11 +0100
Subject: [PATCH 2/4] Move Connection tests out of ConnectionsState tests into
 separate file.

Motivation:

Clean up of code to address issue #234 - here we move away
connection tests to separate files outside of ConnectionsState tests
so we will be able to work on the ConnectionsState in focussed mode.

Modifications:

Connection tests moved to separate files.

Result:

No observable changes.
---
 .../ConnectionPoolTests+XCTest.swift          |   5 -
 .../ConnectionPoolTests.swift                 | 246 ------------------
 .../ConnectionTests+XCTest.swift              |  35 +++
 .../ConnectionTests.swift                     | 229 ++++++++++++++++
 Tests/LinuxMain.swift                         |   1 +
 5 files changed, 265 insertions(+), 251 deletions(-)
 create mode 100644 Tests/AsyncHTTPClientTests/ConnectionTests+XCTest.swift
 create mode 100644 Tests/AsyncHTTPClientTests/ConnectionTests.swift

diff --git a/Tests/AsyncHTTPClientTests/ConnectionPoolTests+XCTest.swift b/Tests/AsyncHTTPClientTests/ConnectionPoolTests+XCTest.swift
index 04cfa6421..0fea77411 100644
--- a/Tests/AsyncHTTPClientTests/ConnectionPoolTests+XCTest.swift
+++ b/Tests/AsyncHTTPClientTests/ConnectionPoolTests+XCTest.swift
@@ -57,11 +57,6 @@ extension ConnectionPoolTests {
             ("testTimeoutAvailableConnection", testTimeoutAvailableConnection),
             ("testRemoteClosedLeasedConnection", testRemoteClosedLeasedConnection),
             ("testRemoteClosedAvailableConnection", testRemoteClosedAvailableConnection),
-            ("testConnectionReleaseActive", testConnectionReleaseActive),
-            ("testConnectionReleaseInactive", testConnectionReleaseInactive),
-            ("testConnectionRemoteCloseRelease", testConnectionRemoteCloseRelease),
-            ("testConnectionTimeoutRelease", testConnectionTimeoutRelease),
-            ("testAcquireAvailableBecomesUnavailable", testAcquireAvailableBecomesUnavailable),
             ("testShutdownOnPendingAndSuccess", testShutdownOnPendingAndSuccess),
             ("testShutdownOnPendingAndError", testShutdownOnPendingAndError),
             ("testShutdownTimeout", testShutdownTimeout),
diff --git a/Tests/AsyncHTTPClientTests/ConnectionPoolTests.swift b/Tests/AsyncHTTPClientTests/ConnectionPoolTests.swift
index e46e9d841..a42d6198e 100644
--- a/Tests/AsyncHTTPClientTests/ConnectionPoolTests.swift
+++ b/Tests/AsyncHTTPClientTests/ConnectionPoolTests.swift
@@ -1219,176 +1219,6 @@ class ConnectionPoolTests: XCTestCase {
         self.http1ConnectionProvider.execute(action, logger: HTTPClient.loggingDisabled)
     }
 
-    // MARK: - Connection Tests
-
-    func testConnectionReleaseActive() throws {
-        let channel = ActiveChannel(eventLoop: self.eventLoop)
-        var snapshot = self.http1ConnectionProvider.state.testsOnly_getInternalState()
-
-        let connection = Connection(channel: channel, provider: self.http1ConnectionProvider)
-        snapshot.openedConnectionsCount = 1
-        snapshot.leasedConnections.insert(ConnectionKey(connection))
-
-        self.http1ConnectionProvider.state.testsOnly_setInternalState(snapshot)
-
-        XCTAssertEqual(0, snapshot.availableConnections.count)
-        XCTAssertEqual(1, snapshot.leasedConnections.count)
-        XCTAssertEqual(0, snapshot.waiters.count)
-        XCTAssertEqual(0, snapshot.pending)
-        XCTAssertEqual(1, snapshot.openedConnectionsCount)
-
-        connection.release(closing: false, logger: HTTPClient.loggingDisabled)
-
-        // XCTAssertFalse(connection.isInUse)
-        snapshot = self.http1ConnectionProvider.state.testsOnly_getInternalState()
-        XCTAssertEqual(1, snapshot.availableConnections.count)
-        XCTAssertEqual(0, snapshot.leasedConnections.count)
-        XCTAssertEqual(0, snapshot.waiters.count)
-        XCTAssertEqual(0, snapshot.pending)
-        XCTAssertEqual(1, snapshot.openedConnectionsCount)
-
-        // cleanup
-        // this cleanup code needs to go and use HTTP1ConnectionProvider's API instead
-        // (https://github.com/swift-server/async-http-client/issues/234)
-        connection.remoteClosed(logger: HTTPClient.loggingDisabled)
-    }
-
-    func testConnectionReleaseInactive() throws {
-        let channel = EmbeddedChannel()
-
-        var snapshot = self.http1ConnectionProvider.state.testsOnly_getInternalState()
-
-        let connection = Connection(channel: channel, provider: self.http1ConnectionProvider)
-        snapshot.openedConnectionsCount = 1
-        snapshot.leasedConnections.insert(ConnectionKey(connection))
-
-        self.http1ConnectionProvider.state.testsOnly_setInternalState(snapshot)
-
-        XCTAssertEqual(0, snapshot.availableConnections.count)
-        XCTAssertEqual(1, snapshot.leasedConnections.count)
-        XCTAssertEqual(0, snapshot.waiters.count)
-        XCTAssertEqual(0, snapshot.pending)
-        XCTAssertEqual(1, snapshot.openedConnectionsCount)
-
-        connection.release(closing: true, logger: HTTPClient.loggingDisabled)
-
-        snapshot = self.http1ConnectionProvider.state.testsOnly_getInternalState()
-        XCTAssertEqual(0, snapshot.availableConnections.count)
-        XCTAssertEqual(0, snapshot.leasedConnections.count)
-        XCTAssertEqual(0, snapshot.waiters.count)
-        XCTAssertEqual(0, snapshot.pending)
-        XCTAssertEqual(0, snapshot.openedConnectionsCount)
-    }
-
-    func testConnectionRemoteCloseRelease() throws {
-        let channel = EmbeddedChannel()
-
-        var snapshot = self.http1ConnectionProvider.state.testsOnly_getInternalState()
-
-        let connection = Connection(channel: channel, provider: self.http1ConnectionProvider)
-        snapshot.availableConnections.append(connection)
-        snapshot.openedConnectionsCount = 1
-
-        self.http1ConnectionProvider.state.testsOnly_setInternalState(snapshot)
-
-        XCTAssertEqual(1, snapshot.availableConnections.count)
-        XCTAssertEqual(0, snapshot.leasedConnections.count)
-        XCTAssertEqual(0, snapshot.waiters.count)
-        XCTAssertEqual(0, snapshot.pending)
-        XCTAssertEqual(1, snapshot.openedConnectionsCount)
-
-        connection.remoteClosed(logger: HTTPClient.loggingDisabled)
-
-        snapshot = self.http1ConnectionProvider.state.testsOnly_getInternalState()
-        XCTAssertEqual(0, snapshot.availableConnections.count)
-        XCTAssertEqual(0, snapshot.leasedConnections.count)
-        XCTAssertEqual(0, snapshot.waiters.count)
-        XCTAssertEqual(0, snapshot.pending)
-        XCTAssertEqual(0, snapshot.openedConnectionsCount)
-    }
-
-    func testConnectionTimeoutRelease() throws {
-        let channel = EmbeddedChannel()
-
-        var snapshot = self.http1ConnectionProvider.state.testsOnly_getInternalState()
-
-        let connection = Connection(channel: channel, provider: self.http1ConnectionProvider)
-        snapshot.availableConnections.append(connection)
-        snapshot.openedConnectionsCount = 1
-
-        self.http1ConnectionProvider.state.testsOnly_setInternalState(snapshot)
-
-        XCTAssertEqual(1, snapshot.availableConnections.count)
-        XCTAssertEqual(0, snapshot.leasedConnections.count)
-        XCTAssertEqual(0, snapshot.waiters.count)
-        XCTAssertEqual(0, snapshot.pending)
-        XCTAssertEqual(1, snapshot.openedConnectionsCount)
-
-        connection.timeout(logger: HTTPClient.loggingDisabled)
-
-        snapshot = self.http1ConnectionProvider.state.testsOnly_getInternalState()
-        XCTAssertEqual(0, snapshot.availableConnections.count)
-        XCTAssertEqual(0, snapshot.leasedConnections.count)
-        XCTAssertEqual(0, snapshot.waiters.count)
-        XCTAssertEqual(0, snapshot.pending)
-        XCTAssertEqual(0, snapshot.openedConnectionsCount)
-    }
-
-    func testAcquireAvailableBecomesUnavailable() throws {
-        let channel = ActiveChannel(eventLoop: self.eventLoop)
-        var snapshot = self.http1ConnectionProvider.state.testsOnly_getInternalState()
-
-        let connection = Connection(channel: channel, provider: self.http1ConnectionProvider)
-        snapshot.availableConnections.append(connection)
-        snapshot.openedConnectionsCount = 1
-
-        self.http1ConnectionProvider.state.testsOnly_setInternalState(snapshot)
-
-        XCTAssertEqual(1, snapshot.availableConnections.count)
-        XCTAssertEqual(0, snapshot.leasedConnections.count)
-        XCTAssertEqual(0, snapshot.waiters.count)
-        XCTAssertEqual(0, snapshot.pending)
-        XCTAssertEqual(1, snapshot.openedConnectionsCount)
-
-        XCTAssertTrue(self.http1ConnectionProvider.enqueue())
-
-        let action = self.http1ConnectionProvider.state.acquire(waiter: .init(promise: self.eventLoop.makePromise(), setupComplete: self.eventLoop.makeSucceededFuture(()), preference: .indifferent))
-        switch action {
-        case .lease(let connection, let waiter):
-            // Since this connection is already in use, this should be a no-op and state should not have changed from normal lease
-            connection.timeout(logger: HTTPClient.loggingDisabled)
-
-            snapshot = self.http1ConnectionProvider.state.testsOnly_getInternalState()
-            XCTAssertTrue(connection.isActiveEstimation)
-            XCTAssertTrue(snapshot.leasedConnections.contains(ConnectionKey(connection)))
-            XCTAssertEqual(0, snapshot.availableConnections.count)
-            XCTAssertEqual(1, snapshot.leasedConnections.count)
-            XCTAssertEqual(0, snapshot.waiters.count)
-            XCTAssertEqual(0, snapshot.pending)
-            XCTAssertEqual(1, snapshot.openedConnectionsCount)
-
-            // This is unrecoverable, but in this case we create a new connection, so state again should not change, even though release will be called
-            // This is important to preventself.http1ConnectionProvider deletion since connection is released and there could be 0 waiters
-            connection.remoteClosed(logger: HTTPClient.loggingDisabled)
-
-            snapshot = self.http1ConnectionProvider.state.testsOnly_getInternalState()
-            XCTAssertTrue(snapshot.leasedConnections.contains(ConnectionKey(connection)))
-            XCTAssertEqual(0, snapshot.availableConnections.count)
-            XCTAssertEqual(1, snapshot.leasedConnections.count)
-            XCTAssertEqual(0, snapshot.waiters.count)
-            XCTAssertEqual(0, snapshot.pending)
-            XCTAssertEqual(1, snapshot.openedConnectionsCount)
-
-            // cleanup
-            // this cleanup code needs to go and use HTTP1ConnectionProvider's API instead
-            // (https://github.com/swift-server/async-http-client/issues/234)
-            waiter.promise.fail(TempError())
-            self.http1ConnectionProvider.release(connection: connection, closing: true, logger: HTTPClient.loggingDisabled)
-        default:
-            XCTFail("Unexpected action: \(action)")
-        }
-    }
-
     // MARK: - Shutdown tests
 
     func testShutdownOnPendingAndSuccess() {
@@ -1555,79 +1385,3 @@ class ConnectionPoolTests: XCTestCase {
         self.http1ConnectionProvider = nil
     }
 }
-
-class ActiveChannel: Channel, ChannelCore {
-    struct NotImplementedError: Error {}
-
-    func localAddress0() throws -> SocketAddress {
-        throw NotImplementedError()
-    }
-
-    func remoteAddress0() throws -> SocketAddress {
-        throw NotImplementedError()
-    }
-
-    func register0(promise: EventLoopPromise<Void>?) {
-        promise?.fail(NotImplementedError())
-    }
-
-    func bind0(to: SocketAddress, promise: EventLoopPromise<Void>?) {
-        promise?.fail(NotImplementedError())
-    }
-
-    func connect0(to: SocketAddress, promise: EventLoopPromise<Void>?) {
-        promise?.fail(NotImplementedError())
-    }
-
-    func write0(_ data: NIOAny, promise: EventLoopPromise<Void>?) {
-        promise?.fail(NotImplementedError())
-    }
-
-    func flush0() {}
-
-    func read0() {}
-
-    func close0(error: Error, mode: CloseMode, promise: EventLoopPromise<Void>?) {
-        promise?.succeed(())
-    }
-
-    func triggerUserOutboundEvent0(_ event: Any, promise: EventLoopPromise<Void>?) {
-        promise?.fail(NotImplementedError())
-    }
-
-    func channelRead0(_: NIOAny) {}
-
-    func errorCaught0(error: Error) {}
-
-    var allocator: ByteBufferAllocator
-    var closeFuture: EventLoopFuture<Void>
-    var eventLoop: EventLoop
-
-    var localAddress: SocketAddress?
-    var remoteAddress: SocketAddress?
-    var parent: Channel?
-    var isWritable: Bool = true
-    var isActive: Bool = true
-
-    init(eventLoop: EmbeddedEventLoop) {
-        self.allocator = ByteBufferAllocator()
-        self.eventLoop = eventLoop
-        self.closeFuture = self.eventLoop.makeSucceededFuture(())
-    }
-
-    var _channelCore: ChannelCore {
-        return self
-    }
-
-    var pipeline: ChannelPipeline {
-        return ChannelPipeline(channel: self)
-    }
-
-    func setOption<Option>(_ option: Option, value: Option.Value) -> EventLoopFuture<Void> where Option: ChannelOption {
-        preconditionFailure("Not implemented")
-    }
-
-    func getOption<Option>(_: Option) -> EventLoopFuture<Option.Value> where Option: ChannelOption {
-        preconditionFailure("Not implemented")
-    }
-}
diff --git a/Tests/AsyncHTTPClientTests/ConnectionTests+XCTest.swift b/Tests/AsyncHTTPClientTests/ConnectionTests+XCTest.swift
new file mode 100644
index 000000000..87837a2ca
--- /dev/null
+++ b/Tests/AsyncHTTPClientTests/ConnectionTests+XCTest.swift
@@ -0,0 +1,35 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the AsyncHTTPClient open source project
+//
+// Copyright (c) 2018-2019 Apple Inc. and the AsyncHTTPClient project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of AsyncHTTPClient project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+//
+// ConnectionTests+XCTest.swift
+//
+import XCTest
+
+///
+/// NOTE: This file was generated by generate_linux_tests.rb
+///
+/// Do NOT edit this file directly as it will be regenerated automatically when needed.
+///
+
+extension ConnectionTests {
+    static var allTests: [(String, (ConnectionTests) -> () throws -> Void)] {
+        return [
+            ("testConnectionReleaseActive", testConnectionReleaseActive),
+            ("testConnectionReleaseInactive", testConnectionReleaseInactive),
+            ("testConnectionRemoteCloseRelease", testConnectionRemoteCloseRelease),
+            ("testConnectionTimeoutRelease", testConnectionTimeoutRelease),
+            ("testAcquireAvailableBecomesUnavailable", testAcquireAvailableBecomesUnavailable),
+        ]
+    }
+}
diff --git a/Tests/AsyncHTTPClientTests/ConnectionTests.swift b/Tests/AsyncHTTPClientTests/ConnectionTests.swift
new file mode 100644
index 000000000..e01cefc99
--- /dev/null
+++ b/Tests/AsyncHTTPClientTests/ConnectionTests.swift
@@ -0,0 +1,229 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the AsyncHTTPClient open source project
+//
+// Copyright (c) 2018-2019 Apple Inc. and the AsyncHTTPClient project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of AsyncHTTPClient project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+@testable import AsyncHTTPClient
+import NIO
+import XCTest
+
+class ConnectionTests: XCTestCase {
+    var eventLoop: EmbeddedEventLoop!
+    var http1ConnectionProvider: HTTP1ConnectionProvider!
+
+    func buildState(connection: Connection, release: Bool) {
+        XCTAssertTrue(self.http1ConnectionProvider.state.enqueue())
+        let action = self.http1ConnectionProvider.state.acquire(waiter: .init(promise: self.eventLoop.makePromise(), setupComplete: self.eventLoop.makeSucceededFuture(()), preference: .indifferent))
+
+        switch action {
+        case .create(let waiter):
+            waiter.promise.succeed(connection)
+        default:
+            XCTFail("Unexpected action: \(action)")
+        }
+
+        // We offer the connection to the pool so that it can be tracked
+        _ = self.http1ConnectionProvider.state.offer(connection: connection)
+
+        if release {
+            _ = self.http1ConnectionProvider.state.release(connection: connection, closing: false)
+        }
+    }
+
+    // MARK: - Connection Tests
+
+    func testConnectionReleaseActive() throws {
+        let channel = ActiveChannel(eventLoop: self.eventLoop)
+        let connection = Connection(channel: channel, provider: self.http1ConnectionProvider)
+
+        self.buildState(connection: connection, release: false)
+
+        XCTAssertState(self.http1ConnectionProvider.state, available: 0, leased: 1, waiters: 0, pending: 0, opened: 1)
+
+        connection.release(closing: false, logger: HTTPClient.loggingDisabled)
+
+        // XCTAssertFalse(connection.isInUse)
+        XCTAssertState(self.http1ConnectionProvider.state, available: 1, leased: 0, waiters: 0, pending: 0, opened: 1)
+
+        // cleanup
+        // this cleanup code needs to go and use HTTP1ConnectionProvider's API instead
+        // (https://github.com/swift-server/async-http-client/issues/234)
+        connection.remoteClosed(logger: HTTPClient.loggingDisabled)
+    }
+
+    func testConnectionReleaseInactive() throws {
+        let channel = EmbeddedChannel()
+        let connection = Connection(channel: channel, provider: self.http1ConnectionProvider)
+
+        self.buildState(connection: connection, release: false)
+
+        XCTAssertState(self.http1ConnectionProvider.state, available: 0, leased: 1, waiters: 0, pending: 0, opened: 1)
+
+        connection.release(closing: true, logger: HTTPClient.loggingDisabled)
+        XCTAssertState(self.http1ConnectionProvider.state, available: 0, leased: 0, waiters: 0, pending: 0, opened: 0)
+    }
+
+    func testConnectionRemoteCloseRelease() throws {
+        let channel = ActiveChannel(eventLoop: self.eventLoop)
+        let connection = Connection(channel: channel, provider: self.http1ConnectionProvider)
+
+        self.buildState(connection: connection, release: true)
+
+        XCTAssertState(self.http1ConnectionProvider.state, available: 1, leased: 0, waiters: 0, pending: 0, opened: 1)
+
+        connection.remoteClosed(logger: HTTPClient.loggingDisabled)
+
+        XCTAssertState(self.http1ConnectionProvider.state, available: 0, leased: 0, waiters: 0, pending: 0, opened: 0)
+    }
+
+    func testConnectionTimeoutRelease() throws {
+        let channel = ActiveChannel(eventLoop: self.eventLoop)
+        let connection = Connection(channel: channel, provider: self.http1ConnectionProvider)
+
+        self.buildState(connection: connection, release: true)
+
+        XCTAssertState(self.http1ConnectionProvider.state, available: 1, leased: 0, waiters: 0, pending: 0, opened: 1)
+
+        connection.timeout(logger: HTTPClient.loggingDisabled)
+
+        XCTAssertState(self.http1ConnectionProvider.state, available: 0, leased: 0, waiters: 0, pending: 0, opened: 0)
+    }
+
+    func testAcquireAvailableBecomesUnavailable() throws {
+        let channel = ActiveChannel(eventLoop: self.eventLoop)
+        let connection = Connection(channel: channel, provider: self.http1ConnectionProvider)
+
+        self.buildState(connection: connection, release: true)
+
+        XCTAssertState(self.http1ConnectionProvider.state, available: 1, leased: 0, waiters: 0, pending: 0, opened: 1)
+
+        XCTAssertTrue(self.http1ConnectionProvider.enqueue())
+        let action = self.http1ConnectionProvider.state.acquire(waiter: .init(promise: self.eventLoop.makePromise(), setupComplete: self.eventLoop.makeSucceededFuture(()), preference: .indifferent))
+        switch action {
+        case .lease(let connection, let waiter):
+            // Since this connection is already in use, this should be a no-op and state should not have changed from normal lease
+            connection.timeout(logger: HTTPClient.loggingDisabled)
+
+            XCTAssertTrue(connection.isActiveEstimation)
+            XCTAssertState(self.http1ConnectionProvider.state, available: 0, leased: 1, waiters: 0, pending: 0, opened: 1, isLeased: connection)
+
+            // This is unrecoverable, but in this case we create a new connection, so state again should not change, even though release will be called
+            // This is important to preventself.http1ConnectionProvider deletion since connection is released and there could be 0 waiters
+            connection.remoteClosed(logger: HTTPClient.loggingDisabled)
+
+            XCTAssertState(self.http1ConnectionProvider.state, available: 0, leased: 1, waiters: 0, pending: 0, opened: 1, isLeased: connection)
+
+            // cleanup
+            waiter.promise.fail(TempError())
+            connection.release(closing: true, logger: HTTPClient.loggingDisabled)
+        default:
+            XCTFail("Unexpected action: \(action)")
+        }
+    }
+
+    override func setUp() {
+        XCTAssertNil(self.eventLoop)
+        XCTAssertNil(self.http1ConnectionProvider)
+        self.eventLoop = EmbeddedEventLoop()
+        XCTAssertNoThrow(self.http1ConnectionProvider = try HTTP1ConnectionProvider(key: .init(.init(url: "http://some.test")),
+                                                                                    eventLoop: self.eventLoop,
+                                                                                    configuration: .init(),
+                                                                                    pool: .init(configuration: .init(),
+                                                                                                backgroundActivityLogger: HTTPClient.loggingDisabled),
+                                                                                    backgroundActivityLogger: HTTPClient.loggingDisabled))
+    }
+
+    override func tearDown() {
+        XCTAssertNotNil(self.eventLoop)
+        XCTAssertNotNil(self.http1ConnectionProvider)
+        XCTAssertNoThrow(try self.http1ConnectionProvider.close().wait())
+        XCTAssertNoThrow(try self.eventLoop.syncShutdownGracefully())
+        self.eventLoop = nil
+        self.http1ConnectionProvider = nil
+    }
+}
+
+class ActiveChannel: Channel, ChannelCore {
+    struct NotImplementedError: Error {}
+
+    func localAddress0() throws -> SocketAddress {
+        throw NotImplementedError()
+    }
+
+    func remoteAddress0() throws -> SocketAddress {
+        throw NotImplementedError()
+    }
+
+    func register0(promise: EventLoopPromise<Void>?) {
+        promise?.fail(NotImplementedError())
+    }
+
+    func bind0(to: SocketAddress, promise: EventLoopPromise<Void>?) {
+        promise?.fail(NotImplementedError())
+    }
+
+    func connect0(to: SocketAddress, promise: EventLoopPromise<Void>?) {
+        promise?.fail(NotImplementedError())
+    }
+
+    func write0(_ data: NIOAny, promise: EventLoopPromise<Void>?) {
+        promise?.fail(NotImplementedError())
+    }
+
+    func flush0() {}
+
+    func read0() {}
+
+    func close0(error: Error, mode: CloseMode, promise: EventLoopPromise<Void>?) {
+        promise?.succeed(())
+    }
+
+    func triggerUserOutboundEvent0(_ event: Any, promise: EventLoopPromise<Void>?) {
+        promise?.fail(NotImplementedError())
+    }
+
+    func channelRead0(_: NIOAny) {}
+
+    func errorCaught0(error: Error) {}
+
+    var allocator: ByteBufferAllocator
+    var closeFuture: EventLoopFuture<Void>
+    var eventLoop: EventLoop
+
+    var localAddress: SocketAddress?
+    var remoteAddress: SocketAddress?
+    var parent: Channel?
+    var isWritable: Bool = true
+    var isActive: Bool = true
+
+    init(eventLoop: EmbeddedEventLoop) {
+        self.allocator = ByteBufferAllocator()
+        self.eventLoop = eventLoop
+        self.closeFuture = self.eventLoop.makeSucceededFuture(())
+    }
+
+    var _channelCore: ChannelCore {
+        return self
+    }
+
+    var pipeline: ChannelPipeline {
+        return ChannelPipeline(channel: self)
+    }
+
+    func setOption<Option>(_ option: Option, value: Option.Value) -> EventLoopFuture<Void> where Option: ChannelOption {
+        preconditionFailure("Not implemented")
+    }
+
+    func getOption<Option>(_: Option) -> EventLoopFuture<Option.Value> where Option: ChannelOption {
+        preconditionFailure("Not implemented")
+    }
+}
diff --git a/Tests/LinuxMain.swift b/Tests/LinuxMain.swift
index ad5e377c0..c6dfc8291 100644
--- a/Tests/LinuxMain.swift
+++ b/Tests/LinuxMain.swift
@@ -27,6 +27,7 @@ import XCTest
 
     XCTMain([
         testCase(ConnectionPoolTests.allTests),
+        testCase(ConnectionTests.allTests),
         testCase(HTTPClientCookieTests.allTests),
         testCase(HTTPClientInternalTests.allTests),
         testCase(HTTPClientNIOTSTests.allTests),

From f348d89cfaffce149b030f54221d66accd4eccbf Mon Sep 17 00:00:00 2001
From: Ilya Teterin <iteterin@apple.com>
Date: Sat, 26 Sep 2020 13:25:23 +0100
Subject: [PATCH 3/4] Gather Connection code into Connection.swift

Motivation:

For tests we will need a simple version of Connection, so here I gather
Connection code in one place and will generify ConnectionsState on next
commit.

Modifications:

Code of Connection is moved from multiple files into single
Connections.swift.

Result:

All tests are passing, no observable behaviour change.
---
 Sources/AsyncHTTPClient/Connection.swift      | 174 ++++++++++++++++++
 Sources/AsyncHTTPClient/ConnectionPool.swift  | 128 -------------
 .../StringConvertibleInstances.swift          |   6 -
 3 files changed, 174 insertions(+), 134 deletions(-)
 create mode 100644 Sources/AsyncHTTPClient/Connection.swift

diff --git a/Sources/AsyncHTTPClient/Connection.swift b/Sources/AsyncHTTPClient/Connection.swift
new file mode 100644
index 000000000..347800aff
--- /dev/null
+++ b/Sources/AsyncHTTPClient/Connection.swift
@@ -0,0 +1,174 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the AsyncHTTPClient open source project
+//
+// Copyright (c) 2019-2020 Apple Inc. and the AsyncHTTPClient project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of AsyncHTTPClient project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import Foundation
+import Logging
+import NIO
+import NIOConcurrencyHelpers
+import NIOHTTP1
+import NIOHTTPCompression
+import NIOTLS
+import NIOTransportServices
+
+/// A `Connection` represents a `Channel` in the context of the connection pool
+///
+/// In the `ConnectionPool`, each `Channel` belongs to a given `HTTP1ConnectionProvider`
+/// and has a certain "lease state" (see the `inUse` property).
+/// The role of `Connection` is to model this by storing a `Channel` alongside its associated properties
+/// so that they can be passed around together and correct provider can be identified when connection is released.
+class Connection {
+    /// The provider this `Connection` belongs to.
+    ///
+    /// This enables calling methods like `release()` directly on a `Connection` instead of
+    /// calling `provider.release(connection)`. This gives a more object oriented feel to the API
+    /// and can avoid having to keep explicit references to the pool at call site.
+    private let provider: HTTP1ConnectionProvider
+
+    /// The `Channel` of this `Connection`
+    ///
+    /// - Warning: Requests that lease connections from the `ConnectionPool` are responsible
+    /// for removing the specific handlers they added to the `Channel` pipeline before releasing it to the pool.
+    let channel: Channel
+
+    init(channel: Channel, provider: HTTP1ConnectionProvider) {
+        self.channel = channel
+        self.provider = provider
+    }
+}
+
+extension Connection {
+    /// Release this `Connection` to its associated `HTTP1ConnectionProvider`.
+    ///
+    /// - Warning: This only releases the connection and doesn't take care of cleaning handlers in the `Channel` pipeline.
+    func release(closing: Bool, logger: Logger) {
+        self.channel.eventLoop.assertInEventLoop()
+        self.provider.release(connection: self, closing: closing, logger: logger)
+    }
+
+    /// Called when channel exceeds idle time in pool.
+    func timeout(logger: Logger) {
+        self.channel.eventLoop.assertInEventLoop()
+        self.provider.timeout(connection: self, logger: logger)
+    }
+
+    /// Called when channel goes inactive while in the pool.
+    func remoteClosed(logger: Logger) {
+        self.channel.eventLoop.assertInEventLoop()
+        self.provider.remoteClosed(connection: self, logger: logger)
+    }
+
+    /// Called from `HTTP1ConnectionProvider.close` when client is shutting down.
+    func close() -> EventLoopFuture<Void> {
+        return self.channel.close()
+    }
+}
+
+/// Methods of Connection which are used in ConnectionsState extracted as protocol
+/// to facilitate test of ConnectionsState.
+protocol PoolManageableConnection: AnyObject {
+    func cancel() -> EventLoopFuture<Void>
+    var eventLoop: EventLoop { get }
+    var isActiveEstimation: Bool { get }
+}
+
+/// Implementation of methods used by ConnectionsState and its tests to manage Connection
+extension Connection: PoolManageableConnection {
+    /// Convenience property indicating whether the underlying `Channel` is active or not.
+    var isActiveEstimation: Bool {
+        return self.channel.isActive
+    }
+
+    var eventLoop: EventLoop {
+        return self.channel.eventLoop
+    }
+
+    func cancel() -> EventLoopFuture<Void> {
+        return self.channel.triggerUserOutboundEvent(TaskCancelEvent())
+    }
+}
+
+extension Connection {
+    /// Sets idle timeout handler and channel inactivity listener.
+    func setIdleTimeout(timeout: TimeAmount?, logger: Logger) {
+        _ = self.channel.pipeline.addHandler(IdleStateHandler(writeTimeout: timeout), position: .first).flatMap { _ in
+            self.channel.pipeline.addHandler(IdlePoolConnectionHandler(connection: self, logger: logger))
+        }
+    }
+
+    /// Removes idle timeout handler and channel inactivity listener
+    func cancelIdleTimeout() -> EventLoopFuture<Void> {
+        return self.removeHandler(IdleStateHandler.self).flatMap { _ in
+            self.removeHandler(IdlePoolConnectionHandler.self)
+        }
+    }
+}
+
+class IdlePoolConnectionHandler: ChannelInboundHandler, RemovableChannelHandler {
+    typealias InboundIn = NIOAny
+
+    let connection: Connection
+    var eventSent: Bool
+    let logger: Logger
+
+    init(connection: Connection, logger: Logger) {
+        self.connection = connection
+        self.eventSent = false
+        self.logger = logger
+    }
+
+    // this is needed to detect when remote end closes connection while connection is in the pool idling
+    func channelInactive(context: ChannelHandlerContext) {
+        if !self.eventSent {
+            self.eventSent = true
+            self.connection.remoteClosed(logger: self.logger)
+        }
+    }
+
+    func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) {
+        if let idleEvent = event as? IdleStateHandler.IdleStateEvent, idleEvent == .write {
+            if !self.eventSent {
+                self.eventSent = true
+                self.connection.timeout(logger: self.logger)
+            }
+        } else {
+            context.fireUserInboundEventTriggered(event)
+        }
+    }
+}
+
+extension Connection: CustomStringConvertible {
+    var description: String {
+        return "\(self.channel)"
+    }
+}
+
+struct ConnectionKey: Hashable {
+    let connection: Connection
+
+    init(_ connection: Connection) {
+        self.connection = connection
+    }
+
+    static func == (lhs: ConnectionKey, rhs: ConnectionKey) -> Bool {
+        return ObjectIdentifier(lhs.connection) == ObjectIdentifier(rhs.connection)
+    }
+
+    func hash(into hasher: inout Hasher) {
+        hasher.combine(ObjectIdentifier(self.connection))
+    }
+
+    func cancel() -> EventLoopFuture<Void> {
+        return self.connection.cancel()
+    }
+}
diff --git a/Sources/AsyncHTTPClient/ConnectionPool.swift b/Sources/AsyncHTTPClient/ConnectionPool.swift
index e482ceba6..c7df079e4 100644
--- a/Sources/AsyncHTTPClient/ConnectionPool.swift
+++ b/Sources/AsyncHTTPClient/ConnectionPool.swift
@@ -163,101 +163,6 @@ final class ConnectionPool {
     }
 }
 
-/// A `Connection` represents a `Channel` in the context of the connection pool
-///
-/// In the `ConnectionPool`, each `Channel` belongs to a given `HTTP1ConnectionProvider`
-/// and has a certain "lease state" (see the `inUse` property).
-/// The role of `Connection` is to model this by storing a `Channel` alongside its associated properties
-/// so that they can be passed around together and correct provider can be identified when connection is released.
-class Connection {
-    /// The provider this `Connection` belongs to.
-    ///
-    /// This enables calling methods like `release()` directly on a `Connection` instead of
-    /// calling `provider.release(connection)`. This gives a more object oriented feel to the API
-    /// and can avoid having to keep explicit references to the pool at call site.
-    let provider: HTTP1ConnectionProvider
-
-    /// The `Channel` of this `Connection`
-    ///
-    /// - Warning: Requests that lease connections from the `ConnectionPool` are responsible
-    /// for removing the specific handlers they added to the `Channel` pipeline before releasing it to the pool.
-    let channel: Channel
-
-    init(channel: Channel, provider: HTTP1ConnectionProvider) {
-        self.channel = channel
-        self.provider = provider
-    }
-
-    /// Convenience property indicating wether the underlying `Channel` is active or not.
-    var isActiveEstimation: Bool {
-        return self.channel.isActive
-    }
-
-    /// Release this `Connection` to its associated `HTTP1ConnectionProvider`.
-    ///
-    /// - Warning: This only releases the connection and doesn't take care of cleaning handlers in the `Channel` pipeline.
-    func release(closing: Bool, logger: Logger) {
-        self.channel.eventLoop.assertInEventLoop()
-        self.provider.release(connection: self, closing: closing, logger: logger)
-    }
-
-    /// Called when channel exceeds idle time in pool.
-    func timeout(logger: Logger) {
-        self.channel.eventLoop.assertInEventLoop()
-        self.provider.timeout(connection: self, logger: logger)
-    }
-
-    /// Called when channel goes inactive while in the pool.
-    func remoteClosed(logger: Logger) {
-        self.channel.eventLoop.assertInEventLoop()
-        self.provider.remoteClosed(connection: self, logger: logger)
-    }
-
-    func cancel() -> EventLoopFuture<Void> {
-        return self.channel.triggerUserOutboundEvent(TaskCancelEvent())
-    }
-
-    /// Called from `HTTP1ConnectionProvider.close` when client is shutting down.
-    func close() -> EventLoopFuture<Void> {
-        return self.channel.close()
-    }
-
-    /// Sets idle timeout handler and channel inactivity listener.
-    func setIdleTimeout(timeout: TimeAmount?, logger: Logger) {
-        _ = self.channel.pipeline.addHandler(IdleStateHandler(writeTimeout: timeout), position: .first).flatMap { _ in
-            self.channel.pipeline.addHandler(IdlePoolConnectionHandler(connection: self,
-                                                                       logger: logger))
-        }
-    }
-
-    /// Removes idle timeout handler and channel inactivity listener
-    func cancelIdleTimeout() -> EventLoopFuture<Void> {
-        return self.removeHandler(IdleStateHandler.self).flatMap { _ in
-            self.removeHandler(IdlePoolConnectionHandler.self)
-        }
-    }
-}
-
-struct ConnectionKey: Hashable {
-    let connection: Connection
-
-    init(_ connection: Connection) {
-        self.connection = connection
-    }
-
-    static func == (lhs: ConnectionKey, rhs: ConnectionKey) -> Bool {
-        return ObjectIdentifier(lhs.connection) == ObjectIdentifier(rhs.connection)
-    }
-
-    func hash(into hasher: inout Hasher) {
-        hasher.combine(ObjectIdentifier(self.connection))
-    }
-
-    func cancel() -> EventLoopFuture<Void> {
-        return self.connection.cancel()
-    }
-}
-
 /// A connection provider of `HTTP/1.1` connections with a given `Key` (host, scheme, port)
 ///
 /// On top of enabling connection reuse this provider it also facilitates the creation
@@ -586,39 +491,6 @@ class HTTP1ConnectionProvider {
     }
 }
 
-class IdlePoolConnectionHandler: ChannelInboundHandler, RemovableChannelHandler {
-    typealias InboundIn = NIOAny
-
-    let connection: Connection
-    var eventSent: Bool
-    let logger: Logger
-
-    init(connection: Connection, logger: Logger) {
-        self.connection = connection
-        self.eventSent = false
-        self.logger = logger
-    }
-
-    // this is needed to detect when remote end closes connection while connection is in the pool idling
-    func channelInactive(context: ChannelHandlerContext) {
-        if !self.eventSent {
-            self.eventSent = true
-            self.connection.remoteClosed(logger: self.logger)
-        }
-    }
-
-    func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) {
-        if let idleEvent = event as? IdleStateHandler.IdleStateEvent, idleEvent == .write {
-            if !self.eventSent {
-                self.eventSent = true
-                self.connection.timeout(logger: self.logger)
-            }
-        } else {
-            context.fireUserInboundEventTriggered(event)
-        }
-    }
-}
-
 extension CircularBuffer {
     mutating func swap(at index: Index, with value: Element) -> Element {
         let tmp = self[index]
diff --git a/Sources/AsyncHTTPClient/StringConvertibleInstances.swift b/Sources/AsyncHTTPClient/StringConvertibleInstances.swift
index 1039dc3f1..9cbb0afad 100644
--- a/Sources/AsyncHTTPClient/StringConvertibleInstances.swift
+++ b/Sources/AsyncHTTPClient/StringConvertibleInstances.swift
@@ -12,12 +12,6 @@
 //
 //===----------------------------------------------------------------------===//
 
-extension Connection: CustomStringConvertible {
-    var description: String {
-        return "\(self.channel)"
-    }
-}
-
 extension HTTP1ConnectionProvider.Waiter: CustomStringConvertible {
     var description: String {
         return "HTTP1ConnectionProvider.Waiter(\(self.preference))"

From 5cb8902a075565834269245b3fb9134439705c4c Mon Sep 17 00:00:00 2001
From: Ilya Teterin <iteterin@apple.com>
Date: Sat, 26 Sep 2020 13:40:26 +0100
Subject: [PATCH 4/4] Introduce generic type ConnectionType into
 ConnectionsState

Motivation:

To rework tests of ConnectionsState we want to have a "simpler" version
of Connection to be used, therefore here we convert ConnectionsState to
support generic type ConnectionType. We will substitute current
Connection with a test version in follow up commit.

Modifications:

ConnectionsState is altered to work on generic type ConnectionType
instead of solid type Connection.

Users of ConnectionsState are modified to provide type Connection into
ConnectionType in this commit.

Result:

Test are passing, no observable behaviour change.
---
 Sources/AsyncHTTPClient/Connection.swift      |  8 +--
 Sources/AsyncHTTPClient/ConnectionPool.swift  | 14 ++---
 .../AsyncHTTPClient/ConnectionsState.swift    | 62 +++++++++----------
 .../ConnectionPoolTests.swift                 | 18 +++---
 .../ConnectionPoolTestsSupport.swift          | 12 ++--
 5 files changed, 57 insertions(+), 57 deletions(-)

diff --git a/Sources/AsyncHTTPClient/Connection.swift b/Sources/AsyncHTTPClient/Connection.swift
index 347800aff..3923603ee 100644
--- a/Sources/AsyncHTTPClient/Connection.swift
+++ b/Sources/AsyncHTTPClient/Connection.swift
@@ -153,14 +153,14 @@ extension Connection: CustomStringConvertible {
     }
 }
 
-struct ConnectionKey: Hashable {
-    let connection: Connection
+struct ConnectionKey<ConnectionType>: Hashable where ConnectionType: PoolManageableConnection {
+    let connection: ConnectionType
 
-    init(_ connection: Connection) {
+    init(_ connection: ConnectionType) {
         self.connection = connection
     }
 
-    static func == (lhs: ConnectionKey, rhs: ConnectionKey) -> Bool {
+    static func == (lhs: ConnectionKey<ConnectionType>, rhs: ConnectionKey<ConnectionType>) -> Bool {
         return ObjectIdentifier(lhs.connection) == ObjectIdentifier(rhs.connection)
     }
 
diff --git a/Sources/AsyncHTTPClient/ConnectionPool.swift b/Sources/AsyncHTTPClient/ConnectionPool.swift
index c7df079e4..ed9bcbd8c 100644
--- a/Sources/AsyncHTTPClient/ConnectionPool.swift
+++ b/Sources/AsyncHTTPClient/ConnectionPool.swift
@@ -191,7 +191,7 @@ class HTTP1ConnectionProvider {
 
     var closePromise: EventLoopPromise<Void>
 
-    var state: ConnectionsState
+    var state: ConnectionsState<Connection>
 
     private let backgroundActivityLogger: Logger
 
@@ -222,7 +222,7 @@ class HTTP1ConnectionProvider {
         self.state.assertInvariants()
     }
 
-    func execute(_ action: Action, logger: Logger) {
+    func execute(_ action: Action<Connection>, logger: Logger) {
         switch action {
         case .lease(let connection, let waiter):
             // if connection is became inactive, we create a new one.
@@ -297,7 +297,7 @@ class HTTP1ConnectionProvider {
     func getConnection(preference: HTTPClient.EventLoopPreference,
                        setupComplete: EventLoopFuture<Void>,
                        logger: Logger) -> EventLoopFuture<Connection> {
-        let waiter = Waiter(promise: self.eventLoop.makePromise(), setupComplete: setupComplete, preference: preference)
+        let waiter = Waiter<Connection>(promise: self.eventLoop.makePromise(), setupComplete: setupComplete, preference: preference)
 
         let action: Action = self.lock.withLock {
             self.state.acquire(waiter: waiter)
@@ -309,10 +309,10 @@ class HTTP1ConnectionProvider {
     }
 
     func connect(_ result: Result<Channel, Error>,
-                 waiter: Waiter,
+                 waiter: Waiter<Connection>,
                  replacing closedConnection: Connection? = nil,
                  logger: Logger) {
-        let action: Action
+        let action: Action<Connection>
         switch result {
         case .success(let channel):
             logger.trace("successfully created connection",
@@ -478,9 +478,9 @@ class HTTP1ConnectionProvider {
     ///
     /// `Waiter`s are created when `maximumConcurrentConnections` is reached
     /// and we cannot create new connections anymore.
-    struct Waiter {
+    struct Waiter<ConnectionType: PoolManageableConnection> {
         /// The promise to complete once a connection is available
-        let promise: EventLoopPromise<Connection>
+        let promise: EventLoopPromise<ConnectionType>
 
         /// Future that will be succeeded when request timeout handler and `TaskHandler` are added to the pipeline.
         let setupComplete: EventLoopFuture<Void>
diff --git a/Sources/AsyncHTTPClient/ConnectionsState.swift b/Sources/AsyncHTTPClient/ConnectionsState.swift
index d99416a49..2ceace465 100644
--- a/Sources/AsyncHTTPClient/ConnectionsState.swift
+++ b/Sources/AsyncHTTPClient/ConnectionsState.swift
@@ -15,29 +15,29 @@
 import NIO
 
 extension HTTP1ConnectionProvider {
-    enum Action {
-        case lease(Connection, Waiter)
-        case create(Waiter)
-        case replace(Connection, Waiter)
+    enum Action<ConnectionType: PoolManageableConnection> {
+        case lease(ConnectionType, Waiter<ConnectionType>)
+        case create(Waiter<ConnectionType>)
+        case replace(ConnectionType, Waiter<ConnectionType>)
         case closeProvider
-        case park(Connection)
+        case park(ConnectionType)
         case none
-        case fail(Waiter, Error)
-        indirect case closeAnd(Connection, Action)
-        indirect case parkAnd(Connection, Action)
+        case fail(Waiter<ConnectionType>, Error)
+        indirect case closeAnd(ConnectionType, Action<ConnectionType>)
+        indirect case parkAnd(ConnectionType, Action<ConnectionType>)
     }
 
-    struct ConnectionsState {
+    struct ConnectionsState<ConnectionType: PoolManageableConnection> {
         enum State {
             case active
             case closed
         }
 
-        struct Snapshot {
+        struct Snapshot<ConnectionType: PoolManageableConnection> {
             var state: State
-            var availableConnections: CircularBuffer<Connection>
-            var leasedConnections: Set<ConnectionKey>
-            var waiters: CircularBuffer<Waiter>
+            var availableConnections: CircularBuffer<ConnectionType>
+            var leasedConnections: Set<ConnectionKey<ConnectionType>>
+            var waiters: CircularBuffer<Waiter<ConnectionType>>
             var openedConnectionsCount: Int
             var pending: Int
         }
@@ -48,16 +48,16 @@ extension HTTP1ConnectionProvider {
         private var state: State = .active
 
         /// Opened connections that are available.
-        private var availableConnections: CircularBuffer<Connection> = .init(initialCapacity: 8)
+        private var availableConnections: CircularBuffer<ConnectionType> = .init(initialCapacity: 8)
 
         /// Opened connections that are leased to the user.
-        private var leasedConnections: Set<ConnectionKey> = .init()
+        private var leasedConnections: Set<ConnectionKey<ConnectionType>> = .init()
 
         /// Consumers that weren't able to get a new connection without exceeding
         /// `maximumConcurrentConnections` get a `Future<Connection>`
         /// whose associated promise is stored in `Waiter`. The promise is completed
         /// as soon as possible by the provider, in FIFO order.
-        private var waiters: CircularBuffer<Waiter> = .init(initialCapacity: 8)
+        private var waiters: CircularBuffer<Waiter<ConnectionType>> = .init(initialCapacity: 8)
 
         /// Number of opened or opening connections, used to keep track of all connections and enforcing `maximumConcurrentConnections` limit.
         private var openedConnectionsCount: Int = 0
@@ -70,11 +70,11 @@ extension HTTP1ConnectionProvider {
             self.eventLoop = eventLoop
         }
 
-        func testsOnly_getInternalState() -> Snapshot {
+        func testsOnly_getInternalState() -> Snapshot<ConnectionType> {
             return Snapshot(state: self.state, availableConnections: self.availableConnections, leasedConnections: self.leasedConnections, waiters: self.waiters, openedConnectionsCount: self.openedConnectionsCount, pending: self.pending)
         }
 
-        mutating func testsOnly_setInternalState(_ snapshot: Snapshot) {
+        mutating func testsOnly_setInternalState(_ snapshot: Snapshot<ConnectionType>) {
             self.state = snapshot.state
             self.availableConnections = snapshot.availableConnections
             self.leasedConnections = snapshot.leasedConnections
@@ -109,7 +109,7 @@ extension HTTP1ConnectionProvider {
             return self.openedConnectionsCount == 0 && self.pending == 0
         }
 
-        mutating func acquire(waiter: Waiter) -> Action {
+        mutating func acquire(waiter: Waiter<ConnectionType>) -> Action<ConnectionType> {
             switch self.state {
             case .active:
                 self.pending -= 1
@@ -117,7 +117,7 @@ extension HTTP1ConnectionProvider {
                 let (eventLoop, required) = self.resolvePreference(waiter.preference)
                 if required {
                     // If there is an opened connection on the same EL - use it
-                    if let found = self.availableConnections.firstIndex(where: { $0.channel.eventLoop === eventLoop }) {
+                    if let found = self.availableConnections.firstIndex(where: { $0.eventLoop === eventLoop }) {
                         let connection = self.availableConnections.remove(at: found)
                         self.leasedConnections.insert(ConnectionKey(connection))
                         return .lease(connection, waiter)
@@ -151,7 +151,7 @@ extension HTTP1ConnectionProvider {
             }
         }
 
-        mutating func release(connection: Connection, closing: Bool) -> Action {
+        mutating func release(connection: ConnectionType, closing: Bool) -> Action<ConnectionType> {
             switch self.state {
             case .active:
                 assert(self.leasedConnections.contains(ConnectionKey(connection)))
@@ -161,12 +161,12 @@ extension HTTP1ConnectionProvider {
                         let (eventLoop, required) = self.resolvePreference(waiter.preference)
 
                         // If returned connection is on same EL or we do not require special EL - lease it
-                        if connection.channel.eventLoop === eventLoop || !required {
+                        if connection.eventLoop === eventLoop || !required {
                             return .lease(connection, waiter)
                         }
 
                         // If there is an opened connection on the same loop, lease it and park returned
-                        if let found = self.availableConnections.firstIndex(where: { $0.channel.eventLoop === eventLoop }) {
+                        if let found = self.availableConnections.firstIndex(where: { $0.eventLoop === eventLoop }) {
                             self.leasedConnections.remove(ConnectionKey(connection))
                             let replacement = self.availableConnections.swap(at: found, with: connection)
                             self.leasedConnections.insert(ConnectionKey(replacement))
@@ -203,7 +203,7 @@ extension HTTP1ConnectionProvider {
             }
         }
 
-        mutating func offer(connection: Connection) -> Action {
+        mutating func offer(connection: ConnectionType) -> Action<ConnectionType> {
             switch self.state {
             case .active:
                 self.leasedConnections.insert(ConnectionKey(connection))
@@ -214,7 +214,7 @@ extension HTTP1ConnectionProvider {
             }
         }
 
-        mutating func drop(connection: Connection) {
+        mutating func drop(connection: ConnectionType) {
             switch self.state {
             case .active:
                 self.leasedConnections.remove(ConnectionKey(connection))
@@ -223,7 +223,7 @@ extension HTTP1ConnectionProvider {
             }
         }
 
-        mutating func connectFailed() -> Action {
+        mutating func connectFailed() -> Action<ConnectionType> {
             switch self.state {
             case .active:
                 self.openedConnectionsCount -= 1
@@ -239,7 +239,7 @@ extension HTTP1ConnectionProvider {
             }
         }
 
-        mutating func remoteClosed(connection: Connection) -> Action {
+        mutating func remoteClosed(connection: ConnectionType) -> Action<ConnectionType> {
             switch self.state {
             case .active:
                 // Connection can be closed remotely while we wait for `.lease` action to complete.
@@ -260,7 +260,7 @@ extension HTTP1ConnectionProvider {
             }
         }
 
-        mutating func timeout(connection: Connection) -> Action {
+        mutating func timeout(connection: ConnectionType) -> Action<ConnectionType> {
             switch self.state {
             case .active:
                 // We can get timeout and inUse = true when we decided to lease the connection, but this action is not executed yet.
@@ -285,12 +285,12 @@ extension HTTP1ConnectionProvider {
             }
         }
 
-        mutating func processNextWaiter() -> Action {
+        mutating func processNextWaiter() -> Action<ConnectionType> {
             if let waiter = self.waiters.popFirst() {
                 let (eventLoop, required) = self.resolvePreference(waiter.preference)
 
                 // If specific EL is required, we have only two options - find open one or create a new one
-                if required, let found = self.availableConnections.firstIndex(where: { $0.channel.eventLoop === eventLoop }) {
+                if required, let found = self.availableConnections.firstIndex(where: { $0.eventLoop === eventLoop }) {
                     let connection = self.availableConnections.remove(at: found)
                     self.leasedConnections.insert(ConnectionKey(connection))
                     return .lease(connection, waiter)
@@ -313,7 +313,7 @@ extension HTTP1ConnectionProvider {
             return .none
         }
 
-        mutating func close() -> (CircularBuffer<Waiter>, CircularBuffer<Connection>, Set<ConnectionKey>, Bool)? {
+        mutating func close() -> (CircularBuffer<Waiter<ConnectionType>>, CircularBuffer<ConnectionType>, Set<ConnectionKey<ConnectionType>>, Bool)? {
             switch self.state {
             case .active:
                 let waiters = self.waiters
diff --git a/Tests/AsyncHTTPClientTests/ConnectionPoolTests.swift b/Tests/AsyncHTTPClientTests/ConnectionPoolTests.swift
index a42d6198e..62a8d1db0 100644
--- a/Tests/AsyncHTTPClientTests/ConnectionPoolTests.swift
+++ b/Tests/AsyncHTTPClientTests/ConnectionPoolTests.swift
@@ -28,7 +28,7 @@ class ConnectionPoolTests: XCTestCase {
     var http1ConnectionProvider: HTTP1ConnectionProvider!
 
     func testPending() {
-        var state = HTTP1ConnectionProvider.ConnectionsState(eventLoop: self.eventLoop)
+        var state = HTTP1ConnectionProvider.ConnectionsState<Connection>(eventLoop: self.eventLoop)
 
         var snapshot = state.testsOnly_getInternalState()
         XCTAssertEqual(0, snapshot.availableConnections.count)
@@ -50,7 +50,7 @@ class ConnectionPoolTests: XCTestCase {
     // MARK: - Acquire Tests
 
     func testAcquireWhenEmpty() {
-        var state = HTTP1ConnectionProvider.ConnectionsState(eventLoop: self.eventLoop)
+        var state = HTTP1ConnectionProvider.ConnectionsState<Connection>(eventLoop: self.eventLoop)
 
         var snapshot = state.testsOnly_getInternalState()
         XCTAssertEqual(0, snapshot.availableConnections.count)
@@ -155,7 +155,7 @@ class ConnectionPoolTests: XCTestCase {
     // MARK: - Acquire on Specific EL Tests
 
     func testAcquireWhenEmptySpecificEL() {
-        var state = HTTP1ConnectionProvider.ConnectionsState(eventLoop: self.eventLoop)
+        var state = HTTP1ConnectionProvider.ConnectionsState<Connection>(eventLoop: self.eventLoop)
         var snapshot = state.testsOnly_getInternalState()
 
         XCTAssertEqual(0, snapshot.availableConnections.count)
@@ -301,7 +301,7 @@ class ConnectionPoolTests: XCTestCase {
     // MARK: - Acquire Errors Tests
 
     func testAcquireWhenClosed() {
-        var state = HTTP1ConnectionProvider.ConnectionsState(eventLoop: self.eventLoop)
+        var state = HTTP1ConnectionProvider.ConnectionsState<Connection>(eventLoop: self.eventLoop)
         _ = state.close()
 
         XCTAssertFalse(state.enqueue())
@@ -317,7 +317,7 @@ class ConnectionPoolTests: XCTestCase {
     }
 
     func testConnectFailedWhenClosed() {
-        var state = HTTP1ConnectionProvider.ConnectionsState(eventLoop: self.eventLoop)
+        var state = HTTP1ConnectionProvider.ConnectionsState<Connection>(eventLoop: self.eventLoop)
         _ = state.close()
 
         let action = state.connectFailed()
@@ -1222,7 +1222,7 @@ class ConnectionPoolTests: XCTestCase {
     // MARK: - Shutdown tests
 
     func testShutdownOnPendingAndSuccess() {
-        var state = HTTP1ConnectionProvider.ConnectionsState(eventLoop: self.eventLoop)
+        var state = HTTP1ConnectionProvider.ConnectionsState<Connection>(eventLoop: self.eventLoop)
 
         XCTAssertTrue(state.enqueue())
 
@@ -1261,7 +1261,7 @@ class ConnectionPoolTests: XCTestCase {
     }
 
     func testShutdownOnPendingAndError() {
-        var state = HTTP1ConnectionProvider.ConnectionsState(eventLoop: self.eventLoop)
+        var state = HTTP1ConnectionProvider.ConnectionsState<Connection>(eventLoop: self.eventLoop)
 
         XCTAssertTrue(state.enqueue())
 
@@ -1298,7 +1298,7 @@ class ConnectionPoolTests: XCTestCase {
     }
 
     func testShutdownTimeout() {
-        var state = HTTP1ConnectionProvider.ConnectionsState(eventLoop: self.eventLoop)
+        var state = HTTP1ConnectionProvider.ConnectionsState<Connection>(eventLoop: self.eventLoop)
 
         var snapshot = self.http1ConnectionProvider.state.testsOnly_getInternalState()
 
@@ -1333,7 +1333,7 @@ class ConnectionPoolTests: XCTestCase {
     }
 
     func testShutdownRemoteClosed() {
-        var state = HTTP1ConnectionProvider.ConnectionsState(eventLoop: self.eventLoop)
+        var state = HTTP1ConnectionProvider.ConnectionsState<Connection>(eventLoop: self.eventLoop)
 
         var snapshot = self.http1ConnectionProvider.state.testsOnly_getInternalState()
 
diff --git a/Tests/AsyncHTTPClientTests/ConnectionPoolTestsSupport.swift b/Tests/AsyncHTTPClientTests/ConnectionPoolTestsSupport.swift
index dc5281e60..70edd1426 100644
--- a/Tests/AsyncHTTPClientTests/ConnectionPoolTestsSupport.swift
+++ b/Tests/AsyncHTTPClientTests/ConnectionPoolTestsSupport.swift
@@ -17,10 +17,10 @@ import NIO
 import XCTest
 
 extension ConnectionPoolTests {
-    func buildState(count: Int, release: Bool = true, eventLoop: EventLoop? = nil) -> (HTTP1ConnectionProvider.ConnectionsState, [Connection]) {
+    func buildState(count: Int, release: Bool = true, eventLoop: EventLoop? = nil) -> (HTTP1ConnectionProvider.ConnectionsState<Connection>, [Connection]) {
         let eventLoop = eventLoop ?? self.eventLoop!
 
-        var state = HTTP1ConnectionProvider.ConnectionsState(eventLoop: eventLoop)
+        var state = HTTP1ConnectionProvider.ConnectionsState<Connection>(eventLoop: eventLoop)
         var items: [Connection] = []
 
         if count == 0 {
@@ -60,7 +60,7 @@ extension ConnectionPoolTests {
     }
 }
 
-func XCTAssertState(_ state: HTTP1ConnectionProvider.ConnectionsState, available: Int, leased: Int, waiters: Int, pending: Int, opened: Int) {
+func XCTAssertState<ConnectionType>(_ state: HTTP1ConnectionProvider.ConnectionsState<ConnectionType>, available: Int, leased: Int, waiters: Int, pending: Int, opened: Int) {
     let snapshot = state.testsOnly_getInternalState()
     XCTAssertEqual(available, snapshot.availableConnections.count)
     XCTAssertEqual(leased, snapshot.leasedConnections.count)
@@ -69,7 +69,7 @@ func XCTAssertState(_ state: HTTP1ConnectionProvider.ConnectionsState, available
     XCTAssertEqual(opened, snapshot.openedConnectionsCount)
 }
 
-func XCTAssertState(_ state: HTTP1ConnectionProvider.ConnectionsState, available: Int, leased: Int, waiters: Int, pending: Int, opened: Int, isLeased connection: Connection) {
+func XCTAssertState<ConnectionType>(_ state: HTTP1ConnectionProvider.ConnectionsState<ConnectionType>, available: Int, leased: Int, waiters: Int, pending: Int, opened: Int, isLeased connection: ConnectionType) {
     let snapshot = state.testsOnly_getInternalState()
     XCTAssertEqual(available, snapshot.availableConnections.count)
     XCTAssertEqual(leased, snapshot.leasedConnections.count)
@@ -79,7 +79,7 @@ func XCTAssertState(_ state: HTTP1ConnectionProvider.ConnectionsState, available
     XCTAssertTrue(snapshot.leasedConnections.contains(ConnectionKey(connection)))
 }
 
-func XCTAssertState(_ state: HTTP1ConnectionProvider.ConnectionsState, available: Int, leased: Int, waiters: Int, pending: Int, opened: Int, isNotLeased connection: Connection) {
+func XCTAssertState<ConnectionType>(_ state: HTTP1ConnectionProvider.ConnectionsState<ConnectionType>, available: Int, leased: Int, waiters: Int, pending: Int, opened: Int, isNotLeased connection: ConnectionType) {
     let snapshot = state.testsOnly_getInternalState()
     XCTAssertEqual(available, snapshot.availableConnections.count)
     XCTAssertEqual(leased, snapshot.leasedConnections.count)
@@ -100,7 +100,7 @@ func XCTUnwrap<T>(_ value: T?) throws -> T {
 
 struct TempError: Error {}
 
-func XCTAssertStateClose(_ state: HTTP1ConnectionProvider.ConnectionsState, available: Int, leased: Int, waiters: Int, clean: Bool) throws {
+func XCTAssertStateClose<ConnectionType>(_ state: HTTP1ConnectionProvider.ConnectionsState<ConnectionType>, available: Int, leased: Int, waiters: Int, clean: Bool) throws {
     var state = state
 
     let (foundWaiters, foundAvailable, foundLeased, foundClean) = try XCTUnwrap(state.close())