IJPL-8 IJPL-9 explicitly support propagation in ProgressRunner

This allows to install newly obtained modality state into the context.

GitOrigin-RevId: 37c65317448103bdcc27cd183d8593bdc44b5140

platform/core-impl/src/com/intellij/openapi/progress/impl/ProgressRunner.java

@@ -2,7 +2,10 @@
 package com.intellij.openapi.progress.impl;
 
 import com.intellij.codeWithMe.ClientId;
+import com.intellij.concurrency.ContextAwareRunnable;
+import com.intellij.concurrency.ThreadContext;
 import com.intellij.openapi.Disposable;
+import com.intellij.openapi.application.AccessToken;
 import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.application.ModalityState;
 import com.intellij.openapi.application.ex.ApplicationManagerEx;
@@ -13,8 +16,14 @@ import com.intellij.openapi.progress.*;
 import com.intellij.openapi.util.Disposer;
 import com.intellij.openapi.wm.ex.ProgressIndicatorEx;
 import com.intellij.util.concurrency.AppExecutorUtil;
+import com.intellij.util.concurrency.Propagation;
 import com.intellij.util.concurrency.Semaphore;
 import com.intellij.util.ui.EDT;
+import kotlin.Pair;
+import kotlin.Unit;
+import kotlin.coroutines.CoroutineContext;
+import kotlin.coroutines.EmptyCoroutineContext;
+import kotlinx.coroutines.CompletableJob;
 import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
@@ -23,6 +32,8 @@ import org.jetbrains.annotations.Nullable;
 import java.util.concurrent.*;
 import java.util.function.Function;
 
+import static com.intellij.openapi.application.ModalityKt.asContextElement;
+
 /**
  * <p>
  * A builder-like API for running tasks with {@link ProgressIndicator}.
@@ -438,6 +449,27 @@ public final class ProgressRunner<R> {
     @NotNull CompletableFuture<? extends @NotNull ProgressIndicator> progressIndicatorFuture
   ) {
     CompletableFuture<R> resultFuture = new CompletableFuture<>();
+    Pair<CoroutineContext, CompletableJob> childContextAndJob = Propagation.createChildContext();
+    CoroutineContext childContext = childContextAndJob.getFirst();
+    CompletableJob childJob = childContextAndJob.getSecond();
+    if (childJob != null) {
+      // cancellation of the Job cancels the future
+      childJob.invokeOnCompletion(true, true, (throwable) -> {
+        if (throwable != null) {
+          resultFuture.completeExceptionally(throwable);
+        }
+        return Unit.INSTANCE;
+      });
+      // completion of the future completes the Job
+      resultFuture.whenComplete((result, throwable) -> {
+        if (throwable == null) {
+          childJob.complete();
+        }
+        else {
+          childJob.completeExceptionally(throwable);
+        }
+      });
+    }
     progressIndicatorFuture.whenComplete((progressIndicator, throwable) -> {
       if (throwable != null) {
         resultFuture.completeExceptionally(throwable);
@@ -451,13 +483,19 @@ public final class ProgressRunner<R> {
           resultFuture.completeExceptionally(e);
         }
       };
+      Runnable contextRunnable = childContext.equals(EmptyCoroutineContext.INSTANCE) ? runnable : (ContextAwareRunnable)() -> {
+        CoroutineContext effectiveContext = childContext.plus(asContextElement(progressIndicator.getModalityState()));
+        try (AccessToken ignored = ThreadContext.installThreadContext(effectiveContext, false)) {
+          runnable.run();
+        }
+      };
       switch (myThreadToUse) {
         case POOLED:
-          AppExecutorUtil.getAppExecutorService().execute(runnable);
+          AppExecutorUtil.getAppExecutorService().execute(contextRunnable);
           break;
         case WRITE:
           ModalityState processModality = progressIndicator.getModalityState();
-          ApplicationManager.getApplication().invokeLaterOnWriteThread(runnable, processModality);
+          ApplicationManager.getApplication().invokeLaterOnWriteThread(contextRunnable, processModality);
           break;
         default:
           throw new IllegalStateException("Unexpected value: " + myThreadToUse);

platform/platform-tests/testSrc/com/intellij/util/concurrency/ThreadContextPropagationTest.kt

@@ -5,22 +5,22 @@ import com.intellij.concurrency.*
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.EDT
 import com.intellij.openapi.application.ModalityState
-import com.intellij.openapi.progress.blockingContext
-import com.intellij.openapi.progress.timeoutWaitUp
+import com.intellij.openapi.application.currentThreadContextModality
+import com.intellij.openapi.progress.*
 import com.intellij.openapi.util.Conditions
+import com.intellij.testFramework.junit5.SystemProperty
 import com.intellij.testFramework.junit5.TestApplication
 import com.intellij.util.timeoutRunBlocking
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.launch
-import kotlinx.coroutines.suspendCancellableCoroutine
-import kotlinx.coroutines.withContext
+import kotlinx.coroutines.*
 import org.junit.jupiter.api.Assertions
+import org.junit.jupiter.api.Assertions.assertFalse
 import org.junit.jupiter.api.Assertions.assertSame
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.extension.ExtendWith
 import org.junit.jupiter.api.extension.ExtensionContext
 import org.junit.jupiter.api.extension.InvocationInterceptor
 import org.junit.jupiter.api.extension.ReflectiveInvocationContext
+import java.lang.Runnable
 import java.lang.reflect.Method
 import java.util.concurrent.ExecutorService
 import java.util.concurrent.Future
@@ -182,4 +182,37 @@ class ThreadContextPropagationTest {
       }
     }
   }
+
+  @SystemProperty("intellij.progress.task.ignoreHeadless", "true")
+  @Test
+  fun `Task Modal`(): Unit = timeoutRunBlocking {
+    doTest {
+      object : Task.Modal(null, "", true) {
+        override fun run(indicator: ProgressIndicator) {
+          it()
+        }
+      }.queue()
+    }
+  }
+
+  @SystemProperty("intellij.progress.task.ignoreHeadless", "true")
+  @Test
+  fun `Task Modal receives newly entered modality state in the context`(): Unit = timeoutRunBlocking {
+    val finished = CompletableDeferred<Unit>()
+    withModalProgress(ModalTaskOwner.guess(), "", TaskCancellation.cancellable()) {
+      blockingContext {
+        assertSame(currentThreadContextModality(), ModalityState.defaultModalityState())
+        object : Task.Modal(null, "", true) {
+          override fun run(indicator: ProgressIndicator) {
+            finished.completeWith(runCatching {
+              assertFalse(currentThreadContextModality() == ModalityState.NON_MODAL)
+              assertSame(currentThreadContextModality(), indicator.modalityState)
+              assertSame(currentThreadContextModality(), ModalityState.defaultModalityState())
+            })
+          }
+        }.queue()
+      }
+    }
+    finished.await()
+  }
 }

platform/util/src/com/intellij/util/concurrency/propagation.kt

@@ -62,7 +62,8 @@ class BlockingJob(val blockingJob: Job) : AbstractCoroutineContextElement(Blocki
   companion object : CoroutineContext.Key<BlockingJob>
 }
 
-private fun createChildContext(): Pair<CoroutineContext, CompletableJob?> {
+@Internal
+fun createChildContext(): Pair<CoroutineContext, CompletableJob?> {
   val currentThreadContext = currentThreadContext()
 
   // Problem: a task may infinitely reschedule itself