Merge pull request #2331 from github/aeisenberg/fix-flakes

Fix various test flakiness

extensions/ql-vscode/src/extension.ts

@@ -177,7 +177,13 @@ function getCommands(
         cliServer.restartCliServer();
         await Promise.all([
           queryRunner.restartQueryServer(progress, token),
-          ideServer.restart(),
+          async () => {
+            if (ideServer.isRunning()) {
+              await ideServer.restart();
+            } else {
+              await ideServer.start();
+            }
+          },
         ]);
         void showAndLogInformationMessage("CodeQL Query Server restarted.", {
           outputLogger: queryServerLogger,

extensions/ql-vscode/src/json-rpc-server.ts

@@ -23,6 +23,7 @@ export class ServerProcess implements Disposable {
   dispose(): void {
     void this.logger.log(`Stopping ${this.name}...`);
     this.connection.dispose();
+    this.connection.end();
     this.child.stdin!.end();
     this.child.stderr!.destroy();
     // TODO kill the process if it doesn't terminate after a certain time limit.

extensions/ql-vscode/src/local-databases.ts

@@ -1029,7 +1029,19 @@ export class DatabaseManager extends DisposableObject {
     token: vscode.CancellationToken,
     dbItem: DatabaseItem,
   ) {
+    try {
       await this.qs.deregisterDatabase(progress, token, dbItem);
+    } catch (e) {
+      const message = getErrorMessage(e);
+      if (message === "Connection is disposed.") {
+        // This is expected if the query server is not running.
+        void extLogger.log(
+          `Could not de-register database '${dbItem.name}' because query server is not running.`,
+        );
+        return;
+      }
+      throw e;
+    }
   }
   private async registerDatabase(
     progress: ProgressCallback,

extensions/ql-vscode/src/query-server/queryserver-client.ts

@@ -1,6 +1,6 @@
 import { ensureFile } from "fs-extra";
 
-import { DisposableObject } from "../pure/disposable-object";
+import { DisposableObject, DisposeHandler } from "../pure/disposable-object";
 import { CancellationToken } from "vscode";
 import { createMessageConnection, RequestType } from "vscode-jsonrpc/node";
 import * as cli from "../cli";
@@ -224,4 +224,10 @@ export class QueryServerClient extends DisposableObject {
       delete this.progressCallbacks[id];
     }
   }
+
+  public dispose(disposeHandler?: DisposeHandler | undefined): void {
+    this.progressCallbacks = {};
+    this.stopQueryServer();
+    super.dispose(disposeHandler);
+  }
 }