add commands for warming the overlay-base cache

extensions/ql-vscode/package.json

@@ -515,6 +515,14 @@
         "command": "codeQL.runQueryContextEditor",
         "title": "CodeQL: Run Query on Selected Database"
       },
+      {
+        "command": "codeQL.runWarmOverlayBaseCacheForQuery",
+        "title": "CodeQL: Warm Overlay-Base Cache for Query"
+      },
+      {
+        "command": "codeQL.runWarmOverlayBaseCacheForQueryContextEditor",
+        "title": "CodeQL: Warm Overlay-Base Cache for Query"
+      },
       {
         "command": "codeQL.debugQuery",
         "title": "CodeQL: Debug Query"
@@ -571,10 +579,18 @@
         "command": "codeQL.runQueries",
         "title": "CodeQL: Run Queries in Selected Files"
       },
+      {
+        "command": "codeQL.runWarmOverlayBaseCacheForQueries",
+        "title": "CodeQL: Warm Overlay-Base Cache for Queries in Selected Files"
+      },
       {
         "command": "codeQL.runQuerySuite",
         "title": "CodeQL: Run Selected Query Suite"
       },
+      {
+        "command": "codeQL.runWarmOverlayBaseCacheForQuerySuite",
+        "title": "CodeQL: Warm Overlay-Base Cache for Query Suite"
+      },
       {
         "command": "codeQL.quickEval",
         "title": "CodeQL: Quick Evaluation"
@@ -1378,11 +1394,21 @@
           "group": "9_qlCommands",
           "when": "resourceScheme != codeql-zip-archive"
         },
+        {
+          "command": "codeQL.runWarmOverlayBaseCacheForQueries",
+          "group": "9_qlCommands",
+          "when": "resourceScheme != codeql-zip-archive && config.codeQL.canary"
+        },
         {
           "command": "codeQL.runQuerySuite",
           "group": "9_qlCommands",
           "when": "resourceScheme != codeql-zip-archive && resourceExtname == .qls && !explorerResourceIsFolder && !listMultiSelection && config.codeQL.canary"
         },
+        {
+          "command": "codeQL.runWarmOverlayBaseCacheForQuerySuite",
+          "group": "9_qlCommands",
+          "when": "resourceScheme != codeql-zip-archive && resourceExtname == .qls && !explorerResourceIsFolder && !listMultiSelection && config.codeQL.canary"
+        },
         {
           "command": "codeQL.runVariantAnalysisContextExplorer",
           "group": "9_qlCommands",
@@ -1408,6 +1434,10 @@
           "command": "codeQL.runQuery",
           "when": "resourceLangId == ql && resourceExtname == .ql"
         },
+        {
+          "command": "codeQL.runWarmOverlayBaseCacheForQuery",
+          "when": "resourceLangId == ql && resourceExtname == .ql && config.codeQL.canary"
+        },
         {
           "command": "codeQLQueries.runLocalQueryFromQueriesPanel",
           "when": "false"
@@ -1428,6 +1458,10 @@
           "command": "codeQL.runQueryContextEditor",
           "when": "false"
         },
+        {
+          "command": "codeQL.runWarmOverlayBaseCacheForQueryContextEditor",
+          "when": "false"
+        },
         {
           "command": "codeQL.debugQuery",
           "when": "config.codeQL.canary && editorLangId == ql && resourceExtname == .ql && !inDebugMode"
@@ -1480,10 +1514,18 @@
           "command": "codeQL.runQueries",
           "when": "false"
         },
+        {
+          "command": "codeQL.runWarmOverlayBaseCacheForQueries",
+          "when": "false"
+        },
         {
           "command": "codeQL.runQuerySuite",
           "when": "false"
         },
+        {
+          "command": "codeQL.runWarmOverlayBaseCacheForQuerySuite",
+          "when": "false"
+        },
         {
           "command": "codeQL.quickEval",
           "when": "editorLangId == ql"
@@ -1832,6 +1874,10 @@
           "command": "codeQL.runQueryContextEditor",
           "when": "editorLangId == ql && resourceExtname == .ql && !inDebugMode"
         },
