Implement cancellation of test run

Author: Dave Bartolomeo

common/config/rush/shrinkwrap.yaml

@@ -56,6 +56,7 @@ dependencies:
   style-loader: 0.23.1
   through2: 3.0.1
   tmp: 0.1.0
+  tree-kill: 1.2.2
   ts-loader: 5.4.5
   ts-node: 8.3.0
   ts-protoc-gen: 0.9.0
@@ -5690,6 +5691,11 @@ packages:
     dev: false
     resolution:
       integrity: sha1-cXuPIgzAu3tE5AUUwisui7xw2Lk=
+  /tree-kill/1.2.2:
+    dev: false
+    hasBin: true
+    resolution:
+      integrity: sha512-L0Orpi8qGpRG//Nd+H90vFB+3iHnue1zSSGmNOOCh1GLJ7rUKVwV2HvijphGQS2UmhUZewS9VgvxYIdgr+fG1A==
   /true-case-path/2.2.1:
     dev: false
     resolution:
@@ -6566,6 +6572,7 @@ packages:
       style-loader: 0.23.1
       through2: 3.0.1
       tmp: 0.1.0
+      tree-kill: 1.2.2
       ts-loader: /ts-loader/5.4.5/typescript@3.7.2
       ts-node: /ts-node/8.3.0/typescript@3.7.2
       ts-protoc-gen: 0.9.0
@@ -6585,7 +6592,7 @@ packages:
     dev: false
     name: '@rush-temp/vscode-codeql'
     resolution:
-      integrity: sha512-1zj1CO19sanU2olFV3Kwf9OrG+92KO25WR6K1vnVHmScMj8H7aegRrjwyNeHi/IzxXDCAp68CZK+d4vEWBEa+g==
+      integrity: sha512-pHdIjv8YEIRicmnIzv2OB5cqdCRoKcF2ZbylLcaR+VP28TzoyYZK5q2anIYayibWa3Ik4la4pVGnSa344IatvA==
       tarball: 'file:projects/vscode-codeql.tgz'
     version: 0.0.0
 registry: 'https://registry.npmjs.org/'
@@ -6649,6 +6656,7 @@ specifiers:
   style-loader: ~0.23.1
   through2: ^3.0.1
   tmp: ^0.1.0
+  tree-kill: ~1.2.2
   ts-loader: ^5.4.5
   ts-node: ^8.3.0
   ts-protoc-gen: ^0.9.0

extensions/ql-vscode/package.json

@@ -371,6 +371,7 @@
     "semmle-io-node": "^0.0.1",
     "semmle-vscode-utils": "^0.0.1",
     "tmp": "^0.1.0",
+    "tree-kill": "~1.2.2",
     "unzipper": "~0.10.5",
     "vscode-jsonrpc": "^4.0.0",
     "vscode-languageclient": "^5.2.1",

extensions/ql-vscode/src/cli.ts

@@ -3,9 +3,10 @@ import * as cpp from 'child-process-promise';
 import * as fs from 'fs-extra';
 import * as path from 'path';
 import * as sarif from 'sarif';
+import * as tk from 'tree-kill';
 import * as util from 'util';
 import { Logger, ProgressReporter } from './logging';
-import { Disposable } from 'vscode';
+import { Disposable, CancellationToken } from 'vscode';
 import { DistributionProvider } from './distribution';
 import { SortDirection } from './interface-types';
 import { assertNever } from './helpers-pure';
@@ -86,6 +87,7 @@ export type ResolvedTests = string[];
  * Options for `codeql test run`.
  */
 export interface TestRunOptions {
+  cancellationToken?: CancellationToken;
   logger?: Logger;
 }
 
