IJPL-162718 WorkspaceEntitiesLifecycleActivity - get rid of WriteAction.runAndWait

(cherry picked from commit 4879c07b3e8f4db3b4663a1bbfd098a4fe5ff0d5)

IJ-CR-146879

GitOrigin-RevId: eba1be90aed6c7059a29e890361171a15465ed9e

platform/lang-impl/src/com/intellij/ide/scratch/workspace/ScratchWorkspaceStartupActivity.kt

@@ -2,27 +2,32 @@
 package com.intellij.ide.scratch.workspace
 
 import com.intellij.ide.scratch.RootType
-import com.intellij.openapi.Disposable
-import com.intellij.openapi.components.Service
-import com.intellij.openapi.components.serviceAsync
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.extensions.ExtensionNotApplicableException
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.startup.ProjectActivity
 import com.intellij.platform.backend.workspace.WorkspaceEntityLifecycleSupporter
 import com.intellij.workspaceModel.ide.impl.WorkspaceEntityLifecycleSupporterUtils
+import kotlinx.coroutines.awaitCancellation
+import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.launch
 
 private class ScratchWorkspaceStartupActivity : ProjectActivity {
-  override suspend fun execute(project: Project) {
-    val listener = {
-      val provider = WorkspaceEntityLifecycleSupporter.EP_NAME.findExtensionOrFail(ScratchEntityLifecycleSupporter::class.java)
-      WorkspaceEntityLifecycleSupporterUtils.ensureEntitiesInWorkspaceAreAsProviderDefined(project, provider)
+  init {
+    if (ApplicationManager.getApplication().isUnitTestMode) {
+      throw ExtensionNotApplicableException.create()
     }
-    RootType.ROOT_EP.addChangeListener(listener, project.serviceAsync<ScratchDisposableService>())
   }
-}
 
-@Service(Service.Level.PROJECT)
-private class ScratchDisposableService : Disposable {
-  override fun dispose() {
+  override suspend fun execute(project: Project) {
+    coroutineScope {
+      RootType.ROOT_EP.addChangeListener(coroutineScope = this) {
+        val provider = WorkspaceEntityLifecycleSupporter.EP_NAME.findExtensionOrFail(ScratchEntityLifecycleSupporter::class.java)
+        launch {
+          WorkspaceEntityLifecycleSupporterUtils.ensureEntitiesInWorkspaceAreAsProviderDefined(project, provider)
+        }
+      }
+      awaitCancellation()
+    }
   }
-}
-
+}

platform/projectModel-impl/src/com/intellij/workspaceModel/ide/impl/WorkspaceEntitiesLifecycleActivity.kt

@@ -1,29 +1,37 @@
 // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.workspaceModel.ide.impl
 
-import com.intellij.openapi.Disposable
-import com.intellij.openapi.components.Service
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.extensions.ExtensionPointListener
 import com.intellij.openapi.extensions.PluginDescriptor
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.startup.ProjectActivity
 import com.intellij.platform.backend.workspace.WorkspaceEntityLifecycleSupporter
 import com.intellij.platform.workspace.storage.WorkspaceEntity
+import kotlinx.coroutines.awaitCancellation
+import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.launch
 
-internal class WorkspaceEntitiesLifecycleActivity : ProjectActivity {
+private class WorkspaceEntitiesLifecycleActivity : ProjectActivity {
   override suspend fun execute(project: Project) {
+    if (ApplicationManager.getApplication().isUnitTestMode) {
+      //todo GoGutterMarkerTest failed without it - ideally, we should setup project in tests without this activity
+      WorkspaceEntityLifecycleSupporterUtils.ensureAllEntitiesInWorkspaceAreAsProvidersDefined(project)
+      return
+    }
 
-    WorkspaceEntityLifecycleSupporter.EP_NAME.addExtensionPointListener(object : ExtensionPointListener<WorkspaceEntityLifecycleSupporter<out WorkspaceEntity, out WorkspaceEntity.Builder<out WorkspaceEntity>>> {
-      override fun extensionAdded(extension: WorkspaceEntityLifecycleSupporter<out WorkspaceEntity, out WorkspaceEntity.Builder<out WorkspaceEntity>>, pluginDescriptor: PluginDescriptor) {
-        WorkspaceEntityLifecycleSupporterUtils.ensureEntitiesInWorkspaceAreAsProviderDefined(project, extension)
-      }
-    }, project.getService(ConstantEntitiesDisposableService::class.java))
+    coroutineScope {
+      WorkspaceEntityLifecycleSupporter.EP_NAME.addExtensionPointListener(this, object : ExtensionPointListener<WorkspaceEntityLifecycleSupporter<out WorkspaceEntity, out WorkspaceEntity.Builder<out WorkspaceEntity>>> {
+        override fun extensionAdded(extension: WorkspaceEntityLifecycleSupporter<out WorkspaceEntity, out WorkspaceEntity.Builder<out WorkspaceEntity>>, pluginDescriptor: PluginDescriptor) {
+          launch {
+            WorkspaceEntityLifecycleSupporterUtils.ensureEntitiesInWorkspaceAreAsProviderDefined(project, extension)
+          }
+        }
+      })
 
-    WorkspaceEntityLifecycleSupporterUtils.ensureAllEntitiesInWorkspaceAreAsProvidersDefined(project)
+      WorkspaceEntityLifecycleSupporterUtils.ensureAllEntitiesInWorkspaceAreAsProvidersDefined(project)
+
+      awaitCancellation()
+    }
   }
-}
-
-@Service(Service.Level.PROJECT)
-internal class ConstantEntitiesDisposableService : Disposable {
-  override fun dispose() {}
-}
+}

platform/projectModel-impl/src/com/intellij/workspaceModel/ide/impl/WorkspaceEntityLifecycleSupporterUtils.kt

@@ -1,8 +1,8 @@
 // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.workspaceModel.ide.impl
 
-import com.intellij.openapi.application.WriteAction
 import com.intellij.openapi.application.backgroundWriteAction
+import com.intellij.openapi.components.serviceAsync
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Ref
 import com.intellij.platform.backend.workspace.WorkspaceEntityLifecycleSupporter
@@ -15,39 +15,38 @@ import org.jetbrains.annotations.ApiStatus
 
 @ApiStatus.Internal
 object WorkspaceEntityLifecycleSupporterUtils {
-  fun ensureEntitiesInWorkspaceAreAsProviderDefined(project: Project, provider: WorkspaceEntityLifecycleSupporter<*, *>) {
-    val workspaceModel = WorkspaceModel.getInstance(project)
-    val builderRef: Ref<MutableEntityStorage?> = Ref.create()
+  suspend fun ensureEntitiesInWorkspaceAreAsProviderDefined(project: Project, provider: WorkspaceEntityLifecycleSupporter<*, *>) {
+    val workspaceModel = project.serviceAsync<WorkspaceModel>()
+    val builderRef = Ref<MutableEntityStorage>()
     val snapshot = workspaceModel.currentSnapshot
+    ensureInitialized(project = project, provider = provider, snapshot = snapshot, builderRef = builderRef)
 
-    ensureInitialized(project, provider, snapshot, builderRef)
-
-    builderRef.get()?.also { builder -> WriteAction.runAndWait<RuntimeException> { writeBuilder(workspaceModel, builder) } }
+    builderRef.get()?.let { backgroundWriteAction { writeBuilder(workspaceModel, it) } }
   }
 
   suspend fun ensureAllEntitiesInWorkspaceAreAsProvidersDefined(project: Project) {
-    val workspaceModel = WorkspaceModel.getInstance(project)
+    val workspaceModel = project.serviceAsync<WorkspaceModel>()
     val snapshot = workspaceModel.currentSnapshot
-    val builderRef: Ref<MutableEntityStorage?> = Ref.create()
+    val builderRef = Ref<MutableEntityStorage>()
 
     WorkspaceEntityLifecycleSupporter.EP_NAME.forEachExtensionSafe { provider ->
-      ensureInitialized(project, provider, snapshot, builderRef)
+      ensureInitialized(project = project, provider = provider, snapshot = snapshot, builderRef = builderRef)
     }
-    builderRef.get()?.also { builder -> backgroundWriteAction { writeBuilder(workspaceModel, builder) } }
+    builderRef.get()?.let { backgroundWriteAction { writeBuilder(workspaceModel, it) } }
   }
 
-
-  private fun writeBuilder(workspaceModel: WorkspaceModel,
-                           builder: MutableEntityStorage) {
-    workspaceModel.updateProjectModel("ConstantEntitiesCheckActivity") { tempBuilder: MutableEntityStorage ->
-      tempBuilder.applyChangesFrom(builder)
+  private fun writeBuilder(workspaceModel: WorkspaceModel, builder: MutableEntityStorage) {
+    workspaceModel.updateProjectModel("ConstantEntitiesCheckActivity") {
+      it.applyChangesFrom(builder)
     }
   }
 
-  private fun <E : WorkspaceEntity, M : WorkspaceEntity.Builder<E>> ensureInitialized(project: Project,
-                                                                                      provider: WorkspaceEntityLifecycleSupporter<E, M>,
-                                                                                      snapshot: ImmutableEntityStorage,
-                                                                                      builderRef: Ref<MutableEntityStorage?>) {
+  private fun <E : WorkspaceEntity, M : WorkspaceEntity.Builder<E>> ensureInitialized(
+    project: Project,
+    provider: WorkspaceEntityLifecycleSupporter<E, M>,
+    snapshot: ImmutableEntityStorage,
+    builderRef: Ref<MutableEntityStorage>,
+  ) {
     val expectedEntity = provider.createSampleEntity(project)
     val actualEntities = snapshot.entities(provider.getEntityClass()).toList()
 
@@ -59,10 +58,11 @@ object WorkspaceEntityLifecycleSupporterUtils {
       return
     }
 
-    if (builderRef.isNull) {
-      builderRef.set(snapshot.toBuilder())
+    var builder = builderRef.get()
+    if (builder == null) {
+      builder = snapshot.toBuilder()
+      builderRef.set(builder)
     }
-    val builder = builderRef.get()!!
 
     for (oldEntity in actualEntities) {
       builder.removeEntity(oldEntity)