Avoid resource leaks for database tracker

Author: simolus3

core/build.gradle.kts

@@ -211,6 +211,11 @@ kotlin {
             dependsOn(commonTest.get())
         }
 
+        val commonJava by creating {
+            kotlin.srcDir("commonJava")
+            dependsOn(commonMain.get())
+        }
+
         commonMain.dependencies {
             implementation(libs.uuid)
             implementation(libs.kotlin.stdlib)
@@ -226,13 +231,18 @@ kotlin {
             api(libs.kermit)
         }
 
-        androidMain.dependencies {
-            implementation(libs.ktor.client.okhttp)
+        androidMain {
+            dependsOn(commonJava)
+            dependencies.implementation(libs.ktor.client.okhttp)
         }
 
-        jvmMain.dependencies {
-            implementation(libs.ktor.client.okhttp)
-            implementation(libs.sqlite.jdbc)
+        jvmMain {
+            dependsOn(commonJava)
+
+            dependencies {
+                implementation(libs.ktor.client.okhttp)
+                implementation(libs.sqlite.jdbc)
+            }
         }
 
         iosMain.dependencies {

core/src/commonJava/kotlin/com/powersync/db/ActiveInstanceStore.commonJava.kt

@@ -0,0 +1,18 @@
+package com.powersync.db
+
+import java.lang.ref.Cleaner
+
+internal actual fun disposeWhenDeallocated(resource: ActiveDatabaseResource): Any {
+    // Note: It's important that the returned object does not reference the resource directly
+    val wrapper = CleanableWrapper()
+    CleanableWrapper.cleaner.register(wrapper, resource::dispose)
+    return wrapper
+}
+
+private class CleanableWrapper {
+    var cleanable: Cleaner.Cleanable? = null
+
+    companion object {
+        val cleaner: Cleaner = Cleaner.create()
+    }
+}

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

@@ -1,23 +1,81 @@
 package com.powersync.db
 
-import com.powersync.PowerSyncDatabase
-import com.powersync.utils.ExclusiveMethodProvider
-
-internal class ActiveInstanceStore : ExclusiveMethodProvider() {
-    private val instances = mutableListOf<PowerSyncDatabase>()
-
-    /**
-     * Registers an instance. Returns true if multiple instances with the same identifier are
-     * present.
-     */
-    suspend fun registerAndCheckInstance(db: PowerSyncDatabase) =
-        exclusiveMethod("instances") {
-            instances.add(db)
-            return@exclusiveMethod instances.filter { it.identifier == db.identifier }.size > 1
+import co.touchlab.kermit.Logger
+import kotlinx.atomicfu.atomic
+import kotlinx.atomicfu.locks.SynchronizedObject
+import kotlinx.atomicfu.locks.synchronized
+import kotlinx.coroutines.sync.Mutex
+
+/**
+ * Returns an object that, when deallocated, calls [ActiveDatabaseResource.dispose].
+ */
+internal expect fun disposeWhenDeallocated(resource: ActiveDatabaseResource): Any
+
+/**
+ * An collection of PowerSync databases with the same path / identifier.
+ *
+ * We expect that each group will only ever have one database because we encourage users to write their databases as
+ * singletons. We print a warning when two databases are part of the same group.
+ * Additionally, we want to avoid two databases in the same group having a sync stream open at the same time to avoid
+ * duplicate resources being used. For this reason, each active database group has a coroutine mutex guarding the
+ * sync job.
+ */
+internal class ActiveDatabaseGroup(val identifier: String) {
+    internal var refCount = 0 // Guarded by companion object
+    internal val syncMutex = Mutex()
+
+    fun removeUsage() {
+        synchronized(ActiveDatabaseGroup) {
+            if (--refCount == 0) {
+                allGroups.remove(this)
+            }
+        }
+    }
+
+    companion object: SynchronizedObject() {
+        internal val multipleInstancesMessage =
+            """
+            Multiple PowerSync instances for the same database have been detected.
+            This can cause unexpected results.
+            Please check your PowerSync client instantiation logic if this is not intentional.
+            """.trimIndent()
+
+        private val allGroups = mutableListOf<ActiveDatabaseGroup>()
+
+        private fun findGroup(warnOnDuplicate: Logger, identifier: String): ActiveDatabaseGroup {
+            return synchronized(this) {
+                val existing = allGroups.asSequence().firstOrNull { it.identifier == identifier }
+                val resolvedGroup = if (existing == null) {
+                    val added = ActiveDatabaseGroup(identifier)
+                    allGroups.add(added)
+                    added
+                } else {
+                    existing
+                }
+
+                if (resolvedGroup.refCount++ != 0) {
+                    warnOnDuplicate.w { multipleInstancesMessage }
+                }
+
+                resolvedGroup
+            }
         }
 
-    suspend fun removeInstance(db: PowerSyncDatabase) =
-        exclusiveMethod("instances") {
-            instances.remove(db)
+        internal fun referenceDatabase(warnOnDuplicate: Logger, identifier: String): Pair<ActiveDatabaseResource, Any> {
+            val group = findGroup(warnOnDuplicate, identifier)
+            val resource = ActiveDatabaseResource(group)
+
+            return resource to disposeWhenDeallocated(resource)
+        }
+    }
+}
+
+internal class ActiveDatabaseResource(val group: ActiveDatabaseGroup) {
+    var disposed = atomic(false)
+
+    fun dispose() {
+        if (!disposed.getAndSet(true)) {
+            group.removeUsage()
         }
+    }
 }

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

@@ -20,7 +20,6 @@ import com.powersync.sync.PriorityStatusEntry
 import com.powersync.sync.SyncStatus
 import com.powersync.sync.SyncStatusData
 import com.powersync.sync.SyncStream
-import com.powersync.utils.ExclusiveMethodProvider
 import com.powersync.utils.JsonParam
 import com.powersync.utils.JsonUtil
 import com.powersync.utils.throttle
@@ -57,8 +56,7 @@ internal class PowerSyncDatabaseImpl(
     private val dbFilename: String,
     val logger: Logger = Logger,
     driver: PsSqlDriver = factory.createDriver(scope, dbFilename),
-) : ExclusiveMethodProvider(),
-    PowerSyncDatabase {
+): PowerSyncDatabase {
     companion object {
         internal val streamConflictMessage =
             """
@@ -68,21 +66,15 @@ internal class PowerSyncDatabaseImpl(
             This connection attempt will be queued and will only be executed after
             currently connecting clients are disconnected.
             """.trimIndent()
-
-        internal val multipleInstancesMessage =
-            """
-            Multiple PowerSync instances for the same database have been detected.
-            This can cause unexpected results.
-            Please check your PowerSync client instantiation logic if this is not intentional.
-            """.trimIndent()
-
-        internal val instanceStore = ActiveInstanceStore()
     }
 
     override val identifier = dbFilename
 
     private val internalDb = InternalDatabaseImpl(driver, scope)
     internal val bucketStorage: BucketStorage = BucketStorageImpl(internalDb, logger)
+    private val resource: ActiveDatabaseResource
+    private val clearResourceWhenDisposed: Any
+
     var closed = false
 
     /**
@@ -97,12 +89,11 @@ internal class PowerSyncDatabaseImpl(
     private var uploadJob: Job? = null
 
     init {
-        val db = this
+        val res = ActiveDatabaseGroup.referenceDatabase(logger, identifier)
+        resource = res.first
+        clearResourceWhenDisposed = res.second
+
         runBlocking {
-            val isMultiple = instanceStore.registerAndCheckInstance(db)
-            if (isMultiple) {
-                logger.w { multipleInstancesMessage }
-            }
             val sqliteVersion = internalDb.queries.sqliteVersion().executeAsOne()
             logger.d { "SQLiteVersion: $sqliteVersion" }
             checkVersion()
@@ -154,8 +145,7 @@ internal class PowerSyncDatabaseImpl(
         syncJob =
             scope.launch {
                 // Get a global lock for checking mutex maps
-                val streamMutex =
-                    globalMutexFor("streaming-$identifier")
+                val streamMutex = resource.group.syncMutex
 
                 // Poke the streaming mutex to see if another client is using it
                 var obtainedLock = false
@@ -437,7 +427,7 @@ internal class PowerSyncDatabaseImpl(
             }
             disconnect()
             internalDb.close()
-            instanceStore.removeInstance(this)
+            resource.dispose()
             closed = true
         }
 

core/src/commonMain/kotlin/com/powersync/utils/ExclusiveMethodProvider.kt

@@ -0,0 +1,9 @@
+package com.powersync.db
+
+import kotlin.experimental.ExperimentalNativeApi
+import kotlin.native.ref.createCleaner
+
+@OptIn(ExperimentalNativeApi::class)
+internal actual fun disposeWhenDeallocated(resource: ActiveDatabaseResource): Any {
+    return createCleaner(resource) { it.dispose() }
+}