Add Benchmarks for ConnectionPool (#545)

Author: fabianfett

Diff for: Benchmarks/.gitignore

@@ -0,0 +1 @@
+.build

Diff for: Benchmarks/Benchmarks/ConnectionPoolBenchmarks/ConnectionPoolBenchmarks.swift

@@ -0,0 +1,51 @@
+import _ConnectionPoolModule
+import _ConnectionPoolTestUtils
+import Benchmark
+
+let benchmarks: @Sendable () -> Void = {
+    Benchmark("Minimal benchmark", configuration: .init(scalingFactor: .kilo)) { benchmark in
+        let clock = MockClock()
+        let factory = MockConnectionFactory<MockClock>(autoMaxStreams: 1)
+        var configuration = ConnectionPoolConfiguration()
+        configuration.maximumConnectionSoftLimit = 50
+        configuration.maximumConnectionHardLimit = 50
+
+        let pool = ConnectionPool(
+            configuration: configuration,
+            idGenerator: ConnectionIDGenerator(),
+            keepAliveBehavior: MockPingPongBehavior(keepAliveFrequency: nil, connectionType: MockConnection.self),
+            observabilityDelegate: NoOpConnectionPoolMetrics(connectionIDType: MockConnection.ID.self),
+            clock: clock
+        ) {
+            try await factory.makeConnection(id: $0, for: $1)
+        }
+
+        await withTaskGroup { taskGroup in
+
+            taskGroup.addTask {
+                await pool.run()
+            }
+
+            let sequential = benchmark.scaledIterations.upperBound / configuration.maximumConnectionSoftLimit
+
+            for parallel in 0..<configuration.maximumConnectionSoftLimit {
+                taskGroup.addTask {
+                    for _ in 0..<sequential {
+                        do {
+                            try await pool.withConnection { connection in
+                                blackHole(connection)
+                            }
+                        } catch {
+                            fatalError()
+                        }
+                    }
+                }
+            }
+
+            for i in 0..<configuration.maximumConnectionSoftLimit {
+                await taskGroup.next()
+            }
+            taskGroup.cancelAll()
+        }
+    }
+}

Diff for: Benchmarks/Package.swift

@@ -0,0 +1,28 @@
+// swift-tools-version: 6.0
+
+import PackageDescription
+
+let package = Package(
+    name: "benchmarks",
+    platforms: [
+        .macOS("14")
+    ],
+    dependencies: [
+        .package(path: "../"),
+        .package(url: "https://github.com/ordo-one/package-benchmark.git", from: "1.29.0"),
+    ],
+    targets: [
+        .executableTarget(
+            name: "ConnectionPoolBenchmarks",
+            dependencies: [
+                .product(name: "_ConnectionPoolModule", package: "postgres-nio"),
+                .product(name: "_ConnectionPoolTestUtils", package: "postgres-nio"),
+                .product(name: "Benchmark", package: "package-benchmark"),
+            ],
+            path: "Benchmarks/ConnectionPoolBenchmarks",
+            plugins: [
+                .plugin(name: "BenchmarkPlugin", package: "package-benchmark")
+            ]
+        ),
+    ]
+)

Diff for: Package.swift

@@ -16,6 +16,7 @@ let package = Package(
     products: [
         .library(name: "PostgresNIO", targets: ["PostgresNIO"]),
         .library(name: "_ConnectionPoolModule", targets: ["_ConnectionPoolModule"]),
+        .library(name: "_ConnectionPoolTestUtils", targets: ["_ConnectionPoolTestUtils"]),
     ],
     dependencies: [
         .package(url: "https://github.com/apple/swift-atomics.git", from: "1.2.0"),
@@ -58,11 +59,13 @@ let package = Package(
             swiftSettings: swiftSettings
         ),
         .target(
-            name: "ConnectionPoolTestUtils",
+            name: "_ConnectionPoolTestUtils",
             dependencies: [
                 "_ConnectionPoolModule",
                 .product(name: "NIOConcurrencyHelpers", package: "swift-nio"),
-            ]
+            ],
+            path: "Sources/ConnectionPoolTestUtils",
+            swiftSettings: swiftSettings
         ),
         .testTarget(
             name: "PostgresNIOTests",
@@ -77,7 +80,7 @@ let package = Package(
             name: "ConnectionPoolModuleTests",
             dependencies: [
                 .target(name: "_ConnectionPoolModule"),
-                .target(name: "ConnectionPoolTestUtils"),
+                .target(name: "_ConnectionPoolTestUtils"),
                 .product(name: "DequeModule", package: "swift-collections"),
                 .product(name: "NIOCore", package: "swift-nio"),
                 .product(name: "NIOConcurrencyHelpers", package: "swift-nio"),

Diff for: Sources/ConnectionPoolTestUtils/MockConnectionFactory.swift

@@ -21,7 +21,11 @@ public final class MockConnectionFactory<Clock: _Concurrency.Clock>: Sendable wh
         var runningConnections = [ConnectionID: Connection]()
     }
 
-    public init() {}
+    let autoMaxStreams: UInt16?
+
+    public init(autoMaxStreams: UInt16? = nil) {
+        self.autoMaxStreams = autoMaxStreams
+    }
 
     public var pendingConnectionAttemptsCount: Int {
         self.stateBox.withLockedValue { $0.attempts.count }
@@ -35,6 +39,15 @@ public final class MockConnectionFactory<Clock: _Concurrency.Clock>: Sendable wh
         id: Int,
         for pool: ConnectionPool<MockConnection, Int, ConnectionIDGenerator, some ConnectionRequestProtocol, Int, MockPingPongBehavior<MockConnection>, NoOpConnectionPoolMetrics<Int>, Clock>
     ) async throws -> ConnectionAndMetadata<MockConnection> {
+        if let autoMaxStreams = self.autoMaxStreams {
+            let connection = MockConnection(id: id)
+            Task {
+                try? await connection.signalToClose
+                connection.closeIfClosing()
+            }
+            return .init(connection: connection, maximalStreamsOnConnection: autoMaxStreams)
+        }
+
         // we currently don't support cancellation when creating a connection
         let result = try await withCheckedThrowingContinuation { (checkedContinuation: CheckedContinuation<(MockConnection, UInt16), any Error>) in
             let waiter = self.stateBox.withLockedValue { state -> (CheckedContinuation<(ConnectionID, CheckedContinuation<(MockConnection, UInt16), any Error>), Never>)? in

Diff for: Sources/ConnectionPoolTestUtils/MockPingPongBehaviour.swift

@@ -1,5 +1,6 @@
-@testable import _ConnectionPoolModule
+import _ConnectionPoolModule
 import DequeModule
+import NIOConcurrencyHelpers
 
 @available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)
 public final class MockPingPongBehavior<Connection: PooledConnection>: ConnectionKeepAliveBehavior {

Diff for: Tests/ConnectionPoolModuleTests/ConnectionPoolTests.swift

@@ -1,6 +1,6 @@
 @testable import _ConnectionPoolModule
+import _ConnectionPoolTestUtils
 import Atomics
-import ConnectionPoolTestUtils
 import NIOEmbedded
 import XCTest
 

Diff for: Tests/ConnectionPoolModuleTests/ConnectionRequestTests.swift

@@ -1,5 +1,5 @@
 @testable import _ConnectionPoolModule
-import ConnectionPoolTestUtils
+import _ConnectionPoolTestUtils
 import XCTest
 
 final class ConnectionRequestTests: XCTestCase {

Diff for: Tests/ConnectionPoolModuleTests/NoKeepAliveBehaviorTests.swift

@@ -1,5 +1,5 @@
 import _ConnectionPoolModule
-import ConnectionPoolTestUtils
+import _ConnectionPoolTestUtils
 import XCTest
 
 @available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)

Diff for: Tests/ConnectionPoolModuleTests/PoolStateMachine+ConnectionGroupTests.swift

@@ -1,5 +1,5 @@
 @testable import _ConnectionPoolModule
-import ConnectionPoolTestUtils
+import _ConnectionPoolTestUtils
 import XCTest
 
 @available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)

Diff for: Tests/ConnectionPoolModuleTests/PoolStateMachine+ConnectionStateTests.swift

@@ -1,5 +1,5 @@
 @testable import _ConnectionPoolModule
-import ConnectionPoolTestUtils
+import _ConnectionPoolTestUtils
 import XCTest
 
 @available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)

Diff for: Tests/ConnectionPoolModuleTests/PoolStateMachine+RequestQueueTests.swift

@@ -1,5 +1,5 @@
 @testable import _ConnectionPoolModule
-import ConnectionPoolTestUtils
+import _ConnectionPoolTestUtils
 import XCTest
 
 @available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)

Diff for: Tests/ConnectionPoolModuleTests/PoolStateMachineTests.swift

@@ -1,6 +1,6 @@
-import ConnectionPoolTestUtils
-import XCTest
 @testable import _ConnectionPoolModule
+import _ConnectionPoolTestUtils
+import XCTest
 
 @available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)
 typealias TestPoolStateMachine = PoolStateMachine<