Build: Rework test workers calculation with respect to available memory

#KTI-787

buildSrc/src/main/kotlin/tasks.kt

@@ -7,6 +7,7 @@
 // usages in build scripts are not tracked properly
 @file:Suppress("unused")
 
+import com.sun.management.OperatingSystemMXBean
 import org.gradle.api.Project
 import org.gradle.api.Task
 import org.gradle.api.artifacts.ProjectDependency
@@ -20,6 +21,7 @@ import org.gradle.kotlin.dsl.support.serviceOf
 import java.io.File
 import java.lang.Character.isLowerCase
 import java.lang.Character.isUpperCase
+import java.lang.management.ManagementFactory
 import java.nio.file.Files
 import java.nio.file.Path
 
@@ -99,6 +101,9 @@ fun Project.projectTest(
     parallel: Boolean = false,
     shortenTempRootName: Boolean = false,
     jUnitMode: JUnitMode = JUnitMode.JUnit4,
+    maxHeapSizeMb: Int? = null,
+    minHeapSizeMb: Int? = null,
+    reservedCodeCacheSizeMb: Int = 256,
     body: Test.() -> Unit = {}
 ): TaskProvider<Test> {
     val shouldInstrument = project.providers.gradleProperty("kotlin.test.instrumentation.disable")
@@ -170,11 +175,24 @@ fun Project.projectTest(
             "-ea",
             "-XX:+HeapDumpOnOutOfMemoryError",
             "-XX:+UseCodeCacheFlushing",
-            "-XX:ReservedCodeCacheSize=256m",
+            "-XX:ReservedCodeCacheSize=${reservedCodeCacheSizeMb}m",
             "-Djna.nosys=true"
         )
 
-        maxHeapSize = "1600m"
+        val junit5ParallelTestWorkers =
+            project.kotlinBuildProperties.junit5NumberOfThreadsForParallelExecution ?: Runtime.getRuntime().availableProcessors()
+
+        val memoryPerTestProcessMb = maxHeapSizeMb ?: if (jUnitMode == JUnitMode.JUnit5)
+            totalMaxMemoryForTestsMb.coerceAtMost(defaultMaxMemoryPerTestWorkerMb * junit5ParallelTestWorkers)
+        else
+            defaultMaxMemoryPerTestWorkerMb
+
+        maxHeapSize = "${memoryPerTestProcessMb}m"
+
+        if (minHeapSizeMb != null) {
+            minHeapSize = "${minHeapSizeMb}m"
+        }
+
         systemProperty("idea.is.unit.test", "true")
         systemProperty("idea.home.path", project.ideaHomePathForTests().canonicalPath)
         systemProperty("idea.use.native.fs.for.win", false)
@@ -222,13 +240,23 @@ fun Project.projectTest(
         }
 
         if (parallel && jUnitMode != JUnitMode.JUnit5) {
+            val forks = (totalMaxMemoryForTestsMb / memoryPerTestProcessMb).coerceAtMost(16)
             maxParallelForks =
                 project.providers.gradleProperty("kotlin.test.maxParallelForks").forUseAtConfigurationTime().orNull?.toInt()
-                    ?: (Runtime.getRuntime().availableProcessors() / if (project.kotlinBuildProperties.isTeamcityBuild) 2 else 4).coerceAtLeast(1)
+                    ?: forks.coerceAtMost(Runtime.getRuntime().availableProcessors())
         }
     }.apply { configure(body) }
 }
 
+val defaultMaxMemoryPerTestWorkerMb = 1600
+val reservedMemoryMb = 9000 // system processes, gradle daemon, kotlin daemon, etc ...
+
+val totalMaxMemoryForTestsMb: Int
+    get() {
+        val mxbean = ManagementFactory.getOperatingSystemMXBean() as OperatingSystemMXBean
+        return (mxbean.totalPhysicalMemorySize / 1048576 - reservedMemoryMb).toInt()
+    }
+
 val Test.commandLineIncludePatterns: Set<String>
     get() = (filter as? DefaultTestFilter)?.commandLineIncludePatterns.orEmpty()
 

compiler/fir/analysis-tests/build.gradle.kts

@@ -59,9 +59,6 @@ if (kotlinBuildProperties.isInJpsBuildIdeaSync) {
 projectTest(jUnitMode = JUnitMode.JUnit5) {
     dependsOn(":dist")
     workingDir = rootDir
-    jvmArgs!!.removeIf { it.contains("-Xmx") }
-    maxHeapSize = "3g"
-
     useJUnitPlatform()
 }
 

compiler/fir/analysis-tests/legacy-fir-tests/build.gradle.kts

@@ -47,11 +47,9 @@ if (kotlinBuildProperties.isInJpsBuildIdeaSync) {
     }
 }
 
