Use file for storing distribution state instead of globalState

Author: koesie10

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

@@ -3,6 +3,7 @@ import {
   createWriteStream,
   mkdtemp,
   pathExists,
+  readJson,
   remove,
   writeJson,
 } from "fs-extra";
@@ -25,7 +26,9 @@ import {
   InvocationRateLimiter,
   InvocationRateLimiterResultKind,
 } from "../common/invocation-rate-limiter";
+import type { NotificationLogger } from "../common/logging";
 import {
+  showAndLogExceptionWithTelemetry,
   showAndLogErrorMessage,
   showAndLogWarningMessage,
 } from "../common/logging";
@@ -35,6 +38,10 @@ 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
@@ -60,6 +67,11 @@ const NIGHTLY_DISTRIBUTION_REPOSITORY_NWO = "dsp-testing/codeql-cli-nightlies";
  */
 export const DEFAULT_DISTRIBUTION_VERSION_RANGE: Range = new Range("2.x");
 
+interface DistributionState {
+  folderIndex: number;
+  release: Release | null;
+}
+
 export interface DistributionProvider {
   getCodeQlPathWithoutVersionCheck(): Promise<string | undefined>;
   onDidChangeDistribution?: Event<void>;
@@ -71,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,
@@ -87,6 +101,10 @@ export class DistributionManager implements DistributionProvider {
     );
   }
 
+  public async initialize(): Promise<void> {
+    await this.extensionSpecificDistributionManager.initialize();
+  }
+
   /**
    * Look up a CodeQL launcher binary.
    */
@@ -287,14 +305,53 @@ 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() {
+    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(
+            ExtensionSpecificDistributionManager._currentDistributionFolderIndexStateKey,
+            0,
+          ),
+          release: (this.extensionContext.globalState.get(
+            ExtensionSpecificDistributionManager._installedReleaseStateKey,
+          ) ?? null) as Release | null,
+        };
+
+        // This may result in a race condition, but when this happens both processes should write the same file.
+        await writeJson(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.
@@ -357,14 +414,7 @@ class ExtensionSpecificDistributionManager {
     release: Release,
     progressCallback?: ProgressCallback,
   ): Promise<void> {
-    const distributionStatePath = join(
-      this.extensionContext.globalStorageUri.fsPath,
-      ExtensionSpecificDistributionManager._distributionStateFilename,
-    );
-    if (!(await pathExists(distributionStatePath))) {
-      // This may result in a race condition, but when this happens both processes should write the same file.
-      await writeJson(distributionStatePath, {});
-    }
+    const distributionStatePath = this.getDistributionStatePath();
 
     await withDistributionUpdateLock(
       // .lock will be appended to this filename
@@ -586,23 +636,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 + 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 +
@@ -617,19 +663,50 @@ 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 writeJson(distributionStatePath, newState);
   }
 
   private static readonly _currentDistributionFolderBaseName = "distribution";

extensions/ql-vscode/src/extension.ts

@@ -362,7 +362,9 @@ export async function activate(
     distributionConfigListener,
     codeQlVersionRange,
     ctx,
+    app.logger,
   );
+  await distributionManager.initialize();
 
   registerErrorStubs([checkForUpdatesCommand], (command) => async () => {
     void showAndLogErrorMessage(

extensions/ql-vscode/test/vscode-tests/no-workspace/codeql-cli/distribution.test.ts

@@ -13,6 +13,7 @@ import type {
   showAndLogErrorMessage,
   showAndLogWarningMessage,
 } from "../../../../src/common/logging";
+import { createMockLogger } from "../../../__mocks__/loggerMock";
 
 jest.mock("os", () => {
   const original = jest.requireActual("os");
@@ -108,6 +109,7 @@ describe("Launcher path", () => {
       { customCodeQlPath: pathToCmd } as any,
       {} as any,
       {} as any,
+      createMockLogger(),
     );
 
     const result = await manager.getCodeQlPathWithoutVersionCheck();
@@ -126,6 +128,7 @@ describe("Launcher path", () => {
       { customCodeQlPath: pathToCmd } as any,
       {} as any,
       {} as any,
+      createMockLogger(),
     );
 
     const result = await manager.getCodeQlPathWithoutVersionCheck();
@@ -141,6 +144,7 @@ describe("Launcher path", () => {
       { customCodeQlPath: pathToCmd } as any,
       {} as any,
       {} as any,
+      createMockLogger(),
     );
 
     const result = await manager.getCodeQlPathWithoutVersionCheck();