Adds MultiCancellationToken

This is a cancellation token that cancels when any of its constituent
cancellation tokens are cancelled.

This token is used to fix a bug in Find Definitions. Previously, when
clicking `CTRL` (or `CMD` on macs) inside a source file in an archive
and hovering over a token, this will automatically invoke the
definitions finder (in preparation for navigating to the definition).
The only way to cancel is to move down to the message popup and click
cancel there.

However, this is a bug. What _should_ happen is that if a user moves
their mouse away from the token, the operation should cancel.

The underlying problem is that the extension was only listening to the
cancellation token from inside `getLocationsForUriString` the
cancellation token used by the Language Server protocol to cancel
operations in flight was being ignored.

This fix will ensure we are listening to _both_ cancellation tokens
and cancel the query if either are cancelled.

Author: aeisenberg

extensions/ql-vscode/src/common/multi-cancellation-token.ts

@@ -0,0 +1,28 @@
+import { CancellationToken, Disposable } from "vscode";
+
+/**
+ * A cancellation token that cancels when any of its constituent
+ * cancellation tokens are cancelled.
+ */
+export class MultiCancellationToken implements CancellationToken {
+  private readonly tokens: CancellationToken[];
+
+  constructor(...tokens: CancellationToken[]) {
+    this.tokens = tokens;
+  }
+
+  get isCancellationRequested(): boolean {
+    return this.tokens.some((t) => t.isCancellationRequested);
+  }
+
+  onCancellationRequested<T>(listener: (e: T) => any): Disposable {
+    this.tokens.forEach((t) => t.onCancellationRequested(listener));
+    return {
+      dispose: () => {
+        this.tokens.forEach((t) =>
+          t.onCancellationRequested(listener).dispose(),
+        );
+      },
+    };
+  }
+}

extensions/ql-vscode/src/language-support/contextual/template-provider.ts

@@ -36,13 +36,15 @@ import {
 import { CoreCompletedQuery, QueryRunner } from "../../query-server";
 import { AstBuilder } from "../ast-viewer/ast-builder";
 import { qlpackOfDatabase } from "../../local-queries";
+import { MultiCancellationToken } from "../../common/multi-cancellation-token";
 
 /**
  * Runs templated CodeQL queries to find definitions in
  * source-language files. We may eventually want to find a way to
  * generalize this to other custom queries, e.g. showing dataflow to
  * or from a selected identifier.
  */
+
 export class TemplateQueryDefinitionProvider implements DefinitionProvider {
   private cache: CachedOperation<LocationLink[]>;
 
@@ -60,11 +62,11 @@ export class TemplateQueryDefinitionProvider implements DefinitionProvider {
   async provideDefinition(
     document: TextDocument,
     position: Position,
-    _token: CancellationToken,
+    token: CancellationToken,
   ): Promise<LocationLink[]> {
     const fileLinks = this.shouldUseCache()
-      ? await this.cache.get(document.uri.toString())
-      : await this.getDefinitions(document.uri.toString());
+      ? await this.cache.get(document.uri.toString(), token)
+      : await this.getDefinitions(document.uri.toString(), token);
 
     const locLinks: LocationLink[] = [];
     for (const link of fileLinks) {
@@ -79,9 +81,13 @@ export class TemplateQueryDefinitionProvider implements DefinitionProvider {
     return !(isCanary() && NO_CACHE_CONTEXTUAL_QUERIES.getValue<boolean>());
   }
 
-  private async getDefinitions(uriString: string): Promise<LocationLink[]> {
+  private async getDefinitions(
+    uriString: string,
+    token: CancellationToken,
+  ): Promise<LocationLink[]> {
     return withProgress(
-      async (progress, token) => {
+      async (progress, tokenInner) => {
+        const multiToken = new MultiCancellationToken(token, tokenInner);
         return getLocationsForUriString(
           this.cli,
           this.qs,
@@ -90,7 +96,7 @@ export class TemplateQueryDefinitionProvider implements DefinitionProvider {
           KeyType.DefinitionQuery,
           this.queryStorageDir,
           progress,
-          token,
+          multiToken,
           (src, _dest) => src === uriString,
         );
       },
@@ -126,11 +132,11 @@ export class TemplateQueryReferenceProvider implements ReferenceProvider {
     document: TextDocument,
     position: Position,
     _context: ReferenceContext,
-    _token: CancellationToken,
+    token: CancellationToken,
   ): Promise<Location[]> {
     const fileLinks = this.shouldUseCache()
-      ? await this.cache.get(document.uri.toString())
-      : await this.getReferences(document.uri.toString());
+      ? await this.cache.get(document.uri.toString(), token)
+      : await this.getReferences(document.uri.toString(), token);
 
     const locLinks: Location[] = [];
     for (const link of fileLinks) {
@@ -148,9 +154,14 @@ export class TemplateQueryReferenceProvider implements ReferenceProvider {
     return !(isCanary() && NO_CACHE_CONTEXTUAL_QUERIES.getValue<boolean>());
   }
 
-  private async getReferences(uriString: string): Promise<FullLocationLink[]> {
+  private async getReferences(
+    uriString: string,
+    token: CancellationToken,
+  ): Promise<FullLocationLink[]> {
     return withProgress(
-      async (progress, token) => {
+      async (progress, tokenInner) => {
+        const multiToken = new MultiCancellationToken(token, tokenInner);
+
         return getLocationsForUriString(
           this.cli,
           this.qs,
@@ -159,7 +170,7 @@ export class TemplateQueryReferenceProvider implements ReferenceProvider {
           KeyType.DefinitionQuery,
           this.queryStorageDir,
           progress,
-          token,
+          multiToken,
           (src, _dest) => src === uriString,
         );
       },