Add tests

Author: fabianfett

Diff for: Sources/AsyncHTTPClient/ConnectionPool/State Machine/HTTPConnectionPool+HTTP1Connections.swift

@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 import NIOCore
+import XCTest
 
 extension HTTPConnectionPool {
     /// Represents the state of a single HTTP/1.1 connection
@@ -219,7 +220,7 @@ extension HTTPConnectionPool {
 
         var startingGeneralPurposeConnections: Int {
             var connecting = 0
-            for connectionState in self.connections {
+            for connectionState in self.connections[0..<self.overflowIndex] {
                 if connectionState.isConnecting || connectionState.isBackingOff {
                     connecting += 1
                 }
@@ -236,7 +237,18 @@ extension HTTPConnectionPool {
             }
         }
 
-        // MARK: - Internal functions -
+        // MARK: - Mutations -
+
+        enum NextConnectionAction<Request> {
+            /// Directly lease the connection again
+            case lease(Connection, Request)
+            /// Start an idle timeout timer and make connection available to new requests
+            case park(Connection.ID)
+            /// Close the connection and remove
+            case close(Connection)
+        }
+
+        // MARK: Connection creation
 
         mutating func createNewConnection(on eventLoop: EventLoop) -> Connection.ID {
             let connection = HTTP1ConnectionState(connectionID: self.generator.next(), eventLoop: eventLoop)
@@ -251,13 +263,16 @@ extension HTTPConnectionPool {
             return connection.connectionID
         }
 
-        mutating func newHTTP1ConnectionCreated(_ connection: Connection) -> Connection.ID {
+        mutating func newHTTP1ConnectionCreated<RequestType>(
+            _ connection: Connection,
+            _ closure: (IdleConnectionContext) -> (IdleConnectionAction<RequestType>)
+        ) -> NextConnectionAction<RequestType> {
             guard let index = self.connections.firstIndex(where: { $0.connectionID == connection.id }) else {
                 preconditionFailure("There is a new connection that we didn't request!")
             }
-
+            assert(connection.eventLoop === self.connections[index].eventLoop, "Expected the new connection to be on EL")
             self.connections[index].started(connection)
-            return connection.id
+            return self.withIdleConnection(index: index, closure: closure)
         }
 
         mutating func backoffNextConnectionAttempt(_ connectionID: Connection.ID) -> EventLoop {
@@ -269,19 +284,114 @@ extension HTTPConnectionPool {
             return self.connections[index].eventLoop
         }
 
-        /// Information around the failed/closed connection.
-        struct FailedConnectionContext {
-            /// The failed connection's use. Did it serve in the pool or was it specialized for an `EventLoop`?
-            enum PreviousUse {
-                case generalPurpose
-                case eventLoop
+        // MARK: Leasing and releasing
+
+        mutating func leaseConnection(onPreferred preferredEL: EventLoop) -> Connection? {
+            guard let index = self.findAvailableConnection(onPreferred: preferredEL) else {
+                return nil
+            }
+
+            return self.connections[index].lease()
+        }
+
+        mutating func leaseConnection(onRequired preferredEL: EventLoop) -> Connection? {
+            guard let index = self.findAvailableConnection(onRequired: preferredEL) else {
+                return nil
+            }
+
+            return self.connections[index].lease()
+        }
+
+        mutating func releaseConnection<RequestType>(
+            _ connectionID: Connection.ID,
+            _ closure: (IdleConnectionContext) -> (IdleConnectionAction<RequestType>)
+        ) -> NextConnectionAction<RequestType> {
+            guard let index = self.connections.firstIndex(where: { $0.connectionID == connectionID }) else {
+                preconditionFailure("A connection that we don't know was released? Something is very wrong...")
+            }
+
+            self.connections[index].release()
+            return self.withIdleConnection(index: index, closure: closure)
+        }
+
+        struct IdleConnectionContext {
+            /// The `EventLoop` the connection runs on.
+            var eventLoop: EventLoop
+            /// The connection's use. Either general purpose or for requests with `EventLoop`
+            /// requirements.
+            var use: ConnectionUse
+        }
+
+        enum IdleConnectionAction<R> {
+            /// Directly lease the connection again
+            case lease(for: R)
+            /// Start an idle timeout timer and make connection available to new requests
+            case park
+            /// Close the connection and remove
+            case close
+        }
+
+        private mutating func withIdleConnection<RequestType>(
+            index: Int,
+            closure: (IdleConnectionContext) -> (IdleConnectionAction<RequestType>)
+        ) -> NextConnectionAction<RequestType> {
+            let eventLoop = self.connections[index].eventLoop
+            let use: ConnectionUse
+            if index < self.overflowIndex {
+                use = .generalPurpose
+            } else {
+                use = .eventLoop(eventLoop)
+            }
+
+            let context = IdleConnectionContext(eventLoop: eventLoop, use: use)
+            switch closure(context) {
+            case .close:
+                if index < self.overflowIndex {
+                    self.overflowIndex -= 1
+                }
+                var connectionState = self.connections.remove(at: index)
+                return .close(connectionState.close())
+
+            case .lease(for: let request):
+                return .lease(self.connections[index].lease(), request)
+
+            case .park:
+                return .park(self.connections[index].connectionID)
+            }
+        }
+
+        // MARK: Connection close/removal
+
+        mutating func closeConnectionIfIdle(_ connectionID: Connection.ID) -> Connection? {
+            guard let index = self.connections.firstIndex(where: { $0.connectionID == connectionID }) else {
+                // because of a race this connection (connection close runs against trigger of timeout)
+                // was already removed from the state machine.
+                return .none
+            }
+
+            guard self.connections[index].isAvailable else {
+                // connection is not available anymore, we may have just leased it for a request
+                return nil
             }
 
+            return self.closeConnection(at: index)
+        }
+
+        // MARK: Connection failure
+
+        /// A connection's use. Did it serve in the pool or was it specialized for an `EventLoop`?
+        enum ConnectionUse {
+            case generalPurpose
+            case eventLoop(EventLoop)
+        }
+
+        /// Information around the failed/closed connection.
+        struct FailedConnectionContext {
             /// The eventLoop the connection ran on.
             var eventLoop: EventLoop
             /// The failed connection's use.
-            var use: PreviousUse
-            /// Connections that we start up for this usecase
+            var use: ConnectionUse
+            /// Connections that we start up for this use-case
             var connectionsStartingForUseCase: Int
         }
 
@@ -292,6 +402,15 @@ extension HTTPConnectionPool {
             case remove
         }
 
+        /// Remove or replace a connection.
+        ///
+        /// Call this function after the connection start backoff is done, or a live connection is closed.
+        ///
+        /// - Parameters:
+        ///   - connectionID: The failed connection's ID
+        ///   - closure: A closure in which a decision is made how to progress with this connection spot
+        ///
+        /// - Returns: A new connectionID and eventLoop to create a new connection, if the action is `.replace`
         mutating func failConnection(
             _ connectionID: Connection.ID,
             _ closure: (FailedConnectionContext) -> (FailedConnectionAction)
@@ -300,15 +419,17 @@ extension HTTPConnectionPool {
                 preconditionFailure("We tried to create a new connection that we know nothing about?")
             }
 
-            let use: FailedConnectionContext.PreviousUse
+            let use: ConnectionUse
             let eventLoop = self.connections[index].eventLoop
             let starting: Int
             if index < self.overflowIndex {
                 use = .generalPurpose
-                starting = self.startingGeneralPurposeConnections
+                // we need to subtract 1, since we dont want this connection to count
+                starting = self.startingGeneralPurposeConnections - 1
             } else {
-                use = .eventLoop
-                starting = self.startingEventLoopConnections(on: eventLoop)
+                use = .eventLoop(eventLoop)
+                // we need to subtract 1, since we dont want this connection to count
+                starting = self.startingEventLoopConnections(on: eventLoop) - 1
             }
 
             let context = FailedConnectionContext(
@@ -335,67 +456,18 @@ extension HTTPConnectionPool {
             }
         }
 
-        mutating func leaseConnection(onPreferred preferredEL: EventLoop) -> Connection? {
-            guard let index = self.findAvailableConnection(onPreferred: preferredEL) else {
-                return nil
-            }
-
-            return self.connections[index].lease()
-        }
-
-        mutating func leaseConnection(onRequired preferredEL: EventLoop) -> Connection? {
-            guard let index = self.findAvailableConnection(onRequired: preferredEL) else {
-                return nil
-            }
-
-            return self.connections[index].lease()
-        }
-
-        mutating func releaseConnection(_ connectionID: Connection.ID) {
-            guard let index = self.connections.firstIndex(where: { $0.connectionID == connectionID }) else {
-                preconditionFailure("A connection that we don't know was released? Something is very wrong...")
-            }
-
-            self.connections[index].release()
-        }
-
-        /// Remove deletes an already closed connection from the connection pool.
-        mutating func removeConnection(_ connectionID: Connection.ID) {
-            guard let index = self.connections.firstIndex(where: { $0.connectionID == connectionID }) else {
-                preconditionFailure("We tried to create a new connection, that we know nothing about?")
-            }
-            self.connections.remove(at: index)
-            self.overflowIndex = self.connections.index(before: self.overflowIndex)
-        }
-
-        mutating func closeConnectionIfIdle(_ connectionID: Connection.ID) -> Connection? {
-            guard let index = self.connections.firstIndex(where: { $0.connectionID == connectionID }) else {
-                // because of a race this connection (connection close runs against trigger of timeout)
-                // was already removed from the state machine.
-                return .none
-            }
-
-            guard self.connections[index].isAvailable else {
-                // connection is not available anymore, we may have just leased it for a request
-                return nil
-            }
-
-            return self.closeConnection(at: index)
-        }
-
-        mutating func closeConnection(_ connectionID: Connection.ID) -> Connection {
-            guard let index = self.connections.firstIndex(where: { $0.connectionID == connectionID }) else {
-                preconditionFailure("We tried to create a new connection, that we know nothing about?")
-            }
-            return self.closeConnection(at: index)
-        }
+        // MARK: Shutdown
 
         mutating func shutdown() -> CleanupContext {
             var cleanupContext = CleanupContext()
+            let initialOverflowIndex = self.overflowIndex
 
-            self.connections = self.connections.compactMap { connectionState in
+            self.connections = self.connections.enumerated().compactMap { index, connectionState in
                 var connectionState = connectionState
                 if connectionState.cleanup(&cleanupContext) {
+                    if index < initialOverflowIndex {
+                        self.overflowIndex -= 1
+                    }
                     return nil
                 }
                 return connectionState