Merge pull request #550 from modelix/perf/reduce-memory-usage-for-import-into-mps

Oleksandr Dzhychko · web-flow · commit da90ba80878b · 2024-02-29T15:48:48.000+01:00
reduce memory usage for import into mps
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -116,7 +116,7 @@ subprojects {
         if (!name.lowercase().contains("test")) {
             this.kotlinOptions {
                 jvmTarget = "11"
-                freeCompilerArgs += listOf("-Xjvm-default=all-compatibility")
+                freeCompilerArgs += listOf("-Xjvm-default=all-compatibility", "-Xexpect-actual-classes")
                 apiVersion = kotlinApiVersion.version
             }
         }
diff --git a/bulk-model-sync-lib/build.gradle.kts b/bulk-model-sync-lib/build.gradle.kts
@@ -21,6 +21,12 @@ kotlin {
             }
         }
 
+        val jvmMain by getting {
+            dependencies {
+                implementation(libs.trove4j)
+            }
+        }
+
         val commonTest by getting {
             dependencies {
                 implementation(project(":model-api"))
diff --git a/bulk-model-sync-lib/src/commonMain/kotlin/org/modelix/model/sync/bulk/MemoryEfficientMap.kt b/bulk-model-sync-lib/src/commonMain/kotlin/org/modelix/model/sync/bulk/MemoryEfficientMap.kt
@@ -0,0 +1,34 @@
+/*
+ * 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.sync.bulk
+
+/**
+ * Built-in maps like [HashMap] are not the most memory efficient way of map.
+ * A common issue is that entry objects for every item in the table are created.
+ * [MemoryEfficientMap] is an internal implementation that we can use
+ * when the memory overhead becomes too big.
+ *
+ * Java implementation is optimized to not create entry objects by using a map implementation from another library.
+ * The JS implementation is not optimized yet because we did not invest time in finding a suitable library.
+ *
+ * [MemoryEfficientMap] is an internal abstraction.
+ * The API is therefore kept minimal
+ */
+expect class MemoryEfficientMap<KeyT, ValueT>() {
+    operator fun set(key: KeyT, value: ValueT)
+    operator fun get(key: KeyT): ValueT?
+}
diff --git a/bulk-model-sync-lib/src/commonMain/kotlin/org/modelix/model/sync/bulk/ModelImporter.kt b/bulk-model-sync-lib/src/commonMain/kotlin/org/modelix/model/sync/bulk/ModelImporter.kt
@@ -25,6 +25,7 @@ import org.modelix.model.api.SerializedNodeReference
 import org.modelix.model.api.getDescendants
 import org.modelix.model.api.isChildRoleOrdered
 import org.modelix.model.api.remove
+import org.modelix.model.api.resolveInCurrentContext
 import org.modelix.model.data.ModelData
 import org.modelix.model.data.NodeData
 import kotlin.jvm.JvmName
@@ -48,11 +49,17 @@ class ModelImporter(
     private val continueOnError: Boolean,
     private val childFilter: (INode) -> Boolean = { true },
 ) {
+    // We have seen imports where the `originalIdToExisting` had a dozen ten million entries.
+    // 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 val originalIdToExisting by lazy(::buildExistingIndex)
 
-    private val originalIdToExisting: MutableMap<String, INode> = mutableMapOf()
+    // 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 {}
 
@@ -96,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.clear()
-            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 }
-            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)
         }
     }
 
