Clean up old distributions

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

@@ -19,6 +19,7 @@ import {
   InvocationRateLimiter,
   InvocationRateLimiterResultKind,
 } from "../common/invocation-rate-limiter";
+import type { NotificationLogger } from "../common/logging";
 import {
   showAndLogErrorMessage,
   showAndLogWarningMessage,
@@ -28,6 +29,7 @@ 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 { ExtensionManagedDistributionCleaner } from "./distribution/cleaner";
 
 /**
  * distribution.ts
@@ -64,6 +66,7 @@ export class DistributionManager implements DistributionProvider {
     public readonly config: DistributionConfig,
     private readonly versionRange: Range,
     extensionContext: ExtensionContext,
+    logger: NotificationLogger,
   ) {
     this._onDidChangeDistribution = config.onDidChangeConfiguration;
     this.extensionSpecificDistributionManager =
@@ -78,6 +81,12 @@ export class DistributionManager implements DistributionProvider {
       () =>
         this.extensionSpecificDistributionManager.checkForUpdatesToDistribution(),
     );
+    this.extensionManagedDistributionCleaner =
+      new ExtensionManagedDistributionCleaner(
+        extensionContext,
+        logger,
+        this.extensionSpecificDistributionManager,
+      );
   }
 
   /**
@@ -255,6 +264,10 @@ export class DistributionManager implements DistributionProvider {
     );
   }
 
+  public startCleanup() {
+    this.extensionManagedDistributionCleaner.start();
+  }
+
   public get onDidChangeDistribution(): Event<void> | undefined {
     return this._onDidChangeDistribution;
   }
@@ -276,6 +289,7 @@ export class DistributionManager implements DistributionProvider {
 
   private readonly extensionSpecificDistributionManager: ExtensionSpecificDistributionManager;
   private readonly updateCheckRateLimiter: InvocationRateLimiter<DistributionUpdateCheckResult>;
+  private readonly extensionManagedDistributionCleaner: ExtensionManagedDistributionCleaner;
   private readonly _onDidChangeDistribution: Event<void> | undefined;
 }
 
@@ -610,6 +624,19 @@ class ExtensionSpecificDistributionManager {
     );
   }
 
+  public get folderIndex() {
+    return (
+      this.extensionContext.globalState.get(
+        ExtensionSpecificDistributionManager._currentDistributionFolderIndexStateKey,
+        0,
+      ) ?? 0
+    );
+  }
+
+  public get distributionFolderPrefix() {
+    return ExtensionSpecificDistributionManager._currentDistributionFolderBaseName;
+  }
+
   private static readonly _currentDistributionFolderBaseName = "distribution";
   private static readonly _currentDistributionFolderIndexStateKey =
     "distributionFolderIndex";

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

@@ -0,0 +1,113 @@
+import type { ExtensionContext } from "vscode";
+import { getDirectoryNamesInsidePath } from "../../common/files";
+import { sleep } from "../../common/time";
+import type { BaseLogger } from "../../common/logging";
+import { join } from "path";
+import { getErrorMessage } from "../../common/helpers-pure";
+import { pathExists, remove } from "fs-extra";
+
+interface ExtensionManagedDistributionManager {
+  folderIndex: number;
+  distributionFolderPrefix: string;
+}
+
+interface DistributionDirectory {
+  directoryName: string;
+  folderIndex: number;
+}
+
+/**
+ * This class is responsible for cleaning up old distributions that are no longer needed. In normal operation, this
+ * should not be necessary as the old distribution is deleted when the distribution is updated. However, in some cases
+ * the extension may leave behind old distribution which can result in a significant amount of space (> 100 GB) being
+ * taking up by unused distributions.
+ */
+export class ExtensionManagedDistributionCleaner {
+  constructor(
+    private readonly extensionContext: ExtensionContext,
+    private readonly logger: BaseLogger,
+    private readonly manager: ExtensionManagedDistributionManager,
+  ) {}
+
+  public start() {
+    // Intentionally starting this without waiting for it
+    void this.cleanup().catch((e: unknown) => {
+      void this.logger.log(
+        `Failed to clean up old versions of the CLI: ${getErrorMessage(e)}`,
+      );
+    });
+  }
+
+  public async cleanup() {
+    if (!(await pathExists(this.extensionContext.globalStorageUri.fsPath))) {
+      return;
+    }
+
+    const currentFolderIndex = this.manager.folderIndex;
+
+    const distributionDirectoryRegex = new RegExp(
+      `^${this.manager.distributionFolderPrefix}(\\d+)$`,
+    );
+
+    const existingDirectories = await getDirectoryNamesInsidePath(
+      this.extensionContext.globalStorageUri.fsPath,
+    );
+    const distributionDirectories = existingDirectories
+      .map((dir): DistributionDirectory | null => {
+        const match = dir.match(distributionDirectoryRegex);
+        if (!match) {
+          // When the folderIndex is 0, the distributionFolderPrefix is used as the directory name
+          if (dir === this.manager.distributionFolderPrefix) {
+            return {
+              directoryName: dir,
+              folderIndex: 0,
+            };
+          }
+
+          return null;
+        }
+
+        return {
+          directoryName: dir,
+          folderIndex: parseInt(match[1]),
+        };
+      })
+      .filter((dir) => dir !== null);
+
+    // Clean up all directories that are older than the current one
+    const cleanableDirectories = distributionDirectories.filter(
+      (dir) => dir.folderIndex < currentFolderIndex,
+    );
+
+    if (cleanableDirectories.length === 0) {
+      return;
+    }
+
+    void this.logger.log(
+      `Cleaning up ${cleanableDirectories.length} old versions of the CLI.`,
+    );
+
+    for (const cleanableDirectory of cleanableDirectories) {
+      // Wait 60 seconds between each cleanup to avoid overloading the system (even though the remove call should be async)
+      await sleep(10_000);
+
+      const path = join(
+        this.extensionContext.globalStorageUri.fsPath,
+        cleanableDirectory.directoryName,
+      );
+
+      // Delete this directory
+      try {
+        await remove(path);
+      } catch (e) {
+        void this.logger.log(
+          `Tried to clean up an old version of the CLI at ${path} but encountered an error: ${getErrorMessage(e)}.`,
+        );
+      }
+    }
+
+    void this.logger.log(
+      `Cleaned up ${cleanableDirectories.length} old versions of the CLI.`,
+    );
+  }
+}

extensions/ql-vscode/src/extension.ts

@@ -362,6 +362,7 @@ export async function activate(
     distributionConfigListener,
     codeQlVersionRange,
     ctx,
+    app.logger,
   );
 
   registerErrorStubs([checkForUpdatesCommand], (command) => async () => {
@@ -1123,6 +1124,8 @@ async function activateWithInstalledDistribution(
   void extLogger.log("Reading query history");
   await qhm.readQueryHistory();
 
+  distributionManager.startCleanup();
+
   void extLogger.log("Successfully finished extension initialization.");
 
   return {

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();

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

@@ -0,0 +1,146 @@
+import { ExtensionManagedDistributionCleaner } from "../../../../../src/codeql-cli/distribution/cleaner";
+import { mockedObject } from "../../../../mocked-object";
+import type { ExtensionContext } from "vscode";
+import { Uri } from "vscode";
+import { createMockLogger } from "../../../../__mocks__/loggerMock";
+import type { DirectoryResult } from "tmp-promise";
+import { dir } from "tmp-promise";
+import { outputFile, pathExists } from "fs-extra";
+import { join } from "path";
+import { codeQlLauncherName } from "../../../../../src/common/distribution";
+import { getDirectoryNamesInsidePath } from "../../../../../src/common/files";
+
+describe("ExtensionManagedDistributionCleaner", () => {
+  let globalStorageDirectory: DirectoryResult;
+
+  let manager: ExtensionManagedDistributionCleaner;
+
+  beforeEach(async () => {
+    globalStorageDirectory = await dir({
+      unsafeCleanup: true,
+    });
+
+    manager = new ExtensionManagedDistributionCleaner(
+      mockedObject<ExtensionContext>({
+        globalStorageUri: Uri.file(globalStorageDirectory.path),
+      }),
+      createMockLogger(),
+      {
+        folderIndex: 768,
+        distributionFolderPrefix: "distribution",
+      },
+    );
+
+    // Mock setTimeout to call the callback immediately
+    jest.spyOn(global, "setTimeout").mockImplementation((callback) => {
+      callback();
+      return 0 as unknown as ReturnType<typeof setTimeout>;
+    });
+  });
+
+  afterEach(async () => {
+    await globalStorageDirectory.cleanup();
+  });
+
+  it("does nothing when no distributions exist", async () => {
+    await manager.cleanup();
+  });
+
+  it("does nothing when only the current distribution exists", async () => {
+    await outputFile(
+      join(
+        globalStorageDirectory.path,
+        "distribution768",
+        "codeql",
+        "bin",
+        codeQlLauncherName(),
+      ),
+      "launcher!",
+    );
+
+    await manager.cleanup();
+
+    expect(
+      await pathExists(
+        join(
+          globalStorageDirectory.path,
+          "distribution768",
+          "codeql",
+          "bin",
+          codeQlLauncherName(),
+        ),
+      ),
+    ).toBe(true);
+  });
+
+  it("removes old distributions", async () => {
+    await outputFile(
+      join(
+        globalStorageDirectory.path,
+        "distribution",
+        "codeql",
+        "bin",
+        codeQlLauncherName(),
+      ),
+      "launcher!",
+    );
+    await outputFile(
+      join(
+        globalStorageDirectory.path,
+        "distribution12",
+        "codeql",
+        "bin",
+        codeQlLauncherName(),
+      ),
+      "launcher!",
+    );
+    await outputFile(
+      join(
+        globalStorageDirectory.path,
+        "distribution244",
+        "codeql",
+        "bin",
+        codeQlLauncherName(),
+      ),
+      "launcher!",
+    );
+    await outputFile(
+      join(
+        globalStorageDirectory.path,
+        "distribution637",
+        "codeql",
+        "bin",
+        codeQlLauncherName(),
+      ),
+      "launcher!",
+    );
+    await outputFile(
+      join(
+        globalStorageDirectory.path,
+        "distribution768",
+        "codeql",
+        "bin",
+        codeQlLauncherName(),
+      ),
+      "launcher!",
+    );
+    await outputFile(
+      join(
+        globalStorageDirectory.path,
+        "distribution890",
+        "codeql",
+        "bin",
+        codeQlLauncherName(),
+      ),
+      "launcher!",
+    );
+
+    const promise = manager.cleanup();
+
+    await promise;
+
+    expect(
+      (await getDirectoryNamesInsidePath(globalStorageDirectory.path)).sort(),
+    ).toEqual(["distribution768", "distribution890"]);
+  });
+});