feat(model-server): new environment variables for warming up the ModelQL cache

It ensures that the latest model version on a branch is loaded into memory so that also the first
ModelQL query is fast.

Author: slisson

model-datastructure/build.gradle.kts

@@ -55,6 +55,7 @@ kotlin {
 //                implementation("org.json:json:20230618")
                 implementation(libs.trove4j)
                 implementation(libs.apache.commons.collections)
+                implementation(libs.kotlin.coroutines.core)
 //
 //                implementation("com.google.oauth-client:google-oauth-client:1.34.1")
 //                implementation("com.google.oauth-client:google-oauth-client-jetty:1.34.1")

model-datastructure/src/jvmMain/kotlin/org/modelix/model/InMemoryModel.kt

@@ -18,6 +18,10 @@ package org.modelix.model
 
 import gnu.trove.map.TLongObjectMap
 import gnu.trove.map.hash.TLongObjectHashMap
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.launch
 import org.modelix.model.api.ConceptReference
 import org.modelix.model.api.IBranch
 import org.modelix.model.api.IConcept
@@ -44,6 +48,42 @@ import kotlin.time.DurationUnit
 
 private val LOG = mu.KotlinLogging.logger { }
 
+class InMemoryModelLoader(val model: IncrementalInMemoryModel) {
+    private val coroutineScope = CoroutineScope(Dispatchers.IO)
+    private var modelLoadingJob: Job? = null
+
+    /**
+     * Should be called repeatedly by a readiness probe until it returns true.
+     *
+     * @return true if the model is done loading
+     */
+    @Synchronized
+    fun loadModelAsync(tree: CLTree): Boolean {
+        if (model.getLoadedModel()?.loadedMapRef?.getHash() == tree.nodesMap!!.hash) return true
+        if (modelLoadingJob?.isActive != true) {
+            modelLoadingJob = coroutineScope.launch {
+                try {
+                    model.getModel(tree)
+                } catch (ex: Throwable) {
+                    LOG.error(ex) { "Failed loading model ${tree.hash}" }
+                }
+            }
+        }
+        return false
+    }
+}
+
+class InMemoryModels {
+    private val models = HashMap<String, InMemoryModelLoader>()
+
+    @Synchronized
+    fun getModel(id: String) = models.getOrPut(id) { InMemoryModelLoader(IncrementalInMemoryModel()) }
+
+    fun getModel(tree: CLTree) = getModel(tree.getId()).model.getModel(tree)
+
+    fun loadModelAsync(tree: CLTree) = getModel(tree.getId()).loadModelAsync(tree)
+}
+
 class IncrementalInMemoryModel {
     private var lastModel: InMemoryModel? = null
 
@@ -58,6 +98,8 @@ class IncrementalInMemoryModel {
         lastModel = newModel
         return newModel
     }
+
+    fun getLoadedModel() = lastModel
 }
 
 class InMemoryModel private constructor(

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

@@ -44,6 +44,7 @@ import org.modelix.authorization.getUserName
 import org.modelix.authorization.requiresPermission
 import org.modelix.authorization.toKeycloakScope
 import org.modelix.model.lazy.BranchReference
+import org.modelix.model.lazy.RepositoryId
 import org.modelix.model.persistent.HashUtil
 import org.modelix.model.server.store.IStoreClient
 import org.modelix.model.server.store.pollEntry
@@ -93,6 +94,20 @@ class KeyValueLikeModelServer(val repositoriesManager: RepositoriesManager) {
     private fun Application.modelServerModule() {
         routing {
             get<Paths.getHealth> {
+                // eagerly load model into memory to speed up ModelQL queries
+                val branchRef = System.getenv("MODELIX_SERVER_MODELQL_WARMUP_REPOSITORY")?.let { RepositoryId(it) }
+                    ?.getBranchReference(System.getenv("MODELIX_SERVER_MODELQL_WARMUP_BRANCH"))
+                if (branchRef != null) {
+                    val version = repositoriesManager.getVersion(branchRef)
+                    if (!repositoriesManager.inMemoryModels.loadModelAsync(version!!.getTree())) {
+                        call.respondText(
+                            status = HttpStatusCode.ServiceUnavailable,
+                            text = "Waiting for version $version to be loaded into memory",
+                        )
+                        return@get
+                    }
+                }
+
                 if (isHealthy()) {
                     call.respondText(text = "healthy", contentType = ContentType.Text.Plain, status = HttpStatusCode.OK)
                 } else {

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

@@ -47,7 +47,6 @@ import kotlinx.serialization.encodeToString
 import kotlinx.serialization.json.Json
 import org.modelix.api.public.Paths
 import org.modelix.authorization.getUserName
-import org.modelix.model.IncrementalInMemoryModel
 import org.modelix.model.api.ITree
 import org.modelix.model.api.PBranch
 import org.modelix.model.api.TreePointer
@@ -266,25 +265,22 @@ class ModelReplicationServer(val repositoriesManager: RepositoriesManager) {
             TODO()
         }
 
-        val inMemoryModel = IncrementalInMemoryModel()
-        post<Paths.postRepositoryBranchQuery> {
-            fun ApplicationCall.repositoryId() = RepositoryId(parameters["repository"]!!)
-            fun PipelineContext<Unit, ApplicationCall>.repositoryId() = call.repositoryId()
-
-            fun ApplicationCall.branchRef() = repositoryId().getBranchReference(parameters["branch"]!!)
-            fun PipelineContext<Unit, ApplicationCall>.branchRef() = call.branchRef()
-
-            val branchRef = branchRef()
+        post<Paths.postRepositoryBranchQuery> { parameters ->
+            val branchRef = RepositoryId(parameters.repository).getBranchReference(parameters.branch)
             val version = repositoriesManager.getVersion(branchRef)
             LOG.trace("Running query on {} @ {}", branchRef, version)
             val initialTree = version!!.getTree()
-            val branch = OTBranch(PBranch(initialTree, repositoriesManager.client.idGenerator), repositoriesManager.client.idGenerator, repositoriesManager.client.storeCache)
+            val branch = OTBranch(
+                PBranch(initialTree, repositoriesManager.client.idGenerator),
+                repositoriesManager.client.idGenerator,
+                repositoriesManager.client.storeCache,
+            )
 
             ModelQLServer.handleCall(call, { writeAccess ->
                 if (writeAccess) {
                     branch.getRootNode() to branch.getArea()
                 } else {
-                    val model = inMemoryModel.getModel(initialTree)
+                    val model = repositoriesManager.inMemoryModels.getModel(initialTree)
                     model.getNode(ITree.ROOT_ID) to model.getArea()
                 }
             }, {

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

@@ -21,6 +21,7 @@ import kotlinx.coroutines.flow.map
 import kotlinx.datetime.Clock
 import org.apache.commons.collections4.map.LRUMap
 import org.modelix.model.IKeyValueStore
+import org.modelix.model.InMemoryModels
 import org.modelix.model.VersionMerger
 import org.modelix.model.api.IBranch
 import org.modelix.model.api.IReadTransaction
@@ -51,6 +52,7 @@ class RepositoriesManager(val client: LocalModelClient) {
     private val store: IStoreClient get() = client.store
     private val kvStore: IKeyValueStore get() = client.asyncStore
     private val objectStore: IDeserializingKeyValueStore get() = client.storeCache
+    val inMemoryModels = InMemoryModels()
 
     fun generateClientId(repositoryId: RepositoryId): Long {
         return client.store.generateId("$KEY_PREFIX:${repositoryId.id}:clientId")