add state to HTTP2StateMaschine and add testExecuteOnShuttingDownPool

Author: dnadoba

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

@@ -21,6 +21,12 @@ extension HTTPConnectionPool {
         typealias RequestAction = HTTPConnectionPool.StateMachine.RequestAction
         typealias ConnectionAction = HTTPConnectionPool.StateMachine.ConnectionAction
 
+        private enum State: Equatable {
+            case running
+            case shuttingDown(unclean: Bool)
+            case shutDown
+        }
+
         private var lastConnectFailure: Error?
         private var failedConsecutiveConnectionAttempts = 0
 
@@ -31,6 +37,8 @@ extension HTTPConnectionPool {
 
         private let idGenerator: Connection.ID.Generator
 
+        private var state: State = .running
+
         init(
             idGenerator: Connection.ID.Generator
         ) {
@@ -68,10 +76,24 @@ extension HTTPConnectionPool {
         }
 
         mutating func executeRequest(_ request: Request) -> Action {
-            if let eventLoop = request.requiredEventLoop {
-                return self.executeRequest(request, onRequired: eventLoop)
-            } else {
-                return self.executeRequest(request, onPreferred: request.preferredEventLoop)
+            switch self.state {
+            case .running:
+                if let eventLoop = request.requiredEventLoop {
+                    return self.executeRequest(request, onRequired: eventLoop)
+                } else {
+                    return self.executeRequest(request, onPreferred: request.preferredEventLoop)
+                }
+            case .shutDown, .shuttingDown:
+                // it is fairly unlikely that this condition is met, since the ConnectionPoolManager
+                // also fails new requests immediately, if it is shutting down. However there might
+                // be race conditions in which a request passes through a running connection pool
+                // manager, but hits a connection pool that is already shutting down.
+                //
+                // (Order in one lock does not guarantee order in the next lock!)
+                return .init(
+                    request: .failRequest(request, HTTPClientError.alreadyShutdown, cancelTimeout: false),
+                    connection: .none
+                )
             }
         }
 
@@ -149,36 +171,57 @@ extension HTTPConnectionPool {
             at index: Int,
             context: HTTP2Connections.AvailableConnectionContext
         ) -> Action {
-            // We prioritise requests with a required event loop over those without a requirement.
-            // This can cause starvation for request without a required event loop.
-            // We should come up with a better algorithm in the future.
-
-            var requestsToExecute = self.requests.popFirst(max: context.availableStreams, for: context.eventLoop)
-            let remainingAvailableStreams = context.availableStreams - requestsToExecute.count
-            // use the remaining available streams for requests without a required event loop
-            requestsToExecute += self.requests.popFirst(max: remainingAvailableStreams, for: nil)
-            let connection = self.connections.leaseStreams(at: index, count: requestsToExecute.count)
-
-            let requestAction = { () -> RequestAction in
-                if requestsToExecute.isEmpty {
-                    return .none
-                } else {
-                    return .executeRequestsAndCancelTimeouts(requestsToExecute, connection)
-                }
-            }()
+            switch self.state {
+            case .running:
+                // We prioritise requests with a required event loop over those without a requirement.
+                // This can cause starvation for request without a required event loop.
+                // We should come up with a better algorithm in the future.
+
+                var requestsToExecute = self.requests.popFirst(max: context.availableStreams, for: context.eventLoop)
+                let remainingAvailableStreams = context.availableStreams - requestsToExecute.count
+                // use the remaining available streams for requests without a required event loop
+                requestsToExecute += self.requests.popFirst(max: remainingAvailableStreams, for: nil)
+                let connection = self.connections.leaseStreams(at: index, count: requestsToExecute.count)
+
+                let requestAction = { () -> RequestAction in
+                    if requestsToExecute.isEmpty {
+                        return .none
+                    } else {
+                        return .executeRequestsAndCancelTimeouts(requestsToExecute, connection)
+                    }
+                }()
+
+                let connectionAction = { () -> ConnectionAction in
+                    if context.isIdle, requestsToExecute.isEmpty {
+                        return .scheduleTimeoutTimer(connection.id, on: context.eventLoop)
+                    } else {
+                        return .none
+                    }
+                }()
 
-            let connectionAction = { () -> ConnectionAction in
-                if context.isIdle, requestsToExecute.isEmpty {
-                    return .scheduleTimeoutTimer(connection.id, on: context.eventLoop)
-                } else {
+                return .init(
+                    request: requestAction,
+                    connection: connectionAction
+                )
+            case .shuttingDown(let unclean):
+                guard context.isIdle else {
                     return .none
                 }
-            }()
 
-            return .init(
-                request: requestAction,
-                connection: connectionAction
-            )
+                let connection = self.connections.closeConnection(at: index)
+                if self.connections.isEmpty {
+                    return .init(
+                        request: .none,
+                        connection: .closeConnection(connection, isShutdown: .yes(unclean: unclean))
+                    )
+                }
+                return .init(
+                    request: .none,
+                    connection: .closeConnection(connection, isShutdown: .no)
+                )
+            case .shutDown:
+                preconditionFailure("It the pool is already shutdown, all connections must have been torn down.")
+            }
         }
 
         mutating func newHTTP2MaxConcurrentStreamsReceived(_ connectionID: Connection.ID, newMaxStreams: Int) -> Action {
@@ -199,32 +242,53 @@ extension HTTPConnectionPool {
         }
 
         private mutating func nextActionForFailedConnection(at index: Int, on eventLoop: EventLoop) -> Action {
-            let hasPendingRequest = !self.requests.isEmpty(for: eventLoop) || !self.requests.isEmpty(for: nil)
-            guard hasPendingRequest else {
-                return .none
-            }
+            switch self.state {
+            case .running:
+                let hasPendingRequest = !self.requests.isEmpty(for: eventLoop) || !self.requests.isEmpty(for: nil)
+                guard hasPendingRequest else {
+                    return .none
+                }
 
-            let (newConnectionID, previousEventLoop) = self.connections.createNewConnectionByReplacingClosedConnection(at: index)
-            precondition(previousEventLoop === eventLoop)
+                let (newConnectionID, previousEventLoop) = self.connections.createNewConnectionByReplacingClosedConnection(at: index)
+                precondition(previousEventLoop === eventLoop)
 
-            return .init(
-                request: .none,
-                connection: .createConnection(newConnectionID, on: eventLoop)
-            )
+                return .init(
+                    request: .none,
+                    connection: .createConnection(newConnectionID, on: eventLoop)
+                )
+            case .shuttingDown(let unclean):
+                assert(self.requests.isEmpty)
+                self.connections.removeConnection(at: index)
+                if self.connections.isEmpty {
+                    return .init(
+                        request: .none,
+                        connection: .cleanupConnections(.init(), isShutdown: .yes(unclean: unclean))
+                    )
+                }
+                return .none
+
+            case .shutDown:
+                preconditionFailure("If the pool is already shutdown, all connections must have been torn down.")
+            }
         }
 
         private mutating func nextActionForClosingConnection(on eventLoop: EventLoop) -> Action {
-            let hasPendingRequest = !self.requests.isEmpty(for: eventLoop) || !self.requests.isEmpty(for: nil)
-            guard hasPendingRequest else {
-                return .none
-            }
+            switch self.state {
+            case .running:
+                let hasPendingRequest = !self.requests.isEmpty(for: eventLoop) || !self.requests.isEmpty(for: nil)
+                guard hasPendingRequest else {
+                    return .none
+                }
 
-            let newConnectionID = self.connections.createNewConnection(on: eventLoop)
+                let newConnectionID = self.connections.createNewConnection(on: eventLoop)
 
-            return .init(
-                request: .none,
-                connection: .createConnection(newConnectionID, on: eventLoop)
-            )
+                return .init(
+                    request: .none,
+                    connection: .createConnection(newConnectionID, on: eventLoop)
+                )
+            case .shutDown, .shuttingDown:
+                return .none
+            }
         }
 
         mutating func http2ConnectionStreamClosed(_ connectionID: Connection.ID) -> Action {
@@ -248,7 +312,7 @@ extension HTTPConnectionPool {
             guard let (index, context) = self.connections.failConnection(connectionID) else {
                 preconditionFailure("Backing off a connection that is unknown to us?")
             }
-            return nextActionForFailedConnection(at: index, on: context.eventLoop)
+            return self.nextActionForFailedConnection(at: index, on: context.eventLoop)
         }
 
         mutating func timeoutRequest(_ requestID: Request.ID) -> Action {
@@ -289,8 +353,13 @@ extension HTTPConnectionPool {
 
         mutating func connectionIdleTimeout(_ connectionID: Connection.ID) -> Action {
             guard let connection = connections.closeConnectionIfIdle(connectionID) else {
+                // because of a race this connection (connection close runs against trigger of timeout)
+                // was already removed from the state machine.
                 return .none
             }
+
+            precondition(self.state == .running, "If we are shutting down, we must not have any idle connections")
+
             return .init(
                 request: .none,
                 connection: .closeConnection(connection, isShutdown: .no)
@@ -329,8 +398,10 @@ extension HTTPConnectionPool {
             let unclean = !(cleanupContext.cancel.isEmpty && waitingRequests.isEmpty)
             if self.connections.isEmpty {
                 isShutdown = .yes(unclean: unclean)
+                self.state = .shutDown
             } else {
                 isShutdown = .no
+                self.state = .shuttingDown(unclean: unclean)
             }
             return .init(
                 request: requestAction,

Diff for: Tests/AsyncHTTPClientTests/HTTPConnectionPool+HTTP2StateMachineTests+XCTest.swift

@@ -26,6 +26,9 @@ extension HTTPConnectionPool_HTTP2StateMachineTests {
     static var allTests: [(String, (HTTPConnectionPool_HTTP2StateMachineTests) -> () throws -> Void)] {
         return [
             ("testCreatingOfConnection", testCreatingOfConnection),
+            ("testConnectionFailureBackoff", testConnectionFailureBackoff),
+            ("testCancelRequestWorks", testCancelRequestWorks),
+            ("testExecuteOnShuttingDownPool", testExecuteOnShuttingDownPool),
         ]
     }
 }

Diff for: Tests/AsyncHTTPClientTests/HTTPConnectionPool+HTTP2StateMachineTests.swift

@@ -114,7 +114,7 @@ class HTTPConnectionPool_HTTP2StateMachineTests: XCTestCase {
             isShutdown: .yes(unclean: false)
         ))
     }
-    
+
     func testConnectionFailureBackoff() {
         let elg = EmbeddedEventLoopGroup(loops: 4)
         defer { XCTAssertNoThrow(try elg.syncShutdownGracefully()) }
@@ -170,7 +170,7 @@ class HTTPConnectionPool_HTTP2StateMachineTests: XCTestCase {
         // 4. retry connection, but no more queued requests.
         XCTAssertEqual(state.connectionCreationBackoffDone(newConnectionID), .none)
     }
-    
+
     func testCancelRequestWorks() {
         let elg = EmbeddedEventLoopGroup(loops: 4)
         defer { XCTAssertNoThrow(try elg.syncShutdownGracefully()) }
@@ -207,4 +207,58 @@ class HTTPConnectionPool_HTTP2StateMachineTests: XCTestCase {
         XCTAssertEqual(connectedAction.request, .none, "Request must not be executed")
         XCTAssertEqual(connectedAction.connection, .scheduleTimeoutTimer(connectionID, on: connectionEL))
     }
+
+    func testExecuteOnShuttingDownPool() {
+        let elg = EmbeddedEventLoopGroup(loops: 4)
+        defer { XCTAssertNoThrow(try elg.syncShutdownGracefully()) }
+
+        var state = HTTPConnectionPool.HTTP2StateMaschine(
+            idGenerator: .init()
+        )
+
+        let mockRequest = MockHTTPRequest(eventLoop: elg.next())
+        let request = HTTPConnectionPool.Request(mockRequest)
+
+        let executeAction = state.executeRequest(request)
+        XCTAssertEqual(.scheduleRequestTimeout(for: request, on: mockRequest.eventLoop), executeAction.request)
+
+        // 1. connection attempt
+        guard case .createConnection(let connectionID, on: let connectionEL) = executeAction.connection else {
+            return XCTFail("Unexpected connection action: \(executeAction.connection)")
+        }
+        XCTAssert(connectionEL === mockRequest.eventLoop) // XCTAssertIdentical not available on Linux
+
+        // 2. connection succeeds
+        let connection: HTTPConnectionPool.Connection = .__testOnly_connection(id: connectionID, eventLoop: connectionEL)
+        let connectedAction = state.newHTTP2ConnectionEstablished(connection, maxConcurrentStreams: 100)
+        guard case .executeRequestsAndCancelTimeouts([request], connection) = connectedAction.request else {
+            return XCTFail("Unexpected request action: \(connectedAction.request)")
+        }
+        XCTAssert(request.__testOnly_wrapped_request() === mockRequest) // XCTAssertIdentical not available on Linux
+        XCTAssertEqual(connectedAction.connection, .none)
+
+        // 3. shutdown
+        let shutdownAction = state.shutdown()
+        XCTAssertEqual(.none, shutdownAction.request)
+        guard case .cleanupConnections(let cleanupContext, isShutdown: .no) = shutdownAction.connection else {
+            return XCTFail("Unexpected connection action: \(shutdownAction.connection)")
+        }
+
+        XCTAssertEqual(cleanupContext.cancel.count, 1)
+        XCTAssertEqual(cleanupContext.cancel.first?.id, connectionID)
+        XCTAssertEqual(cleanupContext.close, [])
+        XCTAssertEqual(cleanupContext.connectBackoff, [])
+
+        // 4. execute another request
+        let finalMockRequest = MockHTTPRequest(eventLoop: elg.next())
+        let finalRequest = HTTPConnectionPool.Request(finalMockRequest)
+        let failAction = state.executeRequest(finalRequest)
+        XCTAssertEqual(failAction.connection, .none)
+        XCTAssertEqual(failAction.request, .failRequest(finalRequest, HTTPClientError.alreadyShutdown, cancelTimeout: false))
+
+        // 5. close open connection
+        let closeAction = state.connectionClosed(connectionID)
+        XCTAssertEqual(closeAction.connection, .cleanupConnections(.init(), isShutdown: .yes(unclean: true)))
+        XCTAssertEqual(closeAction.request, .none)
+    }
 }