@@ -145,15 +169,25 @@ class ModelImporter(
         for (role in allRoles) {
             val expectedNodes = expectedParent.children.filter { it.role == role }
             val existingNodes = existingParent.getChildren(role).filter(childFilter).toList()
+            val allExpectedNodesDoNotExist by lazy {
+                expectedNodes.all { expectedNode ->
+                    val originalId = expectedNode.originalId()
+                    checkNotNull(originalId) { "Specified node '$expectedNode' has no ID." }
+                    originalIdToExisting[originalId] == null
+                }
+            }
 
             // optimization that uses the bulk operation .addNewChildren
-            if (existingNodes.isEmpty() && expectedNodes.all { originalIdToExisting[it.originalId()] == null }) {
-                existingParent.addNewChildren(role, -1, expectedNodes.map { it.concept?.let { ConceptReference(it) } }).zip(expectedNodes).forEach { (newChild, expected) ->
-                    val expectedId = checkNotNull(expected.originalId()) { "Specified node '$expected' has no id" }
-                    newChild.setPropertyValue(NodeData.idPropertyKey, expectedId)
-                    originalIdToExisting[expectedId] = newChild
-                    syncNode(newChild, expected, progressReporter)
-                }
+            if (existingNodes.isEmpty() && allExpectedNodesDoNotExist) {
+                existingParent.addNewChildren(role, -1, expectedNodes.map { it.concept?.let { ConceptReference(it) } })
+                    .zip(expectedNodes)
+                    .forEach { (newChild, expected) ->
+                        val expectedId = expected.originalId()
+                        checkNotNull(expectedId) { "Specified node '$expected' has no ID." }
+                        newChild.setPropertyValue(NodeData.idPropertyKey, expectedId)
+                        originalIdToExisting[expectedId] = newChild.reference
+                        syncNode(newChild, expected, progressReporter)
+                    }
                 continue
             }
 
@@ -186,13 +220,18 @@ class ModelImporter(
                 val nodeAtIndex = existingChildren.getOrNull(newIndex)
                 val expectedConcept = expected.concept?.let { s -> ConceptReference(s) }
                 val childNode = if (nodeAtIndex?.originalId() != expectedId) {
-                    val existingNode = originalIdToExisting[expectedId]
-                    if (existingNode == null) {
+                    val existingNodeReference = originalIdToExisting[expectedId]
+                    if (existingNodeReference == null) {
                         val newChild = existingParent.addNewChild(role, newIndex, expectedConcept)
                         newChild.setPropertyValue(NodeData.idPropertyKey, expectedId)
-                        originalIdToExisting[expectedId] = newChild
+                        originalIdToExisting[expectedId] = newChild.reference
                         newChild
                     } else {
+                        val existingNode = existingNodeReference.resolveInCurrentContext()
+                        checkNotNull(existingNode) {
+                            // This reference should always be resolvable because the node existed or was created before.
+                            "Could not resolve $existingNodeReference."
+                        }
                         // The existing child node is not only moved to a new index,
                         // it is potentially moved to a new parent and role.
                         existingParent.moveChild(role, newIndex, existingNode)
@@ -217,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 }
+            node.originalId()?.let { localOriginalIdToExisting[it] = node.reference }
         }
+        return localOriginalIdToExisting
     }
 
     private fun syncProperties(node: INode, nodeData: NodeData) {
@@ -268,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,
+)
diff --git a/bulk-model-sync-lib/src/jsMain/kotlin/org/modelix/model/sync/bulk/MemoryEfficientMap.kt b/bulk-model-sync-lib/src/jsMain/kotlin/org/modelix/model/sync/bulk/MemoryEfficientMap.kt
@@ -0,0 +1,25 @@
+/*
+ * 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.sync.bulk
+
+actual class MemoryEfficientMap<KeyT, ValueT> {
+    private val map: MutableMap<KeyT, ValueT> = mutableMapOf()
+
+    actual operator fun set(key: KeyT, value: ValueT) = map.set(key, value)
+
+    actual operator fun get(key: KeyT) = map[key]
+}
diff --git a/bulk-model-sync-lib/src/jvmMain/kotlin/org/modelix/model/sync/bulk/MemoryEfficientMap.kt b/bulk-model-sync-lib/src/jvmMain/kotlin/org/modelix/model/sync/bulk/MemoryEfficientMap.kt
@@ -0,0 +1,28 @@
+/*
+ * 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.sync.bulk
+
+import gnu.trove.map.TMap
+import gnu.trove.map.hash.THashMap
+
+actual class MemoryEfficientMap<KeyT, ValueT> {
+    private val map: TMap<KeyT, ValueT> = THashMap()
+
+    actual operator fun set(key: KeyT, value: ValueT) = map.set(key, value)
+
+    actual operator fun get(key: KeyT) = map[key]
+}
diff --git a/bulk-model-sync-mps/build.gradle.kts b/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 {
diff --git a/bulk-model-sync-mps/src/main/kotlin/org/modelix/mps/model/sync/bulk/MPSBulkSynchronizer.kt b/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()
diff --git a/mps-model-adapters/src/main/kotlin/org/modelix/model/mpsadapters/MPSArea.kt b/mps-model-adapters/src/main/kotlin/org/modelix/model/mpsadapters/MPSArea.kt

Original file line number	Diff line number	Diff line change
`@@ -116,7 +116,7 @@ subprojects {`
`116`	`116`	`if (!name.lowercase().contains("test")) {`
`117`	`117`	`this.kotlinOptions {`
`118`	`118`	`jvmTarget = "11"`
`119`		`- freeCompilerArgs += listOf("-Xjvm-default=all-compatibility")`
	`119`	`+ freeCompilerArgs += listOf("-Xjvm-default=all-compatibility", "-Xexpect-actual-classes")`
`120`	`120`	`apiVersion = kotlinApiVersion.version`
`121`	`121`	`}`
`122`	`122`	`}`
Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,12 @@ kotlin {`
`21`	`21`	`}`
`22`	`22`	`}`
`23`	`23`
	`24`	`+ val jvmMain by getting {`
	`25`	`+ dependencies {`
	`26`	`+ implementation(libs.trove4j)`
	`27`	`+ }`
	`28`	`+ }`
	`29`	`+`
`24`	`30`	`val commonTest by getting {`
`25`	`31`	`dependencies {`
`26`	`32`	`implementation(project(":model-api"))`
Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,7 @@ dependencies {`
`21`	`21`
`22`	`22`	`implementation(kotlin("stdlib"))`
`23`	`23`	`implementation(libs.kotlin.logging)`
	`24`	`+ implementation(libs.kotlin.serialization.json)`
`24`	`25`	`}`
`25`	`26`
`26`	`27`	`publishing {`