Fixed locking granularity in ScriptDefinitionsManager

#KT-34552 Fixed
#KT-39547 Fixed

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDefinitionsManager.kt

@@ -16,7 +16,6 @@ import com.intellij.openapi.diagnostic.ControlFlowException
 import com.intellij.openapi.extensions.ExtensionPointName
 import com.intellij.openapi.extensions.Extensions
 import com.intellij.openapi.fileTypes.FileTypeManager
-import com.intellij.openapi.progress.runBackgroundableTask
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.JavaSdk
 import com.intellij.openapi.projectRoots.ex.PathUtilEx
@@ -76,6 +75,8 @@ class LoadScriptDefinitionsStartupActivity : StartupActivity {
 
 class ScriptDefinitionsManager(private val project: Project) : LazyScriptDefinitionProvider() {
     private var definitionsBySource = mutableMapOf<ScriptDefinitionsSource, List<ScriptDefinition>>()
+
+    @Volatile
     private var definitions: List<ScriptDefinition>? = null
 
     private val failedContributorsHashes = HashSet<Int>()
@@ -121,21 +122,26 @@ class ScriptDefinitionsManager(private val project: Project) : LazyScriptDefinit
 
     override fun findScriptDefinition(fileName: String): KotlinScriptDefinition? = findDefinition(File(fileName).toScriptSource())?.legacyDefinition
 
-    fun reloadDefinitionsBy(source: ScriptDefinitionsSource) = lock.write {
+    fun reloadDefinitionsBy(source: ScriptDefinitionsSource) {
         if (definitions == null) return // not loaded yet
 
-        if (source !in definitionsBySource) error("Unknown script definition source: $source")
+        lock.write {
+            if (source !in definitionsBySource) error("Unknown script definition source: $source")
+        }
 
-        definitionsBySource[source] = source.safeGetDefinitions()
+        val safeGetDefinitions = source.safeGetDefinitions()
+        lock.write {
+            definitionsBySource[source] = safeGetDefinitions
 
-        definitions = definitionsBySource.values.flattenTo(mutableListOf())
+            definitions = definitionsBySource.values.flattenTo(mutableListOf())
 
-        updateDefinitions()
+            updateDefinitions()
+        }
     }
 
     override val currentDefinitions
         get() =
-            (definitions ?: kotlin.run {
+            (definitions ?: run {
                 reloadScriptDefinitions()
                 definitions!!
             }).asSequence().filter { KotlinScriptingSettings.getInstance(project).isScriptDefinitionEnabled(it) }
@@ -149,45 +155,46 @@ class ScriptDefinitionsManager(private val project: Project) : LazyScriptDefinit
         return fromNewEp.dropLast(1) + fromDeprecatedEP + fromNewEp.last()
     }
 
-    fun reloadScriptDefinitionsIfNeeded() = lock.write {
-        if (definitions == null) {
-            loadScriptDefinitions()
-        }
+    fun reloadScriptDefinitionsIfNeeded() {
+        definitions ?: loadScriptDefinitions()
     }
 
-    fun reloadScriptDefinitions() = lock.write {
-        loadScriptDefinitions()
-    }
+    fun reloadScriptDefinitions() = loadScriptDefinitions()
 
     private fun loadScriptDefinitions() {
-        for (source in getSources()) {
-            val definitions = source.safeGetDefinitions()
-            definitionsBySource[source] = definitions
+        val newDefinitionsBySource = getSources().map { it to it.safeGetDefinitions() }.toMap()
+
+        lock.write {
+            definitionsBySource.putAll(newDefinitionsBySource)
+            definitions = definitionsBySource.values.flattenTo(mutableListOf())
+
+            updateDefinitions()
         }
-
-        definitions = definitionsBySource.values.flattenTo(mutableListOf())
-
-        updateDefinitions()
     }
 
-    fun reorderScriptDefinitions() = lock.write {
+    fun reorderScriptDefinitions() {
         definitions?.forEach {
-            it.order = KotlinScriptingSettings.getInstance(project).getScriptDefinitionOrder(it)
+            val order = KotlinScriptingSettings.getInstance(project).getScriptDefinitionOrder(it)
+            lock.write {
+                it.order = order
+            }
         }
-        definitions = definitions?.sortedBy { it.order }
+        lock.write {
+            definitions = definitions?.sortedBy { it.order }
 
-        updateDefinitions()
+            updateDefinitions()
+        }
     }
 
-    fun getAllDefinitions(): List<ScriptDefinition> {
-        return definitions ?: kotlin.run {
-            reloadScriptDefinitions()
-            definitions!!
-        }
+    fun getAllDefinitions(): List<ScriptDefinition> = definitions ?: run {
+        reloadScriptDefinitions()
+        definitions!!
     }
 
     fun isReady(): Boolean {
-        return definitions != null && definitionsBySource.keys.all { source ->
+        if (definitions == null) return false
+        val keys = lock.write { definitionsBySource.keys }
+        return keys.all { source ->
             // TODO: implement another API for readiness checking
             (source as? ScriptDefinitionContributor)?.isReady() != false
         }

idea/src/org/jetbrains/kotlin/idea/script/ScriptTemplatesFromDependenciesProvider.kt

@@ -78,6 +78,7 @@ class ScriptTemplatesFromDependenciesProvider(private val project: Project) : Sc
         }
     }
 
+    @Volatile
     private var _definitions: List<ScriptDefinition>? = null
     private val definitionsLock = ReentrantLock()