perf(bulk-model-sync): load module data into memory only on demand and release it after import

Load module data and consume module data lazily.
Do not hold references to imported module data after importing a module
so that it can be garbage collected.

Author: Oleksandr Dzhychko

bulk-model-sync-lib/src/commonMain/kotlin/org/modelix/model/sync/bulk/ModelImporter.kt

@@ -53,14 +53,13 @@ class ModelImporter(
     // Therefore, choose a map with is optimized for memory usage.
     // For the same reason store `INodeReference`s instead of `INode`s.
     // In a few cases, where we need the `INode` we can resolve it.
-    private var originalIdToExisting = MemoryEfficientMap<String, INodeReference>()
+    private val originalIdToExisting by lazy(::buildExistingIndex)
 
     // Use`INode` instead of `INodeReference` in `postponedReferences` and `nodesToRemove`
     // because we know that we will always need the `INode`s in those cases.
     // Those cases are deleting nodes and adding references to nodes.
     private val postponedReferences = mutableListOf<PostponedReference>()
     private val nodesToRemove = HashSet<INode>()
-    private var numExpectedNodes = 0
     private var currentNodeProgress = 0
     private val logger = KotlinLogging.logger {}
 
@@ -104,34 +103,51 @@ class ModelImporter(
      */
     @JvmName("importData")
     fun import(data: ModelData) {
-        INodeResolutionScope.runWithAdditionalScope(root.getArea()) {
-            logImportSize(data.root, logger)
-            logger.info { "Building indices for import..." }
-            originalIdToExisting = MemoryEfficientMap()
-            postponedReferences.clear()
-            nodesToRemove.clear()
-            numExpectedNodes = countExpectedNodes(data.root)
-            val progressReporter = ProgressReporter(numExpectedNodes.toULong(), logger)
-            currentNodeProgress = 0
-            buildExistingIndex(root)
+        importIntoNodes(sequenceOf(ExistingAndExpectedNode(root, data)))
+    }
 
-            logger.info { "Importing nodes..." }
-            data.root.originalId()?.let { originalIdToExisting[it] = root.reference }
-            syncNode(root, data.root, progressReporter)
+    /**
+     * Incrementally updates existing children of the given with specified data.
+     *
+     * @param nodeCombinationsToImport Combinations of an old existing child and the new expected data.
+     * The combinations are consumed lazily.
+     * Callers can use this to load expected data on demand.
+     */
+    fun importIntoNodes(nodeCombinationsToImport: Sequence<ExistingAndExpectedNode>) {
+        logger.info { "Building indices for import..." }
+        postponedReferences.clear()
+        nodesToRemove.clear()
 
-            logger.info { "Synchronizing references..." }
-            postponedReferences.forEach { it.setPostponedReference() }
+        nodeCombinationsToImport.forEach { nodeCombination ->
+            importIntoNode(nodeCombination.expectedNodeData, nodeCombination.existingNode)
+        }
 
-            logger.info { "Removing extra nodes..." }
-            nodesToRemove.forEach {
-                doAndPotentiallyContinueOnErrors {
-                    if (it.isValid) { // if it's invalid then it's already removed
-                        it.remove()
-                    }
+        logger.info { "Synchronizing references..." }
+        postponedReferences.forEach { it.setPostponedReference() }
+
+        logger.info { "Removing extra nodes..." }
+        nodesToRemove.forEach {
+            doAndPotentiallyContinueOnErrors {
+                if (it.isValid) { // if it's invalid then it's already removed
+                    it.remove()
                 }
             }
+        }
 
-            logger.info { "Synchronization finished." }
+        logger.info { "Synchronization finished." }
+    }
+
+    private fun importIntoNode(expectedNodeData: ModelData, existingNode: INode = root) {
+        INodeResolutionScope.runWithAdditionalScope(existingNode.getArea()) {
+            logImportSize(expectedNodeData.root, logger)
+            logger.info { "Building indices for nodes import..." }
+            currentNodeProgress = 0
+            val numExpectedNodes = countExpectedNodes(expectedNodeData.root)
+            val progressReporter = ProgressReporter(numExpectedNodes.toULong(), logger)
+
+            logger.info { "Importing nodes..." }
+            expectedNodeData.root.originalId()?.let { originalIdToExisting[it] = existingNode.reference }
+            syncNode(existingNode, expectedNodeData.root, progressReporter)
         }
     }
 
@@ -240,10 +256,12 @@ class ModelImporter(
         }
     }
 
-    private fun buildExistingIndex(root: INode) {
+    private fun buildExistingIndex(): MemoryEfficientMap<String, INodeReference> {
+        val localOriginalIdToExisting = MemoryEfficientMap<String, INodeReference>()
         root.getDescendants(true).forEach { node ->
-            node.originalId()?.let { originalIdToExisting[it] = node.reference }
+            node.originalId()?.let { localOriginalIdToExisting[it] = node.reference }
         }
+        return localOriginalIdToExisting
     }
 
     private fun syncProperties(node: INode, nodeData: NodeData) {
@@ -291,3 +309,8 @@ internal fun INode.originalId(): String? {
 internal fun NodeData.originalId(): String? {
     return properties[NodeData.idPropertyKey] ?: id
 }
+
+data class ExistingAndExpectedNode(
+    val existingNode: INode,
+    val expectedNodeData: ModelData,
+)

bulk-model-sync-mps/build.gradle.kts

@@ -21,6 +21,7 @@ dependencies {
 
     implementation(kotlin("stdlib"))
     implementation(libs.kotlin.logging)
+    implementation(libs.kotlin.serialization.json)
 }
 
 publishing {

bulk-model-sync-mps/src/main/kotlin/org/modelix/mps/model/sync/bulk/MPSBulkSynchronizer.kt

@@ -19,14 +19,17 @@ package org.modelix.mps.model.sync.bulk
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.ProjectManager
 import jetbrains.mps.ide.project.ProjectHelper
+import kotlinx.serialization.ExperimentalSerializationApi
+import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.decodeFromStream
+import org.jetbrains.mps.openapi.module.SModule
 import org.jetbrains.mps.openapi.module.SRepository
-import org.modelix.model.api.BuiltinLanguages
-import org.modelix.model.api.INode
+import org.modelix.model.data.ModelData
 import org.modelix.model.mpsadapters.MPSModuleAsNode
 import org.modelix.model.mpsadapters.MPSRepositoryAsNode
+import org.modelix.model.sync.bulk.ExistingAndExpectedNode
 import org.modelix.model.sync.bulk.ModelExporter
 import org.modelix.model.sync.bulk.ModelImporter
-import org.modelix.model.sync.bulk.importFilesAsRootChildren
 import org.modelix.model.sync.bulk.isModuleIncluded
 import java.io.File
 import java.util.concurrent.atomic.AtomicInteger
@@ -37,7 +40,8 @@ object MPSBulkSynchronizer {
     fun exportRepository() {
         val repository = getRepository()
         val includedModuleNames = parseRawPropertySet(System.getProperty("modelix.mps.model.sync.bulk.output.modules"))
-        val includedModulePrefixes = parseRawPropertySet(System.getProperty("modelix.mps.model.sync.bulk.output.modules.prefixes"))
+        val includedModulePrefixes =
+            parseRawPropertySet(System.getProperty("modelix.mps.model.sync.bulk.output.modules.prefixes"))
 
         repository.modelAccess.runReadAction {
             val allModules = repository.modules
@@ -62,6 +66,7 @@ object MPSBulkSynchronizer {
         }
     }
 
+    @OptIn(ExperimentalSerializationApi::class)
     @JvmStatic
     fun importRepository() {
         val repository = getRepository()
@@ -70,26 +75,44 @@ object MPSBulkSynchronizer {
         val inputPath = System.getProperty("modelix.mps.model.sync.bulk.input.path")
         val continueOnError = System.getProperty("modelix.mps.model.sync.bulk.input.continueOnError", "false").toBoolean()
         val jsonFiles = File(inputPath).listFiles()?.filter {
-            it.extension == "json" && isModuleIncluded(it.nameWithoutExtension, includedModuleNames, includedModulePrefixes)
+            it.extension == "json" && isModuleIncluded(
+                it.nameWithoutExtension,
+                includedModuleNames,
+                includedModulePrefixes,
+            )
         }
 
         if (jsonFiles.isNullOrEmpty()) error("no json files found for included modules")
 
         println("Found ${jsonFiles.size} modules to be imported")
         val access = repository.modelAccess
         access.runWriteInEDT {
+            val allModules = repository.modules
+            val includedModules: Iterable<SModule> = allModules.filter {
+                isModuleIncluded(it.moduleName!!, includedModuleNames, includedModulePrefixes)
+            }
+            val numIncludedModules = includedModules.count()
             access.executeCommand {
                 val repoAsNode = MPSRepositoryAsNode(repository)
-
-                // Without the filter MPS would attempt to delete all modules that are not included
-                fun moduleFilter(node: INode): Boolean {
-                    if (node.getConceptReference()?.getUID() != BuiltinLanguages.MPSRepositoryConcepts.Module.getUID()) return true
-                    val moduleName = node.getPropertyValue(BuiltinLanguages.jetbrains_mps_lang_core.INamedConcept.name) ?: return false
-                    return isModuleIncluded(moduleName, includedModuleNames, includedModulePrefixes)
-                }
                 println("Importing modules...")
                 try {
-                    ModelImporter(repoAsNode, continueOnError, childFilter = ::moduleFilter).importFilesAsRootChildren(jsonFiles)
+                    println("Importing modules...")
+                    // `modulesToImport` lazily produces modules to import
+                    // so that loaded model data can be garbage collected.
+                    val modulesToImport = includedModules.asSequence().flatMapIndexed { index, module ->
+                        println("Importing module ${index + 1} of $numIncludedModules: '${module.moduleName}'")
+                        val fileName = inputPath + File.separator + module.moduleName + ".json"
+                        val moduleFile = File(fileName)
+                        if (moduleFile.exists()) {
+                            val expectedData: ModelData = moduleFile.inputStream().use(Json::decodeFromStream)
+                            sequenceOf(ExistingAndExpectedNode(MPSModuleAsNode(module), expectedData))
+                        } else {
+                            println("Skip importing ${module.moduleName}} because $fileName does not exist.")
+                            sequenceOf()
+                        }
+                    }
+                    ModelImporter(repoAsNode, continueOnError).importIntoNodes(modulesToImport)
+                    println("Import finished.")
                 } catch (ex: Exception) {
                     // Exceptions are only visible in the MPS log file by default
                     ex.printStackTrace()