Fix various test flakiness

This commit addresses various test flakiness:

1. Bump timeouts for queries tests
2. Add a dispose handler to queryserver-client. This will help us during
   tests because if there is a test that timesout while a query is
   running, the query's progress callback won't be invoked. We will
   still get a timeout error in the first test, but the second test will
   not get a spurious error.
3. Handle a disposed query server in `deregisterDatabase`. This method
   will remove the database from the currently running query server.
   If there is no query server, then there is nothing to remove. So,
   this error is safe to ignore.
4. Explicitly `end()` a connection `ServerProcess`. I'm not 100% sure if
   this is necessary, but it seems like it prevents responses from being
   handled and erroring out.
5. Better handling of ideServer restarts. Previously, if you quickly
   called `CodeQL: Restart Query Server` twice in a row, you would get
   an error from the ideServer restart. Restart fails if the server is
   not already started. So, in this case just call a start.

Author: aeisenberg

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

@@ -1022,7 +1022,19 @@ export class DatabaseManager extends DisposableObject {
     token: vscode.CancellationToken,
     dbItem: DatabaseItem,
   ) {
-    await this.qs.deregisterDatabase(progress, token, dbItem);
+    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);
+  }
 }

extensions/ql-vscode/test/vscode-tests/cli-integration/queries.test.ts

@@ -30,6 +30,7 @@ import { AllCommands, QueryServerCommands } from "../../../src/common/commands";
 /**
  * Integration tests for queries
  */
+jest.setTimeout(60_000);
 describeWithCodeQL()("Queries", () => {
   let dbItem: DatabaseItem;
   let databaseManager: DatabaseManager;