feat(model-client): lazy loading support for IModelClientV2

Author: Sascha Lisson

model-api/src/commonMain/kotlin/org/modelix/model/api/INode.kt

@@ -314,6 +314,10 @@ interface IReplaceableNode : INode {
 @Deprecated("Use .key(INode), .key(IBranch), .key(ITransaction) or .key(ITree)")
 fun IRole.key(): String = RoleAccessContext.getKey(this)
 fun IRole.key(node: INode): String = if (node.usesRoleIds()) getUID() else getSimpleName()
+fun IChildLink.key(node: INode): String? = when (this) {
+    is NullChildLink -> null
+    else -> (this as IRole).key(node)
+}
 fun INode.usesRoleIds(): Boolean = if (this is INodeEx) this.usesRoleIds() else false
 fun INode.getChildren(link: IChildLink): Iterable<INode> = if (this is INodeEx) getChildren(link) else getChildren(link.key(this))
 fun INode.moveChild(role: IChildLink, index: Int, child: INode): Unit = if (this is INodeEx) moveChild(role, index, child) else moveChild(role.key(this), index, child)
@@ -433,3 +437,5 @@ fun INode.getContainmentLink() = if (this is INodeEx) {
 fun INode.getRoot(): INode = parent?.getRoot() ?: this
 fun INode.isInstanceOf(superConcept: IConcept?): Boolean = concept.let { it != null && it.isSubConceptOf(superConcept) }
 fun INode.isInstanceOfSafe(superConcept: IConcept): Boolean = tryGetConcept()?.isSubConceptOf(superConcept) ?: false
+
+fun INode.addNewChild(role: IChildLink, index: Int) = addNewChild(role, index, null as IConceptReference?)

model-client/build.gradle.kts

@@ -21,6 +21,7 @@ val ktorVersion: String by rootProject
 val kotlinxSerializationVersion: String by rootProject
 
 kotlin {
+    jvmToolchain(11)
     jvm()
     js(IR) {
         browser {

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

@@ -86,4 +86,8 @@ interface IModelClientV2 {
     suspend fun <R> query(branch: BranchReference, body: (IMonoStep<INode>) -> IMonoStep<R>): R
 
     suspend fun <R> query(repositoryId: RepositoryId, versionHash: String, body: (IMonoStep<INode>) -> IMonoStep<R>): R
+
+    suspend fun getObjects(repository: RepositoryId, keys: Sequence<String>): Map<String, String>
+
+    suspend fun pushObjects(repository: RepositoryId, objects: Sequence<Pair<String, String>>)
 }

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

@@ -265,6 +265,29 @@ class ModelClientV2(
         }
     }
 
+    override suspend fun getObjects(repository: RepositoryId, keys: Sequence<String>): Map<String, String> {
+        val response = httpClient.post {
+            url {
+                takeFrom(baseUrl)
+                appendPathSegments("repositories", repository.id, "objects", "getAll")
+            }
+            setBody(keys.joinToString("\n"))
+        }
+
+        val content = response.bodyAsChannel()
+        val objects = HashMap<String, String>()
+        while (true) {
+            val key = checkNotNull(content.readUTF8Line()) { "Empty line expected at the end of the stream" }
+            if (key == "") {
+                check(content.readUTF8Line() == null) { "Empty line is only allowed at the end of the stream" }
+                break
+            }
+            val value = checkNotNull(content.readUTF8Line()) { "Object missing for hash $key" }
+            objects[key] = value
+        }
+        return objects
+    }
+
     override suspend fun push(branch: BranchReference, version: IVersion, baseVersion: IVersion?): IVersion {
         LOG.debug { "${clientId.toString(16)}.push($branch, $version, $baseVersion)" }
         require(version is CLVersion)
@@ -274,7 +297,7 @@ class ModelClientV2(
         HashUtil.checkObjectHashes(objects)
         val delta = if (objects.size > 1000) {
             // large HTTP requests and large Json objects don't scale well
-            uploadObjects(branch.repositoryId, objects.asSequence().map { it.key to it.value })
+            pushObjects(branch.repositoryId, objects.asSequence().map { it.key to it.value })
             VersionDelta(version.getContentHash(), null)
         } else {
             VersionDelta(version.getContentHash(), null, objectsMap = objects)
@@ -292,7 +315,7 @@ class ModelClientV2(
         }
     }
 
-    private suspend fun uploadObjects(repository: RepositoryId, objects: Sequence<Pair<String, String>>) {
+    override suspend fun pushObjects(repository: RepositoryId, objects: Sequence<Pair<String, String>>) {
         LOG.debug { "${clientId.toString(16)}.pushObjects($repository)" }
         objects.chunked(100_000).forEach { unsortedChunk ->
             // Entries are sorted to avoid deadlocks on the server side between transactions.

model-client/src/jvmMain/kotlin/org/modelix/model/client2/LazyLoading.kt

@@ -0,0 +1,71 @@
+/*
+ * Copyright (c) 2024.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.modelix.model.client2
+
+import kotlinx.coroutines.runBlocking
+import org.modelix.model.IKeyListener
+import org.modelix.model.IKeyValueStore
+import org.modelix.model.IVersion
+import org.modelix.model.lazy.BranchReference
+import org.modelix.model.lazy.CLVersion
+import org.modelix.model.lazy.ObjectStoreCache
+import org.modelix.model.lazy.RepositoryId
+
+fun IModelClientV2.lazyLoadVersion(repositoryId: RepositoryId, versionHash: String, cacheSize: Int = 100_000): IVersion {
+    val store = ObjectStoreCache(ModelClientAsStore(this, repositoryId), cacheSize)
+    return CLVersion.loadFromHash(versionHash, store)
+}
+
+suspend fun IModelClientV2.lazyLoadVersion(branchRef: BranchReference, cacheSize: Int = 100_000): IVersion {
+    return lazyLoadVersion(branchRef.repositoryId, pullHash(branchRef), cacheSize)
+}
+
+class ModelClientAsStore(val client: IModelClientV2, val repositoryId: RepositoryId) : IKeyValueStore {
+    override fun get(key: String): String? {
+        return getAll(listOf(key))[key]
+    }
+
+    override fun put(key: String, value: String?) {
+        TODO("Not yet implemented")
+    }
+
+    override fun getAll(keys: Iterable<String>): Map<String, String?> {
+        return runBlocking {
+            client.getObjects(repositoryId, keys.asSequence())
+        }
+    }
+
+    override fun putAll(entries: Map<String, String?>) {
+        TODO("Not yet implemented")
+    }
+
+    override fun prefetch(key: String) {
+        TODO("Not yet implemented")
+    }
+
+    override fun listen(key: String, listener: IKeyListener) {
+        TODO("Not yet implemented")
+    }
+
+    override fun removeListener(key: String, listener: IKeyListener) {
+        TODO("Not yet implemented")
+    }
+
+    override fun getPendingSize(): Int {
+        TODO("Not yet implemented")
+    }
+}

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

@@ -17,9 +17,10 @@ package org.modelix.model.lazy
 
 import org.modelix.model.IKeyValueStore
 import org.modelix.model.createLRUMap
+import kotlin.jvm.JvmOverloads
 
-class ObjectStoreCache(override val keyValueStore: IKeyValueStore) : IDeserializingKeyValueStore {
-    private val cache: MutableMap<String?, Any> = createLRUMap(100000)
+class ObjectStoreCache @JvmOverloads constructor(override val keyValueStore: IKeyValueStore, cacheSize: Int = 100_000) : IDeserializingKeyValueStore {
+    private val cache: MutableMap<String?, Any> = createLRUMap(cacheSize)
 
     override fun <T> getAll(hashes_: Iterable<String>, deserializer: (String, String) -> T): Iterable<T> {
         val hashes = hashes_.toList()

model-server-openapi/specifications/model-server-v2.yaml

@@ -75,6 +75,18 @@ paths:
           $ref: '#/components/responses/200json'
         default:
           $ref: '#/components/responses/GeneralError'
+  /repositories/{repository}/objects/getAll:
+    post:
+      operationId: postRepositoryObjectsGetAll
+      parameters:
+        - name: repository
+          in: "path"
+          required: true
+          schema:
+            type: string
+      responses:
+        "200":
+          $ref: '#/components/responses/200'
   /repositories/{repository}/branches:
     get:
       operationId: getRepositoryBranches

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

@@ -25,6 +25,7 @@ import io.ktor.server.request.receiveStream
 import io.ktor.server.response.respond
 import io.ktor.server.response.respondBytesWriter
 import io.ktor.server.response.respondText
+import io.ktor.server.response.respondTextWriter
 import io.ktor.server.routing.route
 import io.ktor.server.routing.routing
 import io.ktor.util.cio.use
@@ -200,6 +201,23 @@ class ModelReplicationServer(
         }
     }
 
+    override suspend fun PipelineContext<Unit, ApplicationCall>.postRepositoryObjectsGetAll(repository: String) {
+        val keys = call.receiveStream().bufferedReader().use { reader ->
+            reader.lineSequence().toHashSet()
+        }
+        val objects = withContext(Dispatchers.IO) { modelClient.store.getAll(keys) }
+        call.respondTextWriter(contentType = VersionDeltaStream.CONTENT_TYPE) {
+            objects.forEach {
+                append(it.key)
+                append("\n")
+                append(it.value)
+                append("\n")
+            }
+            // additional empty line indicates end of stream and can be used to verify completeness of data transfer
+            append("\n")
+        }
+    }
+
     override suspend fun PipelineContext<Unit, ApplicationCall>.pollRepositoryBranchHash(
         repository: String,
         branch: String,

model-server/src/test/kotlin/org/modelix/model/server/LazyLoadingTest.kt

@@ -0,0 +1,114 @@
+/*
+ * Copyright (c) 2024.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.modelix.model.server
+
+import io.ktor.server.testing.ApplicationTestBuilder
+import io.ktor.server.testing.testApplication
+import org.modelix.authorization.installAuthentication
+import org.modelix.model.api.INode
+import org.modelix.model.api.NullChildLink
+import org.modelix.model.api.TreePointer
+import org.modelix.model.api.addNewChild
+import org.modelix.model.api.getDescendants
+import org.modelix.model.api.getRootNode
+import org.modelix.model.client2.lazyLoadVersion
+import org.modelix.model.client2.runWrite
+import org.modelix.model.lazy.RepositoryId
+import org.modelix.model.persistent.MapBasedStore
+import org.modelix.model.server.api.v2.ObjectHash
+import org.modelix.model.server.api.v2.SerializedObject
+import org.modelix.model.server.handlers.IdsApiImpl
+import org.modelix.model.server.handlers.ModelReplicationServer
+import org.modelix.model.server.store.InMemoryStoreClient
+import kotlin.test.Test
+import kotlin.test.assertTrue
+
+class LazyLoadingTest {
+
+    private lateinit var statistics: StoreClientWithStatistics
+
+    private fun runTest(block: suspend ApplicationTestBuilder.() -> Unit) = testApplication {
+        application {
+            installAuthentication(unitTestMode = true)
+            installDefaultServerPlugins()
+            statistics = StoreClientWithStatistics(InMemoryStoreClient())
+            ModelReplicationServer(statistics).init(this)
+            IdsApiImpl(statistics).init(this)
+        }
+        block()
+    }
+
+    private fun assertRequestCount(atLeast: Long, body: () -> Unit): Long {
+        val requestCount = measureRequests(body)
+        assertTrue(requestCount >= atLeast, "At least $atLeast requests expected, but was $requestCount")
+        return requestCount
+    }
+
+    private fun measureRequests(body: () -> Unit): Long {
+        val before = statistics.getTotalRequests()
+        body()
+        val after = statistics.getTotalRequests()
+        val requestCount = after - before
+        println("Requests: $requestCount")
+        return requestCount
+    }
+
+    @Test
+    fun `model data is loaded on demand`() = runTest {
+        // After optimizing the lazy loading to send less (but bigger) requests, this test might fail.
+        // Just update the model size, cache size and expected request count to fix it.
+
+        val client = createModelClient()
+        val branchRef = RepositoryId("my-repo").getBranchReference()
+        client.runWrite(branchRef) {
+            fun createNodes(parentNode: INode, numberOfNodes: Int) {
+                if (numberOfNodes == 0) return
+                if (numberOfNodes == 1) {
+                    parentNode.addNewChild(NullChildLink, 0)
+                    return
+                }
+                val subtreeSize1 = numberOfNodes / 2
+                val subtreeSize2 = numberOfNodes - subtreeSize1
+                createNodes(parentNode.addNewChild(NullChildLink, 0), subtreeSize1 - 1)
+                createNodes(parentNode.addNewChild(NullChildLink, 1), subtreeSize2 - 1)
+            }
+
+            createNodes(it, 5_000)
+        }
+        val version = client.lazyLoadVersion(branchRef, cacheSize = 500)
+
+        val rootNode = TreePointer(version.getTree()).getRootNode()
+
+        // Traverse to the first leaf node. This should load some data, but not the whole model.
+        assertRequestCount(1) {
+            generateSequence(rootNode) { it.allChildren.firstOrNull() }.count()
+        }
+
+        // Traverse the whole model.
+        val requestCountFirstTraversal = assertRequestCount(10) {
+            rootNode.getDescendants(true).count()
+        }
+
+        // Traverse the whole model a second time. The model doesn't fit into the cache and some parts are already
+        // unloaded during the first traversal. The unloaded parts need to be requested again.
+        // But the navigation to the first leaf is like a warmup of the cache for the whole model traversal.
+        // The previous traversal can benefit from that, but the next one cannot and is expected to need more requests.
+        assertRequestCount(requestCountFirstTraversal + 1) {
+            rootNode.getDescendants(true).count()
+        }
+    }
+}