Add PoolStateMachine.RequestQueue (vapor#424)

Author: fabianfett

Sources/ConnectionPoolModule/ConnectionRequest.swift

@@ -4,11 +4,13 @@ public struct ConnectionRequest<Connection: PooledConnection>: ConnectionRequest
 
     public var id: ID
 
-    private var continuation: CheckedContinuation<Connection, ConnectionPoolError>
+    @usableFromInline
+    private(set) var continuation: CheckedContinuation<Connection, any Error>
 
+    @inlinable
     init(
         id: Int,
-        continuation: CheckedContinuation<Connection, ConnectionPoolError>
+        continuation: CheckedContinuation<Connection, any Error>
     ) {
         self.id = id
         self.continuation = continuation

Sources/ConnectionPoolModule/OneElementFastSequence.swift

@@ -17,7 +17,7 @@ struct OneElementFastSequence<Element>: Sequence {
     }
 
     @inlinable
-    init(_ element: Element) {
+    init(element: Element) {
         self.base = .one(element, reserveCapacity: 1)
     }
 

Sources/ConnectionPoolModule/PoolStateMachine+RequestQueue.swift

@@ -0,0 +1,71 @@
+import DequeModule
+
+@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)
+extension PoolStateMachine {
+
+    /// A request queue, which can enqueue requests in O(1), dequeue requests in O(1) and even cancel requests in O(1).
+    ///
+    /// While enqueueing and dequeueing on O(1) is trivial, cancellation is hard, as it normally requires a removal within the
+    /// underlying Deque. However thanks to having an additional `requests` dictionary, we can remove the cancelled
+    /// request from the dictionary and keep it inside the queue. Whenever we pop a request from the deque, we validate
+    /// that it hasn't been cancelled in the meantime by checking if the popped request is still in the `requests` dictionary.
+    @usableFromInline
+    struct RequestQueue {
+        @usableFromInline
+        private(set) var queue: Deque<RequestID>
+
+        @usableFromInline
+        private(set) var requests: [RequestID: Request]
+
+        @inlinable
+        var count: Int {
+            self.requests.count
+        }
+
+        @inlinable
+        var isEmpty: Bool {
+            self.count == 0
+        }
+
+        @usableFromInline
+        init() {
+            self.queue = .init(minimumCapacity: 256)
+            self.requests = .init(minimumCapacity: 256)
+        }
+
+        @inlinable
+        mutating func queue(_ request: Request) {
+            self.requests[request.id] = request
+            self.queue.append(request.id)
+        }
+
+        @inlinable
+        mutating func pop(max: UInt16) -> OneElementFastSequence<Request> {
+            var result = OneElementFastSequence<Request>()
+            result.reserveCapacity(Int(max))
+            var popped = 0
+            while let requestID = self.queue.popFirst(), popped < max {
+                if let requestIndex = self.requests.index(forKey: requestID) {
+                    popped += 1
+                    result.append(self.requests.remove(at: requestIndex).value)
+                }
+            }
+
+            assert(result.count <= max)
+            return result
+        }
+
+        @inlinable
+        mutating func remove(_ requestID: RequestID) -> Request? {
+            self.requests.removeValue(forKey: requestID)
+        }
+
+        @inlinable
+        mutating func removeAll() -> OneElementFastSequence<Request> {
+            let result = OneElementFastSequence(self.requests.values)
+            self.requests.removeAll()
+            self.queue.removeAll()
+            return result
+        }
+    }
+}

Sources/ConnectionPoolModule/PoolStateMachine.swift

@@ -0,0 +1,74 @@
+#if canImport(Darwin)
+import Darwin
+#else
+import Glibc
+#endif
+
+@usableFromInline
+@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)
+struct PoolConfiguration {
+    /// The minimum number of connections to preserve in the pool.
+    ///
+    /// If the pool is mostly idle and the remote servers closes idle connections,
+    /// the `ConnectionPool` will initiate new outbound connections proactively
+    /// to avoid the number of available connections dropping below this number.
+    @usableFromInline
+    var minimumConnectionCount: Int = 0
+
+    /// The maximum number of connections to for this pool, to be preserved.
+    @usableFromInline
+    var maximumConnectionSoftLimit: Int = 10
+
+    @usableFromInline
+    var maximumConnectionHardLimit: Int = 10
+
+    @usableFromInline
+    var keepAliveDuration: Duration?
+
+    @usableFromInline
+    var idleTimeoutDuration: Duration = .seconds(30)
+}
+
+@usableFromInline
+@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)
+struct PoolStateMachine<
+    Connection: PooledConnection,
+    ConnectionIDGenerator: ConnectionIDGeneratorProtocol,
+    ConnectionID: Hashable & Sendable,
+    Request: ConnectionRequestProtocol,
+    RequestID,
+    TimerCancellationToken
+> where Connection.ID == ConnectionID, ConnectionIDGenerator.ID == ConnectionID, RequestID == Request.ID {
+
+    @usableFromInline
+    struct Timer: Hashable, Sendable {
+        @usableFromInline
+        enum Usecase: Sendable {
+            case backoff
+            case idleTimeout
+            case keepAlive
+        }
+
+        @usableFromInline
+        var connectionID: ConnectionID
+
+        @usableFromInline
+        var timerID: Int
+
+        @usableFromInline
+        var duration: Duration
+
+        @usableFromInline
+        var usecase: Usecase
+
+        @inlinable
+        init(connectionID: ConnectionID, timerID: Int, duration: Duration, usecase: Usecase) {
+            self.connectionID = connectionID
+            self.timerID = timerID
+            self.duration = duration
+            self.usecase = usecase
+        }
+    }
+
+
+}

