Refactor ConnectionPoolImpl for enhanced lifecycle and pooling operations. Introduce PooledConnection for better connection management and thread safety. Expand test coverage with advanced pooling scenarios.

Author: smyrgeorge

sqlx4k/src/commonMain/kotlin/io/github/smyrgeorge/sqlx4k/impl/pool/ConnectionPoolImpl.kt

@@ -7,24 +7,20 @@ import io.github.smyrgeorge.sqlx4k.ConnectionPool
 import io.github.smyrgeorge.sqlx4k.SQLError
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.Channel
-import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.Semaphore
-import kotlinx.coroutines.sync.withLock
 import kotlin.concurrent.atomics.*
 import kotlin.time.Duration.Companion.seconds
-import kotlin.time.TimeMark
-import kotlin.time.TimeSource
 
 class ConnectionPoolImpl(
-    private val options: ConnectionPool.Options,
+    val options: ConnectionPool.Options,
     private val connectionFactory: suspend () -> Connection,
 ) : ConnectionPool {
-    private var closed = AtomicBoolean(false)
-    private val idleCount = AtomicInt(0)
-    private val totalConnections = AtomicInt(0)
+    internal var closed = AtomicBoolean(false)
+    internal val idleCount = AtomicInt(0)
+    internal val totalConnections = AtomicInt(0)
+    internal val semaphore = Semaphore(options.maxConnections)
 
     private val cleanupJob: Job
-    private val semaphore = Semaphore(options.maxConnections)
     private val scope = CoroutineScope(Dispatchers.Default + SupervisorJob())
     private val idleConnections = Channel<PooledConnection>(options.maxConnections)
 
@@ -45,8 +41,8 @@ class ConnectionPoolImpl(
             when (error) {
                 is TimeoutCancellationException -> {
                     SQLError(
-                        SQLError.Code.PoolTimedOut,
-                        "Timed out waiting for connection after ${options.acquireTimeout}"
+                        code = SQLError.Code.PoolTimedOut,
+                        message = "Timed out waiting for connection after ${options.acquireTimeout}"
                     ).ex()
                 }
 
@@ -81,12 +77,6 @@ class ConnectionPoolImpl(
                     val newConnection = connectionFactory()
                     val pooled = PooledConnection(newConnection, this)
                     totalConnections.incrementAndFetch()
-
-                    // Double-check pool wasn't closed during connection creation
-                    if (closed.load()) {
-                        pooled.close()
-                        SQLError(SQLError.Code.PoolClosed, "Connection pool is closed").ex()
-                    }
                     return pooled.acquire()
                 } catch (e: Exception) {
                     // Release semaphore if connection creation failed
@@ -122,30 +112,23 @@ class ConnectionPoolImpl(
     }
 
     override suspend fun close(): Result<Unit> = runCatching {
-        if (closed.load()) return@runCatching
-        closed.store(true)
+        if (!closed.compareAndSet(expectedValue = false, newValue = true)) {
+            return@runCatching
+        }
 
-        // Cancel cleanup job
         cleanupJob.cancel()
-
-        // Close the idle channel to wake any waiters and prevent further sends
         idleConnections.close()
-
-        // Close all idle connections
         while (true) {
             val pooled = idleConnections.tryReceive().getOrNull() ?: break
-            // Adjust idle counter for each drained connection
             idleCount.decrementAndFetch()
             pooled.close()
         }
         // Ensure idleCount cannot go negative due to races; clamp to 0
         idleCount.update { if (it < 0) 0 else it }
-
-        // Cancel scope
         scope.cancel()
     }
 
-    private suspend fun sendToIdle(connection: PooledConnection): Boolean {
+    internal suspend fun sendToIdle(connection: PooledConnection): Boolean {
         return try {
             idleConnections.send(connection)
             idleCount.incrementAndFetch()
@@ -162,34 +145,43 @@ class ConnectionPoolImpl(
     }
 
     private suspend fun warmup(minConnections: Int) {
+        if (closed.load()) return
+
         repeat(minConnections) {
-            var semaphoreAcquired = false
-            var connectionCreated = false
-            var pooledConnection: PooledConnection? = null
+            if (closed.load()) return
+
+            var acquired = false
+            var pooled: PooledConnection? = null
 
             try {
                 semaphore.acquire()
-                semaphoreAcquired = true
-                pooledConnection = PooledConnection(connectionFactory(), this)
-                connectionCreated = true
+                acquired = true
 
+                // Check again after potentially blocking on semaphore
+                if (closed.load()) {
+                    semaphore.release()
+                    return
+                }
+
+                pooled = PooledConnection(connectionFactory(), this)
                 totalConnections.incrementAndFetch()
 
-                // Use atomic send operation - if it fails, rollback
-                if (!sendToIdle(pooledConnection)) {
+                if (!sendToIdle(pooled)) {
                     // Send failed (pool closing), close the connection
-                    pooledConnection.close()
+                    pooled.close()
+                    // No point continuing warmup if pool is closing
+                    return
                 }
             } catch (_: Exception) {
                 // Clean up based on what succeeded
-                if (connectionCreated && pooledConnection != null) {
+                if (pooled != null) {
                     // Connection was created but something failed, close it
                     try {
-                        pooledConnection.close()
+                        pooled.close()
                     } catch (_: Exception) {
                         // Ignore errors during cleanup
                     }
-                } else if (semaphoreAcquired) {
+                } else if (acquired) {
                     // Only semaphore was acquired, release it
                     semaphore.release()
                 }
@@ -199,27 +191,51 @@ class ConnectionPoolImpl(
     }
 
     private suspend fun cleanup() {
-        val temp = mutableListOf<PooledConnection>()
+        // Early exit if pool is closing/closed
+        if (closed.load()) return
+
         val minConnections = options.minConnections ?: 0
+        var processedCount = 0
 
-        // Poll all idle connections using atomic wrapper
-        while (true) {
+        // Process connections incrementally without draining all at once
+        // This allows other coroutines to acquire connections during cleanup
+        val maxBatchSize: Int = options.maxConnections / 2 // Limit how many we process in one cleanup cycle
+
+        while (processedCount < maxBatchSize) {
             val pooled = tryReceiveFromIdle() ?: break
-            temp.add(pooled)
-        }
+            processedCount++
 
-        // Filter and return valid connections, close expired ones
-        for (pooled in temp) {
-            if (pooled.isExpired()) {
-                // Try to close if above minimum (atomically checked)
-                val wasClosed = pooled.closeIfAboveMinimum(minConnections)
-                if (!wasClosed) {
-                    // Couldn't close because we're at or below minimum, keep the connection
-                    sendToIdle(pooled)
+            try {
+                if (pooled.isExpired()) {
+                    // Try to close if above minimum (atomically checked)
+                    val wasClosed = pooled.closeIfAboveMinimum(minConnections)
+                    if (!wasClosed) {
+                        // At or below minimum - keep the expired connection
+                        // It will be used until a new connection can replace it
+                        if (!sendToIdle(pooled)) {
+                            // Pool is closing, cleanup the connection
+                            pooled.close()
+                        }
+                    }
+                } else {
+                    // Connection still valid, return it
+                    if (!sendToIdle(pooled)) {
+                        // Pool is closing, cleanup the connection
+                        pooled.close()
+                    }
                 }
-            } else {
-                // Connection still valid, keep it
-                sendToIdle(pooled)
+            } catch (_: Exception) {
+                // If anything fails, try to close the connection to prevent leaks
+                try {
+                    pooled.close()
+                } catch (_: Exception) {
+                    // Ignore errors during emergency cleanup
+                }
+            }
+
+            // Yield periodically to prevent CPU spinning
+            if (processedCount % 10 == 0) {
+                yield()
             }
         }
     }
@@ -230,122 +246,14 @@ class ConnectionPoolImpl(
                 delay(CLEANUP_INTERVAL)
                 cleanup()
             } catch (_: CancellationException) {
-                // Job was cancelled, exit loop
                 break
             } catch (_: Exception) {
                 // TODO: Log error but continue cleanup loop
             }
         }
     }
 
-    private class PooledConnection(
-        private val connection: Connection,
-        private val pool: ConnectionPoolImpl
-    ) : Connection by connection {
-        private val mutex = Mutex()
-        private val createdAt: TimeMark = TIME_SOURCE.markNow()
-        private var lastUsedAt: TimeMark = createdAt
-        override var status: Connection.Status = connection.status // TODO: Check if this is correct
-        private val released get() = status == Connection.Status.Released
-
-        suspend fun acquire(): Connection {
-            mutex.withLock {
-                if (pool.closed.load()) {
-                    close()
-                    SQLError(SQLError.Code.PoolClosed, "Connection pool is closed").ex()
-                }
-
-                status = Connection.Status.Acquired
-                lastUsedAt = TIME_SOURCE.markNow()
-            }
-            return this
-        }
-
-        override suspend fun release(): Result<Unit> = runCatching {
-            mutex.withLock {
-                if (released) return@runCatching
-                status = Connection.Status.Released
-            }
-
-            // Read closed flag under lock, but perform actions outside to avoid deadlocks
-            if (pool.closed.load()) {
-                // Pool is closed, close the connection outside of the mutex
-                close()
-                return@runCatching
-            }
-
-            lastUsedAt = TIME_SOURCE.markNow()
-
-            // Check if connection is expired
-            if (isExpired()) {
-                // Connection is expired, try to close it only if we're above minimum
-                val minConnections = pool.options.minConnections ?: 0
-                val wasClosed = closeIfAboveMinimum(minConnections)
-                if (!wasClosed) {
-                    // At or below minimum, keep the expired connection (cleanup will replace later)
-                    // Enqueue back to idle using suspending send to properly wake waiters
-                    if (!pool.sendToIdle(this)) {
-                        // Pool closing/closed; ensure the connection is closed
-                        close()
-                    }
-                }
-            } else {
-                // Connection is still valid, return to idle pool using suspending send
-                if (!pool.sendToIdle(this)) {
-                    // Pool closing/closed; ensure the connection is closed
-                    close()
-                }
-            }
-        }
-
-        fun isExpired(): Boolean {
-            pool.options.maxLifetime?.let { maxLifetime -> if (createdAt.elapsedNow() >= maxLifetime) return true }
-            pool.options.idleTimeout?.let { idleTimeout -> if (lastUsedAt.elapsedNow() >= idleTimeout) return true }
-            return false
-        }
-
-        suspend fun close() {
-            try {
-                connection.release().getOrThrow()
-            } catch (_: Exception) {
-                // Ignore errors on close
-            } finally {
-                pool.totalConnections.decrementAndFetch()
-                pool.semaphore.release()
-            }
-        }
-
-        suspend fun closeIfAboveMinimum(minConnections: Int): Boolean {
-            // Atomically check and decide whether to close
-            var shouldClose = false
-            pool.totalConnections.update {
-                if (it > minConnections) {
-                    // We're above minimum, mark for closing by decrementing now
-                    shouldClose = true
-                    it - 1
-                } else {
-                    shouldClose = false
-                    it
-                }
-            }
-
-            if (shouldClose) {
-                // Close the actual connection outside the lock
-                try {
-                    connection.release().getOrThrow()
-                } catch (_: Exception) {
-                    // Ignore errors on close
-                } finally {
-                    pool.semaphore.release()
-                }
-            }
-
-            return shouldClose
-        }
-    }
-
     companion object {
         private val CLEANUP_INTERVAL = 2.seconds
-        private val TIME_SOURCE = TimeSource.Monotonic
     }
 }