update connection pool to use channel implementation

Author: stevensJourney

core/src/commonIntegrationTest/kotlin/com/powersync/DatabaseTest.kt

@@ -10,6 +10,8 @@ import com.powersync.db.PowerSyncDatabaseImpl
 import com.powersync.db.schema.Schema
 import com.powersync.testutils.UserRow
 import com.powersync.testutils.waitFor
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.async
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.runTest
 import kotlin.test.AfterTest

core/src/commonMain/kotlin/com/powersync/db/PowerSyncDatabaseImpl.kt

@@ -55,7 +55,6 @@ internal class PowerSyncDatabaseImpl(
     private val dbFilename: String,
     private val dbDirectory: String? = null,
     val logger: Logger = Logger,
-    driver: PsSqlDriver = factory.createDriver(scope, dbFilename),
 ) : ExclusiveMethodProvider(),
     PowerSyncDatabase {
     companion object {
@@ -81,12 +80,12 @@ internal class PowerSyncDatabaseImpl(
     override val identifier = dbDirectory + dbFilename
 
     private val internalDb =
-    InternalDatabaseImpl(
-        factory = factory,
-        scope = scope,
-        dbFilename = dbFilename,
-        dbDirectory = dbDirectory,
-    )
+        InternalDatabaseImpl(
+            factory = factory,
+            scope = scope,
+            dbFilename = dbFilename,
+            dbDirectory = dbDirectory,
+        )
 
     internal val bucketStorage: BucketStorage = BucketStorageImpl(internalDb, logger)
     var closed = false
@@ -112,13 +111,18 @@ internal class PowerSyncDatabaseImpl(
             if (isMultiple) {
                 logger.w { multipleInstancesMessage }
             }
-            
+
             powerSyncVersion =
                 internalDb.get("SELECT powersync_rs_version()") { it.getString(0)!! }
             logger.d { "SQLiteVersion: $powerSyncVersion" }
-    
+
             checkVersion()
             logger.d { "PowerSyncVersion: ${getPowerSyncVersion()}" }
+
+            internalDb.writeTransaction { tx ->
+                tx.getOptional("SELECT powersync_init()") {}
+            }
+
             applySchema()
             updateHasSynced()
         }

core/src/commonMain/kotlin/com/powersync/db/internal/ConnectionPool.kt

@@ -1,152 +1,67 @@
 package com.powersync.db.internal
 
-import co.touchlab.stately.concurrency.AtomicBoolean
 import com.powersync.PowerSyncException
 import com.powersync.PsSqlDriver
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.channels.ClosedSendChannelException
 import kotlinx.coroutines.joinAll
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.sync.Mutex
-import kotlinx.coroutines.sync.withLock
-
-internal class WrappedConnection(
-    val transactorDriver: TransactorDriver,
-) {
-    var busy: Boolean = false
-}
-
-internal data class DeferredAction<R>(
-    val deferred: CompletableDeferred<R>,
-    val action: suspend (connection: TransactorDriver) -> R,
-)
 
 internal class ConnectionPool(
     factory: () -> PsSqlDriver,
     size: Int = 5,
     private val scope: CoroutineScope,
 ) {
-    val closed = AtomicBoolean(false)
-
-    private val mutex = Mutex()
-    private val connections = List(size) { WrappedConnection(TransactorDriver(factory())) }
-    private val queue = mutableListOf<DeferredAction<Any?>>()
-    private val activeOperations = mutableListOf<CompletableDeferred<Any?>>()
-
-    @Suppress("UNCHECKED_CAST")
-    suspend fun <R> withConnection(action: suspend (connection: TransactorDriver) -> R): R {
-        if (closed.value) {
-            throw PowerSyncException(
-                message = "Cannot process connection pool request",
-                cause = Exception("Pool is closed"),
-            )
-        }
-
-        val wrappedDeferred = DeferredAction(CompletableDeferred<Any?>(), action)
-
-        val connection =
-            mutex.withLock {
-                // check if any connections are free
-                val freeConnection = connections.find { !it.busy }
-                if (freeConnection == null) {
-                    queue.add(wrappedDeferred)
-                    return@withLock null
+    private val available = Channel<Pair<TransactorDriver, CompletableDeferred<Unit>>>()
+    private val connections: List<Job> =
+        buildList {
+            repeat(size) {
+                scope.launch {
+                    val driver = TransactorDriver(factory())
+                    try {
+                        while (true) {
+                            val done = CompletableDeferred<Unit>()
+                            try {
+                                available.send(driver to done)
+                            } catch (_: ClosedSendChannelException) {
+                                break // Pool closed
+                            }
+
+                            done.await()
+                        }
+                    } finally {
+                        driver.driver.close()
+                    }
                 }
-                activeOperations.add(wrappedDeferred.deferred)
-                freeConnection.busy = true
-                return@withLock freeConnection
-            }
-
-        if (connection != null) {
-            // Can process immediately
-            scope.launch {
-                processRequest(wrappedDeferred, connection)
             }
         }
 
-        val result = wrappedDeferred.deferred.await()
-        if (result != null) {
-            val casted =
-                result as? R
-                    ?: throw PowerSyncException(
-                        "Could not cast deferred result",
-                        Exception("Casting error"),
-                    )
-            return casted
-        }
-        return null as R
-    }
+    suspend fun <R> withConnection(action: suspend (connection: TransactorDriver) -> R): R {
+        val (connection, done) =
+            try {
+                available.receive()
+            } catch (e: PoolClosedException) {
+                throw PowerSyncException(
+                    message = "Cannot process connection pool request",
+                    cause = e,
+                )
+            }
 
-    private suspend fun processRequest(
-        request: DeferredAction<Any?>,
-        wrappedConnection: WrappedConnection,
-    ) {
         try {
-            val result = request.action(wrappedConnection.transactorDriver)
-            request.deferred.complete(result)
-        } catch (exception: Exception) {
-            request.deferred.completeExceptionally(exception)
+            return action(connection)
         } finally {
-            mutex.withLock {
-                wrappedConnection.busy = false
-                activeOperations.remove(request.deferred)
-            }
-            scope.launch {
-                processNext()
-            }
+            done.complete(Unit)
         }
     }
 
-    private suspend fun processNext() {
-        val next: Pair<DeferredAction<Any?>, WrappedConnection> =
-            mutex.withLock {
-                if (queue.size == 0) {
-                    return
-                }
-                // check if a connection is open
-                val connection = connections.find { !it.busy }
-                if (connection == null) {
-                    return
-                }
-                val next = queue.first()
-                queue.removeFirst()
-
-                connection.busy = true
-                return@withLock next to connection
-            }
-
-        processRequest(next.first, next.second)
-    }
-
     suspend fun close() {
-        if (closed.value) {
-            return
-        }
-
-        closed.value = true
-        val activeOperations =
-            mutex.withLock {
-                // These should all be pending
-                for (item in queue) {
-                    item.deferred.completeExceptionally(
-                        PowerSyncException(
-                            message = "Pending requests are aborted",
-                            cause = Exception("Pool has been closed"),
-                        ),
-                    )
-                }
-                queue.clear()
-                // Return a copy of active items in order to check them
-                return@withLock activeOperations.toList()
-            }
-
-        // Wait for all pending operations
-        // Alternatively we could cancel and ongoing jobs
-        activeOperations.joinAll()
-
-        // Close all connections
-        for (connection in connections) {
-            connection.transactorDriver.driver.close()
-        }
+        available.cancel(PoolClosedException)
+        connections.joinAll()
     }
 }
+
+internal object PoolClosedException : CancellationException("Pool is closed")