Started work on a pool manager

Author: fabianfett

Sources/ConnectionPoolModule/ConnectionLease.swift

@@ -0,0 +1,22 @@
+//
+//  ConnectionLease.swift
+//  postgres-nio
+//
+//  Created by Fabian Fett on 05.05.25.
+//
+
+struct ConnectionLease<Connection: PooledConnection> {
+
+    var connection: Connection
+    var _release: () -> ()
+
+    init(connection: Connection, release: @escaping () -> Void) {
+        self.connection = connection
+        self._release = release
+    }
+
+    func release() {
+        self._release()
+    }
+
+}

Sources/ConnectionPoolModule/ConnectionPool.swift

@@ -134,9 +134,11 @@ public final class ConnectionPool<
     Connection: PooledConnection,
     ConnectionID: Hashable & Sendable,
     ConnectionIDGenerator: ConnectionIDGeneratorProtocol,
+    ConnectionConfiguration: Equatable & Sendable,
     Request: ConnectionRequestProtocol,
     RequestID: Hashable & Sendable,
     KeepAliveBehavior: ConnectionKeepAliveBehavior,
+    Executor: ConnectionPoolExecutor,
     ObservabilityDelegate: ConnectionPoolObservabilityDelegate,
     Clock: _Concurrency.Clock
 >: Sendable where
