perf(model-server): avoid converting object maps to object flows and vice versa

Author: Oleksandr Dzhychko

model-server/src/main/kotlin/org/modelix/model/server/handlers/ModelReplicationServer.kt

@@ -39,7 +39,6 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.flow
-import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.onEmpty
 import kotlinx.coroutines.flow.withIndex
 import kotlinx.coroutines.withContext
@@ -61,7 +60,6 @@ import org.modelix.model.operations.OTBranch
 import org.modelix.model.persistent.HashUtil
 import org.modelix.model.server.api.v2.VersionDelta
 import org.modelix.model.server.api.v2.VersionDeltaStream
-import org.modelix.model.server.api.v2.toMap
 import org.modelix.model.server.store.IStoreClient
 import org.modelix.model.server.store.LocalModelClient
 import org.modelix.modelql.server.ModelQLServer
@@ -231,9 +229,8 @@ class ModelReplicationServer(val repositoriesManager: RepositoriesManager) {
                     val delta = VersionDelta(
                         newVersionHash,
                         lastVersionHash,
-                        objectsMap = repositoriesManager.computeDelta(newVersionHash, lastVersionHash).toMap(),
+                        objectsMap = repositoriesManager.computeDelta(newVersionHash, lastVersionHash).asMap(),
                     )
-                    delta.checkObjectHashes()
                     send(Json.encodeToString(delta))
                     lastVersionHash = newVersionHash
                 }
