Fix UI freezes, SlowOperations on EDT, and write-unsafe context issues

These issues were happening in the project creator, both from loading
the initial wizard, as well as during actual project creation.

Author: DenWav

src/main/kotlin/creator/JdkComboBoxWithPreference.kt

@@ -20,6 +20,7 @@
 
 package com.demonwav.mcdev.creator
 
+import com.demonwav.mcdev.util.invokeLater
 import com.intellij.ide.util.PropertiesComponent
 import com.intellij.ide.util.projectWizard.ProjectWizardUtil
 import com.intellij.ide.util.projectWizard.WizardContext
@@ -42,6 +43,7 @@ import com.intellij.openapi.util.Disposer
 import com.intellij.ui.dsl.builder.Cell
 import com.intellij.ui.dsl.builder.Row
 import javax.swing.JComponent
+import org.jetbrains.concurrency.runAsync
 
 internal class JdkPreferenceData(
     var jdk: JavaSdkVersion,
@@ -89,15 +91,17 @@ class JdkComboBoxWithPreference internal constructor(
             preferenceData.jdk = version
             reloadModel()
 
-            for (jdkVersion in version.ordinal until JavaSdkVersion.values().size) {
-                val jdk = JavaSdkVersion.values()[jdkVersion]
+            for (jdkVersion in version.ordinal until JavaSdkVersion.entries.size) {
+                val jdk = JavaSdkVersion.entries[jdkVersion]
 
                 val preferredSdkPath = preferenceData.sdkPathByJdk[jdk]
                 if (preferredSdkPath != null) {
                     val sdk = model.sdks.firstOrNull { it.homePath == preferredSdkPath }
                         ?: suggestions.firstOrNull { it.homePath == preferredSdkPath }
                     if (sdk != null) {
-                        setSelectedItem(sdk)
+                        runAsync {
+                            setSelectedItem(sdk)
+                        }
                         return
                     }
                 }
@@ -145,7 +149,7 @@ fun Row.jdkComboBoxWithPreference(
         for (preferenceDataStr in preferenceDataStrs) {
             val parts = preferenceDataStr.split('=', limit = 2)
             val featureVersion = parts.firstOrNull()?.toIntOrNull() ?: continue
-            val knownJdkVersions = JavaSdkVersion.values()
+            val knownJdkVersions = JavaSdkVersion.entries
             if (featureVersion !in knownJdkVersions.indices) {
                 continue
             }
@@ -176,7 +180,9 @@ fun Row.jdkComboBoxWithPreference(
     }
 
     val lastUsedSdk = stateComponent.getValue(selectedJdkProperty)
-    ProjectWizardUtil.preselectJdkForNewModule(project, lastUsedSdk, comboBox) { true }
+    runAsync {
+        ProjectWizardUtil.preselectJdkForNewModule(project, lastUsedSdk, comboBox) { true }
+    }
 
     val windowChild = context.getUserData(AbstractWizard.KEY)!!.contentPanel
     comboBox.loadSuggestions(windowChild, context.disposable)

src/main/kotlin/creator/custom/CreatorTemplateProcessor.kt

@@ -34,6 +34,7 @@ import com.intellij.codeInsight.actions.ReformatCodeProcessor
 import com.intellij.ide.projectView.ProjectView
 import com.intellij.ide.util.projectWizard.WizardContext
 import com.intellij.openapi.application.WriteAction
+import com.intellij.openapi.concurrency.awaitPromise
 import com.intellij.openapi.diagnostic.Attachment
 import com.intellij.openapi.diagnostic.ControlFlowException
 import com.intellij.openapi.diagnostic.thisLogger
@@ -53,11 +54,15 @@ import com.intellij.ui.dsl.builder.Panel
 import com.intellij.ui.dsl.builder.panel
 import com.intellij.util.application
 import java.nio.file.Path
+import java.time.Duration
+import java.util.concurrent.TimeUnit
 import java.util.function.Consumer
 import kotlin.io.path.createDirectories
 import kotlin.io.path.writeText
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.launch
+import org.jetbrains.concurrency.await
+import org.jetbrains.concurrency.runAsync
 
 interface ExternalTemplatePropertyProvider {
 
@@ -268,7 +273,11 @@ class CreatorTemplateProcessor(
                 destPath.parent.createDirectories()
                 destPath.writeText(processedContent)
 
-                val virtualFile = destPath.refreshAndFindVirtualFile()
+                val virtualFile = runCatching {
+                    runAsync {
+                        destPath.refreshAndFindVirtualFile()
+                    }.blockingGet(20, TimeUnit.MILLISECONDS)
+                }.getOrNull()
                 if (virtualFile != null) {
                     generatedFiles.add(file to virtualFile)
                 } else {

src/main/kotlin/creator/custom/CustomPlatformStep.kt

@@ -27,12 +27,18 @@ import com.demonwav.mcdev.creator.custom.providers.LoadedTemplate
 import com.demonwav.mcdev.creator.custom.providers.TemplateProvider
 import com.demonwav.mcdev.creator.modalityState
 import com.demonwav.mcdev.util.getOrLogException
+import com.demonwav.mcdev.util.invokeAndWait
+import com.demonwav.mcdev.util.runWriteTask
+import com.demonwav.mcdev.util.tryWriteSafeContext
 import com.intellij.ide.wizard.AbstractNewProjectWizardStep
 import com.intellij.ide.wizard.GitNewProjectWizardData
 import com.intellij.ide.wizard.NewProjectWizardBaseData
 import com.intellij.ide.wizard.NewProjectWizardStep
 import com.intellij.openapi.application.EDT
+import com.intellij.openapi.application.ModalityState
+import com.intellij.openapi.application.TransactionGuard
 import com.intellij.openapi.application.asContextElement
+import com.intellij.openapi.application.impl.ModalityStateEx
 import com.intellij.openapi.application.runWriteAction
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.observable.properties.GraphProperty
@@ -51,10 +57,12 @@ import com.intellij.ui.dsl.builder.bindText
 import com.intellij.util.application
 import com.intellij.util.ui.AsyncProcessIcon
 import javax.swing.JLabel
+import javax.swing.SwingUtilities
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.launch
+import org.jetbrains.kotlin.descriptors.Modality
 
 /**
  * The step to select a custom template repo.
@@ -248,11 +256,9 @@ class CustomPlatformStep(
         templateLoadingTextProperty.set(MCDevBundle("creator.step.generic.load_template.message"))
         templateLoadingProperty.set(true)
 
-        // For some reason syncRefresh doesn't play nice with writeAction() coroutines so we do it beforehand
-        application.invokeAndWait(
-            { runWriteAction { VirtualFileManager.getInstance().syncRefresh() } },
-            context.modalityState
-        )
+        tryWriteSafeContext(context.modalityState) {
+            VirtualFileManager.getInstance().syncRefresh()
+        }
 
         val dialogCoroutineContext = context.modalityState.asContextElement()
         val uiContext = dialogCoroutineContext + Dispatchers.EDT

src/main/kotlin/creator/custom/providers/RemoteTemplateProvider.kt

@@ -28,14 +28,18 @@ import com.demonwav.mcdev.creator.custom.TemplateDescriptor
 import com.demonwav.mcdev.creator.modalityState
 import com.demonwav.mcdev.creator.selectProxy
 import com.demonwav.mcdev.update.PluginUtil
+import com.demonwav.mcdev.util.asyncIO
 import com.demonwav.mcdev.util.capitalize
+import com.demonwav.mcdev.util.invokeAndWait
 import com.demonwav.mcdev.util.refreshSync
 import com.github.kittinunf.fuel.core.FuelManager
 import com.github.kittinunf.fuel.coroutines.awaitByteArrayResult
 import com.github.kittinunf.result.getOrNull
 import com.github.kittinunf.result.onError
 import com.intellij.ide.util.projectWizard.WizardContext
 import com.intellij.openapi.application.PathManager
+import com.intellij.openapi.application.readAction
+import com.intellij.openapi.application.writeAction
 import com.intellij.openapi.diagnostic.ControlFlowException
 import com.intellij.openapi.diagnostic.thisLogger
 import com.intellij.openapi.observable.properties.PropertyGraph
@@ -44,6 +48,7 @@ import com.intellij.openapi.observable.util.transform
 import com.intellij.openapi.observable.util.trim
 import com.intellij.openapi.progress.ProgressIndicator
 import com.intellij.openapi.progress.ProgressManager
+import com.intellij.openapi.progress.blockingContext
 import com.intellij.openapi.util.NlsContexts
 import com.intellij.openapi.vfs.JarFileSystem
 import com.intellij.ui.CollectionComboBoxModel
@@ -57,16 +62,21 @@ import com.intellij.ui.dsl.builder.bindText
 import com.intellij.ui.dsl.builder.columns
 import com.intellij.ui.dsl.builder.panel
 import com.intellij.ui.dsl.builder.textValidation
+import com.intellij.util.application
 import com.intellij.util.io.createDirectories
 import java.awt.Component
 import java.nio.file.Path
 import javax.swing.JComponent
 import javax.swing.JLabel
 import javax.swing.JList
 import javax.swing.ListCellRenderer
+import javax.swing.SwingUtilities
 import kotlin.io.path.absolutePathString
 import kotlin.io.path.exists
 import kotlin.io.path.writeBytes
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.withContext
 
 open class RemoteTemplateProvider : TemplateProvider {
 
@@ -144,23 +154,27 @@ open class RemoteTemplateProvider : TemplateProvider {
         return doLoadTemplates(context, repo, remoteRepo.innerPath)
     }
 
-    protected fun doLoadTemplates(
+    protected suspend fun doLoadTemplates(
         context: WizardContext,
         repo: MinecraftSettings.TemplateRepo,
         rawInnerPath: String
-    ): List<LoadedTemplate> {
+    ): List<LoadedTemplate> = withContext(Dispatchers.IO) { // don't run on EDT
         val remoteRootPath = RemoteTemplateRepo.getDestinationZip(repo.name)
         if (!remoteRootPath.exists()) {
-            return emptyList()
+            return@withContext emptyList()
         }
 
         val archiveRoot = remoteRootPath.absolutePathString() + JarFileSystem.JAR_SEPARATOR
 
         val fs = JarFileSystem.getInstance()
+
         val rootFile = fs.refreshAndFindFileByPath(archiveRoot)
-            ?: return emptyList()
+            ?: return@withContext emptyList()
         val modalityState = context.modalityState
-        rootFile.refreshSync(modalityState)
+
+        blockingContext {
+            rootFile.refreshSync(modalityState)
+        }
 
         val innerPath = replaceVariables(rawInnerPath)
         val repoRoot = if (innerPath.isNotBlank()) {
@@ -170,10 +184,10 @@ open class RemoteTemplateProvider : TemplateProvider {
         }
 
         if (repoRoot == null) {
-            return emptyList()
+            return@withContext emptyList()
         }
 
-        return TemplateProvider.findTemplates(modalityState, repoRoot)
+        return@withContext TemplateProvider.findTemplates(modalityState, repoRoot)
     }
 
     private fun replaceVariables(originalRepoUrl: String): String =

src/main/kotlin/creator/custom/providers/TemplateProvider.kt

@@ -45,7 +45,9 @@ import com.intellij.serviceContainer.BaseKeyedLazyInstance
 import com.intellij.util.KeyedLazyInstance
 import com.intellij.util.xmlb.annotations.Attribute
 import java.util.ResourceBundle
+import java.util.concurrent.TimeUnit
 import javax.swing.JComponent
+import org.jetbrains.concurrency.runAsync
 
 /**
  * Extensions responsible for creating a [TemplateDescriptor] based on whatever data it is provided in its configuration
@@ -151,8 +153,9 @@ interface TemplateProvider {
             }
 
             try {
-                return file.refreshSync(modalityState)
-                    ?.inputStream?.reader()?.use { TemplateResourceBundle(it, parent) }
+                return runAsync {
+                    file.inputStream.reader().use { TemplateResourceBundle(it, parent) }
+                }.blockingGet(20, TimeUnit.MILLISECONDS)
             } catch (t: Throwable) {
                 if (t is ControlFlowException) {
                     return parent
@@ -171,8 +174,12 @@ interface TemplateProvider {
             tooltip: String? = null,
             bundle: ResourceBundle? = null
         ): VfsLoadedTemplate? {
-            descriptorFile.refreshSync(modalityState)
-            var descriptor = Gson().fromJson<TemplateDescriptor>(descriptorFile.readText())
+            var descriptor = runCatching {
+                runAsync {
+                    descriptorFile.refreshSync(modalityState)
+                    Gson().fromJson<TemplateDescriptor>(descriptorFile.readText())
+                }.blockingGet(100, TimeUnit.MILLISECONDS)
+            }.getOrNull() ?: return null
             if (descriptor.version != TemplateDescriptor.FORMAT_VERSION) {
                 thisLogger().warn("Cannot handle template ${descriptorFile.path} of version ${descriptor.version}")
                 return null

src/main/kotlin/update/PluginUpdater.kt

@@ -23,14 +23,14 @@ package com.demonwav.mcdev.update
 import com.demonwav.mcdev.util.findDeclaredField
 import com.demonwav.mcdev.util.forEachNotNull
 import com.demonwav.mcdev.util.invokeLater
-import com.demonwav.mcdev.util.invokeLaterAny
 import com.intellij.ide.plugins.IdeaPluginDescriptor
 import com.intellij.ide.plugins.PluginManagerCore
 import com.intellij.ide.plugins.PluginManagerMain
 import com.intellij.ide.plugins.PluginNode
 import com.intellij.ide.plugins.RepositoryHelper
 import com.intellij.openapi.application.ApplicationInfo
 import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.ModalityState
 import com.intellij.openapi.progress.ProgressIndicator
 import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.progress.Task
@@ -58,7 +58,7 @@ object PluginUpdater {
                 .forEachNotNull { updateStatus = updateStatus.mergeWith(checkUpdatesInCustomRepo(it)) }
 
             val finalUpdate = updateStatus
-            invokeLaterAny { callback(finalUpdate) }
+            invokeLater(ModalityState.any()) { callback(finalUpdate) }
         } catch (e: Exception) {
             PluginUpdateStatus.CheckFailed("Minecraft Development plugin update check failed")
         }

src/main/kotlin/util/files.kt

@@ -20,18 +20,20 @@
 
 package com.demonwav.mcdev.util
 
-import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.ModalityState
+import com.intellij.openapi.application.TransactionGuard
 import com.intellij.openapi.vfs.LocalFileSystem
 import com.intellij.openapi.vfs.VfsUtilCore
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.newvfs.RefreshQueue
+import com.intellij.util.application
 import java.io.File
 import java.io.IOException
 import java.nio.file.Path
 import java.util.jar.Attributes
 import java.util.jar.JarFile
 import java.util.jar.Manifest
+import javax.swing.SwingUtilities
 
 val VirtualFile.localFile: File
     get() = VfsUtilCore.virtualToIoFile(this)
@@ -79,17 +81,9 @@ operator fun Manifest.get(attribute: String): String? = mainAttributes.getValue(
 operator fun Manifest.get(attribute: Attributes.Name): String? = mainAttributes.getValue(attribute)
 
 fun VirtualFile.refreshSync(modalityState: ModalityState): VirtualFile? {
-    fun refresh() {
+    tryWriteSafeContext(modalityState) {
         RefreshQueue.getInstance().refresh(false, this.isDirectory, null, modalityState, this)
     }
 
-    if (ApplicationManager.getApplication().isWriteAccessAllowed) {
-        refresh()
-    } else {
-        runWriteTask {
-            refresh()
-        }
-    }
-
     return this.parent?.findOrCreateChildData(this, this.name)
 }