feat: Support cancelling getDependencies

extension/src/client/CancellationKeys.ts

@@ -15,3 +15,7 @@ export function getRunTaskCommandCancellationKey(
 ): string {
   return 'runTask' + rootProjectFolder + taskName;
 }
+
+export function getDependenciesCancellationKey(projectDir: string): string {
+  return 'getDependencies' + projectDir;
+}

extension/src/client/GradleClient.ts

@@ -25,6 +25,8 @@ import {
   DependencyItem,
   GetDependenciesRequest,
   GetDependenciesReply,
+  CancelDependenciesRequest,
+  CancelDependenciesReply,
 } from '../proto/gradle_pb';
 
 import { GradleClient as GrpcClient } from '../proto/gradle_grpc_pb';
@@ -38,7 +40,10 @@ import {
   COMMAND_CANCEL_BUILD,
 } from '../commands';
 import { RootProject } from '../rootProject/RootProject';
-import { getBuildCancellationKey } from './CancellationKeys';
+import {
+  getBuildCancellationKey,
+  getDependenciesCancellationKey,
+} from './CancellationKeys';
 import { EventWaiter } from '../util/EventWaiter';
 import { getGradleConfig, getConfigJavaDebug } from '../util/config';
 
@@ -228,38 +233,59 @@ export class GradleClient implements vscode.Disposable {
 
   public async getDependencies(
     projectDir: string,
-    gradleConfig: GradleConfig
+    gradleConfig: GradleConfig,
+    projectName: string
   ): Promise<DependencyItem | undefined> {
     await this.waitForConnect();
-    const request = new GetDependenciesRequest();
-    request.setProjectDir(projectDir);
-    request.setGradleConfig(gradleConfig);
-    try {
-      return await new Promise((resolve, reject) => {
-        this.grpcClient!.getDependencies(
-          request,
-          (
-            err: grpc.ServiceError | null,
-            getDependenciesReply: GetDependenciesReply | undefined
-          ) => {
-            if (err) {
-              reject(err);
-            } else {
-              resolve(getDependenciesReply?.getItem());
-            }
-          }
+    return vscode.window.withProgress(
+      {
+        location: vscode.ProgressLocation.Window,
+        title: 'Gradle',
+        cancellable: true,
+      },
+      async (
+        progress: vscode.Progress<{ message?: string }>,
+        token: vscode.CancellationToken
+      ) => {
+        progress.report({
+          message: `Getting dependencies for ${projectName}`,
+        });
+        const request = new GetDependenciesRequest();
+        request.setProjectDir(projectDir);
+        request.setGradleConfig(gradleConfig);
+        const cancellationKey = getDependenciesCancellationKey(projectDir);
+        request.setCancellationKey(cancellationKey);
+        token.onCancellationRequested(() =>
+          this.cancelDependencies(cancellationKey)
         );
-      });
-    } catch (err) {
-      logger.error(
-        `Error getting dependencies for ${projectDir}: ${
-          err.details || err.message
-        }`
-      );
-      this.statusBarItem.command = COMMAND_SHOW_LOGS;
-      this.statusBarItem.text = '$(warning) Gradle: Get Dependencies Error';
-      this.statusBarItem.show();
-    }
+        try {
+          return await new Promise((resolve, reject) => {
+            this.grpcClient!.getDependencies(
+              request,
+              (
+                err: grpc.ServiceError | null,
+                getDependenciesReply: GetDependenciesReply | undefined
+              ) => {
+                if (err) {
+                  reject(err);
+                } else {
+                  resolve(getDependenciesReply?.getItem());
+                }
+              }
+            );
+          });
+        } catch (err) {
+          logger.error(
+            `Error getting dependencies for ${projectDir}: ${
+              err.details || err.message
+            }`
+          );
+          this.statusBarItem.command = COMMAND_SHOW_LOGS;
+          this.statusBarItem.text = '$(warning) Gradle: Get dependencies Error';
+          this.statusBarItem.show();
+        }
+      }
+    );
   }
 
   public async runBuild(
@@ -431,6 +457,40 @@ export class GradleClient implements vscode.Disposable {
     }
   }
 
+  public async cancelDependencies(cancellationKey: string): Promise<void> {
+    await this.waitForConnect();
+    this.statusBarItem.hide();
+    const request = new CancelDependenciesRequest();
+    request.setCancellationKey(cancellationKey);
+    try {
+      const reply: CancelDependenciesReply | undefined = await new Promise(
+        (resolve, reject) => {
+          this.grpcClient!.cancelDependencies(
+            request,
+            (
+              err: grpc.ServiceError | null,
+              cancelDependenciesReply: CancelDependenciesReply | undefined
+            ) => {
+              if (err) {
+                reject(err);
+              } else {
+                resolve(cancelDependenciesReply);
+              }
+            }
+          );
+        }
+      );
+      if (reply) {
+        logger.info('Cancel dependencies:', reply.getMessage());
+      }
+    } catch (err) {
+      logger.error(
+        'Error cancelling dependencies:',
+        err.details || err.message
+      );
+    }
+  }
+
   public async getDaemonsStatus(
     projectFolder: string,
     cancelToken: vscode.CancellationToken

extension/src/views/gradleTasks/GradleTasksTreeDataProvider.ts

@@ -116,32 +116,7 @@ export class GradleTasksTreeDataProvider
       return element.projects;
     }
     if (element instanceof ProjectTreeItem) {
-      const projectTaskItem = new ProjectTaskTreeItem(
-        'Tasks',
-        vscode.TreeItemCollapsibleState.Collapsed,
-        element
-      );
-      projectTaskItem.setChildren([...element.groups, ...element.tasks]);
-      const results: vscode.TreeItem[] = [projectTaskItem];
-      const resourceUri = element.resourceUri;
-      if (!resourceUri) {
-        return results;
-      }
-      const dependencyItem = await this.client.getDependencies(
-        path.dirname(resourceUri.fsPath),
-        getGradleConfig()
-      );
-      if (!dependencyItem) {
-        return results;
-      }
-      const projectDependencyItem = protocolItem2ProjectDependencyTreeItem(
-        dependencyItem,
-        element
-      );
-      if (projectDependencyItem) {
-        results.push(projectDependencyItem);
-      }
-      return results;
+      return this.getChildrenForProjectTreeItem(element);
     }
     if (element instanceof GroupTreeItem) {
       return element.tasks;
@@ -166,6 +141,39 @@ export class GradleTasksTreeDataProvider
     return [];
   }
 
+  public async getChildrenForProjectTreeItem(
+    element: ProjectTreeItem
+  ): Promise<vscode.TreeItem[]> {
+    const projectTaskItem = new ProjectTaskTreeItem(
+      'Tasks',
+      vscode.TreeItemCollapsibleState.Collapsed,
+      element
+    );
+    projectTaskItem.setChildren([...element.groups, ...element.tasks]);
+    const results: vscode.TreeItem[] = [projectTaskItem];
+    const resourceUri = element.resourceUri;
+    if (!resourceUri) {
+      return results;
+    }
+    const dirname = path.dirname(resourceUri.fsPath);
+    const dependencyItem = await this.client.getDependencies(
+      dirname,
+      getGradleConfig(),
+      element.label || dirname
+    );
+    if (!dependencyItem) {
+      return results;
+    }
+    const projectDependencyItem = protocolItem2ProjectDependencyTreeItem(
+      dependencyItem,
+      element
+    );
+    if (projectDependencyItem) {
+      results.push(projectDependencyItem);
+    }
+    return results;
+  }
+
   // eslint-disable-next-line sonarjs/cognitive-complexity
   public buildItemsTreeFromTasks(
     tasks: vscode.Task[]

gradle-server/src/main/java/com/github/badsyntax/gradle/GradleService.java

@@ -2,6 +2,7 @@
 
 import com.github.badsyntax.gradle.handlers.CancelBuildHandler;
 import com.github.badsyntax.gradle.handlers.CancelBuildsHandler;
+import com.github.badsyntax.gradle.handlers.CancelDependenciesHandler;
 import com.github.badsyntax.gradle.handlers.GetBuildHandler;
 import com.github.badsyntax.gradle.handlers.GetDaemonsStatusHandler;
 import com.github.badsyntax.gradle.handlers.GetDependenciesHandler;
@@ -46,6 +47,14 @@ public void cancelBuilds(
     cancelRunBuildsHandler.run();
   }
 
+  @Override
+  public void cancelDependencies(
+      CancelDependenciesRequest req, StreamObserver<CancelDependenciesReply> responseObserver) {
+    CancelDependenciesHandler cancelDependenciesHandler =
+        new CancelDependenciesHandler(req, responseObserver);
+    cancelDependenciesHandler.run();
+  }
+
   @Override
   public void getDaemonsStatus(
       GetDaemonsStatusRequest req, StreamObserver<GetDaemonsStatusReply> responseObserver) {

gradle-server/src/main/java/com/github/badsyntax/gradle/handlers/CancelDependenciesHandler.java

@@ -0,0 +1,58 @@
+/*******************************************************************************
+ * Copyright (c) 2021 Microsoft Corporation and others.
+ * All rights reserved. This program and the accompanying materials
+ * are made available under the terms of the Eclipse Public License v1.0
+ * which accompanies this distribution, and is available at
+ * http://www.eclipse.org/legal/epl-v10.html
+ *
+ * Contributors:
+ *    Microsoft Corporation - initial API and implementation
+ *******************************************************************************/
+
+package com.github.badsyntax.gradle.handlers;
+
+import com.github.badsyntax.gradle.CancelDependenciesReply;
+import com.github.badsyntax.gradle.CancelDependenciesRequest;
+import com.github.badsyntax.gradle.GradleBuildCancellation;
+import com.github.badsyntax.gradle.exceptions.GradleCancellationException;
+import io.grpc.stub.StreamObserver;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class CancelDependenciesHandler {
+  private static final Logger logger =
+      LoggerFactory.getLogger(CancelDependenciesHandler.class.getName());
+
+  private CancelDependenciesRequest req;
+  private StreamObserver<CancelDependenciesReply> responseObserver;
+
+  public CancelDependenciesHandler(
+      CancelDependenciesRequest req, StreamObserver<CancelDependenciesReply> responseObserver) {
+    this.req = req;
+    this.responseObserver = responseObserver;
+  }
+
+  public void run() {
+    try {
+      GradleBuildCancellation.cancelBuild(req.getCancellationKey());
+      replyWithCancelledSuccess();
+    } catch (GradleCancellationException e) {
+      logger.error(e.getMessage());
+      replyWithCancelError(e);
+    } finally {
+      responseObserver.onCompleted();
+    }
+  }
+
+  private void replyWithCancelledSuccess() {
+    responseObserver.onNext(
+        CancelDependenciesReply.newBuilder()
+            .setMessage("Cancel getting dependencies requested")
+            .build());
+  }
+
+  private void replyWithCancelError(Exception e) {
+    responseObserver.onNext(
+        CancelDependenciesReply.newBuilder().setMessage(e.getMessage()).build());
+  }
+}

gradle-server/src/main/java/com/github/badsyntax/gradle/handlers/GetDependenciesHandler.java

@@ -15,6 +15,7 @@
 import com.github.badsyntax.gradle.ErrorMessageBuilder;
 import com.github.badsyntax.gradle.GetDependenciesReply;
 import com.github.badsyntax.gradle.GetDependenciesRequest;
+import com.github.badsyntax.gradle.GradleBuildCancellation;
 import com.github.badsyntax.gradle.GradleProjectConnector;
 import com.github.badsyntax.gradle.exceptions.GradleConnectionException;
 import com.microsoft.gradle.api.GradleDependencyNode;
@@ -29,6 +30,8 @@
 import java.util.ArrayList;
 import java.util.List;
 import org.gradle.tooling.BuildActionExecuter;
+import org.gradle.tooling.BuildCancelledException;
+import org.gradle.tooling.CancellationToken;
 import org.gradle.tooling.GradleConnector;
 import org.gradle.tooling.ProjectConnection;
 import org.slf4j.Logger;
@@ -66,6 +69,9 @@ public void run() {
       createPluginJar("/gradle-plugin.jar", pluginFile);
       createTemplateScript(pluginFile, initScript);
       action.withArguments("--init-script", initScript.getAbsolutePath());
+      CancellationToken cancellationToken =
+          GradleBuildCancellation.buildToken(req.getCancellationKey());
+      action.withCancellationToken(cancellationToken);
       GradleToolingModel gradleModel = action.run();
       GradleDependencyNode root = gradleModel.getDependencyNode();
       responseObserver.onNext(
@@ -74,6 +80,10 @@ public void run() {
     } catch (IOException e) {
       logger.error(e.getMessage());
       responseObserver.onError(ErrorMessageBuilder.build(e));
+    } catch (BuildCancelledException e) {
+      replyWithCancelled();
+    } finally {
+      GradleBuildCancellation.clearToken(req.getCancellationKey());
     }
   }
 
@@ -114,4 +124,9 @@ private void createTemplateScript(File pluginFile, File outputFile) throws IOExc
             + "}\n";
     Files.writeString(outputFile.toPath(), template);
   }
+
+  private void replyWithCancelled() {
+    responseObserver.onNext(GetDependenciesReply.newBuilder().build());
+    responseObserver.onCompleted();
+  }
 }

proto/gradle.proto

@@ -12,6 +12,7 @@ service Gradle {
   rpc GetDependencies(GetDependenciesRequest) returns (GetDependenciesReply) {}
   rpc CancelBuild(CancelBuildRequest) returns (CancelBuildReply) {}
   rpc CancelBuilds(CancelBuildsRequest) returns (CancelBuildsReply) {}
+  rpc CancelDependencies(CancelDependenciesRequest) returns (CancelDependenciesReply) {}
   rpc GetDaemonsStatus(GetDaemonsStatusRequest) returns (GetDaemonsStatusReply) {}
   rpc StopDaemons(StopDaemonsRequest) returns (StopDaemonsReply) {}
   rpc StopDaemon(StopDaemonRequest) returns (StopDaemonReply) {}
@@ -42,6 +43,7 @@ message GetBuildResult {
 message GetDependenciesRequest {
   string project_dir = 1;
   GradleConfig gradle_config = 2;
+  string cancellation_key = 3;
 }
 
 message GetDependenciesReply {
@@ -100,6 +102,14 @@ message CancelBuildsReply {
   string message = 1;
 }
 
+message CancelDependenciesRequest {
+  string cancellation_key = 1;
+}
+
+message CancelDependenciesReply {
+  string message = 1;
+}
+
 message GetDaemonsStatusRequest {
   string project_dir = 1;
 }