Merge pull request #2885 from github/aeisenberg/multi-token

Adds `MultiCancellationToken`

extensions/ql-vscode/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## [UNRELEASED]
 
+- Fix a bug where the query to Find Definitions in database source files would not be cancelled appropriately. [#2885](https://github.com/github/vscode-codeql/pull/2885)
 - It is now possible to show the language of query history items using the `%l` specifier in the `codeQL.queryHistory.format` setting. Note that this only works for queries run after this upgrade, and older items will show `unknown` as a language. [#2892](https://github.com/github/vscode-codeql/pull/2892)
 - Increase the required version of VS Code to 1.82.0. [#2877](https://github.com/github/vscode-codeql/pull/2877)
 - Fix a bug where the query server was restarted twice after configuration changes. [#2884](https://github.com/github/vscode-codeql/pull/2884).

extensions/ql-vscode/src/common/disposable-object.ts

@@ -9,10 +9,16 @@ export type DisposeHandler = (disposable: Disposable) => void;
 /**
  * Base class to make it easier to implement a `Disposable` that owns other disposable object.
  */
-export abstract class DisposableObject implements Disposable {
+export class DisposableObject implements Disposable {
   private disposables: Disposable[] = [];
   private tracked?: Set<Disposable> = undefined;
 
+  constructor(...dispoables: Disposable[]) {
+    for (const d of dispoables) {
+      this.push(d);
+    }
+  }
+
   /**
    * Adds `obj` to a list of objects to dispose when `this` is disposed. Objects added by `push` are
    * disposed in reverse order of being added.

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

@@ -0,0 +1,24 @@
+import { CancellationToken, Disposable } from "vscode";
+import { DisposableObject } from "../disposable-object";
+
+/**
+ * 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 {
+    return new DisposableObject(
+      ...this.tokens.map((t) => t.onCancellationRequested(listener)),
+    );
+  }
+}

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

@@ -36,6 +36,7 @@ import {
 import { CoreCompletedQuery, QueryRunner } from "../../query-server";
 import { AstBuilder } from "../ast-viewer/ast-builder";
 import { qlpackOfDatabase } from "../../local-queries";
+import { MultiCancellationToken } from "../../common/vscode/multi-cancellation-token";
 
 /**
  * Runs templated CodeQL queries to find definitions in
@@ -43,6 +44,7 @@ import { qlpackOfDatabase } from "../../local-queries";
  * 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.cache.get(document.uri.toString(), token)
-      : await this.getDefinitions(document.uri.toString());
+      : 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.cache.get(document.uri.toString(), token)
-      : await this.getReferences(document.uri.toString());
+      : 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,
         );
       },