From b11447ea6b66e2addc2cd99518be159619910bf3 Mon Sep 17 00:00:00 2001
From: Roman Golyshev <roman.golyshev@jetbrains.com>
Date: Wed, 31 May 2023 16:06:10 +0200
Subject: [PATCH] KTIJ-24983 [FIR] Ensure safe publication of configured plugin
 factories in FirExtensionRegistrar

Previous approach prevented double registration of the plugins, but
might have caused a race condition: atomic flag was set before the
plugins initialization, not after, so another thread could see garbage
data even if it could read the flag set to "true".

This might or might not be a root cause for the KTIJ-24983 exception,
so we'll have to wait for some time before actually closing the issue.
---
 .../fir/extensions/FirExtensionRegistrar.kt   | 35 ++++++++++++-------
 1 file changed, 22 insertions(+), 13 deletions(-)

diff --git a/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/extensions/FirExtensionRegistrar.kt b/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/extensions/FirExtensionRegistrar.kt
index e9b86062117..df0cb0702cf 100644
--- a/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/extensions/FirExtensionRegistrar.kt
+++ b/compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/extensions/FirExtensionRegistrar.kt
@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.fir.SessionConfiguration
 import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtension
 import org.jetbrains.kotlin.fir.builder.FirScriptConfiguratorExtension
 import org.jetbrains.kotlin.fir.resolve.FirSamConversionTransformerExtension
-import java.util.concurrent.atomic.AtomicBoolean
 import kotlin.reflect.KClass
 
 abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
@@ -184,25 +183,35 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
 
     @OptIn(PluginServicesInitialization::class)
     fun configure(): BunchOfRegisteredExtensions {
-        if (isInitialized.compareAndSet(false, true)) {
-            // Extension registrars can survive FirSession recreation in IDE mode, but we don't want to
-            // call `configurePlugin` more than once, because it will lead to registering all plugins twice.
-            // Please see KT-51444 for the details.
+        return BunchOfRegisteredExtensions(configuredExtensionFactories)
+    }
 
-            ExtensionRegistrarContext().configurePlugin()
+    private val extensionFactories: Map<KClass<out FirExtension>, MutableList<FirExtension.Factory<FirExtension>>> =
+        AVAILABLE_EXTENSIONS.associateWith {
+            mutableListOf()
         }
 
-        return BunchOfRegisteredExtensions(map)
-    }
+    /**
+     * A lazy property which returns the [extensionFactories] map, but calls
+     * [configurePlugin] to make sure that it's correctly configured.
+     *
+     * Extension registrars can survive FirSession recreation in IDE mode, but we don't want to
+     * call [configurePlugin] more than once, because it will lead to registering all plugins twice.
+     * That's why we don't want to call [configurePlugin] directly from the [configure].
+     *
+     * Instead, we use [lazy] to ensure that initialization happens only once, and that the
+     * resulting [extensionFactories] map is visible to all possible callers, so no races occur.
+     */
+    private val configuredExtensionFactories: Map<KClass<out FirExtension>, List<FirExtension.Factory<FirExtension>>> by lazy(
+        LazyThreadSafetyMode.SYNCHRONIZED
+    ) {
+        ExtensionRegistrarContext().configurePlugin()
 
-    private val map: Map<KClass<out FirExtension>, MutableList<FirExtension.Factory<FirExtension>>> = AVAILABLE_EXTENSIONS.associateWith {
-        mutableListOf()
+        extensionFactories
     }
 
-    private var isInitialized: AtomicBoolean = AtomicBoolean(false)
-
     private fun <P : FirExtension> registerExtension(kClass: KClass<out P>, factory: FirExtension.Factory<P>) {
-        val registeredExtensions = map.getValue(kClass)
+        val registeredExtensions = extensionFactories.getValue(kClass)
         registeredExtensions += factory
     }
 }