@@ -370,16 +367,14 @@ class ModelReplicationServer(val repositoriesManager: RepositoriesManager) {
         val delta = VersionDelta(
             versionHash,
             baseVersionHash,
-            objectsMap = repositoriesManager.computeDelta(versionHash, baseVersionHash).toMap(),
+            objectsMap = repositoriesManager.computeDelta(versionHash, baseVersionHash).asMap(),
         )
-        delta.checkObjectHashes()
         respond(delta)
     }
 
     private suspend fun ApplicationCall.respondDeltaAsObjectStream(versionHash: String, baseVersionHash: String?, plainText: Boolean) {
         respondTextWriter(contentType = if (plainText) ContentType.Text.Plain else VersionDeltaStream.CONTENT_TYPE) {
-            repositoriesManager.computeDelta(versionHash, baseVersionHash)
-                .checkObjectHashes()
+            repositoriesManager.computeDelta(versionHash, baseVersionHash).asFlow()
                 .flatten()
                 .withSeparator("\n")
                 .onEmpty { emit(versionHash) }
@@ -410,7 +405,3 @@ private fun Flow<String>.withSeparator(separator: String) = flow {
         emit(it)
     }
 }
-
-private fun <V : String?> Flow<Pair<String, V>>.checkObjectHashes(): Flow<Pair<String, V>> {
-    return onEach { HashUtil.checkObjectHash(it.first, it.second) }
-}

model-server/src/main/kotlin/org/modelix/model/server/handlers/RepositoriesManager.kt

@@ -16,8 +16,8 @@ package org.modelix.model.server.handlers
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.asFlow
 import kotlinx.coroutines.flow.channelFlow
-import kotlinx.coroutines.flow.emptyFlow
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.runBlocking
 import kotlinx.datetime.Clock
 import org.apache.commons.collections4.map.LRUMap
@@ -30,6 +30,7 @@ import org.modelix.model.api.ITree
 import org.modelix.model.api.IdGeneratorDummy
 import org.modelix.model.api.PBranch
 import org.modelix.model.api.runSynchronized
+import org.modelix.model.client2.checkObjectHashes
 import org.modelix.model.lazy.BranchReference
 import org.modelix.model.lazy.CLTree
 import org.modelix.model.lazy.CLVersion
@@ -39,6 +40,8 @@ import org.modelix.model.lazy.RepositoryId
 import org.modelix.model.lazy.computeDelta
 import org.modelix.model.metameta.MetaModelBranch
 import org.modelix.model.persistent.CPVersion
+import org.modelix.model.persistent.HashUtil
+import org.modelix.model.server.api.v2.toMap
 import org.modelix.model.server.store.IStoreClient
 import org.modelix.model.server.store.LocalModelClient
 import org.modelix.model.server.store.pollEntry
@@ -255,39 +258,12 @@ class RepositoriesManager(val client: LocalModelClient) {
             ?: throw IllegalStateException("No version found for branch '${branch.branchName}' in repository '${branch.repositoryId}'")
     }
 
-    private val deltaCache = LRUMap<Pair<String, String?>, SoftReference<Lazy<Map<String, String>>>>(10)
-    fun computeDelta(versionHash: String, baseVersionHash: String?): Flow<Pair<String, String>> {
-        if (versionHash == baseVersionHash) return emptyFlow()
+    private val deltaCache = LRUMap<Pair<String, String?>, SoftReference<Lazy<ObjectDataMap>>>(10)
+    fun computeDelta(versionHash: String, baseVersionHash: String?): ObjectData {
+        if (versionHash == baseVersionHash) return ObjectData.empty
         if (baseVersionHash == null) {
             // no need to cache anything if there is no delta computation happening
-
-            return channelFlow {
-                val version = CLVersion(versionHash, objectStore)
-                // Use a bulk query to make as few request to the underlying store as possible.
-                val bulkQuery = objectStore.newBulkQuery()
-                // It is unsatisfactory that we have to keep already emitted hashes in memory.
-                // But without changing the underlying model,
-                // we have to do this to not emit objects more than once.
-                val seenHashes = mutableSetOf<String>()
-                fun emitObjects(entry: KVEntryReference<*>) {
-                    if (seenHashes.contains(entry.getHash())) return
-                    seenHashes.add(entry.getHash())
-                    bulkQuery.get(entry).onSuccess {
-                        val value = checkNotNull(it) { "No value received for ${entry.getHash()}" }
-                        // Use `send` instead of `trySend`,
-                        // because `trySend` fails if the channel capacity is full.
-                        // This might happen if the data is produced faster than consumed.
-                        // A better solution would be to have bulk queries which itself are asynchronous
-                        // but doing that needs more consideration.
-                        runBlocking { channel.send(entry.getHash() to value.serialize()) }
-                        for (referencedEntry in value.getReferencedEntries()) {
-                            emitObjects(referencedEntry)
-                        }
-                    }
-                }
-                emitObjects(KVEntryReference(versionHash, CPVersion.DESERIALIZER))
-                bulkQuery.process()
-            }
+            return allObjectDataAsFlow(versionHash)
         }
 
         return runSynchronized(deltaCache) {
@@ -297,9 +273,47 @@ class RepositoriesManager(val client: LocalModelClient) {
                 // SoftReference because deltas can be very large
                 val version = CLVersion(versionHash, client.storeCache)
                 val baseVersion = CLVersion(baseVersionHash, client.storeCache)
-                version.computeDelta(baseVersion)
+                val objectsMap = version.computeDelta(baseVersion)
+                ObjectDataMap(objectsMap)
             }.also { deltaCache[key] = SoftReference(it) }
-        }.value.entries.asFlow().map { it.toPair() }
+        }.value
+    }
+
+    private fun allObjectDataAsFlow(versionHash: String): ObjectDataFlow {
+        val hashObjectFlow = channelFlow {
+            val version = CLVersion(versionHash, objectStore)
+            // Use a bulk query to make as few request to the underlying store as possible.
+            val bulkQuery = objectStore.newBulkQuery()
+            // It is unsatisfactory that we have to keep already emitted hashes in memory.
+            // But without changing the underlying model,
+            // we have to do this to not emit objects more than once.
+            val seenHashes = mutableSetOf<String>()
+            fun emitObjects(entry: KVEntryReference<*>) {
+                if (seenHashes.contains(entry.getHash())) return
+                seenHashes.add(entry.getHash())
+                bulkQuery.get(entry).onSuccess {
+                    val value = checkNotNull(it) { "No value received for ${entry.getHash()}" }
+                    // Use `send` instead of `trySend`,
+                    // because `trySend` fails if the channel capacity is full.
+                    // This might happen if the data is produced faster than consumed.
+                    // A better solution would be to have bulk queries which itself are asynchronous
+                    // but doing that needs more consideration.
+                    runBlocking {
+                        // Maybe we should avoid Flow<Pair<String, String>> and use Flow<String>.
+                        // This needs profiling.
+                        channel.send(entry.getHash() to value.serialize())
+                    }
+                    for (referencedEntry in value.getReferencedEntries()) {
+                        emitObjects(referencedEntry)
+                    }
+                }
+            }
+            emitObjects(KVEntryReference(versionHash, CPVersion.DESERIALIZER))
+            bulkQuery.process()
+        }
+        val checkedHashObjectFlow = hashObjectFlow.checkObjectHashes()
+        val objectData = ObjectDataFlow(checkedHashObjectFlow)
+        return objectData
     }
 
     private fun branchKey(branch: BranchReference): String {
@@ -342,3 +356,29 @@ class RepositoriesManager(val client: LocalModelClient) {
 }
 
 class RepositoryAlreadyExistsException(val name: String) : IllegalStateException("Repository '$name' already exists")
+
+sealed interface ObjectData {
+    suspend fun asMap(): Map<String, String>
+    fun asFlow(): Flow<Pair<String, String>>
+
+    companion object {
+        val empty = ObjectDataMap(emptyMap())
+    }
+}
+
+class ObjectDataMap(private val byHashObjects: Map<String, String>) : ObjectData {
+    init {
+        HashUtil.checkObjectHashes(byHashObjects)
+    }
+    override suspend fun asMap(): Map<String, String> = byHashObjects
+    override fun asFlow(): Flow<Pair<String, String>> = byHashObjects.entries.asFlow().map { it.toPair() }
+}
+
+class ObjectDataFlow(private val hashObjectFlow: Flow<Pair<String, String>>) : ObjectData {
+    override suspend fun asMap(): Map<String, String> = hashObjectFlow.toMap()
+    override fun asFlow(): Flow<Pair<String, String>> = hashObjectFlow
+}
+
+private fun Flow<Pair<String, String>>.checkObjectHashes(): Flow<Pair<String, String>> {
+    return onEach { HashUtil.checkObjectHash(it.first, it.second) }
+}