[Gradle] Introduce KotlinBuildStatsServicesRegistry to manage FUS services

KT-57371 introduced a legacy JMX service without proper management.
#KT-62318 Fixed

libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/kotlin/org/jetbrains/kotlin/gradle/BuildFusStatisticsIT.kt

@@ -22,18 +22,24 @@ class BuildFusStatisticsIT : KGPDaemonsBaseTest() {
         ) {
             build("assemble") {
                 //register build service for buildSrc.
-                assertOutputContains("Instantiated class org.jetbrains.kotlin.gradle.plugin.statistics.KotlinBuildStatsService: new instance")
+                val legacyBuildServiceMessagesCount = if (gradleVersion < GradleVersion.version(TestVersions.Gradle.G_8_0)) {
-                assertOutputContains("Instantiated class org.jetbrains.kotlin.gradle.plugin.statistics.KotlinBuildStatsService_v2: new instance")
+                    // until 8.0, Gradle was embedding the Kotlin version that used a slightly different approach to detect build finish,
-                // Since gradle 8 kotlinDsl was updated to 1.8 version
+                    // so the service was unregistered after the finish of the buildSrc build
-                // https://docs.gradle.org/8.0/release-notes.html#kotlin-dsl-updated-to-kotlin-api-level-1.8 ,
+                    // and then registered again in the root build
-                // and it registers old service, so we don't need check with re-registering old version service.
+                    2
-                if (gradleVersion < GradleVersion.version(TestVersions.Gradle.G_8_0)) {
+                } else {
-                    //kotlin 1.4 in kotlinDsl does not create jmx service yet
+                    1
-                    assertOutputContains("Register JMX service for backward compatibility")
                 }
+                assertOutputContainsExactTimes(
+                    "Instantiated class org.jetbrains.kotlin.gradle.plugin.statistics.KotlinBuildStatsService: new instance", // the legacy service for compatibility
+                    legacyBuildServiceMessagesCount
+                )
+                assertOutputContainsExactTimes(
+                    "Instantiated class org.jetbrains.kotlin.gradle.plugin.statistics.KotlinBuildStatsService_v2: new instance", // the current default version of the service
+                    1
+                )
                 assertOutputDoesNotContain("[org.jetbrains.kotlin.gradle.plugin.statistics.KotlinBuildStatHandler] Could not execute")
             }
         }
     }
-
 }

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/statistics/BuildFlowService.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.gradle.plugin.statistics
 
 import org.gradle.api.Project
-import org.gradle.api.invocation.Gradle
 import org.gradle.api.logging.Logging
 import org.gradle.api.provider.Property
 import org.gradle.api.provider.Provider
@@ -108,9 +107,7 @@ internal abstract class BuildFlowService : BuildService<BuildFlowService.Paramet
         if (parameters.fusStatisticsAvailable.get()) {
             recordBuildFinished(null, buildFailed)
         }
-        KotlinBuildStatsService.applyIfInitialised {
+        KotlinBuildStatsService.closeServices()
-            it.close()
-        }
         log.kotlinDebug("Close ${this.javaClass.simpleName}")
     }
 

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/statistics/DaemonReusageCounter.kt

