Always scan root for script templates in background

^KT-29893 Fixed

Provide AsyncScriptDefinitionsContributor.
Show a notification if the highlighting for script is requested when not all script templates are loaded

idea/idea-analysis/src/org/jetbrains/kotlin/idea/highlighter/KotlinHighlightingUtil.kt

@@ -20,6 +20,7 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.caches.project.NotUnderContentRootModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.core.script.IdeScriptReportSink
+import org.jetbrains.kotlin.idea.core.script.ScriptDefinitionsManager
 import org.jetbrains.kotlin.idea.core.script.scriptDependencies
 import org.jetbrains.kotlin.idea.util.ProjectRootsUtil
 import org.jetbrains.kotlin.psi.KtCodeFragment
@@ -69,6 +70,8 @@ object KotlinHighlightingUtil {
             return false
         }
 
+        if (!ScriptDefinitionsManager.getInstance(ktFile.project).isReady()) return false
+
         return ProjectRootsUtil.isInProjectSource(ktFile, includeScriptsOutsideSourceRoots = true)
     }
 }

idea/idea-analysis/src/org/jetbrains/kotlin/idea/highlighter/ScriptExternalHighlightingPass.kt

@@ -31,8 +31,13 @@ import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.progress.ProgressIndicator
 import com.intellij.openapi.project.DumbAware
+import com.intellij.openapi.ui.MessageType
+import com.intellij.openapi.wm.WindowManager
+import com.intellij.openapi.wm.ex.StatusBarEx
 import com.intellij.psi.PsiFile
+import com.intellij.util.ui.UIUtil
 import org.jetbrains.kotlin.idea.core.script.IdeScriptReportSink
+import org.jetbrains.kotlin.idea.core.script.ScriptDefinitionsManager
 import org.jetbrains.kotlin.psi.KtFile
 import kotlin.script.experimental.dependencies.ScriptReport
 
