perf(model-server): transfer model objects as plain text instead of JSON

VersionDelta is basically just a map of objects and their SHA hashes.
This can be serialized more efficiently with just a stream of strings separated by \n.
Using flows/sequences/streams is more memory efficient, because the data can directly be streamed
to the client without having to create this large JSON object. This should also fix the timeout
issues, because the server already starts responding while still iterating the over the model.

Author: slisson

api/model-server.yaml

@@ -107,7 +107,7 @@ paths:
         "404":
           $ref: '#/components/responses/404'
         "200":
-          description: OK
+          $ref: '#/components/responses/versionDelta'
 #          content:
 #            '*/*':
 #              schema:
@@ -137,11 +137,7 @@ paths:
 #        required: true
       responses:
         "200":
-          $ref: '#/components/responses/200'
-#          content:
-#            '*/*':
-#              schema:
-#                type: object
+          $ref: '#/components/responses/versionDelta'
   /v2/repositories/{repository}/branches/{branch}/hash:
     get:
       operationId: getRepositoryBranchHash
@@ -247,11 +243,7 @@ paths:
         "500":
           $ref: '#/components/responses/500'
         "200":
-          description: OK
-          content:
-            '*/*':
-              schema:
-                type: object
+          $ref: '#/components/responses/versionDelta'
   /v2/repositories/{repository}/branches/{branch}/pollHash:
     get:
       operationId: pollRepositoryBranchHash
@@ -321,11 +313,7 @@ paths:
         "500":
           $ref: '#/components/responses/500'
         "200":
-          description: OK
-          content:
-            '*/*':
-              schema:
-                type: object
+          $ref: '#/components/responses/versionDelta'
   /v2/repositories/{repository}/versions/{versionHash}:
     get:
       operationId: getRepositoryVersionHash
@@ -414,11 +402,7 @@ paths:
         "404":
           $ref: '#/components/responses/404'
         "200":
-          description: OK
-          content:
-            '*/*':
-              schema:
-                type: object
+          $ref: '#/components/responses/versionDelta'
   /v2/versions/{versionHash}/history/{oldestVersionHash}:
     get:
       operationId: getOldestVersionHashForVersion
@@ -599,7 +583,18 @@ components:
         text/plain:
           schema:
             type: string
-
+    "versionDelta":
+      description: OK
+      content:
+        'application/x-modelix-objects':
+          schema:
+            type: string
+        'application/json':
+          schema:
+            type: object
+        'text/plain':
+          schema:
+            type: string
   schemas:
     VersionDelta:
       title: VersionDelta

gradle/libs.versions.toml

@@ -73,6 +73,7 @@ ktor-client-websockets = { group = "io.ktor", name = "ktor-client-websockets", v
 ktor-client-js = { group = "io.ktor", name = "ktor-client-js", version.ref = "ktor" }
 ktor-client-auth = { group = "io.ktor", name = "ktor-client-auth", version.ref = "ktor" }
 
+ktor-serialization = { group = "io.ktor", name = "ktor-serialization", version.ref = "ktor" }
 ktor-serialization-json = { group = "io.ktor", name = "ktor-serialization-kotlinx-json", version.ref = "ktor" }
 
 keycloak-authz-client = { group = "org.keycloak", name = "keycloak-authz-client", version = "23.0.3" }

model-client/src/commonMain/kotlin/org/modelix/model/client2/ModelClientV2.kt

@@ -21,21 +21,28 @@ import io.ktor.client.plugins.HttpTimeout
 import io.ktor.client.plugins.ResponseException
 import io.ktor.client.plugins.contentnegotiation.ContentNegotiation
 import io.ktor.client.plugins.expectSuccess
+import io.ktor.client.request.HttpRequestBuilder
 import io.ktor.client.request.get
 import io.ktor.client.request.post
 import io.ktor.client.request.put
 import io.ktor.client.request.setBody
+import io.ktor.client.statement.HttpResponse
+import io.ktor.client.statement.bodyAsChannel
 import io.ktor.client.statement.bodyAsText
 import io.ktor.http.ContentType
+import io.ktor.http.HttpHeaders
 import io.ktor.http.HttpStatusCode
 import io.ktor.http.URLBuilder
 import io.ktor.http.appendPathSegments
 import io.ktor.http.contentType
 import io.ktor.http.takeFrom
 import io.ktor.serialization.kotlinx.json.json
+import io.ktor.utils.io.readUTF8Line
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.emptyFlow
+import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.launch
-import kotlinx.serialization.json.Json
 import org.modelix.kotlin.utils.DeprecationInfo
 import org.modelix.model.IVersion
 import org.modelix.model.api.IIdGenerator
@@ -54,6 +61,8 @@ import org.modelix.model.operations.OTBranch
 import org.modelix.model.persistent.HashUtil
 import org.modelix.model.persistent.MapBasedStore
 import org.modelix.model.server.api.v2.VersionDelta
+import org.modelix.model.server.api.v2.VersionDeltaStream
+import org.modelix.model.server.api.v2.asStream
 import org.modelix.modelql.client.ModelQLClient
 import org.modelix.modelql.core.IMonoStep
 import kotlin.time.Duration
@@ -117,9 +126,9 @@ class ModelClientV2(
                 takeFrom(baseUrl)
                 appendPathSegmentsEncodingSlash("repositories", repository.id, "init")
             }
+            useVersionStreamFormat()
         }