@@ -51,7 +51,7 @@ class DaemonReuseCounter private constructor() : IDaemonReuseCounterMXBean {
         }
 
         fun getOrdinal(): Long {
-            val beanName = ObjectName(KotlinBuildStatsService.JMX_BEAN_NAME)
+            val beanName = ObjectName(JMX_BEAN_NAME)
             val mbs: MBeanServer = ManagementFactory.getPlatformMBeanServer()
             ensureRegistered(beanName, mbs)
             return mbs.invoke(beanName, "getOrdinal", emptyArray(), emptyArray<String>()) as? Long ?: 0

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/statistics/KotlinBuildStatsService.kt

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.gradle.plugin.statistics
 
 import org.gradle.api.Project
 import org.gradle.api.invocation.Gradle
-import org.gradle.api.logging.Logger
 import org.gradle.api.logging.Logging
 import org.gradle.api.provider.ProviderFactory
 import org.gradle.initialization.BuildRequestMetaData
@@ -19,8 +18,6 @@ import org.jetbrains.kotlin.gradle.plugin.internal.ConfigurationTimePropertiesAc
 import org.jetbrains.kotlin.gradle.plugin.internal.configurationTimePropertiesAccessor
 import org.jetbrains.kotlin.gradle.plugin.internal.usedAtConfigurationTime
 import org.jetbrains.kotlin.gradle.plugin.statistics.KotlinBuildStatHandler.Companion.runSafe
-import org.jetbrains.kotlin.gradle.plugin.statistics.old.Pre232IdeaKotlinBuildStatsMXBean
-import org.jetbrains.kotlin.gradle.plugin.statistics.old.Pre232IdeaKotlinBuildStatsService
 import org.jetbrains.kotlin.gradle.utils.loadProperty
 import org.jetbrains.kotlin.gradle.utils.localProperties
 import org.jetbrains.kotlin.statistics.BuildSessionLogger
@@ -31,10 +28,8 @@ import org.jetbrains.kotlin.statistics.metrics.NumericalMetrics
 import org.jetbrains.kotlin.statistics.metrics.StringMetrics
 import java.io.Closeable
 import java.io.File
-import java.lang.management.ManagementFactory
 import javax.management.MBeanServer
 import javax.management.ObjectName
-import javax.management.StandardMBean
 import kotlin.system.measureTimeMillis
 
 /**
@@ -51,28 +46,19 @@ interface KotlinBuildStatsMXBean {
     fun reportString(name: String, value: String, subprojectName: String?, weight: Long?): Boolean
 }
 
-
 internal abstract class KotlinBuildStatsService internal constructor() : IStatisticsValuesConsumer, Closeable {
     companion object {
-        // Do not rename this bean otherwise compatibility with the older Kotlin Gradle Plugins would be lost
-        private const val JMX_BEAN_NAME_BEFORE_232_IDEA = "org.jetbrains.kotlin.gradle.plugin.statistics:type=StatsService"
-
-        //update name when API changed
-        private const val SERVICE_NAME = "v2"
-        const val JMX_BEAN_NAME = "org.jetbrains.kotlin.gradle.plugin.statistics:type=StatsService,name=$SERVICE_NAME"
-
-
         // Property name for disabling saving statistical information
-        const val ENABLE_STATISTICS_PROPERTY_NAME = "enable_kotlin_performance_profile"
+        private const val ENABLE_STATISTICS_PROPERTY_NAME = "enable_kotlin_performance_profile"
 
         // Property used for tests. Build will fail fast if collected value doesn't fit regexp
         const val FORCE_VALUES_VALIDATION = "kotlin_performance_profile_force_validation"
 
         // default state
-        const val DEFAULT_STATISTICS_STATE = true
+        private const val DEFAULT_STATISTICS_STATE = true
 
         // "emergency file" collecting statistics is disabled it the file exists
-        const val DISABLE_STATISTICS_FILE_NAME = "${STATISTICS_FOLDER_NAME}/.disable"
+        private const val DISABLE_STATISTICS_FILE_NAME = "${STATISTICS_FOLDER_NAME}/.disable"
 
         /**
          * Method for getting IStatisticsValuesConsumer for reporting some statistics
@@ -85,23 +71,21 @@ internal abstract class KotlinBuildStatsService internal constructor() : IStatis
             if (statisticsIsEnabled != true) {
                 return null
             }
-            return instance
+            return kotlinBuildStatsServicesRegistry?.getDefaultService()
         }
 
-        private fun getServiceName(): String = "${KotlinBuildStatsService::class.java}_$SERVICE_NAME"
-
         /**
          * Method for creating new instance of IStatisticsValuesConsumer
          * It could be invoked only when applying Kotlin gradle plugin.
-         * When executed, this method checks, whether it is already executed in current build.
+         * When executed, this method checks, whether it is already executed in the current build.
          * If it was not executed, the new instance of IStatisticsValuesConsumer is created
          *
-         * If it was already executed in the same classpath (i.e. with the same version of Kotlin plugin),
+         * If it was already executed in the same classpath (i.e., with the same version of Kotlin plugin),
          * the previously returned instance is returned.
          *
          * If it was already executed in the other classpath, a JXM implementation is returned.
          *
-         * All the created instances are registered as build listeners
+         * [closeServices] must be called at the end of the build in order to release resources.
          */
         @JvmStatic
         @Synchronized
@@ -115,56 +99,35 @@ internal abstract class KotlinBuildStatsService internal constructor() : IStatis
                 if (statisticsIsEnabled != true) {
                     null
                 } else {
-                    val log = getLogger()
+                    val registry = kotlinBuildStatsServicesRegistry ?: KotlinBuildStatsServicesRegistry().also {
-
+                        kotlinBuildStatsServicesRegistry = it
-                    if (instance != null) {
-                        log.debug("${getServiceName()} is already instantiated. Current instance is $instance")
-                    } else {
-                        val beanName = ObjectName(JMX_BEAN_NAME)
-                        val mbs: MBeanServer = ManagementFactory.getPlatformMBeanServer()
-                        if (mbs.isRegistered(beanName)) {
-                            log.debug(
-                                "${getServiceName()} is already instantiated in another classpath. Creating JMX-wrapper"
-                            )
-                            instance = JMXKotlinBuildStatsService(mbs, beanName)
-                        } else {
-                            val newInstance = DefaultKotlinBuildStatsService(
-                                project,
-                                beanName
-                            )
-
-                            instance = newInstance
-                            log.debug("Instantiated ${getServiceName()}: new instance $instance")
-                            mbs.registerMBean(StandardMBean(newInstance, KotlinBuildStatsMXBean::class.java), beanName)
-
-                            registerPre232IdeaStatsBean(mbs, log, project)
-                        }
-
-                        BuildEventsListenerRegistryHolder.getInstance(project).listenerRegistry.onTaskCompletion(project.provider {
-                            OperationCompletionListener { event ->
-                                if (event is TaskFinishEvent) {
-                                    reportTaskIfNeed(event.descriptor.name)
-                                }
-                            }
-                        })
                     }
-                    instance
+
+                    val defaultServiceName = KotlinBuildStatsServicesRegistry.getBeanName(KotlinBuildStatsServicesRegistry.DEFAULT_SERVICE_QUALIFIER)
+                    val instance = kotlinBuildStatsServicesRegistry?.getDefaultService()
+                    if (instance != null) {
+                        registry.logger.debug("$defaultServiceName is already instantiated. Current instance is $instance")
+                        return@runSafe instance
+                    }
+
+                    registry.registerServices(project)
+
+                    BuildEventsListenerRegistryHolder.getInstance(project).listenerRegistry.onTaskCompletion(project.provider {
+                        OperationCompletionListener { event ->
+                            if (event is TaskFinishEvent) {
+                                reportTaskIfNeed(event.descriptor.name)
+                            }
+                        }
+                    })
+
+                    registry.getDefaultService() ?: error("The default kotlin build stats $defaultServiceName service wasn't initialized")
                 }
             }
         }
 
