Merge pull request #558 from modelix/MODELIX-801-2

MODELIX-801 Parallel queries to different branches seem to create a deadlock

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