@@ -45,6 +50,23 @@ class ScriptExternalHighlightingPass(
     override fun doApplyInformationToEditor() {
         val document = document ?: return
 
+        if (!file.isScript()) return
+
+        if (!ScriptDefinitionsManager.getInstance(file.project).isReady()) {
+            UIUtil.invokeLaterIfNeeded {
+                val ideFrame = WindowManager.getInstance().getIdeFrame(file.project)
+                if (ideFrame != null) {
+                    val statusBar = ideFrame.statusBar as StatusBarEx
+                    statusBar.notifyProgressByBalloon(
+                        MessageType.WARNING,
+                        "Highlighting in scripts is not available until all Script Definitions are loaded",
+                        null,
+                        null
+                    )
+                }
+            }
+        }
+
         val reports = file.virtualFile.getUserData(IdeScriptReportSink.Reports) ?: return
 
         val annotations = reports.mapNotNull { (message, severity, position) ->

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

@@ -74,6 +74,7 @@ class ScriptDefinitionsManager(private val project: Project) : LazyScriptDefinit
 
     override fun findScriptDefinition(fileName: String): KotlinScriptDefinition? {
         if (nonScriptFileName(fileName)) return null
+        if (!isReady()) return null
 
         val cached = synchronized(scriptDefinitionsCacheLock) { scriptDefinitionsCache.get(fileName) }
         if (cached != null) return cached
@@ -144,6 +145,12 @@ class ScriptDefinitionsManager(private val project: Project) : LazyScriptDefinit
         }
     }
 
+    fun isReady(): Boolean {
+        return definitionsByContributor.keys.all { contributor ->
+            contributor.isReady()
+        }
+    }
+
     override fun getDefaultScriptDefinition(): KotlinScriptDefinition {
         val standardScriptDefinitionContributor = ScriptDefinitionContributor.find<StandardScriptDefinitionContributor>(project)
             ?: error("StandardScriptDefinitionContributor should be registered is plugin.xml")
@@ -271,6 +278,7 @@ interface ScriptDefinitionContributor {
     val id: String
 
     fun getDefinitions(): List<KotlinScriptDefinition>
+    fun isReady() = true
 
     companion object {
         val EP_NAME: ExtensionPointName<ScriptDefinitionContributor> =

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

@@ -0,0 +1,98 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. 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.idea.script
+
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.progress.ProgressIndicator
+import com.intellij.openapi.progress.Task
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.idea.core.script.ScriptDefinitionContributor
+import org.jetbrains.kotlin.idea.core.script.ScriptDefinitionsManager
+import org.jetbrains.kotlin.script.KotlinScriptDefinition
+import java.util.concurrent.locks.ReentrantReadWriteLock
+import kotlin.concurrent.read
+import kotlin.concurrent.write
+
+abstract class AsyncScriptDefinitionsContributor(protected val project: Project) : ScriptDefinitionContributor {
+    abstract val progressMessage: String
+
+    override fun isReady(): Boolean = definitions != null
+
+    override fun getDefinitions(): List<KotlinScriptDefinition> {
+        definitionsLock.read {
+            if (definitions != null) {
+                return definitions!!
+            }
+        }
+
+        forceStartUpdate = false
+        asyncRunUpdateScriptTemplates()
+        return emptyList()
+    }
+
+    protected fun asyncRunUpdateScriptTemplates() {
+        val backgroundTask = inProgressLock.write {
+            shouldStartNewUpdate = true
+            if (!inProgress) {
+                inProgress = true
+                return@write DefinitionsCollectorBackgroundTask()
+            }
+            return@write null
+        }
+        backgroundTask?.queue()
+    }
+
+    protected abstract fun loadScriptDefinitions(previous: List<KotlinScriptDefinition>?): List<KotlinScriptDefinition>
+
+    private var definitions: List<KotlinScriptDefinition>? = null
+    private val definitionsLock = ReentrantReadWriteLock()
+
+    protected var forceStartUpdate = false
+    protected var shouldStartNewUpdate = false
+
+    private var inProgress = false
+    private val inProgressLock = ReentrantReadWriteLock()
+
+    private inner class DefinitionsCollectorBackgroundTask : Task.Backgroundable(project, progressMessage, true) {
+
+        override fun run(indicator: ProgressIndicator) {
+            while (true) {
+                inProgressLock.read {
+                    if (indicator.isCanceled || !shouldStartNewUpdate) {
+                        inProgressLock.write {
+                            inProgress = false
+                        }
+                        return
+                    }
+                    shouldStartNewUpdate = false
+                }
+
+                val wasRunning = definitionsLock.isWriteLocked
+                val needReload = definitionsLock.write {
+                    if (wasRunning && !forceStartUpdate && definitions != null) return@write false
+                    val newDefinitions = loadScriptDefinitions(definitions)
+                    if (newDefinitions != definitions) {
+                        definitions = newDefinitions
+                        return@write true
+                    }
+                    return@write false
+                }
+
+                if (needReload) {
+                    if (isHeadless) {
+                        // If new script definitions found, then ScriptDefinitionsManager.reloadDefinitionsBy should be called
+                        // This may cause deadlock because Task.Backgroundable.queue executes task synchronously in headless mode
+                        ApplicationManager.getApplication().executeOnPooledThread {
+                            ScriptDefinitionsManager.getInstance(project).reloadDefinitionsBy(this@AsyncScriptDefinitionsContributor)
+                        }
+                    } else {
+                        ScriptDefinitionsManager.getInstance(project).reloadDefinitionsBy(this@AsyncScriptDefinitionsContributor)
+                    }
+                }
+            }
+        }
+    }
+}

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

@@ -6,25 +6,20 @@
 package org.jetbrains.kotlin.idea.script
 
 import com.intellij.ProjectTopics
-import com.intellij.openapi.progress.ProgressIndicator
-import com.intellij.openapi.progress.Task
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ModuleRootEvent
 import com.intellij.openapi.roots.ModuleRootListener
 import com.intellij.openapi.roots.OrderEnumerator
 import com.intellij.openapi.vfs.JarFileSystem
 import com.intellij.openapi.vfs.VirtualFile
-import org.jetbrains.kotlin.idea.core.script.ScriptDefinitionContributor
-import org.jetbrains.kotlin.idea.core.script.ScriptDefinitionsManager
 import org.jetbrains.kotlin.idea.core.script.loadDefinitionsFromTemplates
 import org.jetbrains.kotlin.idea.util.projectStructure.allModules
 import org.jetbrains.kotlin.script.KotlinScriptDefinition
 import java.io.File
 import java.util.concurrent.locks.ReentrantReadWriteLock
-import kotlin.concurrent.read
 import kotlin.concurrent.write
 
-class ScriptTemplatesFromDependenciesProvider(private val project: Project) : ScriptDefinitionContributor {
+class ScriptTemplatesFromDependenciesProvider(project: Project) : AsyncScriptDefinitionsContributor(project) {
 
     private var templates: TemplatesWithCp? = null
     private val templatesLock = ReentrantReadWriteLock()
@@ -34,6 +29,7 @@ class ScriptTemplatesFromDependenciesProvider(private val project: Project) : Sc
         connection.subscribe(ProjectTopics.PROJECT_ROOTS, object : ModuleRootListener {
             override fun rootsChanged(event: ModuleRootEvent) {
                 if (project.isInitialized) {
+                    forceStartUpdate = true
                     asyncRunUpdateScriptTemplates()
                 }
             }
@@ -41,69 +37,27 @@ class ScriptTemplatesFromDependenciesProvider(private val project: Project) : Sc
     }
 
     override val id = "ScriptTemplatesFromDependenciesProvider"
+    override val progressMessage = "Kotlin: scanning dependencies for script definitions..."
 
-    override fun getDefinitions(): List<KotlinScriptDefinition> {
+    override fun loadScriptDefinitions(previous: List<KotlinScriptDefinition>?): List<KotlinScriptDefinition> {
-        val templatesCopy = templatesLock.read {
+        val templatesCopy = templatesLock.write {
-            if (templates == null) {
-                syncUpdateScriptTemplates(false)
-            }
-            templates ?: return emptyList()
-        }
-        return loadDefinitionsFromTemplates(
-            templateClassNames = templatesCopy.templates,
-            templateClasspath = templatesCopy.classpath,
-            environment = mapOf(
-                "projectRoot" to (project.basePath ?: project.baseDir.canonicalPath)?.let(::File)
-            )
-        )
-    }
-
-    private var areRootsChanged = false
-    private var backgroundTask: TemplatesCollectorBackgroundTask? = null
-    private val backgroundTaskLock = ReentrantReadWriteLock()
-
-    private inner class TemplatesCollectorBackgroundTask
-        : Task.Backgroundable(project, "Kotlin: scanning dependencies for script definitions...", true) {
-
-        override fun run(indicator: ProgressIndicator) {
-            while (true) {
-                backgroundTaskLock.read {
-                    if (indicator.isCanceled || !areRootsChanged) {
-                        backgroundTaskLock.write {
-                            backgroundTask = null
-                        }
-                        return
-                    }
-                    areRootsChanged = false
-                }
-                if (syncUpdateScriptTemplates(true)) {
-                    ScriptDefinitionsManager.getInstance(project).reloadDefinitionsBy(this@ScriptTemplatesFromDependenciesProvider)
-                }
-            }
-        }
-    }
-
-    private fun asyncRunUpdateScriptTemplates() {
-        backgroundTaskLock.write {
-            areRootsChanged = true
-            if (backgroundTask == null) {
-                backgroundTask = TemplatesCollectorBackgroundTask()
-                backgroundTask!!.queue()
-            }
-        }
-    }
-
-    private fun syncUpdateScriptTemplates(force: Boolean): Boolean {
-        val wasRunning = templatesLock.isWriteLocked
-        templatesLock.write {
-            if (wasRunning && !force && templates != null) return true
             val newTemplates = scriptDefinitionsFromDependencies(project)
             if (newTemplates != templates) {
                 templates = newTemplates
-                return true
+                return@write newTemplates
             }
-            return false
+            return@write null
         }
+        if (templatesCopy != null) {
+            return loadDefinitionsFromTemplates(
+                templateClassNames = templatesCopy.templates,
+                templateClasspath = templatesCopy.classpath,
+                environment = mapOf(
+                    "projectRoot" to (project.basePath ?: project.baseDir.canonicalPath)?.let(::File)
+                )
+            )
+        }
+        return previous ?: emptyList()
     }
 }