Add cache of query schemas to results view

extensions/ql-vscode/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 - Add a palette command that allows importing all databases directly inside of a parent folder. [#3797](https://github.com/github/vscode-codeql/pull/3797)
 - Only use VS Code telemetry settings instead of using `codeQL.telemetry.enableTelemetry` [#3853](https://github.com/github/vscode-codeql/pull/3853)
+- Improve the performance of the results view with large numbers of results. [#3862](https://github.com/github/vscode-codeql/pull/3862)
 
 ## 1.16.1 - 6 November 2024
 

extensions/ql-vscode/src/language-support/contextual/cached-operation.ts

@@ -7,6 +7,7 @@ export class CachedOperation<S extends unknown[], U> {
   private readonly operation: (t: string, ...args: S) => Promise<U>;
   private readonly cached: Map<string, U>;
   private readonly lru: string[];
+  private generation: number;
   private readonly inProgressCallbacks: Map<
     string,
     Array<[(u: U) => void, (reason?: Error) => void]>
@@ -17,6 +18,7 @@ export class CachedOperation<S extends unknown[], U> {
     private cacheSize = 100,
   ) {
     this.operation = operation;
+    this.generation = 0;
     this.lru = [];
     this.inProgressCallbacks = new Map<
       string,
@@ -46,14 +48,19 @@ export class CachedOperation<S extends unknown[], U> {
         inProgressCallback.push([resolve, reject]);
       });
     }
-
+    const origGeneration = this.generation;
     // Otherwise compute the new value, but leave a callback to allow sharing work
     const callbacks: Array<[(u: U) => void, (reason?: Error) => void]> = [];
     this.inProgressCallbacks.set(t, callbacks);
     try {
       const result = await this.operation(t, ...args);
       callbacks.forEach((f) => f[0](result));
       this.inProgressCallbacks.delete(t);
+      if (this.generation !== origGeneration) {
+        // Cache was reset in the meantime so don't trust this
+        // result enough to cache it.
+        return result;
+      }
       if (this.lru.length > this.cacheSize) {
         const toRemove = this.lru.shift()!;
         this.cached.delete(toRemove);
@@ -69,4 +76,11 @@ export class CachedOperation<S extends unknown[], U> {
       this.inProgressCallbacks.delete(t);
     }
   }
+
+  reset() {
+    this.cached.clear();
+    this.lru.length = 0;
+    this.generation++;
+    this.inProgressCallbacks.clear();
+  }
 }

extensions/ql-vscode/src/local-queries/results-view.ts

@@ -75,6 +75,7 @@ import type { App } from "../common/app";
 import type { Disposable } from "../common/disposable-object";
 import type { RawResultSet } from "../common/raw-result-types";
 import type { BqrsResultSetSchema } from "../common/bqrs-cli-types";
+import { CachedOperation } from "../language-support/contextual/cached-operation";
 
 /**
  * results-view.ts
@@ -177,6 +178,8 @@ export class ResultsView extends AbstractWebview<
   // Event listeners that should be disposed of when the view is disposed.
   private disposableEventListeners: Disposable[] = [];
 
+  private schemaCache: CachedOperation<[], BqrsResultSetSchema[]>;
+
   constructor(
     app: App,
     private databaseManager: DatabaseManager,
@@ -206,6 +209,10 @@ export class ResultsView extends AbstractWebview<
         }
       }),
     );
+
+    this.schemaCache = new CachedOperation(
+      this.getResultSetSchemasImpl.bind(this),
+    );
   }
 
   public getCommands(): ResultsViewCommands {
@@ -420,6 +427,7 @@ export class ResultsView extends AbstractWebview<
       );
       return;
     }
+    this.schemaCache.reset();
     // Notify the webview that it should expect new results.
     await this.postMessage({ t: "resultsUpdating" });
     await this._displayedQuery.completedQuery.updateSortState(
@@ -610,6 +618,12 @@ export class ResultsView extends AbstractWebview<
     selectedTable = "",
   ): Promise<BqrsResultSetSchema[]> {
     const resultsPath = completedQuery.getResultsPath(selectedTable);
+    return this.schemaCache.get(resultsPath);
+  }
+
+  private async getResultSetSchemasImpl(
+    resultsPath: string,
+  ): Promise<BqrsResultSetSchema[]> {
     const schemas = await this.cliServer.bqrsInfo(
       resultsPath,
       PAGE_SIZE.getValue(),

extensions/ql-vscode/src/run-queries-shared.ts

@@ -244,7 +244,7 @@ export class QueryEvaluationInfo extends QueryOutputDir {
    */
   async chooseResultSet(cliServer: CodeQLCliServer) {
     const resultSets = (
-      await cliServer.bqrsInfo(this.resultsPaths.resultsPath, 0)
+      await cliServer.bqrsInfo(this.resultsPaths.resultsPath)
     )["result-sets"];
     if (!resultSets.length) {
       return undefined;