Fix races in new sync client

Author: simolus3

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

@@ -675,7 +675,7 @@ class NewSyncIntegrationTest : BaseSyncIntegrationTest(true) {
                 verifyNoMoreCalls(connector)
 
                 syncLines.send(SyncLine.KeepAlive(tokenExpiresIn = 10))
-                prefetchCalled.complete(Unit)
+                prefetchCalled.await()
                 // Should still be connected before prefetch completes
                 database.currentStatus.connected shouldBe true
 

core/src/commonMain/kotlin/com/powersync/bucket/BucketStorage.kt

@@ -6,6 +6,9 @@ import com.powersync.sync.Instruction
 import com.powersync.sync.LegacySyncImplementation
 import com.powersync.sync.SyncDataBatch
 import com.powersync.sync.SyncLocalDatabaseResult
+import kotlinx.serialization.Serializable
+import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.JsonObject
 
 internal interface BucketStorage {
     fun getMaxOpId(): String
@@ -50,13 +53,27 @@ internal interface BucketStorage {
         partialPriority: BucketPriority? = null,
     ): SyncLocalDatabaseResult
 
-    suspend fun control(
-        op: String,
-        payload: String?,
-    ): List<Instruction>
+    suspend fun control(args: PowerSyncControlArguments): List<Instruction>
+}
 
-    suspend fun control(
-        op: String,
-        payload: ByteArray,
-    ): List<Instruction>
+internal sealed interface PowerSyncControlArguments {
+    @Serializable
+    class Start(
+        val parameters: JsonObject
+    ): PowerSyncControlArguments
+    data object Stop: PowerSyncControlArguments
+
+    data class TextLine(val line: String): PowerSyncControlArguments
+    class BinaryLine(val line: ByteArray): PowerSyncControlArguments {
+        override fun toString(): String {
+            return "BinaryLine"
+        }
+    }
+
+    data object CompletedUpload: PowerSyncControlArguments
 }
+
+@Serializable
+internal class StartSyncIteration(
+    val parameters: JsonObject,
+)

core/src/commonMain/kotlin/com/powersync/bucket/BucketStorageImpl.kt

@@ -365,22 +365,20 @@ internal class BucketStorageImpl(
         return JsonUtil.json.decodeFromString<List<Instruction>>(result)
     }
 
-    override suspend fun control(
-        op: String,
-        payload: String?,
-    ): List<Instruction> =
+    override suspend fun control(args: PowerSyncControlArguments): List<Instruction> =
         db.writeTransaction { tx ->
-            logger.v { "powersync_control($op, $payload)" }
+            logger.v { "powersync_control: $args" }
 
-            tx.get("SELECT powersync_control(?, ?) AS r", listOf(op, payload), ::handleControlResult)
-        }
+            val (op: String, data: Any?) = when (args) {
+                is PowerSyncControlArguments.Start -> "start" to JsonUtil.json.encodeToString(args)
+                PowerSyncControlArguments.Stop -> "stop" to null
 
-    override suspend fun control(
-        op: String,
-        payload: ByteArray,
-    ): List<Instruction> =
-        db.writeTransaction { tx ->
-            logger.v { "powersync_control($op, binary payload)" }
-            tx.get("SELECT powersync_control(?, ?) AS r", listOf(op, payload), ::handleControlResult)
+                PowerSyncControlArguments.CompletedUpload -> "completed_upload" to null
+
+                is PowerSyncControlArguments.BinaryLine -> "line_binary" to args.line
+                is PowerSyncControlArguments.TextLine -> "line_text" to args.line
+            }
+
+            tx.get("SELECT powersync_control(?, ?) AS r", listOf(op, data), ::handleControlResult)
         }
 }

core/src/commonMain/kotlin/com/powersync/sync/SyncStream.kt

@@ -8,6 +8,7 @@ import com.powersync.bucket.BucketChecksum
 import com.powersync.bucket.BucketRequest
 import com.powersync.bucket.BucketStorage
 import com.powersync.bucket.Checkpoint
+import com.powersync.bucket.PowerSyncControlArguments
 import com.powersync.bucket.WriteCheckpointResponse
 import com.powersync.connectors.PowerSyncBackendConnector
 import com.powersync.db.crud.CrudEntry
@@ -39,9 +40,11 @@ import kotlinx.coroutines.NonCancellable
 import kotlinx.coroutines.cancelAndJoin
 import kotlinx.coroutines.channels.BufferOverflow
 import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.channels.consume
 import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.emitAll
 import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.launch
