Use structured concurrency for sync jobs

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## 1.0.0-BETA29 (unreleased)
+
+* Fix potential race condition between jobs in `connect()` and `disconnect()`.
+
 ## 1.0.0-BETA28
 
 * Update PowerSync SQLite core extension to 0.3.12.

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

@@ -27,9 +27,9 @@ import com.powersync.utils.JsonUtil
 import dev.mokkery.answering.returns
 import dev.mokkery.everySuspend
 import dev.mokkery.mock
-import kotlinx.coroutines.CoroutineScope
+import dev.mokkery.verify
+import kotlinx.coroutines.DelicateCoroutinesApi
 import kotlinx.coroutines.channels.Channel
-import kotlinx.coroutines.flow.receiveAsFlow
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.runTest
 import kotlinx.serialization.encodeToString
@@ -38,6 +38,7 @@ import kotlin.test.AfterTest
 import kotlin.test.BeforeTest
 import kotlin.test.Test
 import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
 import kotlin.test.assertFalse
 import kotlin.test.assertNotNull
 import kotlin.test.assertTrue
@@ -99,8 +100,8 @@ class SyncIntegrationTest {
             dbFilename = "testdb",
         ) as PowerSyncDatabaseImpl
 
-    private fun CoroutineScope.syncStream(): SyncStream {
-        val client = MockSyncService.client(this, syncLines.receiveAsFlow())
+    private fun syncStream(): SyncStream {
+        val client = MockSyncService(syncLines)
         return SyncStream(
             bucketStorage = database.bucketStorage,
             connector = connector,
@@ -117,6 +118,68 @@ class SyncIntegrationTest {
         assertEquals(amount, users.size, "Expected $amount users, got $users")
     }
 
+    @Test
+    @OptIn(DelicateCoroutinesApi::class)
+    fun closesResponseStreamOnDatabaseClose() =
+        runTest {
+            val syncStream = syncStream()
+            database.connectInternal(syncStream, 1000L)
+
+            turbineScope(timeout = 10.0.seconds) {
+                val turbine = database.currentStatus.asFlow().testIn(this)
+                turbine.waitFor { it.connected }
+
+                database.close()
+                turbine.waitFor { !it.connected }
+                turbine.cancel()
+            }
+
+            // Closing the database should have closed the channel
+            assertTrue { syncLines.isClosedForSend }
+        }
+
+    @Test
+    @OptIn(DelicateCoroutinesApi::class)
+    fun cleansResourcesOnDisconnect() =
+        runTest {
+            val syncStream = syncStream()
+            database.connectInternal(syncStream, 1000L)
+
+            turbineScope(timeout = 10.0.seconds) {
+                val turbine = database.currentStatus.asFlow().testIn(this)
+                turbine.waitFor { it.connected }
+
+                database.disconnect()
+                turbine.waitFor { !it.connected }
+                turbine.cancel()
+            }
+
+            // Disconnecting should have closed the channel
+            assertTrue { syncLines.isClosedForSend }
+
+            // And called invalidateCredentials on the connector
+            verify { connector.invalidateCredentials() }
+        }
+
+    @Test
+    fun cannotUpdateSchemaWhileConnected() =
+        runTest {
+            val syncStream = syncStream()
+            database.connectInternal(syncStream, 1000L)
+
+            turbineScope(timeout = 10.0.seconds) {
+                val turbine = database.currentStatus.asFlow().testIn(this)
+                turbine.waitFor { it.connected }
+                turbine.cancel()
+            }
+
+            assertFailsWith<PowerSyncException>("Cannot update schema while connected") {
+                database.updateSchema(Schema())
+            }
+
+            database.close()
+        }
+
     @Test
     fun testPartialSync() =
         runTest {

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

@@ -24,10 +24,11 @@ import com.powersync.utils.JsonParam
 import com.powersync.utils.JsonUtil
 import com.powersync.utils.throttle
 import com.powersync.utils.toJsonObject
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.FlowPreview
 import kotlinx.coroutines.Job
-import kotlinx.coroutines.cancelAndJoin
+import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.ensureActive
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.filter
@@ -95,9 +96,7 @@ internal class PowerSyncDatabaseImpl(
     override val currentStatus: SyncStatus = SyncStatus()
 
     private val mutex = Mutex()
-    private var syncStream: SyncStream? = null
-    private var syncJob: Job? = null
-    private var uploadJob: Job? = null
+    private var syncSupervisorJob: Job? = null
 
     // This is set in the init
     private lateinit var powerSyncVersion: String
@@ -123,7 +122,7 @@ internal class PowerSyncDatabaseImpl(
     override suspend fun updateSchema(schema: Schema) =
         runWrappedSuspending {
             mutex.withLock {
-                if (this.syncStream != null) {
+                if (this.syncSupervisorJob != null) {
                     throw PowerSyncException(
                         "Cannot update schema while connected",
                         cause = Exception("PowerSync client is already connected"),
@@ -161,12 +160,11 @@ internal class PowerSyncDatabaseImpl(
         stream: SyncStream,
         crudThrottleMs: Long,
     ) {
-        this.syncStream = stream
-
         val db = this
-
-        syncJob =
-            scope.launch {
+        val job = SupervisorJob(scope.coroutineContext[Job])
+        syncSupervisorJob = job
+        scope.launch(job) {
+            launch {
                 // Get a global lock for checking mutex maps
                 val streamMutex = resource.group.syncMutex
 
@@ -181,7 +179,7 @@ internal class PowerSyncDatabaseImpl(
                         // (The tryLock should throw if this client already holds the lock).
                         logger.w(streamConflictMessage)
                     }
-                } catch (ex: IllegalStateException) {
+                } catch (_: IllegalStateException) {
                     logger.e { "The streaming sync client did not disconnect before connecting" }
                 }
 
@@ -194,40 +192,46 @@ internal class PowerSyncDatabaseImpl(
                 // We have a lock if we reached here
                 try {
                     ensureActive()
-                    syncStream!!.streamingSync()
+                    stream.streamingSync()
                 } finally {
                     streamMutex.unlock(db)
                 }
             }
 
-        scope.launch {
-            syncStream!!.status.asFlow().collect {
-                currentStatus.update(
-                    connected = it.connected,
-                    connecting = it.connecting,
-                    uploading = it.uploading,
-                    downloading = it.downloading,
-                    lastSyncedAt = it.lastSyncedAt,
-                    hasSynced = it.hasSynced,
-                    uploadError = it.uploadError,
-                    downloadError = it.downloadError,
-                    clearDownloadError = it.downloadError == null,
-                    clearUploadError = it.uploadError == null,
-                    priorityStatusEntries = it.priorityStatusEntries,
-                )
+            launch {
+                stream.status.asFlow().collect {
+                    currentStatus.update(
+                        connected = it.connected,
+                        connecting = it.connecting,
+                        uploading = it.uploading,
+                        downloading = it.downloading,
+                        lastSyncedAt = it.lastSyncedAt,
+                        hasSynced = it.hasSynced,
+                        uploadError = it.uploadError,
+                        downloadError = it.downloadError,
+                        clearDownloadError = it.downloadError == null,
+                        clearUploadError = it.uploadError == null,
+                        priorityStatusEntries = it.priorityStatusEntries,
+                    )
+                }
             }
-        }
 
-        uploadJob =
-            scope.launch {
+            launch {
                 internalDb
                     .updatesOnTables()
                     .filter { it.contains(InternalTable.CRUD.toString()) }
                     .throttle(crudThrottleMs)
                     .collect {
-                        syncStream!!.triggerCrudUpload()
+                        stream.triggerCrudUpload()
                     }
             }
+        }
+
+        job.invokeOnCompletion {
+            if (it is DisconnectRequestedException) {
+                stream.invalidateCredentials()
+            }
+        }
     }
 
     override suspend fun getCrudBatch(limit: Int): CrudBatch? {
@@ -364,17 +368,12 @@ internal class PowerSyncDatabaseImpl(
     override suspend fun disconnect() = mutex.withLock { disconnectInternal() }
 
     private suspend fun disconnectInternal() {
-        if (syncJob != null && syncJob!!.isActive) {
-            syncJob?.cancelAndJoin()
-        }
-
-        if (uploadJob != null && uploadJob!!.isActive) {
-            uploadJob?.cancelAndJoin()
-        }
-
-        if (syncStream != null) {
-            syncStream?.invalidateCredentials()
-            syncStream = null
+        val syncJob = syncSupervisorJob
+        if (syncJob != null && syncJob.isActive) {
+            // Using this exception type will also make the sync job invalidate credentials.
+            syncJob.cancel(DisconnectRequestedException)
+            syncJob.join()
+            syncSupervisorJob = null
         }
 
         currentStatus.update(
@@ -470,7 +469,7 @@ internal class PowerSyncDatabaseImpl(
     /**
      * Check that a supported version of the powersync extension is loaded.
      */
-    private suspend fun checkVersion(powerSyncVersion: String) {
+    private fun checkVersion(powerSyncVersion: String) {
         // Parse version
         val versionInts: List<Int> =
             try {
@@ -488,3 +487,5 @@ internal class PowerSyncDatabaseImpl(
         }
     }
 }
+
+internal object DisconnectRequestedException : CancellationException("disconnect() called")

core/src/commonTest/kotlin/com/powersync/sync/SyncStreamTest.kt

@@ -30,7 +30,6 @@ import dev.mokkery.verifySuspend
 import io.ktor.client.engine.mock.MockEngine
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.delay
-import kotlinx.coroutines.flow.receiveAsFlow
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.runTest
 import kotlinx.coroutines.withTimeout
@@ -210,7 +209,7 @@ class SyncStreamTest {
             // TODO: It would be neat if we could use in-memory sqlite instances instead of mocking everything
             // Revisit https://github.com/powersync-ja/powersync-kotlin/pull/117/files at some point
             val syncLines = Channel<SyncLine>()
-            val client = MockSyncService.client(this, syncLines.receiveAsFlow())
+            val client = MockSyncService(syncLines)
 
             syncStream =
                 SyncStream(