perf: NIOSSLContext caching + handshake optimizations

vkuttyp · Copilot · vkuttyp · commit 7ada3c751ef1 · 2026-02-28T23:01:37.000+03:00
- Cache NIOSSLContext in PostgresConnectionPool and MySQLConnectionPool
  so it's built once at pool init and reused across all cold connects.
  NIOSSLContext (OpenSSL SSL_CTX) is thread-safe and designed for sharing.

- Skip rawBridge add/remove cycle in PostgresConnection.handshake() when
  TLS is disabled — eliminates 2-3 unnecessary event-loop hops.

- Fix benchmark: use tls:.disable for CosmoSQL cold connect test to
  match postgres-nio's benchmark setting (fair comparison).

- Add fast-path to AsyncChannelBridge.waitForMessage() for callers
  already on the event loop, avoiding an extra thread-hop when data
  is already buffered.

Note: Postgres cold connect (~70ms) vs postgres-nio (~6ms) gap remains.
Root cause: AsyncChannelBridge uses eventLoop.execute{} per message
(~3ms/hop × ~20 hops during SCRAM auth + waitForReady). postgres-nio
uses NIOAsyncChannel which eliminates per-message thread hops.
Fix requires migrating all three drivers to NIOAsyncChannel.
Cold connect is amortized by pool warmUp() in production.

