Skip to content

Commit 6c794d1

Browse files
koesie10koesie10
committed
Clean up old distributions
1 parent 4413e29 commit 6c794d1

File tree

5 files changed

+293
-0
lines changed

5 files changed

+293
-0
lines changed

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import {
1919
InvocationRateLimiter,
2020
InvocationRateLimiterResultKind,
2121
} from "../common/invocation-rate-limiter";
22+
import type { NotificationLogger } from "../common/logging";
2223
import {
2324
showAndLogErrorMessage,
2425
showAndLogWarningMessage,
@@ -28,6 +29,7 @@ import { reportUnzipProgress } from "../common/vscode/unzip-progress";
2829
import type { Release } from "./distribution/release";
2930
import { ReleasesApiConsumer } from "./distribution/releases-api-consumer";
3031
import { createTimeoutSignal } from "../common/fetch-stream";
32+
import { ExtensionManagedDistributionCleaner } from "./distribution/cleaner";
3133

3234
/**
3335
* distribution.ts
@@ -64,6 +66,7 @@ export class DistributionManager implements DistributionProvider {
6466
public readonly config: DistributionConfig,
6567
private readonly versionRange: Range,
6668
extensionContext: ExtensionContext,
69+
logger: NotificationLogger,
6770
) {
6871
this._onDidChangeDistribution = config.onDidChangeConfiguration;
6972
this.extensionSpecificDistributionManager =
@@ -78,6 +81,12 @@ export class DistributionManager implements DistributionProvider {
7881
() =>
7982
this.extensionSpecificDistributionManager.checkForUpdatesToDistribution(),
8083
);
84+
this.extensionManagedDistributionCleaner =
85+
new ExtensionManagedDistributionCleaner(
86+
extensionContext,
87+
logger,
88+
this.extensionSpecificDistributionManager,
89+
);
8190
}
8291

8392
/**
@@ -255,6 +264,10 @@ export class DistributionManager implements DistributionProvider {
255264
);
256265
}
257266

267+
public startCleanup() {
268+
this.extensionManagedDistributionCleaner.start();
269+
}
270+
258271
public get onDidChangeDistribution(): Event<void> | undefined {
259272
return this._onDidChangeDistribution;
260273
}
@@ -276,6 +289,7 @@ export class DistributionManager implements DistributionProvider {
276289

277290
private readonly extensionSpecificDistributionManager: ExtensionSpecificDistributionManager;
278291
private readonly updateCheckRateLimiter: InvocationRateLimiter<DistributionUpdateCheckResult>;
292+
private readonly extensionManagedDistributionCleaner: ExtensionManagedDistributionCleaner;
279293
private readonly _onDidChangeDistribution: Event<void> | undefined;
280294
}
281295

@@ -610,6 +624,19 @@ class ExtensionSpecificDistributionManager {
610624
);
611625
}
612626

627+
public get folderIndex() {
628+
return (
629+
this.extensionContext.globalState.get(
630+
ExtensionSpecificDistributionManager._currentDistributionFolderIndexStateKey,
631+
0,
632+
) ?? 0
633+
);
634+
}
635+
636+
public get distributionFolderPrefix() {
637+
return ExtensionSpecificDistributionManager._currentDistributionFolderBaseName;
638+
}
639+
613640
private static readonly _currentDistributionFolderBaseName = "distribution";
614641
private static readonly _currentDistributionFolderIndexStateKey =
615642
"distributionFolderIndex";

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

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
import type { ExtensionContext } from "vscode";
2+
import { getDirectoryNamesInsidePath } from "../../common/files";
3+
import { sleep } from "../../common/time";
4+
import type { BaseLogger } from "../../common/logging";
5+
import { join } from "path";
6+
import { getErrorMessage } from "../../common/helpers-pure";
7+
import { pathExists, remove } from "fs-extra";
8+
9+
interface ExtensionManagedDistributionManager {
10+
folderIndex: number;
11+
distributionFolderPrefix: string;
12+
}
13+
14+
interface DistributionDirectory {
15+
directoryName: string;
16+
folderIndex: number;
17+
}
18+
19+
/**
20+
* This class is responsible for cleaning up old distributions that are no longer needed. In normal operation, this
21+
* should not be necessary as the old distribution is deleted when the distribution is updated. However, in some cases
22+
* the extension may leave behind old distribution which can result in a significant amount of space (> 100 GB) being
23+
* taking up by unused distributions.
24+
*/
25+
export class ExtensionManagedDistributionCleaner {
26+
constructor(
27+
private readonly extensionContext: ExtensionContext,
28+
private readonly logger: BaseLogger,
29+
private readonly manager: ExtensionManagedDistributionManager,
30+
) {}
31+
32+
public start() {
33+
// Intentionally starting this without waiting for it
34+
void this.cleanup().catch((e: unknown) => {
35+
void this.logger.log(
36+
`Failed to clean up old versions of the CLI: ${getErrorMessage(e)}`,
37+
);
38+
});
39+
}
40+
41+
public async cleanup() {
42+
if (!(await pathExists(this.extensionContext.globalStorageUri.fsPath))) {
43+
return;
44+
}
45+
46+
const currentFolderIndex = this.manager.folderIndex;
47+
48+
const distributionDirectoryRegex = new RegExp(
49+
`^${this.manager.distributionFolderPrefix}(\\d+)$`,
50+
);
51+
52+
const existingDirectories = await getDirectoryNamesInsidePath(
53+
this.extensionContext.globalStorageUri.fsPath,
54+
);
55+
const distributionDirectories = existingDirectories
56+
.map((dir): DistributionDirectory | null => {
57+
const match = dir.match(distributionDirectoryRegex);
58+
if (!match) {
59+
// When the folderIndex is 0, the distributionFolderPrefix is used as the directory name
60+
if (dir === this.manager.distributionFolderPrefix) {
61+
return {
62+
directoryName: dir,
63+
folderIndex: 0,
64+
};
65+
}
66+
67+
return null;
68+
}
69+
70+
return {
71+
directoryName: dir,
72+
folderIndex: parseInt(match[1]),
73+
};
74+
})
75+
.filter((dir) => dir !== null);
76+
77+
// Clean up all directories that are older than the current one
78+
const cleanableDirectories = distributionDirectories.filter(
79+
(dir) => dir.folderIndex < currentFolderIndex,
80+
);
81+
82+
if (cleanableDirectories.length === 0) {
83+
return;
84+
}
85+
86+
void this.logger.log(
87+
`Cleaning up ${cleanableDirectories.length} old versions of the CLI.`,
88+
);
89+
90+
for (const cleanableDirectory of cleanableDirectories) {
91+
// Wait 60 seconds between each cleanup to avoid overloading the system (even though the remove call should be async)
92+
await sleep(10_000);
93+
94+
const path = join(
95+
this.extensionContext.globalStorageUri.fsPath,
96+
cleanableDirectory.directoryName,
97+
);
98+
99+
// Delete this directory
100+
try {
101+
await remove(path);
102+
} catch (e) {
103+
void this.logger.log(
104+
`Tried to clean up an old version of the CLI at ${path} but encountered an error: ${getErrorMessage(e)}.`,
105+
);
106+
}
107+
}
108+
109+
void this.logger.log(
110+
`Cleaned up ${cleanableDirectories.length} old versions of the CLI.`,
111+
);
112+
}
113+
}

‎extensions/ql-vscode/src/extension.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -362,6 +362,7 @@ export async function activate(
362362
distributionConfigListener,
363363
codeQlVersionRange,
364364
ctx,
365+
app.logger,
365366
);
366367

367368
registerErrorStubs([checkForUpdatesCommand], (command) => async () => {
@@ -1123,6 +1124,8 @@ async function activateWithInstalledDistribution(
11231124
void extLogger.log("Reading query history");
11241125
await qhm.readQueryHistory();
11251126

1127+
distributionManager.startCleanup();
1128+
11261129
void extLogger.log("Successfully finished extension initialization.");
11271130

11281131
return {

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import type {
1313
showAndLogErrorMessage,
1414
showAndLogWarningMessage,
1515
} from "../../../../src/common/logging";
16+
import { createMockLogger } from "../../../__mocks__/loggerMock";
1617

1718
jest.mock("os", () => {
1819
const original = jest.requireActual("os");
@@ -108,6 +109,7 @@ describe("Launcher path", () => {
108109
{ customCodeQlPath: pathToCmd } as any,
109110
{} as any,
110111
{} as any,
112+
createMockLogger(),
111113
);
112114

113115
const result = await manager.getCodeQlPathWithoutVersionCheck();
@@ -126,6 +128,7 @@ describe("Launcher path", () => {
126128
{ customCodeQlPath: pathToCmd } as any,
127129
{} as any,
128130
{} as any,
131+
createMockLogger(),
129132
);
130133

131134
const result = await manager.getCodeQlPathWithoutVersionCheck();
@@ -141,6 +144,7 @@ describe("Launcher path", () => {
141144
{ customCodeQlPath: pathToCmd } as any,
142145
{} as any,
143146
{} as any,
147+
createMockLogger(),
144148
);
145149

146150
const result = await manager.getCodeQlPathWithoutVersionCheck();

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

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
import { ExtensionManagedDistributionCleaner } from "../../../../../src/codeql-cli/distribution/cleaner";
2+
import { mockedObject } from "../../../../mocked-object";
3+
import type { ExtensionContext } from "vscode";
4+
import { Uri } from "vscode";
5+
import { createMockLogger } from "../../../../__mocks__/loggerMock";
6+
import type { DirectoryResult } from "tmp-promise";
7+
import { dir } from "tmp-promise";
8+
import { outputFile, pathExists } from "fs-extra";
9+
import { join } from "path";
10+
import { codeQlLauncherName } from "../../../../../src/common/distribution";
11+
import { getDirectoryNamesInsidePath } from "../../../../../src/common/files";
12+
13+
describe("ExtensionManagedDistributionCleaner", () => {
14+
let globalStorageDirectory: DirectoryResult;
15+
16+
let manager: ExtensionManagedDistributionCleaner;
17+
18+
beforeEach(async () => {
19+
globalStorageDirectory = await dir({
20+
unsafeCleanup: true,
21+
});
22+
23+
manager = new ExtensionManagedDistributionCleaner(
24+
mockedObject<ExtensionContext>({
25+
globalStorageUri: Uri.file(globalStorageDirectory.path),
26+
}),
27+
createMockLogger(),
28+
{
29+
folderIndex: 768,
30+
distributionFolderPrefix: "distribution",
31+
},
32+
);
33+
34+
// Mock setTimeout to call the callback immediately
35+
jest.spyOn(global, "setTimeout").mockImplementation((callback) => {
36+
callback();
37+
return 0 as unknown as ReturnType<typeof setTimeout>;
38+
});
39+
});
40+
41+
afterEach(async () => {
42+
await globalStorageDirectory.cleanup();
43+
});
44+
45+
it("does nothing when no distributions exist", async () => {
46+
await manager.cleanup();
47+
});
48+
49+
it("does nothing when only the current distribution exists", async () => {
50+
await outputFile(
51+
join(
52+
globalStorageDirectory.path,
53+
"distribution768",
54+
"codeql",
55+
"bin",
56+
codeQlLauncherName(),
57+
),
58+
"launcher!",
59+
);
60+
61+
await manager.cleanup();
62+
63+
expect(
64+
await pathExists(
65+
join(
66+
globalStorageDirectory.path,
67+
"distribution768",
68+
"codeql",
69+
"bin",
70+
codeQlLauncherName(),
71+
),
72+
),
73+
).toBe(true);
74+
});
75+
76+
it("removes old distributions", async () => {
77+
await outputFile(
78+
join(
79+
globalStorageDirectory.path,
80+
"distribution",
81+
"codeql",
82+
"bin",
83+
codeQlLauncherName(),
84+
),
85+
"launcher!",
86+
);
87+
await outputFile(
88+
join(
89+
globalStorageDirectory.path,
90+
"distribution12",
91+
"codeql",
92+
"bin",
93+
codeQlLauncherName(),
94+
),
95+
"launcher!",
96+
);
97+
await outputFile(
98+
join(
99+
globalStorageDirectory.path,
100+
"distribution244",
101+
"codeql",
102+
"bin",
103+
codeQlLauncherName(),
104+
),
105+
"launcher!",
106+
);
107+
await outputFile(
108+
join(
109+
globalStorageDirectory.path,
110+
"distribution637",
111+
"codeql",
112+
"bin",
113+
codeQlLauncherName(),
114+
),
115+
"launcher!",
116+
);
117+
await outputFile(
118+
join(
119+
globalStorageDirectory.path,
120+
"distribution768",
121+
"codeql",
122+
"bin",
123+
codeQlLauncherName(),
124+
),
125+
"launcher!",
126+
);
127+
await outputFile(
128+
join(
129+
globalStorageDirectory.path,
130+
"distribution890",
131+
"codeql",
132+
"bin",
133+
codeQlLauncherName(),
134+
),
135+
"launcher!",
136+
);
137+
138+
const promise = manager.cleanup();
139+
140+
await promise;
141+
142+
expect(
143+
(await getDirectoryNamesInsidePath(globalStorageDirectory.path)).sort(),
144+
).toEqual(["distribution768", "distribution890"]);
145+
});
146+
});

0 commit comments

Comments
 (0)