-        val delta = response.body<VersionDelta>()
-        return createVersion(null, delta)
+        return createVersion(null, response.readVersionDelta())
     }
 
     override suspend fun listRepositories(): List<RepositoryId> {
@@ -151,9 +160,9 @@ class ModelClientV2(
                     parameters["lastKnown"] = (baseVersion as CLVersion).getContentHash()
                 }
             }
+            useVersionStreamFormat()
         }
-        val delta = Json.decodeFromString<VersionDelta>(response.bodyAsText())
-        return createVersion(baseVersion as CLVersion?, delta)
+        return createVersion(baseVersion as CLVersion?, response.readVersionDelta())
     }
 
     override suspend fun loadVersion(
@@ -169,9 +178,9 @@ class ModelClientV2(
                     parameters["lastKnown"] = (baseVersion as CLVersion).getContentHash()
                 }
             }
+            useVersionStreamFormat()
         }
-        val delta = Json.decodeFromString<VersionDelta>(response.bodyAsText())
-        return createVersion(baseVersion as CLVersion?, delta)
+        return createVersion(baseVersion as CLVersion?, response.readVersionDelta())
     }
 
     override suspend fun push(branch: BranchReference, version: IVersion, baseVersion: IVersion?): IVersion {
@@ -193,11 +202,11 @@ class ModelClientV2(
                 takeFrom(baseUrl)
                 appendPathSegmentsEncodingSlash("repositories", branch.repositoryId.id, "branches", branch.branchName)
             }
+            useVersionStreamFormat()
             contentType(ContentType.Application.Json)
             setBody(delta)
         }
-        val mergedVersionDelta = response.body<VersionDelta>()
-        return createVersion(version, mergedVersionDelta)
+        return createVersion(version, response.readVersionDelta())
     }
 
     private suspend fun uploadObjects(repository: RepositoryId, objects: Sequence<Pair<String, String>>) {
@@ -224,8 +233,9 @@ class ModelClientV2(
                     parameters["lastKnown"] = lastKnownVersion.hash
                 }
             }
+            useVersionStreamFormat()
         }
-        val receivedVersion = createVersion(lastKnownVersion, response.body())
+        val receivedVersion = createVersion(lastKnownVersion, response.readVersionDelta())
         LOG.debug { "${clientId.toString(16)}.pull($branch, $lastKnownVersion) -> $receivedVersion" }
         return receivedVersion
     }
@@ -237,11 +247,12 @@ class ModelClientV2(
                 takeFrom(baseUrl)
                 appendPathSegmentsEncodingSlash("repositories", branch.repositoryId.id, "branches", branch.branchName)
             }