@@ -297,43 +300,41 @@ internal class SyncStream(
      */
     private inner class ActiveIteration(
         val scope: CoroutineScope,
-        var hadSyncLine: Boolean = false,
         var fetchLinesJob: Job? = null,
         var credentialsInvalidation: Job? = null,
     ) {
-        suspend fun start() {
-            @Serializable
-            class StartParameters(
-                val parameters: JsonObject,
-            )
-
-            control("start", JsonUtil.json.encodeToString(StartParameters(params)))
-            fetchLinesJob?.join()
-        }
+        // Using a channel for control invocations so that they're handled by a single coroutine,
+        // avoiding races between concurrent jobs like fetching credentials.
+        private val controlInvocations = Channel<PowerSyncControlArguments>()
 
-        suspend fun stop() {
-            control("stop")
-            fetchLinesJob?.join()
-        }
-
-        private suspend fun control(
-            op: String,
-            payload: String? = null,
-        ) {
-            val instructions = bucketStorage.control(op, payload)
-            handleInstructions(instructions)
+        private suspend fun invokeControl(args: PowerSyncControlArguments) {
+            val instructions = bucketStorage.control(args)
+            instructions.forEach { handleInstruction(it) }
         }
 
-        private suspend fun control(
-            op: String,
-            payload: ByteArray,
-        ) {
-            val instructions = bucketStorage.control(op, payload)
-            handleInstructions(instructions)
+        suspend fun start() {
+            invokeControl(PowerSyncControlArguments.Start(params))
+
+            var hadSyncLine = false
+            for (line in controlInvocations) {
+                val instructions = bucketStorage.control(line)
+                instructions.forEach { handleInstruction(it) }
+
+                if (!hadSyncLine && (line is PowerSyncControlArguments.TextLine || line is PowerSyncControlArguments.BinaryLine)) {
+                    // Trigger a crud upload when receiving the first sync line: We could have
+                    // pending local writes made while disconnected, so in addition to listening on
+                    // updates to `ps_crud`, we also need to trigger a CRUD upload in some other
+                    // cases. We do this on the first sync line because the client is likely to be
+                    // online in that case.
+                    hadSyncLine = true
+                    triggerCrudUploadAsync()
+                }
+            }
         }
 
-        private suspend fun handleInstructions(instructions: List<Instruction>) {
-            instructions.forEach { handleInstruction(it) }
+        suspend fun stop() {
+            invokeControl(PowerSyncControlArguments.Stop)
+            fetchLinesJob?.join()
         }
 
         private suspend fun handleInstruction(instruction: Instruction) {
@@ -344,20 +345,25 @@ internal class SyncStream(
                         scope.launch {
                             launch {
                                 logger.v { "listening for completed uploads" }
-
                                 for (completion in completedCrudUploads) {
-                                    control("completed_upload")
+                                    controlInvocations.send(PowerSyncControlArguments.CompletedUpload)
                                 }
                             }
 
                             launch {
                                 connect(instruction)
                             }
+                        }.also {
+                            it.invokeOnCompletion {
+                                controlInvocations.close()
+                            }
                         }
                 }
                 Instruction.CloseSyncStream -> {
+                    logger.v { "Closing sync stream connection" }
                     fetchLinesJob!!.cancelAndJoin()
                     fetchLinesJob = null
+                    logger.v { "Sync stream connection shut down" }
                 }
                 Instruction.FlushSileSystem -> {
                     // We have durable file systems, so flushing is not necessary
@@ -389,6 +395,7 @@ internal class SyncStream(
                             val job =
                                 scope.launch {
                                     connector.prefetchCredentials().join()
+                                    logger.v { "Stopping because new credentials are available" }
 
                                     // Token has been refreshed, start another iteration
                                     stop()
@@ -409,36 +416,16 @@ internal class SyncStream(
             }
         }
 
-        /**
-         * Triggers a crud upload when called for the first time.
-         *
-         * We could have pending local writes made while disconnected, so in addition to listening
-         * on updates to `ps_crud`, we also need to trigger a CRUD upload in some other cases. We
-         * do this on the first sync line because the client is likely to be online in that case.
-         */
-        private fun triggerCrudUploadIfFirstLine() {
-            if (!hadSyncLine) {
-                triggerCrudUploadAsync()
-                hadSyncLine = true
-            }
-        }
-
-        private suspend fun line(text: String) {
-            control("line_text", text)
-            triggerCrudUploadIfFirstLine()
-        }
-
-        private suspend fun line(blob: ByteArray) {
-            control("line_binary", blob)
-            triggerCrudUploadIfFirstLine()
-        }
-
         private suspend fun connect(start: Instruction.EstablishSyncStream) {
             when (val method = options.method) {
                 ConnectionMethod.Http ->
-                    connectViaHttp(start.request).collect(this::line)
+                    connectViaHttp(start.request).collect {
+                        controlInvocations.send(PowerSyncControlArguments.TextLine(it))
+                    }
                 is ConnectionMethod.WebSocket ->
-                    connectViaWebSocket(start.request, method).collect(this::line)
+                    connectViaWebSocket(start.request, method).collect {
+                        controlInvocations.send(PowerSyncControlArguments.BinaryLine(it))
+                    }
             }
         }
     }