Add ability to await the discovery refresh promise

extensions/ql-vscode/src/common/discovery.ts

@@ -8,18 +8,28 @@ import { getErrorMessage } from "../pure/helpers-pure";
  * same time.
  */
 export abstract class Discovery<T> extends DisposableObject {
-  private retry = false;
-  private discoveryInProgress = false;
+  private restartWhenFinished = false;
+  private currentDiscoveryPromise: Promise<void> | undefined;
 
   constructor(private readonly name: string) {
     super();
   }
 
+  /**
+   * Returns the promise of the currently running refresh operation, if one is in progress.
+   * Otherwise returns a promise that resolves immediately.
+   */
+  public waitForCurrentRefresh(): Promise<void> {
+    return this.currentDiscoveryPromise ?? Promise.resolve();
+  }
+
   /**
    * Force the discovery process to run. Normally invoked by the derived class when a relevant file
    * system change is detected.
+   *
+   * Returns a promise that resolves when the refresh is complete, including any retries.
    */
-  public refresh(): void {
+  public refresh(): Promise<void> {
     // We avoid having multiple discovery operations in progress at the same time. Otherwise, if we
     // got a storm of refresh requests due to, say, the copying or deletion of a large directory
     // tree, we could potentially spawn a separate simultaneous discovery operation for each
@@ -36,49 +46,48 @@ export abstract class Discovery<T> extends DisposableObject {
     // other change notifications that might be coming along. However, this would create more
     // latency in the common case, in order to save a bit of latency in the uncommon case.
 
-    if (this.discoveryInProgress) {
+    if (this.currentDiscoveryPromise !== undefined) {
       // There's already a discovery operation in progress. Tell it to restart when it's done.
-      this.retry = true;
+      this.restartWhenFinished = true;
     } else {
       // No discovery in progress, so start one now.
-      this.discoveryInProgress = true;
-      this.launchDiscovery();
+      this.currentDiscoveryPromise = this.launchDiscovery().finally(() => {
+        this.currentDiscoveryPromise = undefined;
+      });
     }
+    return this.currentDiscoveryPromise;
   }
 
   /**
    * Starts the asynchronous discovery operation by invoking the `discover` function. When the
    * discovery operation completes, the `update` function will be invoked with the results of the
    * discovery.
    */
-  private launchDiscovery(): void {
-    const discoveryPromise = this.discover();
-    discoveryPromise
-      .then((results) => {
-        if (!this.retry) {
-          // Update any listeners with the results of the discovery.
-          this.discoveryInProgress = false;
-          this.update(results);
-        }
-      })
-
-      .catch((err: unknown) => {
-        void extLogger.log(
-          `${this.name} failed. Reason: ${getErrorMessage(err)}`,
-        );
-      })
+  private async launchDiscovery(): Promise<void> {
+    let results: T | undefined;
+    try {
+      results = await this.discover();
+    } catch (err) {
+      void extLogger.log(
+        `${this.name} failed. Reason: ${getErrorMessage(err)}`,
+      );
+      results = undefined;
+    }
 
-      .finally(() => {
-        if (this.retry) {
-          // Another refresh request came in while we were still running a previous discovery
-          // operation. Since the discovery results we just computed are now stale, we'll launch
-          // another discovery operation instead of updating.
-          // Note that by doing this inside of `finally`, we will relaunch discovery even if the
-          // initial discovery operation failed.
-          this.retry = false;
-          this.launchDiscovery();
-        }
-      });
+    if (this.restartWhenFinished) {
+      // Another refresh request came in while we were still running a previous discovery
+      // operation. Since the discovery results we just computed are now stale, we'll launch
+      // another discovery operation instead of updating.
+      // We want to relaunch discovery regardless of if the initial discovery operation
+      // succeeded or failed.
+      this.restartWhenFinished = false;
+      await this.launchDiscovery();
+    } else {
+      // If the discovery was successful, then update any listeners with the results.
+      if (results !== undefined) {
+        this.update(results);
+      }
+    }
   }
 
   /**

extensions/ql-vscode/src/queries-panel/queries-module.ts

@@ -21,7 +21,7 @@ export class QueriesModule extends DisposableObject {
 
     const queryDiscovery = new QueryDiscovery(app, cliServer);
     this.push(queryDiscovery);
-    queryDiscovery.refresh();
+    void queryDiscovery.refresh();
 
     const queriesPanel = new QueriesPanel(queryDiscovery);
     this.push(queriesPanel);

extensions/ql-vscode/src/query-testing/qltest-discovery.ts

@@ -64,7 +64,7 @@ export class QLTestDiscovery extends Discovery<QLTestDiscoveryResults> {
 
   private handleDidChange(uri: Uri): void {
     if (!QLTestDiscovery.ignoreTestPath(uri.fsPath)) {
-      this.refresh();
+      void this.refresh();
     }
   }
   protected async discover(): Promise<QLTestDiscoveryResults> {

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

@@ -115,7 +115,7 @@ export class QLTestAdapter extends DisposableObject implements TestAdapter {
     this.qlTestDiscovery = this.push(
       new QLTestDiscovery(workspaceFolder, cliServer),
     );
-    this.qlTestDiscovery.refresh();
+    void this.qlTestDiscovery.refresh();
 
     this.push(this.qlTestDiscovery.onDidChangeTests(this.discoverTests, this));
   }

extensions/ql-vscode/src/query-testing/test-manager.ts

@@ -92,7 +92,7 @@ class WorkspaceFolderHandler extends DisposableObject {
     this.push(
       this.testDiscovery.onDidChangeTests(this.handleDidChangeTests, this),
     );
-    this.testDiscovery.refresh();
+    void this.testDiscovery.refresh();
   }
 
   private handleDidChangeTests(): void {

extensions/ql-vscode/test/vscode-tests/minimal-workspace/queries-panel/query-discovery.test.ts

@@ -6,14 +6,10 @@ import {
   workspace,
 } from "vscode";
 import { CodeQLCliServer } from "../../../../src/codeql-cli/cli";
-import {
-  QueryDiscovery,
-  QueryDiscoveryResults,
-} from "../../../../src/queries-panel/query-discovery";
+import { QueryDiscovery } from "../../../../src/queries-panel/query-discovery";
 import { createMockApp } from "../../../__mocks__/appMock";
 import { mockedObject } from "../../utils/mocking.helpers";
 import { basename, join, sep } from "path";
-import { sleep } from "../../../../src/pure/time";
 
 describe("QueryDiscovery", () => {
   beforeEach(() => {
@@ -28,11 +24,10 @@ describe("QueryDiscovery", () => {
       });
 
       const discovery = new QueryDiscovery(createMockApp({}), cli);
-      const results: QueryDiscoveryResults = await (
-        discovery as any
-      ).discover();
+      await discovery.refresh();
+      const queries = discovery.queries;
 
-      expect(results.queries).toEqual([]);
+      expect(queries).toEqual([]);
       expect(resolveQueries).toHaveBeenCalledTimes(1);
     });
 
@@ -49,22 +44,18 @@ describe("QueryDiscovery", () => {
       });
 
       const discovery = new QueryDiscovery(createMockApp({}), cli);
-      const results: QueryDiscoveryResults = await (
-        discovery as any
-      ).discover();
-
-      expect(results.queries[0].children.length).toEqual(3);
-      expect(results.queries[0].children[0].name).toEqual("dir1");
-      expect(results.queries[0].children[0].children.length).toEqual(1);
-      expect(results.queries[0].children[0].children[0].name).toEqual(
-        "query1.ql",
-      );
-      expect(results.queries[0].children[1].name).toEqual("dir2");
-      expect(results.queries[0].children[1].children.length).toEqual(1);
-      expect(results.queries[0].children[1].children[0].name).toEqual(
-        "query2.ql",
-      );
-      expect(results.queries[0].children[2].name).toEqual("query3.ql");
+      await discovery.refresh();
+      const queries = discovery.queries;
+      expect(queries).toBeDefined();
+
+      expect(queries![0].children.length).toEqual(3);
+      expect(queries![0].children[0].name).toEqual("dir1");
+      expect(queries![0].children[0].children.length).toEqual(1);
+      expect(queries![0].children[0].children[0].name).toEqual("query1.ql");
+      expect(queries![0].children[1].name).toEqual("dir2");
+      expect(queries![0].children[1].children.length).toEqual(1);
+      expect(queries![0].children[1].children[0].name).toEqual("query2.ql");
+      expect(queries![0].children[2].name).toEqual("query3.ql");
     });
 
     it("should collapse directories containing only a single element", async () => {
@@ -79,25 +70,21 @@ describe("QueryDiscovery", () => {
       });
 
       const discovery = new QueryDiscovery(createMockApp({}), cli);
-      const results: QueryDiscoveryResults = await (
-        discovery as any
-      ).discover();
-
-      expect(results.queries[0].children.length).toEqual(1);
-      expect(results.queries[0].children[0].name).toEqual("dir1");
-      expect(results.queries[0].children[0].children.length).toEqual(2);
-      expect(results.queries[0].children[0].children[0].name).toEqual(
+      await discovery.refresh();
+      const queries = discovery.queries;
+      expect(queries).toBeDefined();
+
+      expect(queries![0].children.length).toEqual(1);
+      expect(queries![0].children[0].name).toEqual("dir1");
+      expect(queries![0].children[0].children.length).toEqual(2);
+      expect(queries![0].children[0].children[0].name).toEqual(
         "dir2 / dir3 / dir3",
       );
-      expect(
-        results.queries[0].children[0].children[0].children.length,
-      ).toEqual(1);
-      expect(
-        results.queries[0].children[0].children[0].children[0].name,
-      ).toEqual("query2.ql");
-      expect(results.queries[0].children[0].children[1].name).toEqual(
-        "query1.ql",
+      expect(queries![0].children[0].children[0].children.length).toEqual(1);
+      expect(queries![0].children[0].children[0].children[0].name).toEqual(
+        "query2.ql",
       );
+      expect(queries![0].children[0].children[1].name).toEqual("query1.ql");
     });
 
     it("calls resolveQueries once for each workspace folder", async () => {
@@ -128,14 +115,14 @@ describe("QueryDiscovery", () => {
       });
 
       const discovery = new QueryDiscovery(createMockApp({}), cli);
-      const results: QueryDiscoveryResults = await (
-        discovery as any
-      ).discover();
+      await discovery.refresh();
+      const queries = discovery.queries;
+      expect(queries).toBeDefined();
 
-      expect(results.queries.length).toEqual(3);
-      expect(results.queries[0].children[0].name).toEqual("query1.ql");
-      expect(results.queries[1].children[0].name).toEqual("query2.ql");
-      expect(results.queries[2].children[0].name).toEqual("query3.ql");
+      expect(queries!.length).toEqual(3);
+      expect(queries![0].children[0].name).toEqual("query1.ql");
+      expect(queries![1].children[0].name).toEqual("query2.ql");
+      expect(queries![2].children[0].name).toEqual("query3.ql");
 
       expect(resolveQueries).toHaveBeenCalledTimes(3);
     });
@@ -176,16 +163,14 @@ describe("QueryDiscovery", () => {
       const onDidChangeQueriesSpy = jest.fn();
       discovery.onDidChangeQueries(onDidChangeQueriesSpy);
 
-      const results = await (discovery as any).discover();
-      (discovery as any).update(results);
+      await discovery.refresh();
 
       expect(createFileSystemWatcherSpy).toHaveBeenCalledTimes(2);
       expect(onDidChangeQueriesSpy).toHaveBeenCalledTimes(1);
 
       onWatcherDidChangeEvent.fire(workspace.workspaceFolders![0].uri);
 
-      // Wait for refresh to finish
-      await sleep(100);
+      await discovery.waitForCurrentRefresh();
 
       expect(onDidChangeQueriesSpy).toHaveBeenCalledTimes(2);
     });
@@ -209,15 +194,13 @@ describe("QueryDiscovery", () => {
       const onDidChangeQueriesSpy = jest.fn();
       discovery.onDidChangeQueries(onDidChangeQueriesSpy);
 
-      const results = await (discovery as any).discover();
-      (discovery as any).update(results);
+      await discovery.refresh();
 
       expect(onDidChangeQueriesSpy).toHaveBeenCalledTimes(1);
 
       onDidChangeWorkspaceFoldersEvent.fire({ added: [], removed: [] });
 
-      // Wait for refresh to finish
-      await sleep(100);
+      await discovery.waitForCurrentRefresh();
 
       expect(onDidChangeQueriesSpy).toHaveBeenCalledTimes(2);
     });

extensions/ql-vscode/test/vscode-tests/minimal-workspace/query-testing/qltest-discovery.test.ts

@@ -52,12 +52,14 @@ describe("qltest-discovery", () => {
     });
 
     it("should run discovery", async () => {
-      const result = await (qlTestDiscover as any).discover();
-      expect(result.watchPath).toEqualPath(baseDir);
-      expect(result.testDirectory.path).toEqualPath(baseDir);
-      expect(result.testDirectory.name).toBe("My tests");
+      await qlTestDiscover.refresh();
+      const testDirectory = qlTestDiscover.testDirectory;
+      expect(testDirectory).toBeDefined();
 
-      let children = result.testDirectory.children;
+      expect(testDirectory!.path).toEqualPath(baseDir);
+      expect(testDirectory!.name).toBe("My tests");
+
+      let children = testDirectory!.children;
       expect(children.length).toBe(1);
 
       expect(children[0].path).toEqualPath(cDir);
@@ -83,12 +85,14 @@ describe("qltest-discovery", () => {
     it("should avoid discovery if a folder does not exist", async () => {
       await fs.remove(baseDir);
 
-      const result = await (qlTestDiscover as any).discover();
-      expect(result.watchPath).toEqualPath(baseDir);
-      expect(result.testDirectory.path).toEqualPath(baseDir);
-      expect(result.testDirectory.name).toBe("My tests");
+      await qlTestDiscover.refresh();
+      const testDirectory = qlTestDiscover.testDirectory;
+      expect(testDirectory).toBeDefined();
+
+      expect(testDirectory!.path).toEqualPath(baseDir);
+      expect(testDirectory!.name).toBe("My tests");
 
-      expect(result.testDirectory.children.length).toBe(0);
+      expect(testDirectory!.children.length).toBe(0);
     });
   });
 });