add commands for warming the overlay-base cache

Author: ginsbach

extensions/ql-vscode/package.json

@@ -515,6 +515,14 @@
         "command": "codeQL.runQueryContextEditor",
         "title": "CodeQL: Run Query on Selected Database"
       },
+      {
+        "command": "codeQL.runWarmOverlayBaseCacheQuery",
+        "title": "CodeQL: Warm Overlay-Base Cache for Query"
+      },
+      {
+        "command": "codeQL.runWarmOverlayBaseCacheQueryContextEditor",
+        "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.runWarmOverlayBaseCacheQueries",
+        "title": "CodeQL: Warm Overlay-Base Cache for Queries in Selected Files"
+      },
       {
         "command": "codeQL.runQuerySuite",
         "title": "CodeQL: Run Selected Query Suite"
       },
+      {
+        "command": "codeQL.runWarmOverlayBaseCacheQuerySuite",
+        "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.runWarmOverlayBaseCacheQueries",
+          "group": "9_qlCommands",
+          "when": "resourceScheme != codeql-zip-archive"
+        },
         {
           "command": "codeQL.runQuerySuite",
           "group": "9_qlCommands",
           "when": "resourceScheme != codeql-zip-archive && resourceExtname == .qls && !explorerResourceIsFolder && !listMultiSelection && config.codeQL.canary"
         },
+        {
+          "command": "codeQL.runWarmOverlayBaseCacheQuerySuite",
+          "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.runWarmOverlayBaseCacheQuery",
+          "when": "resourceLangId == ql && resourceExtname == .ql"
+        },
         {
           "command": "codeQLQueries.runLocalQueryFromQueriesPanel",
           "when": "false"
@@ -1428,6 +1458,10 @@
           "command": "codeQL.runQueryContextEditor",
           "when": "false"
         },
+        {
+          "command": "codeQL.runWarmOverlayBaseCacheQueryContextEditor",
+          "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.runWarmOverlayBaseCacheQueries",
+          "when": "false"
+        },
         {
           "command": "codeQL.runQuerySuite",
           "when": "false"
         },
+        {
+          "command": "codeQL.runWarmOverlayBaseCacheQuerySuite",
+          "when": "false"
+        },
         {
           "command": "codeQL.quickEval",
           "when": "editorLangId == ql"
@@ -1832,6 +1874,10 @@
           "command": "codeQL.runQueryContextEditor",
           "when": "editorLangId == ql && resourceExtname == .ql && !inDebugMode"
         },
+        {
+          "command": "codeQL.runWarmOverlayBaseCacheQueryContextEditor",
+          "when": "editorLangId == ql && resourceExtname == .ql && !inDebugMode"
+        },
         {
           "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.runWarmOverlayBaseCacheQuery": (uri?: Uri) => Promise<void>;
   "codeQL.runQueryContextEditor": (uri?: Uri) => Promise<void>;
+  "codeQL.runWarmOverlayBaseCacheQueryContextEditor": (
+    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.runWarmOverlayBaseCacheQueries": ExplorerSelectionCommandFunction<Uri>;
   "codeQL.runQuerySuite": ExplorerSelectionCommandFunction<Uri>;
+  "codeQL.runWarmOverlayBaseCacheQuerySuite": 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,10 @@ 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.
+export const queryServerForWarmingOverlayBaseCacheLogger =
+  new OutputChannelLogger("CodeQL Query Server for warming overlay-base cache");
+
 // 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

@@ -118,6 +118,7 @@ export class DatabaseManager extends DisposableObject {
     private readonly ctx: ExtensionContext,
     private readonly app: App,
     private readonly qs: QueryRunner,
+    private readonly qsForWarmingOverlayBaseCache: QueryRunner,
     private readonly cli: CodeQLCliServer,
     private readonly languageContext: LanguageContextStore,
     public logger: Logger,
@@ -739,6 +740,31 @@ export class DatabaseManager extends DisposableObject {
     }
   }
 
+  public async withDatabaseInQsForWarmingOverlayBaseCache(
+    whatToDo: () => Promise<void>,
+  ) {
+    try {
+      if (this._currentDatabaseItem) {
+        const dbItem = this._currentDatabaseItem;
+        await this.qs.deregisterDatabase(dbItem);
+        await this.qsForWarmingOverlayBaseCache.registerDatabase(dbItem);
+        await whatToDo();
+        await this.qsForWarmingOverlayBaseCache.deregisterDatabase(dbItem);
+        await this.qs.registerDatabase(dbItem);
+      } else await whatToDo();
+    } 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/debugger/debugger-ui.ts

@@ -141,6 +141,7 @@ class QLDebugAdapterTracker
             [this.configuration.query, results],
           ]),
         },