Tests/ConnectionPoolModuleTests/ConnectionRequestTests.swift

@@ -0,0 +1,27 @@
+@testable import _ConnectionPoolModule
+import XCTest
+
+final class ConnectionRequestTests: XCTestCase {
+
+    func testHappyPath() async throws {
+        let mockConnection = MockConnection(id: 1)
+        let connection = try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<MockConnection, any Error>) in
+            let request = ConnectionRequest(id: 42, continuation: continuation)
+            XCTAssertEqual(request.id, 42)
+            continuation.resume(with: .success(mockConnection))
+        }
+
+        XCTAssert(connection === mockConnection)
+    }
+
+    func testSadPath() async throws {
+        do {
+            _ = try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<MockConnection, any Error>) in
+                continuation.resume(with: .failure(ConnectionPoolError.requestCancelled))
+            }
+            XCTFail("This point should not be reached")
+        } catch {
+            XCTAssertEqual(error as? ConnectionPoolError, .requestCancelled)
+        }
+    }
+}

Tests/ConnectionPoolModuleTests/Mocks/MockRequest.swift

@@ -0,0 +1,28 @@
+import _ConnectionPoolModule
+
+final class MockRequest: ConnectionRequestProtocol, Hashable, Sendable {
+    typealias Connection = MockConnection
+
+    struct ID: Hashable {
+        var objectID: ObjectIdentifier
+
+        init(_ request: MockRequest) {
+            self.objectID = ObjectIdentifier(request)
+        }
+    }
+
+    var id: ID { ID(self) }
+
+
+    static func ==(lhs: MockRequest, rhs: MockRequest) -> Bool {
+        lhs.id == rhs.id
+    }
+
+    func hash(into hasher: inout Hasher) {
+        hasher.combine(self.id)
+    }
+
+    func complete(with: Result<Connection, ConnectionPoolError>) {
+
+    }
+}

Tests/ConnectionPoolModuleTests/Mocks/MockTimerCancellationToken.swift

@@ -0,0 +1,16 @@
+@testable import _ConnectionPoolModule
+
+@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)
+struct MockTimerCancellationToken: Hashable, Sendable {
+    var connectionID: MockConnection.ID
+    var timerID: Int
+    var duration: Duration
+    var usecase: TestPoolStateMachine.Timer.Usecase
+
+    init(_ timer: TestPoolStateMachine.Timer) {
+        self.connectionID = timer.connectionID
+        self.timerID = timer.timerID
+        self.duration = timer.duration
+        self.usecase = timer.usecase
+    }
+}

Tests/ConnectionPoolModuleTests/OneElementFastSequence.swift

@@ -35,7 +35,7 @@ final class OneElementFastSequenceTests: XCTestCase {
         }
         XCTAssertEqual(array.capacity, 8)
 
-        var oneElemSequence = OneElementFastSequence<Int>(1)
+        var oneElemSequence = OneElementFastSequence<Int>(element: 1)
         oneElemSequence.reserveCapacity(8)
         oneElemSequence.append(2)
         guard case .n(let array) = oneElemSequence.base else {