-        //To support backward compatibility with Idea before 232 version
+        fun closeServices() {
-        private fun registerPre232IdeaStatsBean(
+            kotlinBuildStatsServicesRegistry?.close()
-            mbs: MBeanServer,
+            kotlinBuildStatsServicesRegistry = null
-            log: Logger,
-            project: Project
-        ) {
-            val beanName = ObjectName(JMX_BEAN_NAME_BEFORE_232_IDEA)
-            if (!mbs.isRegistered(beanName)) {
-                val newInstance = Pre232IdeaKotlinBuildStatsService(project, beanName)
-                mbs.registerMBean(StandardMBean(newInstance, Pre232IdeaKotlinBuildStatsMXBean::class.java), beanName)
-                log.debug("Register JMX service for backward compatibility")
-            }
         }
 
         protected fun reportTaskIfNeed(task: String) {
@@ -203,10 +166,7 @@ internal abstract class KotlinBuildStatsService internal constructor() : IStatis
             }
         }
 
-        @JvmStatic
+        private var kotlinBuildStatsServicesRegistry: KotlinBuildStatsServicesRegistry? = null
-        internal fun getLogger() = Logging.getLogger(KotlinBuildStatsService::class.java)
-
-        internal var instance: KotlinBuildStatsService? = null
 
         private var statisticsIsEnabled: Boolean? = null
 
