From 9b84e36c79c77374a1050adebba8bb5fbe2d383e Mon Sep 17 00:00:00 2001
From: Nikita Skvortsov <nikita.skvortsov@jetbrains.com>
Date: Wed, 30 May 2018 13:31:03 +0300
Subject: [PATCH] Allow to interrupt gradle wrapper command.

and report exceptions accordingly. IDEA-192927
---
 .../execution/GradleExecutionHelper.java      |  6 ++-
 .../project/GradleProjectResolver.java        | 51 ++++++++++---------
 2 files changed, 31 insertions(+), 26 deletions(-)

diff --git a/plugins/gradle/src/org/jetbrains/plugins/gradle/service/execution/GradleExecutionHelper.java b/plugins/gradle/src/org/jetbrains/plugins/gradle/service/execution/GradleExecutionHelper.java
index 3251f9175a04..aefb69e92b12 100644
--- a/plugins/gradle/src/org/jetbrains/plugins/gradle/service/execution/GradleExecutionHelper.java
+++ b/plugins/gradle/src/org/jetbrains/plugins/gradle/service/execution/GradleExecutionHelper.java
@@ -233,7 +233,8 @@ public class GradleExecutionHelper {
   public void ensureInstalledWrapper(@NotNull ExternalSystemTaskId id,
                                      @NotNull String projectPath,
                                      @NotNull GradleExecutionSettings settings,
-                                     @NotNull ExternalSystemTaskNotificationListener listener) {
+                                     @NotNull ExternalSystemTaskNotificationListener listener,
+                                     @NotNull CancellationTokenSource cancellationTokenSource) {
 
     if (!settings.getDistributionType().isWrapped()) return;
 
@@ -263,6 +264,7 @@ public class GradleExecutionHelper {
         final File tempFile = writeToFileGradleInitScript(StringUtil.join(lines, SystemProperties.getLineSeparator()));
         settings.withArguments(GradleConstants.INIT_SCRIPT_CMD_OPTION, tempFile.getAbsolutePath());
         BuildLauncher launcher = getBuildLauncher(id, connection, settings, listener);
+        launcher.withCancellationToken(cancellationTokenSource.token());
         launcher.forTasks("wrapper");
         launcher.run();
         String wrapperPropertyFile = FileUtil.loadFile(wrapperPropertyFileLocation);
@@ -274,6 +276,8 @@ public class GradleExecutionHelper {
     }
     catch (Throwable e) {
       LOG.warn("Can't update wrapper", e);
+      Throwable rootCause = ExceptionUtil.getRootCause(e);
+      throw new ExternalSystemException(ExceptionUtil.getMessage(rootCause));
     }
     finally {
       settings.setRemoteProcessIdleTtlInMs(ttlInMs);
diff --git a/plugins/gradle/src/org/jetbrains/plugins/gradle/service/project/GradleProjectResolver.java b/plugins/gradle/src/org/jetbrains/plugins/gradle/service/project/GradleProjectResolver.java
index 2210f32bdf4c..a5800adfeb75 100644
--- a/plugins/gradle/src/org/jetbrains/plugins/gradle/service/project/GradleProjectResolver.java
+++ b/plugins/gradle/src/org/jetbrains/plugins/gradle/service/project/GradleProjectResolver.java
@@ -56,7 +56,6 @@ import org.jetbrains.plugins.gradle.model.data.CompositeBuildData;
 import org.jetbrains.plugins.gradle.model.data.GradleSourceSetData;
 import org.jetbrains.plugins.gradle.remote.impl.GradleLibraryNamesMixer;
 import org.jetbrains.plugins.gradle.service.execution.GradleExecutionHelper;
-import org.jetbrains.plugins.gradle.service.execution.UnsupportedCancellationToken;
 import org.jetbrains.plugins.gradle.settings.ClassHolder;
 import org.jetbrains.plugins.gradle.settings.DistributionType;
 import org.jetbrains.plugins.gradle.settings.GradleBuildParticipant;
@@ -129,20 +128,33 @@ public class GradleProjectResolver implements ExternalSystemProjectResolver<Grad
       return projectDataNode;
     }
 
-    if (settings != null) {
-      myHelper.ensureInstalledWrapper(syncTaskId, projectPath, settings, listener);
+    DefaultProjectResolverContext resolverContext = new DefaultProjectResolverContext(syncTaskId, projectPath, settings, listener, false);
+    final CancellationTokenSource cancellationTokenSource = GradleConnector.newCancellationTokenSource();
+    resolverContext.setCancellationTokenSource(cancellationTokenSource);
+    synchronized (myCancellationMap) {
+      myCancellationMap.putValue(resolverContext.getExternalSystemTaskId(), cancellationTokenSource);
     }
 
-    final GradleProjectResolverExtension projectResolverChain = createProjectResolverChain(settings);
-    DefaultProjectResolverContext resolverContext = new DefaultProjectResolverContext(syncTaskId, projectPath, settings, listener, false);
-    final DataNode<ProjectData> projectDataNode = myHelper.execute(
-      projectPath, settings, getProjectDataFunction(resolverContext, projectResolverChain, false));
+    try {
+      if (settings != null) {
+        myHelper.ensureInstalledWrapper(syncTaskId, projectPath, settings, listener, cancellationTokenSource);
+      }
 
-    // auto-discover buildSrc projects of the main and included builds
-    File gradleUserHome = resolverContext.getUserData(GRADLE_HOME_DIR);
-    new GradleBuildSrcProjectsResolver(this, resolverContext, gradleUserHome, settings, listener, syncTaskId, projectResolverChain)
-      .discoverAndAppendTo(projectDataNode);
-    return projectDataNode;
+      final GradleProjectResolverExtension projectResolverChain = createProjectResolverChain(settings);
+      final DataNode<ProjectData> projectDataNode = myHelper.execute(
+        projectPath, settings, getProjectDataFunction(resolverContext, projectResolverChain, false));
+
+
+      // auto-discover buildSrc projects of the main and included builds
+      File gradleUserHome = resolverContext.getUserData(GRADLE_HOME_DIR);
+      new GradleBuildSrcProjectsResolver(this, resolverContext, gradleUserHome, settings, listener, syncTaskId, projectResolverChain)
+        .discoverAndAppendTo(projectDataNode);
+      return projectDataNode;
+    } finally {
+      synchronized (myCancellationMap) {
+        myCancellationMap.remove(resolverContext.getExternalSystemTaskId(), cancellationTokenSource);
+      }
+    }
   }
 
   @NotNull
@@ -258,19 +270,11 @@ public class GradleProjectResolver implements ExternalSystemProjectResolver<Grad
     GradleExecutionHelper.prepare(buildActionExecutor, resolverCtx.getExternalSystemTaskId(),
                                   executionSettings, resolverCtx.getListener(), resolverCtx.getConnection());
     resolverCtx.checkCancelled();
-
     ProjectImportAction.AllModels allModels;
-    final CancellationTokenSource cancellationTokenSource = GradleConnector.newCancellationTokenSource();
+
     final long startTime = System.currentTimeMillis();
     try {
-      resolverCtx.setCancellationTokenSource(cancellationTokenSource);
-      buildActionExecutor.withCancellationToken(cancellationTokenSource.token());
-      synchronized (myCancellationMap) {
-        myCancellationMap.putValue(resolverCtx.getExternalSystemTaskId(), cancellationTokenSource);
-        if (gradleVersion != null && gradleVersion.compareTo(GradleVersion.version("2.1")) < 0) {
-          myCancellationMap.putValue(resolverCtx.getExternalSystemTaskId(), new UnsupportedCancellationToken());
-        }
-      }
+      buildActionExecutor.withCancellationToken(resolverCtx.getCancellationTokenSource().token());
       allModels = buildActionExecutor.run();
       if (allModels == null) {
         throw new IllegalStateException("Unable to get project model for the project: " + resolverCtx.getProjectPath());
@@ -296,9 +300,6 @@ public class GradleProjectResolver implements ExternalSystemProjectResolver<Grad
     }
     finally {
       final long timeInMs = (System.currentTimeMillis() - startTime);
-      synchronized (myCancellationMap) {
-        myCancellationMap.remove(resolverCtx.getExternalSystemTaskId(), cancellationTokenSource);
-      }
       performanceTrace.logPerformance("Gradle data obtained", timeInMs);
       LOG.debug(String.format("Gradle data obtained in %d ms", timeInMs));
     }