gradle.kts: merge data on failed gradle sync, fix updates scheduling in corner cases

idea/idea-gradle/src/org/jetbrains/kotlin/idea/scripting/gradle/importing/KotlinDslScrtiptModelProcessor.kt

@@ -105,6 +105,7 @@ private fun KotlinDslScriptsModel.toListOfScriptModels(project: Project): List<K
 class KotlinDslGradleBuildSync(val workingDir: String, val taskId: ExternalSystemTaskId) {
     val projectRoots = mutableSetOf<String>()
     val models = mutableListOf<KotlinDslScriptModel>()
+    var failed = false
 }
 
 fun saveScriptModels(project: Project, build: KotlinDslGradleBuildSync) {

idea/idea-gradle/src/org/jetbrains/kotlin/idea/scripting/gradle/importing/KotlinDslSyncListener.kt

@@ -10,6 +10,7 @@ import com.intellij.openapi.externalSystem.model.task.ExternalSystemTaskNotifica
 import com.intellij.openapi.externalSystem.model.task.ExternalSystemTaskType
 import org.jetbrains.kotlin.idea.framework.GRADLE_SYSTEM_ID
 import org.jetbrains.kotlin.idea.scripting.gradle.roots.GradleBuildRootsManager
+import java.lang.Exception
 
 class KotlinDslSyncListener : ExternalSystemTaskNotificationListenerAdapter() {
     override fun onStart(id: ExternalSystemTaskId, workingDir: String?) {
@@ -32,4 +33,26 @@ class KotlinDslSyncListener : ExternalSystemTaskNotificationListenerAdapter() {
 
         saveScriptModels(project, sync)
     }
+
+    override fun onFailure(id: ExternalSystemTaskId, e: Exception) {
+        if (id.type != ExternalSystemTaskType.RESOLVE_PROJECT) return
+        if (id.projectSystemId != GRADLE_SYSTEM_ID) return
+
+        val project = id.findProject() ?: return
+        val sync = project.kotlinGradleDslSync[id] ?: return
+
+        sync.failed = true
+    }
+
+    override fun onCancel(id: ExternalSystemTaskId) {
+        if (id.type != ExternalSystemTaskType.RESOLVE_PROJECT) return
+        if (id.projectSystemId != GRADLE_SYSTEM_ID) return
+
+        val project = id.findProject() ?: return
+
+        val cancelled = project.kotlinGradleDslSync.remove(id)
+        if (cancelled != null) {
+            GradleBuildRootsManager.getInstance(project).markImportingInProgress(cancelled.workingDir, false)
+        }
+    }
 }

idea/idea-gradle/src/org/jetbrains/kotlin/idea/scripting/gradle/roots/GradleBuildRootIndex.kt

@@ -44,10 +44,12 @@ class GradleBuildRootIndex {
     fun getBuildByProjectDir(projectDir: String) = byProjectDir[projectDir]
 
     @Synchronized
-    operator fun set(prefix: String, value: GradleBuildRoot.Linked) {
+    fun add(value: GradleBuildRoot.Linked): GradleBuildRoot.Linked? {
+        val prefix = value.pathPrefix
         val old = byWorkingDir.put(prefix, value)
         rebuildProjectRoots()
         log.info("$prefix: $old -> $value")
+        return old
     }
 
     @Synchronized

idea/idea-gradle/src/org/jetbrains/kotlin/idea/scripting/gradle/roots/GradleBuildRootsManager.kt

@@ -5,16 +5,26 @@
 
 package org.jetbrains.kotlin.idea.scripting.gradle.roots
 
+import com.intellij.codeInsight.daemon.DaemonCodeAnalyzer
+import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.externalSystem.util.ExternalSystemApiUtil
+import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.vfs.LocalFileSystem
 import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiManager
+import com.intellij.ui.EditorNotifications
+import kotlinx.coroutines.GlobalScope
+import kotlinx.coroutines.launch
 import org.jetbrains.kotlin.idea.core.script.ScriptConfigurationManager
 import org.jetbrains.kotlin.idea.core.script.configuration.CompositeScriptConfigurationManager
 import org.jetbrains.kotlin.idea.core.script.configuration.ScriptingSupport
 import org.jetbrains.kotlin.idea.core.script.configuration.utils.ScriptClassRootsCache
+import org.jetbrains.kotlin.idea.core.util.EDT
 import org.jetbrains.kotlin.idea.scripting.gradle.*
 import org.jetbrains.kotlin.idea.scripting.gradle.importing.KotlinDslGradleBuildSync
+import org.jetbrains.kotlin.idea.scripting.gradle.importing.KotlinDslScriptModel
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
 import org.jetbrains.plugins.gradle.config.GradleSettingsListenerAdapter
@@ -44,13 +54,13 @@ import java.nio.file.Paths
  *   - [GradleBuildRoot.Imported] - imported
  */
 class GradleBuildRootsManager(val project: Project) : ScriptingSupport.Provider() {
-    val manager: CompositeScriptConfigurationManager
+    private val manager: CompositeScriptConfigurationManager
         get() = ScriptConfigurationManager.getInstance(project) as CompositeScriptConfigurationManager
 
     private val updater
         get() = manager.updater
 
-    val roots = GradleBuildRootIndex()
+    private val roots = GradleBuildRootIndex()
 
     ////////////
     /// ScriptingSupport.Provider implementation:
@@ -127,9 +137,10 @@ class GradleBuildRootsManager(val project: Project) : ScriptingSupport.Provider(
         return ScriptUnderRoot(GradleBuildRoot.Unlinked())
     }
 
-
     fun markImportingInProgress(workingDir: String, inProgress: Boolean = true) {
         actualizeBuildRoot(workingDir)?.importing = inProgress
+        updateNotifications(workingDir)
+        hideNotificationForProjectImport(project)
     }
 
     fun update(build: KotlinDslGradleBuildSync) {
@@ -139,22 +150,37 @@ class GradleBuildRootsManager(val project: Project) : ScriptingSupport.Provider(
             if (root is GradleBuildRoot.Imported && root.data.models.isEmpty()) return
         }
 
-        val old = actualizeBuildRoot(build.workingDir) ?: return
+        val root = actualizeBuildRoot(build.workingDir) ?: return
-        old.importing = false
+        root.importing = false
 
-        if (old is GradleBuildRoot.Legacy) return
+        if (root is GradleBuildRoot.Legacy) return
 
         val templateClasspath = GradleScriptDefinitionsContributor.getDefinitionsTemplateClasspath(project)
-        val data = GradleBuildRootData(build.projectRoots, templateClasspath, build.models)
+        val newData = GradleBuildRootData(build.projectRoots, templateClasspath, build.models)
-        val newSupport = tryCreateImportedRoot(build.workingDir) { data } ?: return
+        val mergedData = if (build.failed && root is GradleBuildRoot.Imported) merge(root.data, newData) else newData
-        GradleBuildRootDataSerializer.write(newSupport.dir, data)
+
-        roots[build.workingDir] = newSupport
+        val newSupport = tryCreateImportedRoot(build.workingDir) { mergedData } ?: return
-        manager.updater.ensureUpdateScheduled()
+        GradleBuildRootDataSerializer.write(newSupport.dir, mergedData)
+
+        add(newSupport)
+
+        hideNotificationForProjectImport(project)
+    }
+
+    private fun merge(old: GradleBuildRootData, new: GradleBuildRootData): GradleBuildRootData {
+        val roots = old.projectRoots.toMutableSet()
+        roots.addAll(new.projectRoots)
+
+        val models = old.models.associateByTo(mutableMapOf()) { it.file }
+        new.models.associateByTo(models) { it.file }
+
+        return GradleBuildRootData(roots, new.templateClasspath, models.values)
     }
 
     init {
         getGradleProjectSettings(project).forEach {
-            roots[it.externalProjectPath] = loadLinkedRoot(it)
+            // don't call this.add, as we are inside scripting manager initialization
+            roots.add(loadLinkedRoot(it))
         }
 
         // subscribe to linked gradle project modification
@@ -167,23 +193,16 @@ class GradleBuildRootsManager(val project: Project) : ScriptingSupport.Provider(
 
             override fun onProjectsUnlinked(linkedProjectPaths: MutableSet<String>) {
                 linkedProjectPaths.forEach {
-                    val buildRoot = VfsUtil.findFile(Paths.get(it), false)
+                    remove(it)
-                    if (buildRoot != null) {
-                        if (roots.remove(buildRoot.localPath) != null) {
-                            GradleBuildRootDataSerializer.remove(buildRoot)
-                        }
-                    }
                 }
             }
 
             override fun onGradleHomeChange(oldPath: String?, newPath: String?, linkedProjectPath: String) {
-                updateBuildRoot(linkedProjectPath)
+                reloadBuildRoot(linkedProjectPath)
-                manager.updater.ensureUpdateScheduled()
             }
 
             override fun onGradleDistributionTypeChange(currentValue: DistributionType?, linkedProjectPath: String) {
-                updateBuildRoot(linkedProjectPath)
+                reloadBuildRoot(linkedProjectPath)
-                manager.updater.ensureUpdateScheduled()
             }
         }
 
@@ -206,7 +225,7 @@ class GradleBuildRootsManager(val project: Project) : ScriptingSupport.Provider(
 
         return when {
             buildRoot != null -> when {
-                !buildRoot.checkActual(actualSettings) -> updateBuildRoot(workingDir)
+                !buildRoot.checkActual(actualSettings) -> reloadBuildRoot(workingDir)
                 else -> buildRoot
             }
             actualSettings != null -> loadLinkedRoot(actualSettings)
@@ -222,22 +241,23 @@ class GradleBuildRootsManager(val project: Project) : ScriptingSupport.Provider(
         return knownAsSupported == shouldBeSupported
     }
 
-    private fun updateBuildRoot(rootPath: String): GradleBuildRoot.Linked? {
+    private fun reloadBuildRoot(rootPath: String): GradleBuildRoot.Linked? {
         val settings = getGradleProjectSettings(rootPath)
         if (settings == null) {
-            val removed = roots.remove(rootPath)
+            remove(rootPath)
-            if (removed is GradleBuildRoot.Imported) {
-                updater.ensureUpdateScheduled()
-            }
             return null
         } else {
             val newRoot = loadLinkedRoot(settings)
-            roots[newRoot.pathPrefix] = newRoot
+            add(newRoot)
-            updater.ensureUpdateScheduled()
             return newRoot
         }
     }
 
+    private fun removeData(rootPath: String) {
+        val buildRoot = LocalFileSystem.getInstance().findFileByPath(rootPath)
+        if (buildRoot != null) GradleBuildRootDataSerializer.remove(buildRoot)
+    }
+
     private fun loadLinkedRoot(settings: GradleProjectSettings) =
         tryLoadFromFsCache(settings) ?: createOtherLinkedRoot(settings)
 
@@ -264,13 +284,47 @@ class GradleBuildRootsManager(val project: Project) : ScriptingSupport.Provider(
             .getExecutionSettings<GradleExecutionSettings>(project, externalProjectPath, GradleConstants.SYSTEM_ID)
             .javaHome?.let { File(it) }
 
-        val newSupport = GradleBuildRoot.Imported(project, buildRoot, javaHome, data)
+        return GradleBuildRoot.Imported(project, buildRoot, javaHome, data)
-        val oldSupport = roots.getBuildRoot(externalProjectPath)
+    }
-        if (oldSupport != null) {
+
-            hideNotificationForProjectImport(project)
+    private fun add(newRoot: GradleBuildRoot.Linked) {
+        val old = roots.add(newRoot)
+        if (old is GradleBuildRoot.Imported) removeData(old.pathPrefix)
+        if (old is GradleBuildRoot.Imported || newRoot is GradleBuildRoot.Imported) {
+            updater.ensureUpdateScheduled()
         }
 
-        return newSupport
+        updateNotifications(newRoot.pathPrefix)
+    }
+
+    private fun remove(rootPath: String) {
+        val removed = roots.remove(rootPath)
+        if (removed is GradleBuildRoot.Imported) {
+            removeData(rootPath)
+            updater.ensureUpdateScheduled()
+        }
+
+        updateNotifications(rootPath)
+    }
+
+    private fun updateNotifications(dir1: String) {
+        if (!project.isOpen) return
+
+        val openedScripts = FileEditorManager.getInstance(project).openFiles.filter {
+            it.path.startsWith(dir1) && isGradleKotlinScript(it)
+        }
+
+        if (openedScripts.isEmpty()) return
+
+        GlobalScope.launch(EDT(project)) {
+            if (project.isDisposed) return@launch
+
+            openedScripts.forEach {
+                val ktFile = PsiManager.getInstance(project).findFile(it)
+                if (ktFile != null) DaemonCodeAnalyzer.getInstance(project).restart(ktFile)
+                EditorNotifications.getInstance(project).updateAllNotifications()
+            }
+        }
     }
 
     companion object {