@@ -226,7 +186,6 @@ internal abstract class KotlinBuildStatsService internal constructor() : IStatis
     }
 
     override fun close() {
-        instance = null
     }
 
     /**

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/statistics/KotlinBuildStatsServicesRegistry.kt

@@ -0,0 +1,103 @@
+/*
+ * 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.gradle.plugin.statistics
+
+import org.gradle.api.Project
+import org.gradle.api.logging.Logging
+import org.jetbrains.kotlin.gradle.plugin.statistics.old.Pre232IdeaKotlinBuildStatsMXBean
+import org.jetbrains.kotlin.gradle.plugin.statistics.old.Pre232IdeaKotlinBuildStatsService
+import java.io.Closeable
+import java.lang.management.ManagementFactory
+import javax.management.MBeanServer
+import javax.management.ObjectName
+import javax.management.StandardMBean
+
+/**
+ * A registry of Kotlin FUS services scoped to a project.
+ * In general, we use only one instance of a service either it is a real service or a JMX wrapper.
+ * However, we also register legacy variants of the services so previous IDEA versions could access them.
+ * We must properly manage the lifecycle of such services.
+ *
+ * The implementation is not thread-safe!
+ */
+internal class KotlinBuildStatsServicesRegistry : Closeable {
+    internal val logger = Logging.getLogger(this::class.java)
+    private val services = hashMapOf<String, KotlinBuildStatsService>()
+
+    /**
+     * Registers the Kotlin build stats services for the given project.
+     *
+     * After a call of that method, [getDefaultService] is expected to return a non-null value.
+     *
+     * The registry must be closed at the end of the usage.
+     */
+    fun registerServices(project: Project) {
+        val defaultBeanName = getBeanName(DEFAULT_SERVICE_QUALIFIER)
+        val mbs: MBeanServer = ManagementFactory.getPlatformMBeanServer()
+        if (mbs.isRegistered(defaultBeanName)) {
+            // because we use only the default service, and it's already registered in the MBean server, then there's no need in caring of legacy services as they're already registered
+            logger.debug("${KotlinBuildStatsService::class.simpleName} $defaultBeanName is already instantiated in another classpath. Creating JMX wrapper for the main service")
+            services[DEFAULT_SERVICE_QUALIFIER] = JMXKotlinBuildStatsService(mbs, defaultBeanName)
+        } else {
+            registerStatsService(
+                mbs,
+                DefaultKotlinBuildStatsService(project, getBeanName(DEFAULT_SERVICE_QUALIFIER)),
+                KotlinBuildStatsMXBean::class.java,
+                DEFAULT_SERVICE_QUALIFIER
+            )
+            // to support backward compatibility with Idea before version 232
+            registerStatsService(
+                mbs,
+                Pre232IdeaKotlinBuildStatsService(project, getBeanName(LEGACY_SERVICE_QUALIFIER)),
+                Pre232IdeaKotlinBuildStatsMXBean::class.java,
+                LEGACY_SERVICE_QUALIFIER
+            )
+        }
+    }
+
+    private fun <T : KotlinBuildStatsService> registerStatsService(
+        mbs: MBeanServer,
+        service: T,
+        beanInterfaceType: Class<in T>,
+        qualifier: String,
+    ) {
+        val beanName = getBeanName(qualifier)
+        if (!mbs.isRegistered(beanName)) {
+            mbs.registerMBean(StandardMBean(service, beanInterfaceType), beanName)
+            services[qualifier] = service
+            val loggedServiceName = "${KotlinBuildStatsService::class.java}" + if (qualifier.isNotEmpty()) "_$qualifier" else ""
+            logger.debug("Instantiated $loggedServiceName: new instance $service")
+        }
+    }
+
+    /**
+     * Retrieves the default Kotlin build stats service. That's the main service we are working with.
+     */
+    fun getDefaultService() = services[DEFAULT_SERVICE_QUALIFIER]
+
+    /**
+     * Unregisters all the registered JMX services and may release other resources allocated by a service.
+     */
+    override fun close() {
+        for (service in services.values) {
+            service.close()
+        }
+        services.clear()
+    }
+
+    companion object {
+        private const val JXM_BEAN_BASE_NAME = "org.jetbrains.kotlin.gradle.plugin.statistics:type=StatsService"
+
+        // Do not rename this bean otherwise compatibility with the older Kotlin Gradle Plugins would be lost
+        private const val LEGACY_SERVICE_QUALIFIER = ""
+
+        // Update name when API changed
+        internal const val DEFAULT_SERVICE_QUALIFIER = "v2"
+
+        internal fun getBeanName(qualifier: String) =
+            ObjectName(JXM_BEAN_BASE_NAME + if (qualifier.isNotEmpty()) ",name=$qualifier" else "")
+    }
+}

libraries/tools/kotlin-gradle-plugin/src/test/kotlin/org/jetbrains/kotlin/gradle/plugin/statistics/BuildStatServiceTest.kt

@@ -12,7 +12,6 @@ import org.junit.Test
 import java.lang.management.ManagementFactory
 import java.util.concurrent.atomic.AtomicInteger
 import javax.management.MBeanServer
-import javax.management.ObjectName
 import javax.management.StandardMBean
 import kotlin.test.assertEquals
 import kotlin.test.assertTrue
@@ -22,7 +21,7 @@ class BuildStatServiceTest {
     @Test
     fun testJmxDoesNotFail() {
         //Test that JMX service does not throw exception even if JMX beans are not configured
-        val beanName = ObjectName(KotlinBuildStatsService.JMX_BEAN_NAME)
+        val beanName = KotlinBuildStatsServicesRegistry.getBeanName(KotlinBuildStatsServicesRegistry.DEFAULT_SERVICE_QUALIFIER)
         val mbs: MBeanServer = ManagementFactory.getPlatformMBeanServer()
 
         val jmxService = JMXKotlinBuildStatsService(mbs, beanName)
@@ -52,7 +51,7 @@ class BuildStatServiceTest {
 
         }
 
-        val beanName = ObjectName(KotlinBuildStatsService.JMX_BEAN_NAME)
+        val beanName = KotlinBuildStatsServicesRegistry.getBeanName(KotlinBuildStatsServicesRegistry.DEFAULT_SERVICE_QUALIFIER)
         val mbs: MBeanServer = ManagementFactory.getPlatformMBeanServer()
         mbs.registerMBean(StandardMBean(instance, KotlinBuildStatsMXBean::class.java), beanName)