diff --git a/extensions/ql-vscode/src/codeql-cli/distribution.ts b/extensions/ql-vscode/src/codeql-cli/distribution.ts index 487906429..2efc6a097 100644 --- a/extensions/ql-vscode/src/codeql-cli/distribution.ts +++ b/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 | undefined { return this._onDidChangeDistribution; } @@ -276,6 +289,7 @@ export class DistributionManager implements DistributionProvider { private readonly extensionSpecificDistributionManager: ExtensionSpecificDistributionManager; private readonly updateCheckRateLimiter: InvocationRateLimiter; + private readonly extensionManagedDistributionCleaner: ExtensionManagedDistributionCleaner; private readonly _onDidChangeDistribution: Event | 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"; diff --git a/extensions/ql-vscode/src/codeql-cli/distribution/cleaner.ts b/extensions/ql-vscode/src/codeql-cli/distribution/cleaner.ts new file mode 100644 index 000000000..99a95c294 --- /dev/null +++ b/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.`, + ); + } +} diff --git a/extensions/ql-vscode/src/extension.ts b/extensions/ql-vscode/src/extension.ts index 4f7c8cd99..f9e493beb 100644 --- a/extensions/ql-vscode/src/extension.ts +++ b/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 { diff --git a/extensions/ql-vscode/test/vscode-tests/no-workspace/codeql-cli/distribution.test.ts b/extensions/ql-vscode/test/vscode-tests/no-workspace/codeql-cli/distribution.test.ts index 189e1b12b..dc6f6d724 100644 --- a/extensions/ql-vscode/test/vscode-tests/no-workspace/codeql-cli/distribution.test.ts +++ b/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(); diff --git a/extensions/ql-vscode/test/vscode-tests/no-workspace/codeql-cli/distribution/cleaner.test.ts b/extensions/ql-vscode/test/vscode-tests/no-workspace/codeql-cli/distribution/cleaner.test.ts new file mode 100644 index 000000000..20a7a71c3 --- /dev/null +++ b/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({ + 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; + }); + }); + + 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"]); + }); +});