-projectTest(parallel = true) {
+projectTest(parallel = true, maxHeapSizeMb = 3072) {
     dependsOn(":dist")
     workingDir = rootDir
-    jvmArgs!!.removeIf { it.contains("-Xmx") }
-    maxHeapSize = "3g"
 }
 
 testsJar()

compiler/fir/modularized-tests/build.gradle.kts

@@ -43,12 +43,9 @@ sourceSets {
     "test" { projectDefault() }
 }
 
-projectTest {
+projectTest(minHeapSizeMb = 8192, maxHeapSizeMb = 8192, reservedCodeCacheSizeMb = 512) {
     systemProperties(project.properties.filterKeys { it.startsWith("fir.") })
     workingDir = rootDir
-    jvmArgs!!.removeIf { it.contains("-Xmx") || it.contains("-Xms") || it.contains("ReservedCodeCacheSize") }
-    minHeapSize = "8g"
-    maxHeapSize = "8g"
     dependsOn(":dist")
 
     run {
@@ -58,7 +55,6 @@ projectTest {
             jvmArgs(paramRegex.findAll(argsExt).map { it.groupValues[1] }.toList())
         }
     }
-    jvmArgs("-XX:ReservedCodeCacheSize=512m")
 }
 
 testsJar()

compiler/tests-common-new/build.gradle.kts

@@ -60,9 +60,6 @@ if (kotlinBuildProperties.isInJpsBuildIdeaSync) {
 projectTest(jUnitMode = JUnitMode.JUnit5) {
     dependsOn(":dist")
     workingDir = rootDir
-    jvmArgs!!.removeIf { it.contains("-Xmx") }
-    maxHeapSize = "3g"
-
     useJUnitPlatform()
 }
 

js/js.tests/build.gradle.kts

@@ -274,7 +274,6 @@ fun Test.setUpBoxTests() {
 
 projectTest(parallel = true, jUnitMode = JUnitMode.JUnit5) {
     setUpJsBoxTests(jsEnabled = true, jsIrEnabled = true)
-    maxHeapSize = "3g"
 
     inputs.dir(rootDir.resolve("compiler/cli/cli-common/resources")) // compiler.xml
 
@@ -294,19 +293,16 @@ projectTest(parallel = true, jUnitMode = JUnitMode.JUnit5) {
 
 projectTest("jsTest", parallel = true, jUnitMode = JUnitMode.JUnit5) {
     setUpJsBoxTests(jsEnabled = true, jsIrEnabled = false)
-    maxHeapSize = "3g"
     useJUnitPlatform()
 }
 
 projectTest("jsIrTest", true, jUnitMode = JUnitMode.JUnit5) {
     setUpJsBoxTests(jsEnabled = false, jsIrEnabled = true)
-    maxHeapSize = "3g"
     useJUnitPlatform()
 }
 
 projectTest("quickTest", parallel = true, jUnitMode = JUnitMode.JUnit5) {
     setUpJsBoxTests(jsEnabled = true, jsIrEnabled = false)
-    maxHeapSize = "3g"
     systemProperty("kotlin.js.skipMinificationTest", "true")
     useJUnitPlatform()
 }

libraries/tools/kotlin-gradle-plugin-integration-tests/build.gradle.kts

@@ -208,7 +208,9 @@ if (isTeamcityBuild) {
 }
 
 val KGP_TEST_TASKS_GROUP = "Kotlin Gradle Plugin Verification"
-val maxParallelTestForks = (Runtime.getRuntime().availableProcessors() / 4).coerceAtLeast(1)
+val memoryPerGradleTestWorkerMb = 6000
+val maxParallelTestForks =
+    (totalMaxMemoryForTestsMb / memoryPerGradleTestWorkerMb).coerceAtMost(Runtime.getRuntime().availableProcessors())
 
 val allParallelTestsTask = tasks.register<Test>("kgpAllParallelTests") {
     group = KGP_TEST_TASKS_GROUP

plugins/fir-plugin-prototype/build.gradle.kts

@@ -59,11 +59,10 @@ if (kotlinBuildProperties.isInJpsBuildIdeaSync) {
 
 projectTest(parallel = true, jUnitMode = JUnitMode.JUnit5) {
     dependsOn(":dist")
+    dependsOn(":plugins:fir-plugin-prototype:plugin-annotations:jar")
+
     workingDir = rootDir
     useJUnitPlatform()
-    jvmArgs!!.removeIf { it.contains("-Xmx") }
-    maxHeapSize = "3g"
-    dependsOn(":plugins:fir-plugin-prototype:plugin-annotations:jar")
 }
 
 testsJar()

plugins/fir-plugin-prototype/fir-plugin-ic-test/build.gradle.kts

@@ -45,11 +45,9 @@ if (kotlinBuildProperties.isInJpsBuildIdeaSync) {
     }
 }
 
-projectTest(parallel = true, jUnitMode = JUnitMode.JUnit4) {
+projectTest(parallel = true, jUnitMode = JUnitMode.JUnit4, maxHeapSizeMb = 3072) {
     workingDir = rootDir
     useJUnitPlatform()
-    jvmArgs!!.removeIf { it.contains("-Xmx") }
-    maxHeapSize = "3g"
     dependsOn(":plugins:fir-plugin-prototype:jar")
     dependsOn(":plugins:fir-plugin-prototype:plugin-annotations:jar")
 }