When restarting the query server wait for the old one to stop properly

Author: josefs

extensions/ql-vscode/CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## [UNRELEASED]
 
+- Fix a bug where restarting the query server could intermittently fail because a new server was started before the previous one had fully exited. The extension now waits for the old process to terminate first. [#4457](https://github.com/github/vscode-codeql/pull/4457)
 - Fix a bug where installing or updating the CodeQL CLI could hang indefinitely while extracting the downloaded archive. Extraction now reports an error if a file cannot be written, and aborts with a clear message if no progress is made within the download timeout (for example due to slow or networked storage, or security software). [#4455](https://github.com/github/vscode-codeql/pull/4455)
 - Remove support for CodeQL CLI versions older than 2.23.9. [#4448](https://github.com/github/vscode-codeql/pull/4448)
 - Added support for selection-based result filtering via a checkbox in the result viewer. When enabled, only results from the currently-viewed file are shown. Additionally, if the editor selection is non-empty, only results within the selection range are shown. [#4362](https://github.com/github/vscode-codeql/pull/4362)

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

@@ -113,6 +113,19 @@ export class QueryServerClient extends DisposableObject {
     }
   }
 
+  /**
+   * Stops the query server and waits for its process to fully exit before
+   * returning. Use this (rather than `stopQueryServer`) when a new server is
+   * about to be started, so that the old process has released any file locks
+   * first. This avoids intermittent failures on Windows where the OS keeps
+   * locks until the process has actually terminated.
+   */
+  private async stopQueryServerAndWait(): Promise<void> {
+    const serverProcess = this.serverProcess;
+    this.stopQueryServer();
+    await serverProcess?.waitForExit();
+  }
+
   /**
    * Restarts the query server by disposing of the current server process and then starting a new one.
    * This resets the unexpected termination count. As hopefully it is an indication that the user has fixed the
@@ -154,7 +167,9 @@ export class QueryServerClient extends DisposableObject {
   private async restartQueryServerInternal(
     progress: ProgressCallback,
   ): Promise<void> {
-    this.stopQueryServer();
+    // Wait for the old process to fully exit before starting a new one, so
+    // that it has released any file locks (important on Windows).
+    await this.stopQueryServerAndWait();
     await this.startQueryServer();
 
     // Ensure we await all responses from event handlers so that

extensions/ql-vscode/src/query-server/server-process.ts

@@ -27,10 +27,57 @@ export class ServerProcess implements Disposable {
     this.child.stdin!.end();
     this.child.stderr!.destroy();
     this.child.removeAllListeners();
-    // TODO kill the process if it doesn't terminate after a certain time limit.
+
+    // `dispose()` is synchronous, so we only signal the process to stop (by
+    // closing its streams) and don't wait for it to actually exit. Callers that
+    // need the process to have terminated — e.g. before starting a replacement
+    // server — should await `waitForExit()` afterwards.
 
     // On Windows, we usually have to terminate the process before closing its stdout.
     this.child.stdout!.destroy();
     void this.logger.log(`Stopped ${this.name}.`);
   }
+
+  /**
+   * Waits for the underlying child process to fully exit, forcibly killing it
+   * after `timeoutMs` if it has not exited on its own.
+   *
+   * Call this after `dispose()` when you need the OS to have actually released
+   * the process before continuing. This matters on Windows, where the OS can
+   * keep file locks (for example on the database or disk cache) until the
+   * process has terminated — so starting a replacement server before the old
+   * one has exited can intermittently fail.
+   */
+  async waitForExit(timeoutMs = 5000): Promise<void> {
+    const hasExited = () =>
+      this.child.exitCode !== null || this.child.signalCode !== null;
+
+    if (hasExited()) {
+      return;
+    }
+
+    await new Promise<void>((resolve) => {
+      const timer = setTimeout(() => {
+        void this.logger.log(
+          `${this.name} did not exit within ${timeoutMs}ms; killing it.`,
+        );
+        this.child.kill("SIGKILL");
+        resolve();
+      }, timeoutMs);
+
+      const done = () => {
+        clearTimeout(timer);
+        resolve();
+      };
+
+      this.child.once("exit", done);
+
+      // Guard against the process having exited between the check above and
+      // attaching the listener.
+      if (hasExited()) {
+        this.child.removeListener("exit", done);
+        done();
+      }
+    });
+  }
 }

extensions/ql-vscode/test/unit-tests/query-server/server-process.test.ts

@@ -0,0 +1,65 @@
+import { EventEmitter } from "events";
+import { ServerProcess } from "../../../src/query-server/server-process";
+
+function createFakeChild(): any {
+  const child = new EventEmitter() as any;
+  child.exitCode = null;
+  child.signalCode = null;
+  child.kill = jest.fn((signal?: string) => {
+    child.exitCode = 1;
+    child.emit("exit", null, signal ?? "SIGKILL");
+    return true;
+  });
+  return child;
+}
+
+const fakeConnection = { dispose: jest.fn(), end: jest.fn() } as any;
+const fakeLogger = { log: jest.fn() } as any;
+
+function createServerProcess(child: any): ServerProcess {
+  return new ServerProcess(
+    child,
+    fakeConnection,
+    "test query server",
+    fakeLogger,
+  );
+}
+
+describe("ServerProcess.waitForExit", () => {
+  it("resolves when the process exits", async () => {
+    const child = createFakeChild();
+    const serverProcess = createServerProcess(child);
+
+    const promise = serverProcess.waitForExit();
+    child.exitCode = 0;
+    child.emit("exit", 0, null);
+
+    await expect(promise).resolves.toBeUndefined();
+    expect(child.kill).not.toHaveBeenCalled();
+  });
+
+  it("resolves immediately if the process has already exited", async () => {
+    const child = createFakeChild();
+    child.exitCode = 0;
+    const serverProcess = createServerProcess(child);
+
+    await expect(serverProcess.waitForExit()).resolves.toBeUndefined();
+    expect(child.kill).not.toHaveBeenCalled();
+  });
+
+  it("force-kills the process if it does not exit within the timeout", async () => {
+    jest.useFakeTimers();
+    try {
+      const child = createFakeChild();
+      const serverProcess = createServerProcess(child);
+
+      const promise = serverProcess.waitForExit(1000);
+      jest.advanceTimersByTime(1000);
+
+      await expect(promise).resolves.toBeUndefined();
+      expect(child.kill).toHaveBeenCalledWith("SIGKILL");
+    } finally {
+      jest.useRealTimers();
+    }
+  });
+});