From f2316a5b17030b5b8a28e6c4e23ca8456e36aa47 Mon Sep 17 00:00:00 2001
From: peter <peter@jetbrains.com>
Date: Thu, 10 Oct 2019 15:36:10 +0200
Subject: [PATCH] ProgressIndicatorUtils.withTimeout: make checkCanceled honor
 thread's indicator, add javadoc

GitOrigin-RevId: 7b3827f00df965f399a67780fb37a75ef31957e3
---
 .../openapi/progress/ProgressManager.java     |  4 +++
 .../progress/util/ProgressIndicatorUtils.java | 29 +++++++++++++++----
 .../progress/impl/ProgressIndicatorTest.java  | 11 +++++++
 3 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/platform/core-api/src/com/intellij/openapi/progress/ProgressManager.java b/platform/core-api/src/com/intellij/openapi/progress/ProgressManager.java
index f823b139b900..f5b06987271b 100644
--- a/platform/core-api/src/com/intellij/openapi/progress/ProgressManager.java
+++ b/platform/core-api/src/com/intellij/openapi/progress/ProgressManager.java
@@ -207,6 +207,10 @@ public abstract class ProgressManager extends ProgressIndicatorProvider {
     getInstance().indicatorCanceled(indicator);
   }
 
+  /**
+   * @throws ProcessCanceledException if the progress indicator associated with the current thread has been canceled.
+   * @see ProgressIndicator#checkCanceled()
+   */
   @SuppressWarnings("MethodOverridesStaticMethodOfSuperclass")
   public static void checkCanceled() throws ProcessCanceledException {
     ProgressManager instance = ourInstance;
diff --git a/platform/platform-impl/src/com/intellij/openapi/progress/util/ProgressIndicatorUtils.java b/platform/platform-impl/src/com/intellij/openapi/progress/util/ProgressIndicatorUtils.java
index f57fecf91825..6d996b546c70 100644
--- a/platform/platform-impl/src/com/intellij/openapi/progress/util/ProgressIndicatorUtils.java
+++ b/platform/platform-impl/src/com/intellij/openapi/progress/util/ProgressIndicatorUtils.java
@@ -1,6 +1,7 @@
 // Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.intellij.openapi.progress.util;
 
+import com.intellij.concurrency.SensitiveProgressWrapper;
 import com.intellij.openapi.Disposable;
 import com.intellij.openapi.application.Application;
 import com.intellij.openapi.application.ApplicationListener;
@@ -11,6 +12,7 @@ import com.intellij.openapi.application.ex.ApplicationManagerEx;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.progress.ProcessCanceledException;
 import com.intellij.openapi.progress.ProgressIndicator;
+import com.intellij.openapi.progress.ProgressIndicatorProvider;
 import com.intellij.openapi.progress.ProgressManager;
 import com.intellij.openapi.util.*;
 import com.intellij.util.concurrency.AppExecutorUtil;
@@ -19,6 +21,7 @@ import org.jetbrains.annotations.Nullable;
 import org.jetbrains.ide.PooledThreadExecutor;
 
 import java.util.concurrent.*;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.Lock;
 
 /**
@@ -249,16 +252,32 @@ public class ProgressIndicatorUtils {
     application.invokeAndWait(EmptyRunnable.INSTANCE, ModalityState.any());
   }
 
+  /**
+   * Run the given computation with its execution time restricted to the given amount of time in milliseconds.<p></p>
+   *
+   * Internally, it creates a new {@link ProgressIndicator}, runs the computation with that indicator and cancels it after the the timeout.
+   * The computation should call {@link ProgressManager#checkCanceled()} frequently enough, so that after the timeout has been exceeded
+   * it can stop the execution by throwing {@link ProcessCanceledException}, which will be caught by this {@code withTimeout}.<p></p>
+   *
+   * If a {@link ProcessCanceledException} happens due to any other reason (e.g. a thread's progress indicator got canceled),
+   * it'll be thrown out of this method.
+   * @return the computation result or {@code null} if timeout has been exceeded.
+   */
   @Nullable
-  public static <T> T withTimeout(long timeoutMs, @NotNull Computable<T> computable) {
+  public static <T> T withTimeout(long timeoutMs, @NotNull Computable<T> computation) {
     ProgressManager.checkCanceled();
-    ProgressIndicatorBase progress = new ProgressIndicatorBase(false, false);
-    ScheduledFuture<?> cancelProgress = AppExecutorUtil.getAppScheduledExecutorService().schedule(progress::cancel, timeoutMs, TimeUnit.MILLISECONDS);
+    ProgressIndicator outer = ProgressIndicatorProvider.getGlobalProgressIndicator();
+    ProgressIndicator inner = outer != null ? new SensitiveProgressWrapper(outer) : new ProgressIndicatorBase(false, false);
+    AtomicBoolean canceledByTimeout = new AtomicBoolean();
+    ScheduledFuture<?> cancelProgress = AppExecutorUtil.getAppScheduledExecutorService().schedule(() -> {
+      canceledByTimeout.set(true);
+      inner.cancel();
+    }, timeoutMs, TimeUnit.MILLISECONDS);
     try {
-      return ProgressManager.getInstance().runProcess(computable, progress);
+      return ProgressManager.getInstance().runProcess(computation, inner);
     }
     catch (ProcessCanceledException e) {
-      if (progress.isCanceled()) {
+      if (canceledByTimeout.get()) {
         return null;
       }
       throw e; // canceled not by timeout
diff --git a/platform/platform-tests/testSrc/com/intellij/openapi/progress/impl/ProgressIndicatorTest.java b/platform/platform-tests/testSrc/com/intellij/openapi/progress/impl/ProgressIndicatorTest.java
index 844e4c787e07..0500ac71326c 100644
--- a/platform/platform-tests/testSrc/com/intellij/openapi/progress/impl/ProgressIndicatorTest.java
+++ b/platform/platform-tests/testSrc/com/intellij/openapi/progress/impl/ProgressIndicatorTest.java
@@ -817,6 +817,17 @@ public class ProgressIndicatorTest extends LightPlatformTestCase {
         throw new ProcessCanceledException();
       });
     });
+
+    ProgressIndicatorBase outer = new ProgressIndicatorBase();
+    ProgressManager.getInstance().runProcess(() -> {
+      assertThrows(ProcessCanceledException.class, () -> {
+        ProgressIndicatorUtils.withTimeout(1, () -> {
+          outer.cancel();
+          ProgressManager.checkCanceled();
+          return null;
+        });
+      });
+    }, outer);
   }
 
   private static class MyAbstractProgressIndicator extends AbstractProgressIndicatorBase {