+        {
+          "command": "codeQL.runWarmOverlayBaseCacheForQueryContextEditor",
+          "when": "editorLangId == ql && resourceExtname == .ql && !inDebugMode && config.codeQL.canary"
+        },
         {
           "command": "codeQL.runQueryOnMultipleDatabasesContextEditor",
           "when": "editorLangId == ql && resourceExtname == .ql"

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

@@ -126,7 +126,11 @@ export type QueryEditorCommands = {
 // Commands used for running local queries
 export type LocalQueryCommands = {
   "codeQL.runQuery": (uri?: Uri) => Promise<void>;
+  "codeQL.runWarmOverlayBaseCacheForQuery": (uri?: Uri) => Promise<void>;
   "codeQL.runQueryContextEditor": (uri?: Uri) => Promise<void>;
+  "codeQL.runWarmOverlayBaseCacheForQueryContextEditor": (
+    uri?: Uri,
+  ) => Promise<void>;
   "codeQL.runQueryOnMultipleDatabases": (uri?: Uri) => Promise<void>;
   "codeQL.runQueryOnMultipleDatabasesContextEditor": (
     uri?: Uri,
@@ -138,7 +142,9 @@ export type LocalQueryCommands = {
   "codeQLQueries.createQuery": () => Promise<void>;
   "codeQL.runLocalQueryFromFileTab": (uri: Uri) => Promise<void>;
   "codeQL.runQueries": ExplorerSelectionCommandFunction<Uri>;
+  "codeQL.runWarmOverlayBaseCacheForQueries": ExplorerSelectionCommandFunction<Uri>;
   "codeQL.runQuerySuite": ExplorerSelectionCommandFunction<Uri>;
+  "codeQL.runWarmOverlayBaseCacheForQuerySuite": ExplorerSelectionCommandFunction<Uri>;
   "codeQL.quickEval": (uri: Uri) => Promise<void>;
   "codeQL.quickEvalCount": (uri: Uri) => Promise<void>;
   "codeQL.quickEvalContextEditor": (uri: Uri) => Promise<void>;

extensions/ql-vscode/src/common/logging/vscode/loggers.ts

@@ -10,6 +10,20 @@ export const extLogger = new OutputChannelLogger("CodeQL Extension Log");
 // Logger for messages from the query server.
 export const queryServerLogger = new OutputChannelLogger("CodeQL Query Server");
 
+// Logger for messages from the query server for warming overlay-base cache.
+let queryServerForWarmingOverlayBaseCacheLogger:
+  | OutputChannelLogger
+  | undefined;
+
+// construct queryServerForWarmingOverlayBaseCacheLogger lazily, as most users don't need it
+export function getQueryServerForWarmingOverlayBaseCacheLogger(): OutputChannelLogger {
+  if (!queryServerForWarmingOverlayBaseCacheLogger)
+    queryServerForWarmingOverlayBaseCacheLogger = new OutputChannelLogger(
+      "CodeQL Query Server for warming overlay-base cache",
+    );
+  return queryServerForWarmingOverlayBaseCacheLogger;
+}
+
 // Logger for messages from the language server.
 export const languageServerLogger = new OutputChannelLogger(
   "CodeQL Language Server",

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

@@ -739,6 +739,35 @@ export class DatabaseManager extends DisposableObject {
     }
   }
 
+  public async runWithDatabaseInSeparateQueryRunner(
+    queryRunner: QueryRunner,
+    whatToDo: () => Promise<void>,
+  ) {
+    try {
+      if (this._currentDatabaseItem) {
+        const dbItem = this._currentDatabaseItem;
+        await this.qs.deregisterDatabase(dbItem);
+        await queryRunner.registerDatabase(dbItem);
+      }
+      await whatToDo();
+      if (this._currentDatabaseItem) {
+        const dbItem = this._currentDatabaseItem;
+        await queryRunner.deregisterDatabase(dbItem);
+        await this.qs.registerDatabase(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 use database for warming overlay-base cache because query server is not running.`,
+        );
+        return;
+      }
+      throw e;
+    }
+  }
+
   private async deregisterDatabase(dbItem: DatabaseItem) {
     try {
       await this.qs.deregisterDatabase(dbItem);

extensions/ql-vscode/src/extension.ts

@@ -86,6 +86,7 @@ import {
   extLogger,
   languageServerLogger,
   queryServerLogger,
+  getQueryServerForWarmingOverlayBaseCacheLogger,
 } from "./common/logging/vscode";
 import { QueryHistoryManager } from "./query-history/query-history-manager";
 import type { CompletedLocalQueryInfo } from "./query-results";
@@ -172,6 +173,7 @@ function getCommands(
   app: App,
   cliServer: CodeQLCliServer,
   queryRunner: QueryRunner,
+  queryRunnerForWarmingOverlayBaseCache: QueryRunner | undefined,
   languageClient: LanguageClient,
 ): BaseCommands {
   const getCliVersion = async () => {
@@ -189,6 +191,9 @@ function getCommands(
         cliServer.restartCliServer();
         await Promise.all([
           queryRunner.restartQueryServer(progress),
+          queryRunnerForWarmingOverlayBaseCache
+            ? queryRunnerForWarmingOverlayBaseCache.restartQueryServer(progress)
+            : {},
           async () => {
             if (languageClient.isRunning()) {
               await languageClient.restart();
@@ -201,6 +206,12 @@ function getCommands(
           queryServerLogger,
           "CodeQL Query Server restarted.",
         );
+        if (queryRunnerForWarmingOverlayBaseCache) {
+          void showAndLogErrorMessage(
+            getQueryServerForWarmingOverlayBaseCacheLogger(),
+            "CodeQL Query Server for warming overlay-base cache restarted.",
+          );
+        }
       },
       {
         title: "Restarting Query Server",
@@ -281,6 +292,7 @@ export interface CodeQLExtensionInterface {
   readonly ctx: ExtensionContext;
   readonly cliServer: CodeQLCliServer;
   readonly qs: QueryRunner;
+  readonly qsForWarmingOverlayBaseCache: QueryRunner | undefined;
   readonly distributionManager: DistributionManager;
   readonly databaseManager: DatabaseManager;
   readonly databaseUI: DatabaseUI;
@@ -795,8 +807,28 @@ async function activateWithInstalledDistribution(
     qlConfigurationListener,
     cliServer,
     ctx,
+    false,
   );
 
+  let qsForWarmingOverlayBaseCache: QueryRunner | undefined;
+
+  // construct qsForWarmingOverlayBaseCache lazily, as most users don't need it
+  async function getQsForWarmingOverlayBaseCache(): Promise<QueryRunner> {
+    if (!qsForWarmingOverlayBaseCache) {
+      void extLogger.log(
+        "Initializing base cache warming query server client.",
+      );
+      qsForWarmingOverlayBaseCache = await createQueryServer(
+        app,
+        qlConfigurationListener,
+        cliServer,
+        ctx,
+        true,
+      );
+    }
+    return qsForWarmingOverlayBaseCache;
+  }
+
   for (const glob of CLEAR_PACK_CACHE_ON_EDIT_GLOBS) {
     const fsWatcher = workspace.createFileSystemWatcher(glob);
     ctx.subscriptions.push(fsWatcher);
@@ -998,6 +1030,7 @@ async function activateWithInstalledDistribution(
   const localQueries = new LocalQueries(
     app,
     qs,
+    getQsForWarmingOverlayBaseCache,
     qhm,
     dbm,
     databaseFetcher,
@@ -1062,7 +1095,13 @@ async function activateWithInstalledDistribution(
   void extLogger.log("Registering top-level command palette commands.");
 
   const allCommands: AllExtensionCommands = {
-    ...getCommands(app, cliServer, qs, languageClient),
+    ...getCommands(
+      app,
+      cliServer,
+      qs,
+      qsForWarmingOverlayBaseCache,
+      languageClient,
+    ),
     ...getQueryEditorCommands({
       commandManager: app.commands,
       queryRunner: qs,
@@ -1165,6 +1204,7 @@ async function activateWithInstalledDistribution(
     cliServer,
     localQueries,
     qs,
+    qsForWarmingOverlayBaseCache,
     distributionManager,
     databaseManager: dbm,
     databaseUI,
@@ -1278,9 +1318,12 @@ async function createQueryServer(
   qlConfigurationListener: QueryServerConfigListener,
   cliServer: CodeQLCliServer,
   ctx: ExtensionContext,
+  warmOverlayBaseCache: boolean,
 ): Promise<QueryRunner> {
   const qsOpts = {
-    logger: queryServerLogger,
+    logger: warmOverlayBaseCache
+      ? getQueryServerForWarmingOverlayBaseCacheLogger()
+      : queryServerLogger,
     contextStoragePath: getContextStoragePath(ctx),
   };
   const progressCallback = (
@@ -1290,7 +1333,12 @@ async function createQueryServer(
     ) => Thenable<void>,
   ) =>
     Window.withProgress(
-      { title: "CodeQL query server", location: ProgressLocation.Window },
+      {
+        title: warmOverlayBaseCache
+          ? "CodeQL query server for warming overlay-base cache"
+          : "CodeQL query server",
+        location: ProgressLocation.Window,
+      },
       task,
     );
 
@@ -1300,6 +1348,7 @@ async function createQueryServer(
     cliServer,
     qsOpts,
     progressCallback,
+    warmOverlayBaseCache,
   );
   ctx.subscriptions.push(qs);
   await qs.startQueryServer();

extensions/ql-vscode/src/language-support/ast-viewer/ast-cfg-commands.ts

@@ -59,6 +59,7 @@ export function getAstCfgCommands({
             progress,
             token,
             undefined,
+            false,
             undefined,
             res[1],
           );