Merge pull request #3763 from github/koesie10/cleanup-distributions

Clean up old distributions

extensions/ql-vscode/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 - 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)
+- Clean up old extension-managed CodeQL CLI distributions. [#3763](https://github.com/github/vscode-codeql/pull/3763)
 
 ## 1.16.0 - 10 October 2024
 

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

@@ -42,6 +42,7 @@ import { asError, getErrorMessage } from "../common/helpers-pure";
 import { isIOError } from "../common/files";
 import { telemetryListener } from "../common/vscode/telemetry";
 import { redactableError } from "../common/errors";
+import { ExtensionManagedDistributionCleaner } from "./distribution/cleaner";
 
 /**
  * distribution.ts
@@ -99,6 +100,12 @@ export class DistributionManager implements DistributionProvider {
       () =>
         this.extensionSpecificDistributionManager.checkForUpdatesToDistribution(),
     );
+    this.extensionManagedDistributionCleaner =
+      new ExtensionManagedDistributionCleaner(
+        extensionContext,
+        logger,
+        this.extensionSpecificDistributionManager,
+      );
   }
 
   public async initialize(): Promise<void> {
@@ -280,6 +287,10 @@ export class DistributionManager implements DistributionProvider {
     );
   }
 
+  public startCleanup() {
+    this.extensionManagedDistributionCleaner.start();
+  }
+
   public get onDidChangeDistribution(): Event<void> | undefined {
     return this._onDidChangeDistribution;
   }
@@ -301,6 +312,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;
 }
 
@@ -718,6 +730,16 @@ class ExtensionSpecificDistributionManager {
     await outputJson(distributionStatePath, newState);
   }
 
+  public get folderIndex() {
+    const distributionState = this.getDistributionState();
+
+    return distributionState.folderIndex;
+  }
+
+  public get distributionFolderPrefix() {
+    return ExtensionSpecificDistributionManager._currentDistributionFolderBaseName;
+  }
+
   private static readonly _currentDistributionFolderBaseName = "distribution";
   private static readonly _codeQlExtractedFolderName = "codeql";
   private static readonly _distributionStateFilename = "distribution.json";

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

@@ -0,0 +1,127 @@
+import type { ExtensionContext } from "vscode";
+import { getDirectoryNamesInsidePath, isIOError } 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;
+    }
+
+    // Shuffle the array so that multiple VS Code processes don't all try to clean up the same directory at the same time
+    for (let i = cleanableDirectories.length - 1; i > 0; i--) {
+      const j = Math.floor(Math.random() * (i + 1));
+      [cleanableDirectories[i], cleanableDirectories[j]] = [
+        cleanableDirectories[j],
+        cleanableDirectories[i],
+      ];
+    }
+
+    void this.logger.log(
+      `Cleaning up ${cleanableDirectories.length} old versions of the CLI.`,
+    );
+
+    for (const cleanableDirectory of cleanableDirectories) {
+      // Wait 10 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) {
+        if (isIOError(e) && e.code === "ENOENT") {
+          // If the directory doesn't exist, that's fine
+          continue;
+        }
+
+        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

@@ -1125,6 +1125,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/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"]);
+  });
+});