added cancel handling to code lens requests

headless-services/commons/commons-language-server/src/main/java/org/springframework/ide/vscode/commons/languageserver/util/CodeLensHandler.java

@@ -1,5 +1,5 @@
 /*******************************************************************************
- * Copyright (c) 2017, 2018 Pivotal, Inc.
+ * Copyright (c) 2017, 2021 Pivotal, Inc.
  * 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
@@ -14,10 +14,11 @@ import java.util.List;
 
 import org.eclipse.lsp4j.CodeLens;
 import org.eclipse.lsp4j.CodeLensParams;
+import org.eclipse.lsp4j.jsonrpc.CancelChecker;
 
 @FunctionalInterface
 public interface CodeLensHandler {
 
-	List<? extends CodeLens> handle(CodeLensParams params);
+	List<? extends CodeLens> handle(CancelChecker cancelToken, CodeLensParams params);
 
 }

headless-services/commons/commons-language-server/src/main/java/org/springframework/ide/vscode/commons/languageserver/util/SimpleTextDocumentService.java

@@ -422,7 +422,7 @@ public class SimpleTextDocumentService implements TextDocumentService, DocumentE
 
 		if (handler != null) {
 			return CompletableFutures.computeAsync(cancelToken -> {
-				return handler.handle(params);
+				return handler.handle(cancelToken, params);
 			});
 		}
 		return CompletableFuture.completedFuture(Collections.emptyList());

headless-services/spring-boot-language-server/src/main/java/org/springframework/ide/vscode/boot/java/handlers/BootJavaCodeLensEngine.java

@@ -13,10 +13,12 @@ package org.springframework.ide.vscode.boot.java.handlers;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 
 import org.eclipse.lsp4j.CodeLens;
 import org.eclipse.lsp4j.CodeLensParams;
+import org.eclipse.lsp4j.jsonrpc.CancelChecker;
 import org.springframework.ide.vscode.boot.java.BootJavaLanguageServerComponents;
 import org.springframework.ide.vscode.commons.languageserver.util.CodeLensHandler;
 import org.springframework.ide.vscode.commons.languageserver.util.SimpleTextDocumentService;
@@ -36,7 +38,7 @@ public class BootJavaCodeLensEngine implements CodeLensHandler {
 	}
 
 	@Override
-	public List<? extends CodeLens> handle(CodeLensParams params) {
+	public List<? extends CodeLens> handle(CancelChecker cancelToken, CodeLensParams params) {
 		SimpleTextDocumentService documents = server.getTextDocumentService();
 		String docURI = params.getTextDocument().getUri();
 
@@ -45,11 +47,14 @@ public class BootJavaCodeLensEngine implements CodeLensHandler {
 			// Spring Boot LS get events from boot properties files as well, so filter them out
 			if (server.getInterestingLanguages().contains(doc.getLanguageId())) {
 				try {
-					List<? extends CodeLens> codeLensesResult = provideCodeLenses(doc);
+					List<? extends CodeLens> codeLensesResult = provideCodeLenses(cancelToken, doc);
 					if (codeLensesResult != null) {
 						return codeLensesResult;
 					}
 				}
+				catch (CancellationException e) {
+					throw e;
+				}
 				catch (Exception e) {
 				}
 			}
@@ -58,13 +63,13 @@ public class BootJavaCodeLensEngine implements CodeLensHandler {
 		return SimpleTextDocumentService.NO_CODELENS;
 	}
 
-	private List<? extends CodeLens> provideCodeLenses(TextDocument document) {
+	private List<? extends CodeLens> provideCodeLenses(CancelChecker cancelToken, TextDocument document) {
 		return server.getCompilationUnitCache().withCompilationUnit(document, cu -> {
 
 			if (cu != null) {
 				List<CodeLens> result = new ArrayList<>();
 				for (CodeLensProvider codeLensProvider : codelensProviders) {
-					codeLensProvider.provideCodeLenses(document, cu, result);
+					codeLensProvider.provideCodeLenses(cancelToken, document, cu, result);
 				}
 
 				if (result.size() > 0) {

headless-services/spring-boot-language-server/src/main/java/org/springframework/ide/vscode/boot/java/handlers/CodeLensProvider.java

@@ -1,5 +1,5 @@
 /*******************************************************************************
- * Copyright (c) 2017, 2018 Pivotal, Inc.
+ * Copyright (c) 2017, 2021 Pivotal, Inc.
  * 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
@@ -14,6 +14,7 @@ import java.util.List;
 
 import org.eclipse.jdt.core.dom.CompilationUnit;
 import org.eclipse.lsp4j.CodeLens;
+import org.eclipse.lsp4j.jsonrpc.CancelChecker;
 import org.springframework.ide.vscode.commons.util.text.TextDocument;
 
 /**
@@ -21,6 +22,6 @@ import org.springframework.ide.vscode.commons.util.text.TextDocument;
  */
 public interface CodeLensProvider {
 
-	public void provideCodeLenses(TextDocument document, CompilationUnit cu, List<CodeLens> resultAccumulator);
+	public void provideCodeLenses(CancelChecker cancelToken, TextDocument document, CompilationUnit cu, List<CodeLens> resultAccumulator);
 
 }