@@ -148,14 +150,19 @@ public final class ConnectionPool<
     ObservabilityDelegate.ConnectionID == ConnectionID,
     Clock.Duration == Duration
 {
-    public typealias ConnectionFactory = @Sendable (ConnectionID, ConnectionPool<Connection, ConnectionID, ConnectionIDGenerator, Request, RequestID, KeepAliveBehavior, ObservabilityDelegate, Clock>) async throws -> ConnectionAndMetadata<Connection>
+    public typealias ConnectionFactory = @Sendable (ConnectionID, ConnectionConfiguration, ConnectionPool<Connection, ConnectionID, ConnectionIDGenerator, ConnectionConfiguration, Request, RequestID, KeepAliveBehavior, Executor, ObservabilityDelegate, Clock>) async throws -> ConnectionAndMetadata<Connection>
 
     @usableFromInline
     typealias StateMachine = PoolStateMachine<Connection, ConnectionIDGenerator, ConnectionID, Request, Request.ID, CheckedContinuation<Void, Never>>
 
     @usableFromInline
     let factory: ConnectionFactory
 
+    public let executor: Executor
+
+    @usableFromInline
+    let connectionConfiguration: ConnectionConfiguration
+
     @usableFromInline
     let keepAliveBehavior: KeepAliveBehavior
 
@@ -188,18 +195,22 @@ public final class ConnectionPool<
 
     public init(
         configuration: ConnectionPoolConfiguration,
+        connectionConfiguration: ConnectionConfiguration,
         idGenerator: ConnectionIDGenerator,
         requestType: Request.Type,
         keepAliveBehavior: KeepAliveBehavior,
+        executor: Executor,
         observabilityDelegate: ObservabilityDelegate,
         clock: Clock,
         connectionFactory: @escaping ConnectionFactory
     ) {
+        self.executor = executor
         self.clock = clock
         self.factory = connectionFactory
         self.keepAliveBehavior = keepAliveBehavior
         self.observabilityDelegate = observabilityDelegate
         self.configuration = configuration
+        self.connectionConfiguration = connectionConfiguration
         var stateMachine = StateMachine(
             configuration: .init(configuration, keepAliveBehavior: keepAliveBehavior),
             generator: idGenerator,
@@ -271,6 +282,13 @@ public final class ConnectionPool<
         }
     }
 
+    @inlinable
+    public func updateConfiguration(_ configuration: ConnectionConfiguration, forceReconnection: Bool) {
+        // TODO: Implement connection will close correctly
+        // If the forceReconnection flag is set, we should gracefully close the connection once they
+        // are returned the next time.
+    }
+
     @inlinable
     public func run() async {
         await withTaskCancellationHandler {
@@ -419,11 +437,11 @@ public final class ConnectionPool<
 
     @inlinable
     /*private*/ func makeConnection(for request: StateMachine.ConnectionRequest, in taskGroup: inout some TaskGroupProtocol) {
-        taskGroup.addTask_ {
+        self.addTask(into: &taskGroup) {
             self.observabilityDelegate.startedConnecting(id: request.connectionID)
 
             do {
-                let bundle = try await self.factory(request.connectionID, self)
+                let bundle = try await self.factory(request.connectionID, self.connectionConfiguration, self)
                 self.connectionEstablished(bundle)
 
                 // after the connection has been established, we keep the task open. This ensures
@@ -468,7 +486,7 @@ public final class ConnectionPool<
     /*private*/ func runKeepAlive(_ connection: Connection, in taskGroup: inout some TaskGroupProtocol) {
         self.observabilityDelegate.keepAliveTriggered(id: connection.id)
 
-        taskGroup.addTask_ {
+        self.addTask(into: &taskGroup) {
             do {
                 try await self.keepAliveBehavior.runKeepAlive(for: connection)
 
@@ -502,8 +520,8 @@ public final class ConnectionPool<
     }
 
     @inlinable
-    /*private*/ func runTimer(_ timer: StateMachine.Timer, in poolGroup: inout some TaskGroupProtocol) {
-        poolGroup.addTask_ { () async -> () in
+    /*private*/ func runTimer(_ timer: StateMachine.Timer, in taskGroup: inout some TaskGroupProtocol) {
+        self.addTask(into: &taskGroup) { () async -> () in
             await withTaskGroup(of: TimerRunResult.self, returning: Void.self) { taskGroup in
                 taskGroup.addTask {
                     do {
@@ -554,6 +572,15 @@ public final class ConnectionPool<
             token.resume()
         }
     }
+
+    @inlinable
+    func addTask(into taskGroup: inout some TaskGroupProtocol, operation: @escaping @Sendable () async -> Void) {
+        if #available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, *), let executor = self.executor as? TaskExecutor {
+            taskGroup.addTask_(executorPreference: executor, operation: operation)
+        } else {
+            taskGroup.addTask_(operation: operation)
+        }
+    }
 }
 
 @available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)
@@ -573,6 +600,9 @@ protocol TaskGroupProtocol {
     // under exactly this name and others have different attributes. So let's pick
     // a name that doesn't clash anywhere and implement it using the standard `addTask`.
     mutating func addTask_(operation: @escaping @Sendable () async -> Void)
+
+    @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, *)
+    mutating func addTask_(executorPreference: ((any TaskExecutor)?), operation: @escaping @Sendable () async -> Void)
 }
 
 @available(macOS 14.0, iOS 17.0, tvOS 17.0, watchOS 10.0, *)
@@ -581,11 +611,23 @@ extension DiscardingTaskGroup: TaskGroupProtocol {
     mutating func addTask_(operation: @escaping @Sendable () async -> Void) {
         self.addTask(priority: nil, operation: operation)
     }
+
+    @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, *)
+    @inlinable
+    mutating func addTask_(executorPreference: (any TaskExecutor)?, operation: @escaping @Sendable () async -> Void) {
+        self.addTask(executorPreference: executorPreference, operation: operation)
+    }
 }
 
 extension TaskGroup<Void>: TaskGroupProtocol {
     @inlinable
     mutating func addTask_(operation: @escaping @Sendable () async -> Void) {
         self.addTask(priority: nil, operation: operation)
     }
+
+    @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, *)
+    @inlinable
+    mutating func addTask_(executorPreference: (any TaskExecutor)?, operation: @escaping @Sendable () async -> Void) {
+        self.addTask(executorPreference: executorPreference, operation: operation)
+    }
 }

Sources/ConnectionPoolModule/ConnectionPoolExecutor.swift

@@ -0,0 +1,15 @@
+public protocol ConnectionPoolExecutor: AnyObject, Sendable {
+    associatedtype ID: Hashable, Sendable
+
+    var id: ID { get }
+
+    static func getExecutorID() -> Self.ID?
+}
+
+public final class NothingConnectionPoolExecutor: ConnectionPoolExecutor {
+    public typealias ID = ObjectIdentifier
+
+    public var id: ObjectIdentifier { ObjectIdentifier(self) }
+
+    public static func getExecutorID() -> ObjectIdentifier? { nil }
+}