IJPL-158: implement proper modality handling in tests instead of suspicious DumbServiceImpl.executeImmediatelyOrScheduleOnEDT

GitOrigin-RevId: ebed32ab3214bcf3ab045dd40d674da1c3064e2f

platform/platform-impl/src/com/intellij/openapi/project/DumbServiceImpl.kt

@@ -195,22 +195,26 @@ open class DumbServiceImpl @NonInjectable @VisibleForTesting constructor(private
   @TestOnly
   fun <T> computeInDumbModeSynchronously(computable: ThrowableComputable<T, in Throwable>): T {
     val trace = Throwable()
-    application.invokeAndWait {
-      val old = myState.getAndUpdate { it.incrementDumbCounter() }
-      if (old.isSmart) {
-        dumbModeStartTrace = trace
-        runCatching { myPublisher.enteredDumbMode() }
+    if (myState.getAndUpdate { it.tryIncrementDumbCounter() }.incrementWillChangeDumbState) {
+      application.invokeAndWait {
+        val old = myState.getAndUpdate { it.incrementDumbCounter() }
+        if (old.isSmart) {
+          dumbModeStartTrace = trace
+          runCatching { myPublisher.enteredDumbMode() }
+        }
       }
     }
     try {
       return computable.compute()
     }
     finally {
-      application.invokeAndWait {
-        val new = myState.updateAndGet { it.decrementDumbCounter() }
-        if (new.isSmart) {
-          dumbModeStartTrace = null
-          runCatching { myPublisher.exitDumbMode() }
+      if (myState.getAndUpdate { it.tryDecrementDumbCounter() }.decrementWillChangeDumbState) {
+        application.invokeAndWait {
+          val new = myState.updateAndGet { it.decrementDumbCounter() }
+          if (new.isSmart) {
+            dumbModeStartTrace = null
+            runCatching { myPublisher.exitDumbMode() }
+          }
         }
       }
     }
@@ -225,39 +229,37 @@ open class DumbServiceImpl @NonInjectable @VisibleForTesting constructor(private
   @TestOnly
   suspend fun <T> runInDumbMode(block: suspend () -> T): T {
     val trace = Throwable()
-    executeImmediatelyOrScheduleOnEDT {
-      val old = myState.getAndUpdate { it.incrementDumbCounter() }
-      if (old.isSmart) {
-        dumbModeStartTrace = trace
-        runCatching { myPublisher.enteredDumbMode() }
+    if (myState.getAndUpdate { it.tryIncrementDumbCounter() }.incrementWillChangeDumbState) {
+      // If already dumb - just increment the counter. We don't need a write action (to not interrupt NBRA), neither we need EDT.
+      // Otherwise, increment the counter on EDT under write action (write action requirement is not implemented in tests yet),
+      //   because this will change dumb state
+      withContext(Dispatchers.EDT) {
+        val old = myState.getAndUpdate { it.incrementDumbCounter() }
+        if (old.isSmart) {
+          dumbModeStartTrace = trace
+          runCatching { myPublisher.enteredDumbMode() }
+        }
       }
     }
     try {
       return block()
     }
     finally {
-      executeImmediatelyOrScheduleOnEDT {
-        val new = myState.updateAndGet { it.decrementDumbCounter() }
-        if (new.isSmart) {
-          dumbModeStartTrace = null
-          runCatching { myPublisher.exitDumbMode() }
+      // If there are other dumb tasks - just decrement the counter. We don't need a write action (to not interrupt NBRA), neither we need EDT.
+      // Otherwise, decrement the counter on EDT under write action (write action requirement is not implemented in tests yet)
+      //   because this will change dumb state
+      if (myState.getAndUpdate { it.tryDecrementDumbCounter() }.decrementWillChangeDumbState) {
+        withContext(Dispatchers.EDT) {
+          val new = myState.updateAndGet { it.decrementDumbCounter() }
+          if (new.isSmart) {
+            dumbModeStartTrace = null
+            runCatching { myPublisher.exitDumbMode() }
+          }
         }
       }
     }
   }
 
-  private suspend fun executeImmediatelyOrScheduleOnEDT(block: suspend () -> Unit) {
-    //Dispatchers.EDT, Dispatchers.Main, and even Dispatchers.Main.immediate may never execute if already on EDT. See SwiftAttributeCompletionTest
-    if (application.isDispatchThread) {
-      block()
-    }
-    else {
-      withContext(Dispatchers.EDT) {
-        block()
-      }
-    }
-  }
-
   override fun runWhenSmart(@Async.Schedule runnable: Runnable) {
     myProject.getService(SmartModeScheduler::class.java).runWhenSmart(runnable)
   }
@@ -541,10 +543,12 @@ open class DumbServiceImpl @NonInjectable @VisibleForTesting constructor(private
       }
     }
 
+    val incrementWillChangeDumbState: Boolean = isSmart
+    val decrementWillChangeDumbState: Boolean = dumbCounter == 1
     fun incrementDumbCounter(): DumbState = nextCounterState(dumbCounter + 1)
     fun decrementDumbCounter(): DumbState = nextCounterState(dumbCounter - 1)
-    fun tryIncrementDumbCounter(): DumbState = if (isDumb) incrementDumbCounter() else this
-    fun tryDecrementDumbCounter(): DumbState = if (dumbCounter > 1) decrementDumbCounter() else this
+    fun tryIncrementDumbCounter(): DumbState = if (incrementWillChangeDumbState) this else incrementDumbCounter()
+    fun tryDecrementDumbCounter(): DumbState = if (decrementWillChangeDumbState) this else decrementDumbCounter()
     val isSmart: Boolean get() = !isDumb
   }
 

platform/testFramework/src/com/intellij/testFramework/DumbModeTestUtils.kt

@@ -1,6 +1,8 @@
 // Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.testFramework
 
+import com.intellij.openapi.progress.runBlockingMaybeCancellable
+import com.intellij.openapi.progress.runWithModalProgressBlocking
 import com.intellij.openapi.project.DumbService
 import com.intellij.openapi.project.DumbServiceImpl
 import com.intellij.openapi.project.Project
@@ -8,8 +10,6 @@ import com.intellij.util.application
 import kotlinx.coroutines.*
 import org.junit.Assert.assertFalse
 import org.junit.Assert.assertTrue
-import org.junit.jupiter.api.fail
-import kotlin.time.Duration
 import kotlin.time.DurationUnit.SECONDS
 import kotlin.time.toDuration
 
@@ -27,22 +27,36 @@ object DumbModeTestUtils {
    */
   @JvmStatic
   fun startEternalDumbModeTask(project: Project): EternalTaskShutdownToken {
-    var dumbModeJob: Job? = null
-    @Suppress("RAW_RUN_BLOCKING")
-    runBlocking {
-      val dumbModeStarted = CompletableDeferred<Boolean>()
+    val finishDumbTask = CompletableDeferred<Boolean>()
+    runModalIfEdt(project) {
+      val dumbTaskStarted = CompletableDeferred<Boolean>()
       withTimeout(10.toDuration(SECONDS)) {
-        dumbModeJob = CoroutineScope(Dispatchers.Main.immediate + Job()).launch {
-          DumbServiceImpl.getInstance(project).runInDumbMode {
-            dumbModeStarted.complete(true)
-            delay(Duration.INFINITE)
+        DumbServiceImpl.getInstance(project).runInDumbMode {
+          CoroutineScope(Job()).launch {
+            // The trick is that we don't need EDT, neither a write action to start dumb task, if we already in dumb mode.
+            DumbServiceImpl.getInstance(project).runInDumbMode {
+              dumbTaskStarted.complete(true)
+              finishDumbTask.await()
+            }
           }
+          // However, we need to make sure that dumb mode does not finish before the second task is started.
+          // Otherwise, the second task will need write action and EDT (to start a new dumb mode). There will be a deadlock,
+          // if startEternalDumbModeTask is invoked from EDT (because CoroutineScope(Job()).launch started with NON_MODAL modality).
+          dumbTaskStarted.await()
         }
-        dumbModeStarted.await()
       }
     }
     assertTrue("Dumb mode didn't start", DumbService.isDumb(project))
-    return EternalTaskShutdownToken(dumbModeJob ?: fail("Could not start dumb mode task"))
+    return EternalTaskShutdownToken(finishDumbTask)
+  }
+
+  private fun runModalIfEdt(project: Project, action: suspend CoroutineScope.() -> Unit) {
+    if (application.isDispatchThread) {
+      runWithModalProgressBlocking(project, "test", action)
+    }
+    else {
+      runBlockingMaybeCancellable(action)
+    }
   }
 
   @JvmStatic