Merge pull request #3762 from github/koesie10/fix-codeql-download

Store state of CodeQL distribution on filesystem instead of in `globalState`

Author: koesie10

extensions/ql-vscode/CHANGELOG.md

@@ -3,6 +3,7 @@
 ## [UNRELEASED]
 
 - Support result columns of type `QlBuiltins::BigInt` in quick evaluations. [#3647](https://github.com/github/vscode-codeql/pull/3647)
+- Fix a bug where the CodeQL CLI would be re-downloaded if you switched to a different filesystem (for example Codespaces or a remote SSH host). [#3762](https://github.com/github/vscode-codeql/pull/3762)
 
 ## 1.16.0 - 10 October 2024
 

extensions/ql-vscode/package-lock.json

@@ -1986,6 +1986,7 @@
     "msw": "^2.2.13",
     "nanoid": "^5.0.7",
     "p-queue": "^8.0.1",
+    "proper-lockfile": "^4.1.2",
     "react": "^18.3.1",
     "react-dom": "^18.3.1",
     "semver": "^7.6.2",
@@ -2040,6 +2041,7 @@
     "@types/js-yaml": "^4.0.6",
     "@types/nanoid": "^3.0.0",
     "@types/node": "20.16.*",
+    "@types/proper-lockfile": "^4.1.4",
     "@types/react": "^18.3.1",
     "@types/react-dom": "^18.3.0",
     "@types/sarif": "^2.1.2",

extensions/ql-vscode/package.json

@@ -1,5 +1,12 @@
 import type { WriteStream } from "fs";
-import { createWriteStream, mkdtemp, pathExists, remove } from "fs-extra";
+import {
+  createWriteStream,
+  mkdtemp,
+  outputJson,
+  pathExists,
+  readJson,
+  remove,
+} from "fs-extra";
 import { tmpdir } from "os";
 import { delimiter, dirname, join } from "path";
 import { Range, satisfies } from "semver";
@@ -19,7 +26,9 @@ import {
   InvocationRateLimiter,
   InvocationRateLimiterResultKind,
 } from "../common/invocation-rate-limiter";
+import type { NotificationLogger } from "../common/logging";
 import {
+  showAndLogExceptionWithTelemetry,
   showAndLogErrorMessage,
   showAndLogWarningMessage,
 } from "../common/logging";
@@ -28,6 +37,11 @@ import { reportUnzipProgress } from "../common/vscode/unzip-progress";
 import type { Release } from "./distribution/release";
 import { ReleasesApiConsumer } from "./distribution/releases-api-consumer";
 import { createTimeoutSignal } from "../common/fetch-stream";
+import { withDistributionUpdateLock } from "./lock";
+import { asError, getErrorMessage } from "../common/helpers-pure";
+import { isIOError } from "../common/files";
+import { telemetryListener } from "../common/vscode/telemetry";
+import { redactableError } from "../common/errors";
 
 /**
  * distribution.ts
@@ -53,6 +67,11 @@ const NIGHTLY_DISTRIBUTION_REPOSITORY_NWO = "dsp-testing/codeql-cli-nightlies";
  */
 export const DEFAULT_DISTRIBUTION_VERSION_RANGE: Range = new Range("2.x");
 
+export interface DistributionState {
+  folderIndex: number;
+  release: Release | null;
+}
+
 export interface DistributionProvider {
   getCodeQlPathWithoutVersionCheck(): Promise<string | undefined>;
   onDidChangeDistribution?: Event<void>;
@@ -64,13 +83,15 @@ export class DistributionManager implements DistributionProvider {
     public readonly config: DistributionConfig,
     private readonly versionRange: Range,
     extensionContext: ExtensionContext,
+    logger: NotificationLogger,
   ) {
     this._onDidChangeDistribution = config.onDidChangeConfiguration;
     this.extensionSpecificDistributionManager =
       new ExtensionSpecificDistributionManager(
         config,
         versionRange,
         extensionContext,
+        logger,
       );
     this.updateCheckRateLimiter = new InvocationRateLimiter(
       extensionContext.globalState,
@@ -80,6 +101,10 @@ export class DistributionManager implements DistributionProvider {
     );
   }
 
+  public async initialize(): Promise<void> {
+    await this.extensionSpecificDistributionManager.initialize();
+  }
+
   /**
    * Look up a CodeQL launcher binary.
    */
@@ -280,14 +305,58 @@ export class DistributionManager implements DistributionProvider {
 }
 
 class ExtensionSpecificDistributionManager {
+  private distributionState: DistributionState | undefined;
+
   constructor(
     private readonly config: DistributionConfig,
     private readonly versionRange: Range,
     private readonly extensionContext: ExtensionContext,
+    private readonly logger: NotificationLogger,
   ) {
     /**/
   }
 
+  public async initialize() {
+    await this.ensureDistributionStateExists();
+  }
+
+  private async ensureDistributionStateExists() {
+    const distributionStatePath = this.getDistributionStatePath();
+    try {
+      this.distributionState = await readJson(distributionStatePath);
+    } catch (e: unknown) {
+      if (isIOError(e) && e.code === "ENOENT") {
+        // If the file doesn't exist, that just means we need to create it
+
+        this.distributionState = {
+          folderIndex:
+            this.extensionContext.globalState.get(
+              "distributionFolderIndex",
+              0,
+            ) ?? 0,
+          release: (this.extensionContext.globalState.get(
+            "distributionRelease",
+          ) ?? null) as Release | null,
+        };
+
+        // This may result in a race condition, but when this happens both processes should write the same file.
+        await outputJson(distributionStatePath, this.distributionState);
+      } else {
+        void showAndLogExceptionWithTelemetry(
+          this.logger,
+          telemetryListener,
+          redactableError(
+            asError(e),
+          )`Failed to read distribution state from ${distributionStatePath}: ${getErrorMessage(e)}`,
+        );
+        this.distributionState = {
+          folderIndex: 0,
+          release: null,
+        };
+      }
+    }
+  }
+
   public async getCodeQlPathWithoutVersionCheck(): Promise<string | undefined> {
     if (this.getInstalledRelease() !== undefined) {
       // An extension specific distribution has been installed.
@@ -350,9 +419,21 @@ class ExtensionSpecificDistributionManager {
     release: Release,
     progressCallback?: ProgressCallback,
   ): Promise<void> {
-    await this.downloadDistribution(release, progressCallback);
-    // Store the installed release within the global extension state.
-    await this.storeInstalledRelease(release);
+    if (!this.distributionState) {
+      await this.ensureDistributionStateExists();
+    }
+
+    const distributionStatePath = this.getDistributionStatePath();
+
+    await withDistributionUpdateLock(
+      // .lock will be appended to this filename
+      distributionStatePath,
+      async () => {
+        await this.downloadDistribution(release, progressCallback);
+        // Store the installed release within the global extension state.
+        await this.storeInstalledRelease(release);
+      },
+    );
   }
 
   private async downloadDistribution(
@@ -564,23 +645,19 @@ class ExtensionSpecificDistributionManager {
   }
 
   private async bumpDistributionFolderIndex(): Promise<void> {
-    const index = this.extensionContext.globalState.get(
-      ExtensionSpecificDistributionManager._currentDistributionFolderIndexStateKey,
-      0,
-    );
-    await this.extensionContext.globalState.update(
-      ExtensionSpecificDistributionManager._currentDistributionFolderIndexStateKey,
-      index + 1,
-    );
+    await this.updateState((oldState) => {
+      return {
+        ...oldState,
+        folderIndex: (oldState.folderIndex ?? 0) + 1,
+      };
+    });
   }
 
   private getDistributionStoragePath(): string {
+    const distributionState = this.getDistributionState();
+
     // Use an empty string for the initial distribution for backwards compatibility.
-    const distributionFolderIndex =
-      this.extensionContext.globalState.get(
-        ExtensionSpecificDistributionManager._currentDistributionFolderIndexStateKey,
-        0,
-      ) || "";
+    const distributionFolderIndex = distributionState.folderIndex || "";
     return join(
       this.extensionContext.globalStorageUri.fsPath,
       ExtensionSpecificDistributionManager._currentDistributionFolderBaseName +
@@ -595,26 +672,55 @@ class ExtensionSpecificDistributionManager {
     );
   }
 
-  private getInstalledRelease(): Release | undefined {
-    return this.extensionContext.globalState.get(
-      ExtensionSpecificDistributionManager._installedReleaseStateKey,
+  private getDistributionStatePath(): string {
+    return join(
+      this.extensionContext.globalStorageUri.fsPath,
+      ExtensionSpecificDistributionManager._distributionStateFilename,
     );
   }
 
+  private getInstalledRelease(): Release | undefined {
+    return this.getDistributionState().release ?? undefined;
+  }
+
   private async storeInstalledRelease(
     release: Release | undefined,
   ): Promise<void> {
-    await this.extensionContext.globalState.update(
-      ExtensionSpecificDistributionManager._installedReleaseStateKey,
-      release,
-    );
+    await this.updateState((oldState) => ({
+      ...oldState,
+      release: release ?? null,
+    }));
+  }
+
+  private getDistributionState(): DistributionState {
+    const distributionState = this.distributionState;
+    if (distributionState === undefined) {
+      throw new Error(
+        "Invariant violation: distribution state not initialized",
+      );
+    }
+    return distributionState;
+  }
+
+  private async updateState(
+    f: (oldState: DistributionState) => DistributionState,
+  ) {
+    const oldState = this.distributionState;
+    if (oldState === undefined) {
+      throw new Error(
+        "Invariant violation: distribution state not initialized",
+      );
+    }
+    const newState = f(oldState);
+    this.distributionState = newState;
+
+    const distributionStatePath = this.getDistributionStatePath();
+    await outputJson(distributionStatePath, newState);
   }
 
   private static readonly _currentDistributionFolderBaseName = "distribution";
-  private static readonly _currentDistributionFolderIndexStateKey =
-    "distributionFolderIndex";
-  private static readonly _installedReleaseStateKey = "distributionRelease";
   private static readonly _codeQlExtractedFolderName = "codeql";
+  private static readonly _distributionStateFilename = "distribution.json";
 }
 
 /*

extensions/ql-vscode/src/codeql-cli/distribution.ts

@@ -0,0 +1,22 @@
+import { lock } from "proper-lockfile";
+
+export async function withDistributionUpdateLock(
+  lockFile: string,
+  f: () => Promise<void>,
+) {
+  const release = await lock(lockFile, {
+    stale: 60_000, // 1 minute. We can take the lock longer than this because that's based on the update interval.
+    update: 10_000, // 10 seconds
+    retries: {
+      minTimeout: 10_000,
+      maxTimeout: 60_000,
+      retries: 100,
+    },
+  });
+
+  try {
+    await f();
+  } finally {
+    await release();
+  }
+}