[FP/MT Test] Document and extract CPU isolation logic to a separate file

compiler/fir/modularized-tests/tests/org/jetbrains/kotlin/fir/FE1FullPipelineModularizedTest.kt

@@ -38,7 +38,7 @@ class FE1FullPipelineModularizedTest : AbstractFullPipelineModularizedTest() {
     }
 
     fun testTotalKotlin() {
-        isolate()
+        pinCurrentThreadToIsolatedCpu()
         for (i in 0 until PASSES) {
             println("Pass $i")
             runTestOnce(i)

compiler/fir/modularized-tests/tests/org/jetbrains/kotlin/fir/FirResolveModularizedTotalKotlinTest.kt

@@ -9,8 +9,6 @@ import com.intellij.openapi.util.Disposer
 import com.intellij.psi.PsiElementFinder
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.search.ProjectScope
-import com.sun.jna.Library
-import com.sun.jna.Native
 import com.sun.management.HotSpotDiagnosticMXBean
 import org.jetbrains.kotlin.KtPsiSourceFile
 import org.jetbrains.kotlin.KtSourceFile
@@ -39,7 +37,6 @@ import java.io.File
 import java.io.FileOutputStream
 import java.io.PrintStream
 import java.lang.management.ManagementFactory
-import java.util.regex.Pattern
 
 
 private const val FAIL_FAST = true
@@ -57,38 +54,7 @@ private val RUN_CHECKERS = System.getProperty("fir.bench.run.checkers", "false")
 private val USE_LIGHT_TREE = System.getProperty("fir.bench.use.light.tree", "true").toBooleanLenient()!!
 private val DUMP_MEMORY = System.getProperty("fir.bench.dump.memory", "false").toBooleanLenient()!!
 
-private interface CLibrary : Library {
-    fun getpid(): Int
-    fun gettid(): Int
 
-    companion object {
-        val INSTANCE = Native.load("c", CLibrary::class.java) as CLibrary
-    }
-}
-
-internal fun isolate() {
-    val isolatedList = System.getenv("DOCKER_ISOLATED_CPUSET")
-    val othersList = System.getenv("DOCKER_CPUSET")
-    println("Trying to set affinity, other: '$othersList', isolated: '$isolatedList'")
-    if (othersList != null) {
-        println("Will move others affinity to '$othersList'")
-        val pidRegex = "[0-9]+".toRegex()
-        File("/proc/").listFiles()?.forEach {
-            if (it.resolve("stat").exists() && it.name.matches(pidRegex)) {
-                ProcessBuilder().command("taskset", "-cap", othersList, it.name).inheritIO().start().waitFor()
-            }
-        }
-    }
-    if (isolatedList != null) {
-        val selfPid = CLibrary.INSTANCE.getpid()
-        val selfTid = CLibrary.INSTANCE.gettid()
-        println("Will pin self affinity, my pid: $selfPid, my tid: $selfTid")
-        ProcessBuilder().command("taskset", "-cp", isolatedList, "$selfTid").inheritIO().start().waitFor()
-    }
-    if (othersList == null && isolatedList == null) {
-        println("No affinity specified")
-    }
-}
 
 class FirResolveModularizedTotalKotlinTest : AbstractFrontendModularizedTest() {
 
@@ -262,7 +228,7 @@ class FirResolveModularizedTotalKotlinTest : AbstractFrontendModularizedTest() {
     }
 
     private fun beforeAllPasses() {
-        isolate()
+        pinCurrentThreadToIsolatedCpu()
     }
 
     fun testTotalKotlin() {

compiler/fir/modularized-tests/tests/org/jetbrains/kotlin/fir/FullPipelineModularizedTest.kt

@@ -40,7 +40,7 @@ class FullPipelineModularizedTest : AbstractFullPipelineModularizedTest() {
     }
 
     fun testTotalKotlin() {
-        isolate()
+        pinCurrentThreadToIsolatedCpu()
         for (i in 0 until PASSES) {
             println("Pass $i")
             runTestOnce(i)

compiler/fir/modularized-tests/tests/org/jetbrains/kotlin/fir/NonFirResolveModularizedTotalKotlinTest.kt

@@ -18,8 +18,6 @@ import org.jetbrains.kotlin.cli.jvm.compiler.KotlinToJVMBytecodeCompiler
 import org.jetbrains.kotlin.config.*
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction
 import org.jetbrains.kotlin.load.kotlin.ModuleVisibilityManager
-import org.jetbrains.kotlin.types.AbstractTypeChecker
-import org.jetbrains.kotlin.types.FlexibleTypeImpl
 import java.io.FileOutputStream
 import java.io.PrintStream
 import kotlin.system.measureNanoTime
@@ -144,7 +142,7 @@ class NonFirResolveModularizedTotalKotlinTest : AbstractFrontendModularizedTest(
 
     fun testTotalKotlin() {
 
-        isolate()
+        pinCurrentThreadToIsolatedCpu()
 
         writeMessageToLog("use_ni: $USE_NI")
 

compiler/fir/modularized-tests/tests/org/jetbrains/kotlin/fir/cpuIsolationUtil.kt

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir
+
+import com.sun.jna.Library
+import com.sun.jna.Native
+import java.io.File
+
+private interface CLibrary : Library {
+    fun getpid(): Int
+    fun gettid(): Int
+
+    companion object {
+        val INSTANCE = Native.load("c", CLibrary::class.java) as CLibrary
+    }
+}
+
+
+/**
+ * Pins the current thread to an isolated CPU.
+ *
+ * This method attempts to pin the current thread to an isolated CPU if the environment variables
+ * 'DOCKER_ISOLATED_CPUSET' and 'DOCKER_CPUSET' are set.
+ *
+ * On benchmark agents, those variables should be assigned to non-overlapping CPUSETs.
+ * It allows us to reduce interference of the other processes and threads.
+ *
+ * Note: CPUSET is in a format compatible with 'taskset' command (e.g., "0-3,8")
+ */
+internal fun pinCurrentThreadToIsolatedCpu() {
+    val isolatedList = System.getenv("DOCKER_ISOLATED_CPUSET")
+    val othersList = System.getenv("DOCKER_CPUSET")
+    println("Trying to set affinity, other: '$othersList', isolated: '$isolatedList'")
+    if (othersList != null) {
+        // Move all processes (including self) to the DOCKER_CPUSET
+        updateAffinityOfAllProcesses(othersList)
+    }
+    if (isolatedList != null) {
+        // Move the current thread to the isolated cpuset
+        // Must be called after the updateAffinityOfAllProcesses, otherwise it wouldn't have an effect
+        updateCurrentThreadAffinity(isolatedList)
+    }
+    if (othersList == null && isolatedList == null) {
+        println("No affinity specified")
+    }
+}
+
+
+/**
+ * Updates the CPU affinity of the current thread.
+ */
+private fun updateCurrentThreadAffinity(cpuList: String) {
+    val selfPid = CLibrary.INSTANCE.getpid()
+    val selfTid = CLibrary.INSTANCE.gettid()
+    println("Will pin self affinity, my pid: $selfPid, my tid: $selfTid")
+    ProcessBuilder().command("taskset", "-cp", cpuList, "$selfTid").inheritIO().start().waitFor()
+}
+
+
+/**
+ * Updates the affinity of all processes in the system to the specified CPU list.
+ *
+ * This method iterates over all processes in the system and modifies their affinity
+ * to the CPUs specified in the 'cpuList' parameter.
+ */
+private fun updateAffinityOfAllProcesses(cpuList: String) {
+    println("Will move others affinity to '$cpuList'")
+    val pidRegex = "[0-9]+".toRegex()
+    File("/proc/").listFiles()?.forEach {
+        if (it.resolve("stat").exists() && it.name.matches(pidRegex)) {
+            ProcessBuilder().command("taskset", "-cap", cpuList, it.name).inheritIO().start().waitFor()
+        }
+    }
+}