+        false,
         (_) => {},
       );
       this.localQueryRun = undefined;

extensions/ql-vscode/src/extension.ts

@@ -86,6 +86,7 @@ import {
   extLogger,
   languageServerLogger,
   queryServerLogger,
+  queryServerForWarmingOverlayBaseCacheLogger,
 } 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,
   languageClient: LanguageClient,
 ): BaseCommands {
   const getCliVersion = async () => {
@@ -189,6 +191,7 @@ function getCommands(
         cliServer.restartCliServer();
         await Promise.all([
           queryRunner.restartQueryServer(progress),
+          queryRunnerForWarmingOverlayBaseCache.restartQueryServer(progress),
           async () => {
             if (languageClient.isRunning()) {
               await languageClient.restart();
@@ -201,6 +204,10 @@ function getCommands(
           queryServerLogger,
           "CodeQL Query Server restarted.",
         );
+        void showAndLogErrorMessage(
+          queryServerForWarmingOverlayBaseCacheLogger,
+          "CodeQL Query Server for warming underlay caches restarted.",
+        );
       },
       {
         title: "Restarting Query Server",
@@ -281,6 +288,7 @@ export interface CodeQLExtensionInterface {
   readonly ctx: ExtensionContext;
   readonly cliServer: CodeQLCliServer;
   readonly qs: QueryRunner;
+  readonly WarmOverlayBaseCacheQs: QueryRunner;
   readonly distributionManager: DistributionManager;
   readonly databaseManager: DatabaseManager;
   readonly databaseUI: DatabaseUI;
@@ -795,6 +803,16 @@ async function activateWithInstalledDistribution(
     qlConfigurationListener,
     cliServer,
     ctx,
+    false,
+  );
+
+  void extLogger.log("Initializing base cache warming query server client.");
+  const WarmOverlayBaseCacheQs = await createQueryServer(
+    app,
+    qlConfigurationListener,
+    cliServer,
+    ctx,
+    true,
   );
 
   for (const glob of CLEAR_PACK_CACHE_ON_EDIT_GLOBS) {
@@ -822,6 +840,7 @@ async function activateWithInstalledDistribution(
     ctx,
     app,
     qs,
+    WarmOverlayBaseCacheQs,
     cliServer,
     languageContext,
     extLogger,
@@ -998,6 +1017,7 @@ async function activateWithInstalledDistribution(
   const localQueries = new LocalQueries(
     app,
     qs,
+    WarmOverlayBaseCacheQs,
     qhm,
     dbm,
     databaseFetcher,
@@ -1062,7 +1082,7 @@ async function activateWithInstalledDistribution(
   void extLogger.log("Registering top-level command palette commands.");
 
   const allCommands: AllExtensionCommands = {
-    ...getCommands(app, cliServer, qs, languageClient),
+    ...getCommands(app, cliServer, qs, WarmOverlayBaseCacheQs, languageClient),
     ...getQueryEditorCommands({
       commandManager: app.commands,
       queryRunner: qs,
@@ -1165,6 +1185,7 @@ async function activateWithInstalledDistribution(
     cliServer,
     localQueries,
     qs,
+    WarmOverlayBaseCacheQs,
     distributionManager,
     databaseManager: dbm,
     databaseUI,
@@ -1278,9 +1299,12 @@ async function createQueryServer(
   qlConfigurationListener: QueryServerConfigListener,
   cliServer: CodeQLCliServer,
   ctx: ExtensionContext,
+  warmOverlayBaseCache: boolean,
 ): Promise<QueryRunner> {
   const qsOpts = {
-    logger: queryServerLogger,
+    logger: warmOverlayBaseCache
+      ? queryServerForWarmingOverlayBaseCacheLogger
+      : queryServerLogger,
     contextStoragePath: getContextStoragePath(ctx),
   };
   const progressCallback = (
@@ -1290,13 +1314,19 @@ async function createQueryServer(
     ) => Thenable<void>,
   ) =>
     Window.withProgress(
-      { title: "CodeQL query server", location: ProgressLocation.Window },
+      {
+        title: warmOverlayBaseCache
+          ? "CodeQL query server for warming underlay caches"
+          : "CodeQL query server",
+        location: ProgressLocation.Window,
+      },
       task,
     );
 
   const qs = new QueryServerClient(
     app,
     qlConfigurationListener,
+    warmOverlayBaseCache,
     cliServer,
     qsOpts,
     progressCallback,

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],
           );