headless-services/spring-boot-language-server/src/main/java/org/springframework/ide/vscode/boot/java/requestmapping/WebfluxHandlerCodeLensProvider.java

@@ -1,5 +1,5 @@
 /*******************************************************************************
- * Copyright (c) 2018, 2019 Pivotal, Inc.
+ * Copyright (c) 2018, 2021 Pivotal, Inc.
  * 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
@@ -18,6 +18,7 @@ import org.eclipse.jdt.core.dom.IMethodBinding;
 import org.eclipse.jdt.core.dom.MethodDeclaration;
 import org.eclipse.lsp4j.CodeLens;
 import org.eclipse.lsp4j.Command;
+import org.eclipse.lsp4j.jsonrpc.CancelChecker;
 import org.springframework.ide.vscode.boot.app.SpringSymbolIndex;
 import org.springframework.ide.vscode.boot.java.handlers.CodeLensProvider;
 import org.springframework.ide.vscode.boot.java.handlers.SymbolAddOnInformation;
@@ -36,17 +37,19 @@ public class WebfluxHandlerCodeLensProvider implements CodeLensProvider {
 	}
 
 	@Override
-	public void provideCodeLenses(TextDocument document, CompilationUnit cu, List<CodeLens> resultAccumulator) {
+	public void provideCodeLenses(CancelChecker cancelToken, TextDocument document, CompilationUnit cu, List<CodeLens> resultAccumulator) {
 		cu.accept(new ASTVisitor() {
 			@Override
 			public boolean visit(MethodDeclaration node) {
-				provideCodeLens(node, document, resultAccumulator);
+				provideCodeLens(cancelToken, node, document, resultAccumulator);
 				return super.visit(node);
 			}
 		});
 	}
 
-	protected void provideCodeLens(MethodDeclaration node, TextDocument document, List<CodeLens> resultAccumulator) {
+	protected void provideCodeLens(CancelChecker cancelToken, MethodDeclaration node, TextDocument document, List<CodeLens> resultAccumulator) {
+		cancelToken.checkCanceled();
+		
 		IMethodBinding methodBinding = node.resolveBinding();
 
 		if (methodBinding != null && methodBinding.getDeclaringClass() != null && methodBinding.getMethodDeclaration() != null
@@ -55,6 +58,8 @@ public class WebfluxHandlerCodeLensProvider implements CodeLensProvider {
 			final String handlerClass = methodBinding.getDeclaringClass().getBinaryName().trim();
 			final String handlerMethod = methodBinding.getMethodDeclaration().toString().trim();
 			
+			cancelToken.checkCanceled();
+
 			List<SymbolAddOnInformation> handlerInfos = this.springIndexer.getAllAdditionalInformation((addon) -> {
 				if (addon instanceof WebfluxHandlerInformation) {
 					WebfluxHandlerInformation handlerInfo = (WebfluxHandlerInformation) addon;