fix(model-server): InMemoryModel was blocking the ktor request handler thread pool

The model was loaded by a synchronized method and all HTTP requests were
blocking the thread while waiting for the result. Since the default pool
is constrained in size it was running out of thread for handling other
request such as the /health endpoint.

Replaced the synchronized method by coroutine features that don't block
the threads.

Author: slisson

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

@@ -18,10 +18,12 @@ package org.modelix.model
 
 import gnu.trove.map.TLongObjectMap
 import gnu.trove.map.hash.TLongObjectHashMap
+import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Deferred
 import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.Job
-import kotlinx.coroutines.launch
+import kotlinx.coroutines.async
+import kotlinx.coroutines.cancel
 import org.modelix.model.api.ConceptReference
 import org.modelix.model.api.IBranch
 import org.modelix.model.api.IConcept
@@ -42,52 +44,66 @@ import org.modelix.model.lazy.NonCachingObjectStore
 import org.modelix.model.persistent.CPHamtNode
 import org.modelix.model.persistent.CPNode
 import org.modelix.model.persistent.CPNodeRef
+import java.util.Collections
 import kotlin.system.measureTimeMillis
 import kotlin.time.Duration.Companion.milliseconds
 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}" }
+class InMemoryModelLoader(val incrementalModel: IncrementalInMemoryModel, val coroutineScope: CoroutineScope) {
+    private val treeHash2modelLoadJob = Collections.synchronizedMap(HashMap<String, Deferred<InMemoryModel>>())
+
+    fun getModel(tree: CLTree): Deferred<InMemoryModel> {
+        val loadedModel = incrementalModel.getLoadedModel()
+        if (loadedModel != null && loadedModel.loadedMapRef.getHash() == tree.nodesMap?.hash) return CompletableDeferred(loadedModel)
+
+        return synchronized(treeHash2modelLoadJob) {
+            val activeJobs = treeHash2modelLoadJob.values.toList()
+            val loadJob = treeHash2modelLoadJob.getOrPut(tree.hash) {
+                coroutineScope.async {
+                    // There should only be one active loading job, because we want to reuse as much data as possible
+                    // from a previously loaded model, so we have to wait for its completion.
+                    // This also limits the number of thread used from the IO dispatcher.
+                    activeJobs.forEach { it.join() }
+
+                    // This is a long-running method that should be executed only once for a new tree version.
+                    // It's executed on the IO dispatcher, because it's not a suspendable function and blocks
+                    // the thread.
+                    incrementalModel.getModel(tree)
                 }
             }
+
+            // cleanup finished jobs
+            treeHash2modelLoadJob -= treeHash2modelLoadJob.entries.filter { !it.value.isActive }.map { it.key }.toSet()
+
+            loadJob
         }
-        return false
     }
 }
 
 class InMemoryModels {
-    private val models = HashMap<String, InMemoryModelLoader>()
+    private val coroutineScope = CoroutineScope(Dispatchers.IO)
+    private val branchId2modelLoader = Collections.synchronizedMap(HashMap<String, InMemoryModelLoader>())
 
-    @Synchronized
-    fun getModel(id: String) = models.getOrPut(id) { InMemoryModelLoader(IncrementalInMemoryModel()) }
+    fun dispose() {
+        coroutineScope.cancel("disposed")
+    }
 
-    fun getModel(tree: CLTree) = getModel(tree.getId()).model.getModel(tree)
+    private fun getModelLoader(branchId: String): InMemoryModelLoader {
+        return synchronized(branchId2modelLoader) {
+            branchId2modelLoader.getOrPut(branchId) { InMemoryModelLoader(IncrementalInMemoryModel(), coroutineScope) }
+        }
+    }
 
-    fun loadModelAsync(tree: CLTree) = getModel(tree.getId()).loadModelAsync(tree)
+    fun getModel(tree: CLTree): Deferred<InMemoryModel> {
+        return getModelLoader(tree.getId()).getModel(tree)
+    }
 }
 
 class IncrementalInMemoryModel {
     private var lastModel: InMemoryModel? = null
 
-    @Synchronized
     fun getModel(tree: CLTree): InMemoryModel {
         val reusable = lastModel?.takeIf { it.branchId == tree.getId() }
         val newModel = if (reusable == null) {

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

@@ -99,7 +99,7 @@ class KeyValueLikeModelServer(val repositoriesManager: RepositoriesManager) {
                     ?.getBranchReference(System.getenv("MODELIX_SERVER_MODELQL_WARMUP_BRANCH"))
                 if (branchRef != null) {
                     val version = repositoriesManager.getVersion(branchRef)
-                    if (!repositoriesManager.inMemoryModels.loadModelAsync(version!!.getTree())) {
+                    if (repositoriesManager.inMemoryModels.getModel(version!!.getTree()).isActive) {
                         call.respondText(
                             status = HttpStatusCode.ServiceUnavailable,
                             text = "Waiting for version $version to be loaded into memory",

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

@@ -280,7 +280,7 @@ class ModelReplicationServer(val repositoriesManager: RepositoriesManager) {
                 if (writeAccess) {
                     branch.getRootNode() to branch.getArea()
                 } else {
-                    val model = repositoriesManager.inMemoryModels.getModel(initialTree)
+                    val model = repositoriesManager.inMemoryModels.getModel(initialTree).await()
                     model.getNode(ITree.ROOT_ID) to model.getArea()
                 }
             }, {

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

@@ -56,6 +56,12 @@ class RepositoriesManager(val client: LocalModelClient) {
     private val objectStore: IDeserializingKeyValueStore get() = client.storeCache
     val inMemoryModels = InMemoryModels()
 
+    fun dispose() {
+        // TODO find instance creations and add a dispose() call if needed. Whoever creates an instance is responsible
+        //      for its lifecycle.
+        inMemoryModels.dispose()
+    }
+
     fun generateClientId(repositoryId: RepositoryId): Long {
         return client.store.generateId("$KEY_PREFIX:${repositoryId.id}:clientId")
     }

modelql-server/src/main/kotlin/org/modelix/modelql/server/ModelQLServer.kt

@@ -80,7 +80,7 @@ class ModelQLServer private constructor(val rootNodeProvider: () -> INode?, val
             handleCall(call, { rootNode to area }, {})
         }
 
-        suspend fun handleCall(call: ApplicationCall, input: (write: Boolean) -> Pair<INode, IArea>, afterQueryExecution: () -> Unit = {}) {
+        suspend fun handleCall(call: ApplicationCall, input: suspend (write: Boolean) -> Pair<INode, IArea>, afterQueryExecution: () -> Unit = {}) {
             try {
                 val serializedQuery = call.receiveText()
                 val json = UntypedModelQL.json