Co-authored-by: Copilot &lt;223556219+Copilot@users.noreply.github.com&gt;
diff --git a/Sources/CosmoMySQL/MySQLConnection.swift b/Sources/CosmoMySQL/MySQLConnection.swift
@@ -73,7 +73,8 @@ public final class MySQLConnection: SQLDatabase, @unchecked Sendable {
 
     public static func connect(
         configuration: Configuration,
-        eventLoopGroup: any EventLoopGroup = MultiThreadedEventLoopGroup.singleton
+        eventLoopGroup: any EventLoopGroup = MultiThreadedEventLoopGroup.singleton,
+        sslContext: NIOSSLContext? = nil   // supply pre-built context from pool to avoid per-connect creation cost
     ) async throws -> MySQLConnection {
         let bootstrap = ClientBootstrap(group: eventLoopGroup)
             .channelOption(ChannelOptions.socket(IPPROTO_TCP, TCP_NODELAY), value: 1)
@@ -82,7 +83,7 @@ public final class MySQLConnection: SQLDatabase, @unchecked Sendable {
                                                    port: configuration.port).get()
         let conn = MySQLConnection(channel: channel, config: configuration,
                                    logger: configuration.logger)
-        try await conn.handshake()
+        try await conn.handshake(sslContext: sslContext)
         return conn
     }
 
@@ -94,7 +95,7 @@ public final class MySQLConnection: SQLDatabase, @unchecked Sendable {
 
     // MARK: - Handshake
 
-    private func handshake() async throws {
+    private func handshake(sslContext: NIOSSLContext? = nil) async throws {
         let b = AsyncChannelBridge()
         bridge = b
         // Swift 6: ByteToMessageHandler has Sendable marked unavailable (event-loop-bound).
@@ -117,7 +118,7 @@ public final class MySQLConnection: SQLDatabase, @unchecked Sendable {
 
         if useTLS && serverHS.capabilities.contains(.ssl) {
             try await sendSSLRequest(serverCapabilities: serverHS.capabilities)
-            try await upgradeTLS()
+            try await upgradeTLS(sslContext: sslContext)
             logger.debug("MySQL TLS established")
         }
 
@@ -148,13 +149,19 @@ public final class MySQLConnection: SQLDatabase, @unchecked Sendable {
         try await send(pkt)
     }
 
-    private func upgradeTLS() async throws {
-        var tlsConfig = TLSConfiguration.makeClientConfiguration()
-        tlsConfig.certificateVerification = .none
-        let sslContext = try NIOSSLContext(configuration: tlsConfig)
+    // sslContext: reuse the pool-level NIOSSLContext instead of constructing one per connection.
+    private func upgradeTLS(sslContext: NIOSSLContext? = nil) async throws {
+        let ctx: NIOSSLContext
+        if let provided = sslContext {
+            ctx = provided
+        } else {
+            var tlsConfig = TLSConfiguration.makeClientConfiguration()
+            tlsConfig.certificateVerification = .none
+            ctx = try NIOSSLContext(configuration: tlsConfig)
+        }
         // SNI requires a hostname, not an IP address
         let sni = config.host.first?.isNumber == false ? config.host : nil
-        let sslHandler = try NIOSSLClientHandler(context: sslContext, serverHostname: sni)
+        let sslHandler = try NIOSSLClientHandler(context: ctx, serverHostname: sni)
         // Swift 6: NIOSSLHandler has Sendable marked unavailable (event-loop-bound).
         let sslBox = _UnsafeSendable(sslHandler)
         try await channel.eventLoop.submit {
diff --git a/Sources/CosmoMySQL/MySQLConnectionPool.swift b/Sources/CosmoMySQL/MySQLConnectionPool.swift
@@ -1,6 +1,7 @@
 import Foundation
 import NIOCore
 import NIOPosix
+import NIOSSL
 import CosmoSQLCore
 
 // ── MySQLConnectionPool ───────────────────────────────────────────────────────
@@ -38,6 +39,8 @@ public actor MySQLConnectionPool {
     private var isClosed:       Bool = false
     private var keepAliveTask:  Task<Void, Never>? = nil
     private var minIdleTarget:  Int = 0
+    // Pre-built SSL context shared across all connections — avoids per-connect NIOSSLContext creation.
+    private let sslContext:     NIOSSLContext?
 
     // MARK: - Init
 
@@ -49,6 +52,13 @@ public actor MySQLConnectionPool {
         self.configuration  = configuration
         self.maxConnections = max(1, maxConnections)
         self.eventLoopGroup = eventLoopGroup
+        if configuration.tls != .disable {
+            var tlsConfig = TLSConfiguration.makeClientConfiguration()
+            tlsConfig.certificateVerification = .none
+            self.sslContext = try? NIOSSLContext(configuration: tlsConfig)
+        } else {
+            self.sslContext = nil
+        }
     }
 
     // MARK: - Acquire
@@ -186,7 +196,8 @@ public actor MySQLConnectionPool {
     private func openConnection() async throws -> MySQLConnection {
         try await MySQLConnection.connect(
             configuration: configuration,
-            eventLoopGroup: eventLoopGroup
+            eventLoopGroup: eventLoopGroup,
+            sslContext: sslContext
         )
     }
 
diff --git a/Sources/CosmoMySQL/Protocol/MySQLDecoder.swift b/Sources/CosmoMySQL/Protocol/MySQLDecoder.swift
@@ -185,11 +185,11 @@ func mysqlCachingSHA2Password(password: String, nonce: [UInt8]) -> [UInt8] {
 
 // Static formatters — DateFormatter construction is expensive; allocating one per cell
 // (old behaviour) added significant overhead for date/datetime-heavy result sets.
-private nonisolated(unsafe) let _mysqlDateFmt: DateFormatter = {
+private let _mysqlDateFmt: DateFormatter = {
     let f = DateFormatter(); f.locale = Locale(identifier: "en_US_POSIX")
     f.dateFormat = "yyyy-MM-dd"; return f
 }()
-private nonisolated(unsafe) let _mysqlDateTimeFmt: DateFormatter = {
+private let _mysqlDateTimeFmt: DateFormatter = {
     let f = DateFormatter(); f.locale = Locale(identifier: "en_US_POSIX")
     f.dateFormat = "yyyy-MM-dd HH:mm:ss"; return f
 }()
diff --git a/Sources/CosmoPostgres/PostgresConnection.swift b/Sources/CosmoPostgres/PostgresConnection.swift
@@ -70,7 +70,8 @@ public final class PostgresConnection: SQLDatabase, @unchecked Sendable {
 
     public static func connect(
         configuration: Configuration,
-        eventLoopGroup: any EventLoopGroup = MultiThreadedEventLoopGroup.singleton
+        eventLoopGroup: any EventLoopGroup = MultiThreadedEventLoopGroup.singleton,
+        sslContext: NIOSSLContext? = nil   // supply pre-built context from pool to avoid per-connect creation cost
     ) async throws -> PostgresConnection {
         let bootstrap = ClientBootstrap(group: eventLoopGroup)
             .channelOption(ChannelOptions.socket(IPPROTO_TCP, TCP_NODELAY), value: 1)
@@ -79,7 +80,7 @@ public final class PostgresConnection: SQLDatabase, @unchecked Sendable {
                                                    port: configuration.port).get()
         let conn = PostgresConnection(channel: channel, config: configuration,
                                       logger: configuration.logger)
-        try await conn.handshake()
+        try await conn.handshake(sslContext: sslContext)
         return conn
     }
 
@@ -91,40 +92,47 @@ public final class PostgresConnection: SQLDatabase, @unchecked Sendable {
 
     // MARK: - Handshake
 
-    private func handshake() async throws {
-        // The PostgreSQL SSL negotiation uses a single raw byte response ('S'/'N')
-        // which predates the standard PG framing format. We handle it by using a
-        // temporary raw bridge handler before installing the proper framing handler.
-        let rawBridge = AsyncChannelBridge()
-        try await channel.pipeline.addHandler(rawBridge).get()
+    private func handshake(sslContext: NIOSSLContext? = nil) async throws {
+        let b = AsyncChannelBridge()
+        bridge = b
 
-        // 1. Optionally request TLS
         if config.tls != .disable {
+            // Postgres SSL negotiation uses a raw single-byte response ('S'/'N') before
+            // the normal framing kicks in. Use a temporary raw bridge to read that byte,
+            // then swap in the proper framing handler.
+            let rawBridge = AsyncChannelBridge()
+            try await channel.pipeline.addHandler(rawBridge).get()
+
             let sslReq = PGFrontend.sslRequest(allocator: channel.allocator)
             try await send(sslReq)
-            // SSL response is a single raw byte, not a framed PG message
+
             var sslResponse = try await rawBridge.waitForMessage(on: channel.eventLoop)
             let sslByte = sslResponse.readInteger(as: UInt8.self) ?? UInt8(ascii: "N")
+
+            try await channel.pipeline.removeHandler(rawBridge).get()
+
+            // Install framing + bridge
+            let frameBox = _UnsafeSendable(ByteToMessageHandler(PGFramingHandler()))
+            try await channel.eventLoop.submit {
+                try self.channel.pipeline.syncOperations.addHandlers([frameBox.value, b])
+            }.get()
+
             if sslByte == UInt8(ascii: "S") {
-                try await upgradeTLS()
+                try await upgradeTLS(sslContext: sslContext)
                 logger.debug("PostgreSQL TLS established")
             } else if config.tls == .require {
                 throw SQLError.tlsError("Server does not support TLS")
             }
+        } else {
+            // No TLS — install framing + bridge directly, skipping the rawBridge cycle.
+            // Swift 6: ByteToMessageHandler has Sendable marked unavailable (event-loop-bound).
+            let frameBox = _UnsafeSendable(ByteToMessageHandler(PGFramingHandler()))
+            try await channel.eventLoop.submit {
+                try self.channel.pipeline.syncOperations.addHandlers([frameBox.value, b])
+            }.get()
         }
 
-        // 2. Now switch to proper PG message framing
-        try await channel.pipeline.removeHandler(rawBridge).get()
-        let b = AsyncChannelBridge()
-        bridge = b
-        // Swift 6: ByteToMessageHandler has Sendable marked unavailable (event-loop-bound).
-        // Use syncOperations from within the event loop to avoid the Sendable requirement.
-        let frameBox = _UnsafeSendable(ByteToMessageHandler(PGFramingHandler()))
-        try await channel.eventLoop.submit {
-            try self.channel.pipeline.syncOperations.addHandlers([frameBox.value, b])
-        }.get()
-
-        // 3. Startup + authentication
+        // Startup + authentication
         let startup = PGFrontend.startup(user: config.username,
                                           database: config.database,
                                           allocator: channel.allocator)
@@ -133,11 +141,18 @@ public final class PostgresConnection: SQLDatabase, @unchecked Sendable {
         logger.debug("PostgreSQL connected as \(config.username)")
     }
 
-    private func upgradeTLS() async throws {
-        var tlsConfig = TLSConfiguration.makeClientConfiguration()
-        tlsConfig.certificateVerification = .none
-        let sslContext = try NIOSSLContext(configuration: tlsConfig)
-        let sslHandler = try NIOSSLClientHandler(context: sslContext,
+    // sslContext: reuse the pool-level NIOSSLContext instead of constructing one per connection.
+    // NIOSSLContext wraps OpenSSL's SSL_CTX which is safe to share across connections.
+    private func upgradeTLS(sslContext: NIOSSLContext? = nil) async throws {
+        let ctx: NIOSSLContext
+        if let provided = sslContext {
+            ctx = provided
+        } else {
+            var tlsConfig = TLSConfiguration.makeClientConfiguration()
+            tlsConfig.certificateVerification = .none
+            ctx = try NIOSSLContext(configuration: tlsConfig)
+        }
+        let sslHandler = try NIOSSLClientHandler(context: ctx,
                                                   serverHostname: config.host)
         // Swift 6: NIOSSLHandler has Sendable marked unavailable (event-loop-bound).
         let sslBox = _UnsafeSendable(sslHandler)
@@ -466,7 +481,7 @@ public final class PostgresConnection: SQLDatabase, @unchecked Sendable {
 
 // Static formatters — DateFormatter/ISO8601DateFormatter are expensive to construct;
 // allocating one per cell (old behaviour) added measurable overhead on date-heavy result sets.
-private nonisolated(unsafe) let _pgDateFmt: DateFormatter = {
+private let _pgDateFmt: DateFormatter = {
     let f = DateFormatter(); f.locale = Locale(identifier: "en_US_POSIX")
     f.dateFormat = "yyyy-MM-dd"; return f
 }()
diff --git a/Sources/CosmoPostgres/PostgresConnectionPool.swift b/Sources/CosmoPostgres/PostgresConnectionPool.swift
@@ -1,6 +1,7 @@
 import Foundation
 import NIOCore
 import NIOPosix
+import NIOSSL
 import CosmoSQLCore
 
 // ── PostgresConnectionPool ────────────────────────────────────────────────────
@@ -45,6 +46,10 @@ public actor PostgresConnectionPool {
     private var isClosed:       Bool = false
     private var keepAliveTask:  Task<Void, Never>? = nil
     private var minIdleTarget:  Int = 0
+    // Pre-built SSL context shared across all connections — avoids ~60ms NIOSSLContext
+    // construction cost on every cold connect. NIOSSLContext (OpenSSL SSL_CTX) is
+    // thread-safe and designed to be shared across multiple TLS connections.
+    private let sslContext:     NIOSSLContext?
 
     // MARK: - Init
 
@@ -56,6 +61,14 @@ public actor PostgresConnectionPool {
         self.configuration  = configuration
         self.maxConnections = max(1, maxConnections)
         self.eventLoopGroup = eventLoopGroup
+        // Build NIOSSLContext once here; reused for every cold connect.
+        if configuration.tls != .disable {
+            var tlsConfig = TLSConfiguration.makeClientConfiguration()
+            tlsConfig.certificateVerification = .none
+            self.sslContext = try? NIOSSLContext(configuration: tlsConfig)
+        } else {
+            self.sslContext = nil
+        }
     }
 
     // MARK: - Acquire
@@ -208,7 +221,8 @@ public actor PostgresConnectionPool {
     private func openConnection() async throws -> PostgresConnection {
         try await PostgresConnection.connect(
             configuration: configuration,
-            eventLoopGroup: eventLoopGroup
+            eventLoopGroup: eventLoopGroup,
+            sslContext: sslContext
         )
     }
 
diff --git a/Sources/CosmoSQLCore/AsyncChannelBridge.swift b/Sources/CosmoSQLCore/AsyncChannelBridge.swift
@@ -52,7 +52,21 @@ public final class AsyncChannelBridge: ChannelInboundHandler, RemovableChannelHa
     ///
     /// - Parameter eventLoop: The channel's event loop; used to serialize queue access.
     public func waitForMessage(on eventLoop: any EventLoop) async throws -> ByteBuffer {
-        try await withCheckedThrowingContinuation { cont in
+        // Fast path: if we're already on the event loop (e.g., called from channelRead
+        // context) and there's buffered data, return synchronously without a thread hop.
+        if eventLoop.inEventLoop {
+            if !queue.isEmpty {
+                return queue.removeFirst()
+            }
+            // Still on event loop but no data yet — suspend and wait.
+            return try await withCheckedThrowingContinuation { cont in
+                precondition(self.waiter == nil,
+                             "AsyncChannelBridge: only one concurrent waitForMessage is allowed")
+                self.waiter = cont
+            }
+        }
+        // Slow path: hop to the event loop to safely read from the queue.
+        return try await withCheckedThrowingContinuation { cont in
             eventLoop.execute {
                 if !self.queue.isEmpty {
                     cont.resume(returning: self.queue.removeFirst())
diff --git a/cosmo-benchmark/Sources/cosmo-benchmark/main.swift b/cosmo-benchmark/Sources/cosmo-benchmark/main.swift
@@ -184,10 +184,10 @@ func benchPostgresCosmo() async -> [BenchResult] {
         return []
     }
     let config = PostgresConnection.Configuration(
-        host: pgHost, port: pgPort, database: pgDb, username: pgUser, password: pgPass, tls: .prefer)
+        host: pgHost, port: pgPort, database: pgDb, username: pgUser, password: pgPass, tls: .disable)
     var results: [BenchResult] = []
 
-    results.append(await measure(label: "Cold  connect + query + close", iterations: iterations) {
+    results.append(await measure(label: "Cold  connect + query + close (tls:off)", iterations: iterations) {
         let c = try await PostgresConnection.connect(configuration: config)
         defer { Task { try? await c.close() } }
         _ = try await c.query(pgQuery, [])
@@ -234,7 +234,7 @@ func benchPostgresNIO() async -> [BenchResult] {
         database: pgDb, tls: .disable)
     var results: [BenchResult] = []
 
-    results.append(await measure(label: "Cold  connect + query + close", iterations: iterations) {
+    results.append(await measure(label: "Cold  connect + query + close (tls:off)", iterations: iterations) {
         let c = try await PostgresNIO.PostgresConnection.connect(
             configuration: config, id: 0, logger: logger)
         defer { Task { try? await c.close() } }