+            useVersionStreamFormat()
         }
 
         val receivedVersion = when (response.status) {
             HttpStatusCode.NotFound -> null
-            HttpStatusCode.OK -> createVersion(null, response.body())
+            HttpStatusCode.OK -> createVersion(null, response.readVersionDelta())
             else -> throw ResponseException(response, response.bodyAsText())
         }
         LOG.debug { "${clientId.toString(16)}.pullIfExists($branch) -> $receivedVersion" }
@@ -284,8 +295,9 @@ class ModelClientV2(
                     parameters["lastKnown"] = lastKnownVersion.hash
                 }
             }
+            useVersionStreamFormat()
         }
-        val receivedVersion = createVersion(lastKnownVersion, response.body())
+        val receivedVersion = createVersion(lastKnownVersion, response.readVersionDelta())
         LOG.debug { "${clientId.toString(16)}.poll($branch, $lastKnownVersion) -> $receivedVersion" }
         return receivedVersion
     }
@@ -310,24 +322,78 @@ class ModelClientV2(
         httpClient.close()
     }
 
-    private fun createVersion(baseVersion: CLVersion?, delta: VersionDelta): CLVersion {
+    private suspend fun HttpResponse.readVersionDelta(): VersionDeltaStream {
+        return if (contentType()?.match(VersionDeltaStream.CONTENT_TYPE) == true) {
+            val content = bodyAsChannel()
+            val versionHash = checkNotNull(content.readUTF8Line()) { "No objects received" }
+            val versionObject = content.readUTF8Line()
+            return if (versionObject == null) {
+                VersionDeltaStream(versionHash, emptyFlow())
+            } else {
+                VersionDeltaStream(
+                    versionHash,
+                    flow {
+                        emit(versionHash to versionObject)
+                        while (true) {
+                            val key = content.readUTF8Line() ?: break
+                            val value = checkNotNull(content.readUTF8Line()) { "Object missing for hash $key" }
+                            emit(key to value)
+                        }
+                    },
+                )
+            }
+        } else {
+            body<VersionDelta>().asStream()
+        }
+    }
+
+    private suspend fun createVersion(baseVersion: CLVersion?, delta: VersionDeltaStream): CLVersion {
+        delta.getObjectsAsFlow().collect {
+            HashUtil.checkObjectHash(it.first, it.second)
+            store.keyValueStore.put(it.first, it.second)
+        }
+        return if (baseVersion == null) {
+            CLVersion(delta.versionHash, store)
+        } else if (delta.versionHash == baseVersion.getContentHash()) {
+            baseVersion
+        } else {
+            require(baseVersion.store == store) { "baseVersion was not created by this client" }
+            CLVersion(delta.versionHash, store)
+        }
+    }
+
+    private suspend fun createVersion(baseVersion: CLVersion?, delta: Flow<String>): CLVersion {
+        var firstHash: String? = null
+        var isHash = true
+        var lastHash: String? = null
+        delta.collect {
+            if (isHash) {
+                lastHash = it
+                if (firstHash == null) {
+                    firstHash = it
+                }
+            } else {
+                val value = it
+                store.keyValueStore.put(lastHash!!, value)
+            }
+            isHash = !isHash
+        }
+        val versionHash = checkNotNull(firstHash) { "No objects received" }
+
         return if (baseVersion == null) {
-            CLVersion(
-                delta.versionHash,
-                store.also { it.keyValueStore.putAll(delta.getAllObjects()) },
-            )
-        } else if (delta.versionHash == baseVersion.hash) {
+            CLVersion(versionHash, store)
+        } else if (versionHash == baseVersion.getContentHash()) {
             baseVersion
         } else {
             require(baseVersion.store == store) { "baseVersion was not created by this client" }
-            store.keyValueStore.putAll(delta.getAllObjects())
-            CLVersion(
-                delta.versionHash,
-                baseVersion.store,
-            )
+            CLVersion(versionHash, store)
         }
     }
 
+    private fun HttpRequestBuilder.useVersionStreamFormat() {
+        headers.set(HttpHeaders.Accept, VersionDeltaStream.CONTENT_TYPE.toString())
+    }
+
     companion object {
         private val LOG = mu.KotlinLogging.logger {}
         fun builder(): ModelClientV2Builder = ModelClientV2PlatformSpecificBuilder()

model-datastructure/src/commonMain/kotlin/org/modelix/model/lazy/CLVersion.kt

@@ -398,7 +398,7 @@ private fun trackAccessedEntries(store: IKeyValueStore, body: (IDeserializingKey
 }
 
 private class AccessTrackingStore(val store: IKeyValueStore) : IKeyValueStore {
-    val accessedEntries: MutableMap<String, String?> = HashMap()
+    val accessedEntries: MutableMap<String, String?> = LinkedHashMap()
 
     override fun newBulkQuery(deserializingCache: IDeserializingKeyValueStore): IBulkQuery {
         return store.newBulkQuery(deserializingCache)

model-datastructure/src/commonMain/kotlin/org/modelix/model/persistent/HashUtil.kt

@@ -40,14 +40,17 @@ object HashUtil {
 
     fun checkObjectHashes(entries: Map<String, String?>) {
         for (entry in entries) {
-            val value = entry.value ?: continue
-            if (!isSha256(entry.key)) continue
-            val computedHash = sha256(value)
-            val providedHash = entry.key
-            require(computedHash == providedHash) {
-                val bytes = value.encodeToByteArray(throwOnInvalidSequence = true)
-                "Provided hash $providedHash doesn't match the computed hash $computedHash for value: $value\n    Value as ByteArray$bytes"
-            }
+            checkObjectHash(entry.key, entry.value)
+        }
+    }
+
+    fun checkObjectHash(providedHash: String, value: String?) {
+        if (value == null) return
+        if (!isSha256(providedHash)) return
+        val computedHash = sha256(value)
+        require(computedHash == providedHash) {
+            val bytes = value.encodeToByteArray(throwOnInvalidSequence = true)
+            "Provided hash $providedHash doesn't match the computed hash $computedHash for value: $value\n    Value as ByteArray$bytes"
         }
     }
 }

model-server-api/build.gradle.kts

@@ -26,6 +26,8 @@ kotlin {
                 implementation(kotlin("stdlib-common"))
                 implementation(libs.kotlin.logging)
                 implementation(libs.kotlin.serialization.json)
+                api(libs.kotlin.coroutines.core)
+                api(libs.ktor.serialization)
 //                implementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:$kotlinCoroutinesVersion")
 
 //                implementation("io.ktor:ktor-client-core:$ktorVersion")

model-server-api/src/commonMain/kotlin/org/modelix/model/server/api/v2/VersionDelta.kt

@@ -13,6 +13,10 @@
  */
 package org.modelix.model.server.api.v2
 
+import io.ktor.http.ContentType
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.asFlow
+import kotlinx.coroutines.flow.emptyFlow
 import kotlinx.serialization.Serializable
 
 @Serializable
@@ -23,3 +27,34 @@ class VersionDelta(
     val objects: Set<String> = emptySet(),
     val objectsMap: Map<String, String> = emptyMap(),
 )
+
+fun VersionDelta.asStream(): VersionDeltaStream {
+    require(objects.isEmpty()) { "Legacy serialization not supported" }
+    return VersionDeltaStream(
+        versionHash = versionHash,
+        objectsFlow = null,
+        objectsSequence = objectsMap.asSequence().map { it.key to it.value },
+    )
+}
+
+class VersionDeltaStream(
+    val versionHash: String,
+    val objectsFlow: Flow<Pair<String, String>>? = null,
+    val objectsSequence: Sequence<Pair<String, String>>? = null,
+) {
+    companion object {
+        val CONTENT_TYPE = ContentType("application", "x-modelix-objects")
+    }
+
+    fun getObjectsAsFlow(): Flow<Pair<String, String>> {
+        return objectsFlow ?: objectsSequence?.asFlow() ?: emptyFlow()
+    }
+}
+
+suspend fun <K, V> Flow<Pair<K, V>>.toMap(): Map<K, V> {
+    val map = LinkedHashMap<K, V>()
+    collect {
+        map[it.first] = it.second
+    }
+    return map
+}