A bit of Sendable work (WIP)

Author: trasch

Package.resolved

@@ -12,9 +12,9 @@ let package = Package(
     ],
     dependencies: [
         .package(url: "https://github.com/apple/swift-atomics.git", from: "1.2.0"),
-        .package(url: "https://github.com/apple/swift-collections", from: "1.1.0"),
-        .package(url: "https://github.com/vapor/postgres-nio.git", from: "1.20.2"),
-        .package(url: "https://github.com/vapor/postgres-kit", from: "2.12.3"),
+        .package(url: "https://github.com/apple/swift-collections", from: "1.1.1"),
+        .package(url: "https://github.com/vapor/postgres-nio.git", from: "1.21.4"),
+        .package(url: "https://github.com/vapor/postgres-kit", from: "2.13.5"),
     ],
     targets: [
         .target(
@@ -30,4 +30,5 @@ let package = Package(
             dependencies: [
                 .target(name: "PostgresConnectionPool"),
             ]),
-    ])
+    ]
+)

Package.swift

@@ -7,15 +7,20 @@ import PostgresNIO
 
 // MARK: CustomStringConvertible
 
-extension PSQLError: CustomStringConvertible {
+extension PSQLError {
 
     /// A short error description.
-    public var description: String {
+    public var pgPoolDescription: String {
         if let serverInfo = self.serverInfo,
            let severity = serverInfo[.severity],
            let message = serverInfo[.message]
         {
-            return "<PSQLError: \(severity): \(message)>"
+            if let code = serverInfo[.sqlState] {
+                return "<PSQLError (code \(code)): \(severity): \(message)>"
+            }
+            else {
+                return "<PSQLError: \(severity): \(message)>"
+            }
         }
 
         return "<PSQLError: \(self.code.description)>"
@@ -25,10 +30,10 @@ extension PSQLError: CustomStringConvertible {
 
 // MARK: - CustomDebugStringConvertible
 
-extension PSQLError: CustomDebugStringConvertible {
+extension PSQLError {
 
     /// A detailed error description suitable for debugging queries and other problems with the server.
-    public var debugDescription: String {
+    public var pgPoolDebugDescription: String {
         var messageElements: [String] = [
             "code: \(self.code)"
         ]

Sources/PostgresConnectionPool/Extensions/PSQLError+Description.swift

@@ -6,7 +6,7 @@ import Foundation
 import PostgresNIO
 
 /// Settings for the pool like connection parameters.
-public struct PoolConfiguration {
+public struct PoolConfiguration: Sendable {
 
     /// PostgreSQL connection parameters.
     @available(*, deprecated, message: "Use `PostgresConnection.Configuration` etc. instead.")
@@ -57,15 +57,15 @@ public struct PoolConfiguration {
     /// Called when new connections to the database are openend.
     ///
     /// Use this to set extra connection options or override the defaults.
-    public var onOpenConnection: ((PostgresConnection, Logger) async throws -> Void)?
+    public let onOpenConnection: (@Sendable (PostgresConnection, Logger) async throws -> Void)?
 
     /// Called before a connection is given to a client.
     ///
     /// Default is to do a quick connection check with "SELECT 1" and close the connection on errors.
-    public var onReturnConnection: ((PostgresConnection, Logger) async throws -> Void)?
+    public let onReturnConnection: (@Sendable (PostgresConnection, Logger) async throws -> Void)?
 
     /// Called just before a connection is being closed.
-    public var onCloseConnection: ((PostgresConnection, Logger) async throws -> Void)?
+    public let onCloseConnection: (@Sendable (PostgresConnection, Logger) async throws -> Void)?
 
     /// Pool configuation.
     ///
@@ -76,14 +76,20 @@ public struct PoolConfiguration {
     ///   - queryTimeout: TImeout for individual database queries
     ///   - poolSize: The maximum number of open connections
     ///   - maxIdleConnections: The maximum number of idle connections
+    ///   - onOpenConnection: Called when new connections to the database are openend
+    ///   - onReturnConnection: Called before a connection is given to a client
+    ///   - onCloseConnection: Called just before a connection is being closed
     public init(
         applicationName: String,
         postgresConfiguration: PostgresConnection.Configuration,
         connectTimeout: TimeInterval = 5.0,
         connectionRetryInterval: TimeInterval = 0.5,
         queryTimeout: TimeInterval? = nil,
         poolSize: Int = 10,
-        maxIdleConnections: Int? = nil)
+        maxIdleConnections: Int? = nil,
+        onOpenConnection: (@Sendable (PostgresConnection, Logger) async throws -> Void)? = nil,
+        onReturnConnection: (@Sendable (PostgresConnection, Logger) async throws -> Void)? = nil,
+        onCloseConnection: (@Sendable (PostgresConnection, Logger) async throws -> Void)? = nil)
     {
         self.applicationName = applicationName
         self.postgresConfiguration = postgresConfiguration
@@ -93,9 +99,11 @@ public struct PoolConfiguration {
         self.poolSize = poolSize.atLeast(1)
         self.maxIdleConnections = maxIdleConnections?.atLeast(0)
 
-        self.onReturnConnection = { connection, logger in
+        self.onOpenConnection = onOpenConnection
+        self.onReturnConnection = onReturnConnection ?? { connection, logger in
             try await connection.query("SELECT 1", logger: logger)
         }
+        self.onCloseConnection = onCloseConnection
     }
 
     @available(*, deprecated, message: "Use `init(applicationName:postgresConfiguration:connectTimeout:queryTimeout:poolSize:maxIdleConnections:)` instead.")

Sources/PostgresConnectionPool/PoolConfiguration.swift

@@ -9,9 +9,9 @@ import PostgresNIO
 final class PoolConnection: Identifiable, Equatable {
 
     // For connection ids
-    private static let atomic: ManagedAtomic<Int> = .init(0)
+    private static let connectionId: ManagedAtomic<Int> = .init(0)
 
-    private(set) var usageCounter = 0
+    let usageCounter: ManagedAtomic<Int> = .init(0)
 
     let id: Int
     var batchId: Int?
@@ -20,7 +20,7 @@ final class PoolConnection: Identifiable, Equatable {
     var state: PoolConnectionState = .connecting {
         didSet {
             if case .active = state {
-                usageCounter += 1
+                usageCounter.wrappingIncrement(by: 1, ordering: .relaxed)
             }
         }
     }
@@ -37,7 +37,7 @@ final class PoolConnection: Identifiable, Equatable {
     }
 
     init() {
-        self.id = PoolConnection.atomic.loadThenWrappingIncrement(by: 1, ordering: .relaxed)
+        self.id = PoolConnection.connectionId.loadThenWrappingIncrement(by: 1, ordering: .relaxed)
     }
 
     static func == (lhs: PoolConnection, rhs: PoolConnection) -> Bool {

Sources/PostgresConnectionPool/PoolConnection.swift

@@ -5,7 +5,7 @@
 import Foundation
 
 /// The possible states of a database connection.
-public enum PoolConnectionState: Equatable {
+public enum PoolConnectionState: Equatable, Sendable {
 
     /// The connection is in use.
     case active(Date)

Sources/PostgresConnectionPool/PoolConnectionState.swift

@@ -6,7 +6,7 @@ import Foundation
 import PostgresNIO
 
 /// Possible errors from the connection pool.
-public enum PoolError: Error {
+public enum PoolError: Error, Sendable {
 
     /// The request was cancelled.
     case cancelled

Sources/PostgresConnectionPool/PoolError.swift

@@ -6,10 +6,10 @@ import Foundation
 import PostgresNIO
 
 /// General information about the pool and its open connections.
-public struct PoolInfo {
+public struct PoolInfo: Sendable {
 
     /// Information about an open connection.
-    public struct ConnectionInfo {
+    public struct ConnectionInfo: Sendable {
         /// The unique connection id.
         public let id: Int
         /// The connection name on the server.

Sources/PostgresConnectionPool/PoolInfo.swift

@@ -263,7 +263,7 @@ public actor PostgresConnectionPool {
             return PoolInfo.ConnectionInfo(
                 id: connection.id,
                 name: nameForConnection(id: connection.id),
-                usageCounter: connection.usageCounter,
+                usageCounter: connection.usageCounter.load(ordering: .relaxed),
                 batchId: connection.batchId,
                 query: connection.query,
                 queryRuntime: connection.queryRuntime,
@@ -323,7 +323,7 @@ public actor PostgresConnectionPool {
                 || (connection.state != .connecting && connection.connection?.isClosed ?? false)
         })
 
-        let usageCounter = connections.reduce(0) { $0 + $1.usageCounter }
+        let usageCounter = connections.reduce(0) { $0 + $1.usageCounter.load(ordering: .relaxed) }
         logger.debug("[\(poolName)] \(connections.count) connections (\(available.count) available, \(usageCounter) queries), \(continuations.count) continuations left")
 
         // Check for waiting continuations and open a new connection if possible

Sources/PostgresConnectionPool/PostgresConnectionPool.swift

@@ -89,12 +89,13 @@ public final class PostgresConnectionWrapper {
     }
 
     public func sql(
-        encoder: PostgresDataEncoder = PostgresDataEncoder(),
-        decoder: PostgresDataDecoder = PostgresDataDecoder())
-        -> SQLDatabase
+        encodingContext: PostgresEncodingContext<some PostgresJSONEncoder> = .default,
+        decodingContext: PostgresDecodingContext<some PostgresJSONDecoder> = .default,
+        queryLogLevel: Logger.Level? = .debug)
+        -> some SQLDatabase
     {
         // TODO: Track the current query
-        postgresConnection.sql(encoder: encoder, decoder: decoder)
+        postgresConnection.sql(encodingContext: encodingContext, decodingContext: decodingContext, queryLogLevel: queryLogLevel)
     }
 
     /// Add a handler for NotificationResponse messages on a certain channel.
@@ -105,11 +106,17 @@ public final class PostgresConnectionWrapper {
     @discardableResult
     public func addListener(
         channel: String,
-        handler notificationHandler: @escaping (PostgresListenContext, PostgresMessage.NotificationResponse) -> Void)
+        handler notificationHandler: @Sendable @escaping (PostgresListenContext, PostgresMessage.NotificationResponse) -> Void)
         -> PostgresListenContext
     {
         poolConnection.query = "LISTEN \(channel)"
         return postgresConnection.addListener(channel: channel, handler: notificationHandler)
     }
 
+    /// Start listening for a channel
+    public func listen(_ channel: String) async throws -> PostgresNotificationSequence {
+        poolConnection.query = "LISTEN \(channel)"
+        return try await postgresConnection.listen(channel)
+    }
+
 }