@@ -286,12 +288,14 @@ export class CodeQLCliServer implements Disposable {
    *
    * @param command The `codeql` command to be run, provided as an array of command/subcommand names.
    * @param commandArgs The arguments to pass to the `codeql` command.
+   * @param cancellationToken CancellationToken to terminate the test process.
    * @param logger Logger to write text output from the command.
    * @returns The sequence of async events produced by the command.
    */
   private async* runAsyncCodeQlCliCommandInternal(
     command: string[],
     commandArgs: string[],
+    cancellationToken?: CancellationToken,
     logger?: Logger
   ): AsyncGenerator<string, void, unknown> {
     // Add format argument first, in case commandArgs contains positional parameters.
@@ -306,16 +310,29 @@ export class CodeQLCliServer implements Disposable {
     const childPromise = cpp.spawn(codeqlPath, args);
     const child = childPromise.childProcess;
 
-    if (logger !== undefined) {
-      // The human-readable output goes to stderr.
-      logStream(child.stderr!, logger);
-    }
+    let cancellationRegistration: Disposable | undefined = undefined;
+    try {
+      if (cancellationToken !== undefined) {
+        cancellationRegistration = cancellationToken.onCancellationRequested(_e => {
+          tk(child.pid);
+        });
+      }
+      if (logger !== undefined) {
+        // The human-readable output goes to stderr.
+        logStream(child.stderr!, logger);
+      }
 
-    for await (const event of await splitStreamAtSeparators(child.stdout!, ['\0'])) {
-      yield event;
-    }
+      for await (const event of await splitStreamAtSeparators(child.stdout!, ['\0'])) {
+        yield event;
+      }
 
-    await childPromise;
+      await childPromise;
+    }
+    finally {
+      if (cancellationRegistration !== undefined) {
+        cancellationRegistration.dispose();
+      }
+    }
   }
 
   /**
@@ -325,16 +342,19 @@ export class CodeQLCliServer implements Disposable {
    * @param command The `codeql` command to be run, provided as an array of command/subcommand names.
    * @param commandArgs The arguments to pass to the `codeql` command.
    * @param description Description of the action being run, to be shown in log and error messages.
+   * @param cancellationToken CancellationToken to terminate the test process.
    * @param logger Logger to write text output from the command.
    * @returns The sequence of async events produced by the command.
    */
   public async* runAsyncCodeQlCliCommand<EventType>(
     command: string[],
     commandArgs: string[],
     description: string,
+    cancellationToken?: CancellationToken,
     logger?: Logger
   ): AsyncGenerator<EventType, void, unknown> {
-    for await (const event of await this.runAsyncCodeQlCliCommandInternal(command, commandArgs, logger)) {
+    for await (const event of await this.runAsyncCodeQlCliCommandInternal(command, commandArgs,
+      cancellationToken, logger)) {
       try {
         yield JSON.parse(event) as EventType;
       } catch (err) {
@@ -451,7 +471,7 @@ export class CodeQLCliServer implements Disposable {
     ];
 
     for await (const event of await this.runAsyncCodeQlCliCommand<TestCompleted>(['test', 'run'],
-      subcommandArgs, 'Run CodeQL Tests', options.logger)) {
+      subcommandArgs, 'Run CodeQL Tests', options.cancellationToken, options.logger)) {
       yield event;
     }
   }

extensions/ql-vscode/src/test-adapter.ts

@@ -14,7 +14,7 @@ import {
 } from 'vscode-test-adapter-api';
 import { TestAdapterRegistrar } from 'vscode-test-adapter-util';
 import { QLTestFile, QLTestNode, QLTestDirectory, QLTestDiscovery } from './qltest-discovery';
-import { Event, EventEmitter, CancellationTokenSource } from 'vscode';
+import { Event, EventEmitter, CancellationTokenSource, CancellationToken } from 'vscode';
 import { DisposableObject } from 'semmle-vscode-utils';
 import { QLPackDiscovery } from './qlpack-discovery';
 import { CodeQLCliServer } from './cli';
@@ -194,7 +194,7 @@ export class QLTestAdapter extends DisposableObject implements TestAdapter {
     const testAdapter = this;
 
     try {
-      await this.runTests(tests);
+      await this.runTests(tests, this.runningTask.token);
     }
     catch (e) {
     }
@@ -218,9 +218,10 @@ export class QLTestAdapter extends DisposableObject implements TestAdapter {
     }
   }
 
-  private async runTests(tests: string[]): Promise<void> {
+  private async runTests(tests: string[], cancellationToken: CancellationToken): Promise<void> {
     const workspacePaths = await getOnDiskWorkspaceFolders();
     for await (const event of await this.cliServer.runTests(tests, workspacePaths, {
+      cancellationToken: cancellationToken,
       logger: testLogger
     })) {
       this._testStates.fire({