From fe75a3e6bf69f3a8a4743d2b70d7cb8688f6334e Mon Sep 17 00:00:00 2001
From: Jinbo Wang <jinbwan@microsoft.com>
Date: Thu, 10 Sep 2020 12:55:51 +0800
Subject: [PATCH] Allow to cancel resolveMainMethod codelens command

---
 src/debugCodeLensProvider.ts | 2 +-
 src/languageServerPlugin.ts  | 6 +++++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/debugCodeLensProvider.ts b/src/debugCodeLensProvider.ts
index e6b3bca3..66f1e9bc 100644
--- a/src/debugCodeLensProvider.ts
+++ b/src/debugCodeLensProvider.ts
@@ -102,7 +102,7 @@ class DebugCodeLensProvider implements vscode.CodeLensProvider {
 
     public async provideCodeLenses(document: vscode.TextDocument, token: vscode.CancellationToken): Promise<vscode.CodeLens[]> {
         try {
-            const mainMethods: IMainMethod[] = await resolveMainMethod(document.uri);
+            const mainMethods: IMainMethod[] = await resolveMainMethod(document.uri, token);
             return _.flatten(mainMethods.map((method) => {
                 return [
                     new vscode.CodeLens(method.range, {
diff --git a/src/languageServerPlugin.ts b/src/languageServerPlugin.ts
index 0e6c6e49..3c542aac 100644
--- a/src/languageServerPlugin.ts
+++ b/src/languageServerPlugin.ts
@@ -33,7 +33,11 @@ export interface ILaunchValidationResponse {
     readonly proposals?: IMainClassOption[];
 }
 
-export async function resolveMainMethod(uri: vscode.Uri): Promise<IMainMethod[]> {
+export async function resolveMainMethod(uri: vscode.Uri, token?: vscode.CancellationToken): Promise<IMainMethod[]> {
+    if (token) {
+        return <IMainMethod[]> await commands.executeJavaLanguageServerCommand(commands.JAVA_RESOLVE_MAINMETHOD, uri.toString(), token);
+    }
+
     return <IMainMethod[]> await commands.executeJavaLanguageServerCommand(commands.JAVA_RESOLVE_MAINMETHOD, uri.toString());
 }