IJPL-162040 Store lock permits in context to allow lock inheritance between coroutines - preliminary steps

Application.invokeLater() now puts modality state into context.

GitOrigin-RevId: bc76fe4c3792a8c352259142cc925983f9f32bb1

platform/platform-impl/src/com/intellij/openapi/application/impl/ApplicationImpl.java

@@ -2,6 +2,7 @@
 package com.intellij.openapi.application.impl;
 
 import com.intellij.CommonBundle;
+import com.intellij.concurrency.ThreadContext;
 import com.intellij.configurationStore.StoreUtil;
 import com.intellij.diagnostic.ActivityCategory;
 import com.intellij.diagnostic.PluginException;
@@ -64,7 +65,9 @@ import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 import static com.intellij.ide.ShutdownKt.cancelAndJoinBlocking;
+import static com.intellij.openapi.application.ModalityKt.asContextElement;
 import static com.intellij.util.concurrency.AppExecutorUtil.propagateContext;
+import static com.intellij.util.concurrency.Propagation.isContextAwareComputation;
 
 @ApiStatus.Internal
 public final class ApplicationImpl extends ClientAwareComponentManager implements ApplicationEx, ReadActionListener, WriteActionListener {
@@ -266,13 +269,19 @@ public final class ApplicationImpl extends ClientAwareComponentManager implement
 
   @Override
   public void invokeLater(@NotNull Runnable runnable, @NotNull ModalityState state, @NotNull Condition<?> expired) {
+    final boolean ctxAware = isContextAwareComputation(runnable);
+    // Start from inner layer: transaction guard
+    final Runnable guarded = myTransactionGuard.wrapLaterInvocation(runnable, state);
+    // Middle layer: lock and modality
+    final Runnable locked = wrapWithRunIntendedWriteActionAndModality(guarded, ctxAware ? null : state);
+    Runnable finalRunnable = locked;
+    // Outer layer, optional: context capture & reset
     if (propagateContext()) {
-      Pair<Runnable, Condition<?>> captured = Propagation.capturePropagationContext(runnable, expired);
-      runnable = captured.getFirst();
+      Pair<Runnable, Condition<?>> captured = Propagation.capturePropagationContext(locked, expired, runnable);
+      finalRunnable = captured.getFirst();
       expired = captured.getSecond();
     }
-    Runnable r = myTransactionGuard.wrapLaterInvocation(runnable, state);
-    LaterInvocator.invokeLater(state, expired, wrapWithRunIntendedWriteAction(r));
+    LaterInvocator.invokeLater(state, expired, finalRunnable);
   }
 
   @ApiStatus.Internal
@@ -387,11 +396,7 @@ public final class ApplicationImpl extends ClientAwareComponentManager implement
 
 
   @Override
-  public void invokeAndWait(@NotNull Runnable runnable, @NotNull ModalityState modalityState) {
-    if (isDispatchThread()) {
-      runIntendedWriteActionOnCurrentThread(runnable);
-      return;
-    }
+  public void invokeAndWait(@NotNull Runnable runnable, @NotNull ModalityState state) {
     if (EDT.isCurrentThreadEdt()) {
       runIntendedWriteActionOnCurrentThread(runnable);
       return;
@@ -401,24 +406,45 @@ public final class ApplicationImpl extends ClientAwareComponentManager implement
       throw new IllegalStateException("Calling invokeAndWait from read-action leads to possible deadlock.");
     }
 
-    Runnable capturingRunnable = AppImplKt.rethrowExceptions(AppScheduledExecutorService::captureContextCancellationForRunnableThatDoesNotOutliveContextScope, runnable);
+    final boolean ctxAware = isContextAwareComputation(runnable);
+    // Start from inner layer: transaction guard
+    final Runnable guarded = myTransactionGuard.wrapLaterInvocation(runnable, state);
+    // Middle layer: lock and modality
+    final Runnable locked = wrapWithRunIntendedWriteActionAndModality(guarded, ctxAware ? null : state);
+    // Outer layer context capture & reset
+    final Runnable finalRunnable = AppImplKt.rethrowExceptions(AppScheduledExecutorService::captureContextCancellationForRunnableThatDoesNotOutliveContextScope, locked);
 
-    Runnable r = myTransactionGuard.wrapLaterInvocation(capturingRunnable, modalityState);
-    LaterInvocator.invokeAndWait(modalityState, wrapWithRunIntendedWriteAction(r));
+    LaterInvocator.invokeAndWait(state, finalRunnable);
   }
 
-  private @NotNull Runnable wrapWithRunIntendedWriteAction(@NotNull Runnable runnable) {
-    return new Runnable() {
-      @Override
-      public void run() {
-        runIntendedWriteActionOnCurrentThread(runnable);
-      }
+  private @NotNull Runnable wrapWithRunIntendedWriteActionAndModality(@NotNull Runnable runnable, @Nullable ModalityState modalityState) {
+    return modalityState != null ?
+           new Runnable() {
+             @Override
+             public void run() {
+               try (AccessToken ignored = ThreadContext.installThreadContext(
+                 ThreadContext.currentThreadContext().plus(asContextElement(modalityState)), true)) {
+                 runIntendedWriteActionOnCurrentThread(runnable);
+               }
+             }
 
-      @Override
-      public String toString() {
-        return runnable.toString();
-      }
-    };
+             @Override
+             public String toString() {
+               return runnable.toString();
+             }
+           }
+                                 :
+           new Runnable() {
+             @Override
+             public void run() {
+               runIntendedWriteActionOnCurrentThread(runnable);
+             }
+
+             @Override
+             public String toString() {
+               return runnable.toString();
+             }
+           };
   }
 
   @Override

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

@@ -68,9 +68,8 @@ class ImplicitBlockingContextTest {
   @Test
   fun invokeLater(): Unit = runBlockingWithCatchingExceptions {
     withContext(E()) {
-      val currentContext = coroutineContext
       ApplicationManager.getApplication().invokeLater {
-        assertContextRemainsOnFreeThread(currentContext)
+        assertContextRemainsOnFreeThread()
       }
     }
   }
@@ -78,9 +77,8 @@ class ImplicitBlockingContextTest {
   @Test
   fun executeOnPooledThread(): Unit = runBlockingWithCatchingExceptions {
     withContext(E()) {
-      val currentContext = coroutineContext
       ApplicationManager.getApplication().executeOnPooledThread {
-        assertContextRemainsOnFreeThread(currentContext)
+        assertContextRemainsOnFreeThread()
       }
     }
   }
@@ -184,9 +182,9 @@ class ImplicitBlockingContextTest {
     assertEquals(context.minusKey(ContinuationInterceptor), currentThreadContext())
   }
 
-  private fun assertContextRemainsOnFreeThread(context: CoroutineContext) {
+  private fun assertContextRemainsOnFreeThread() {
     assertNull(IntellijCoroutines.currentThreadCoroutineContext())
-    val list = currentThreadContext().fold(ArrayList<CoroutineContext.Element>(), { list, elem -> list.apply { add(elem) } })
-    assertEquals(list.single().key, E)
+    val set = currentThreadContext().fold(HashSet<CoroutineContext.Key<*>>(), { list, elem -> list.apply { add(elem.key) } })
+    assertTrue(set.contains(E))
   }
 }

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

@@ -236,7 +236,7 @@ internal fun <V> captureCallableThreadContext(callable: Callable<V>): Callable<V
   return callable
 }
 
-private fun isContextAwareComputation(runnable: Any): Boolean {
+fun isContextAwareComputation(runnable: Any): Boolean {
   return runnable is Continuation<*> || runnable is ContextAwareRunnable || runnable is ContextAwareCallable<*> || runnable is CancellationFutureTask<*>
 }
 
@@ -328,8 +328,8 @@ internal fun capturePropagationContext(r: Runnable, forceUseContextJob : Boolean
 }
 
 @ApiStatus.Internal
-fun capturePropagationContext(r: Runnable, expired: Condition<*>): JBPair<Runnable, Condition<*>> {
-  if (isContextAwareComputation(r)) {
+fun capturePropagationContext(r: Runnable, expired: Condition<*>, signalRunnable: Runnable): JBPair<Runnable, Condition<*>> {
+  if (isContextAwareComputation(signalRunnable)) {
     return JBPair.create(r, expired)
   }
   var command = captureClientIdInRunnable(r)