eng: add assertHeap method for memory assertions (#198334)

This adds an `assertHeap` function that can be used in tests. It
takes a heap snapshot, and asserts the state of classes in memory. This
works in Node and the Electron sandbox, but is a no-op in the browser.
Snapshots are process asynchronously and will report failures at the end
of the suite.

This method should be used sparingly (e.g. once at the end of a suite to
ensure nothing leaked before), as gathering a heap snapshot is fairly
slow, at least until V8 11.5.130 (https://v8.dev/blog/speeding-up-v8-heap-snapshots).

When used, the function will ensure the test has a minimum timeout
duration of 20s to avoid immediate failures.

It takes options containing a mapping of class names, and assertion functions
to run on the number of retained instances of that class. For example:

```ts
assertSnapshot({
	classes: {
		ShouldNeverLeak: count => assert.strictEqual(count, 0),
		SomeSingleton: count => assert(count <= 1),
	}
});
```

Closes #191920

.eslintrc.json

@@ -128,7 +128,8 @@
 						{
 							"message": "Expression must be awaited",
 							"functions": [
-								"assertSnapshot"
+								"assertSnapshot",
+								"assertHeap"
 							]
 						}
 					]

package.json

@@ -140,6 +140,7 @@
     "@vscode/test-cli": "^0.0.3",
     "@vscode/test-electron": "^2.3.5",
     "@vscode/test-web": "^0.0.42",
+    "@vscode/v8-heap-parser": "^0.1.0",
     "@vscode/vscode-perf": "^0.0.14",
     "ansi-colors": "^3.2.3",
     "asar": "^3.0.3",

src/vs/base/test/common/assertHeap.ts

@@ -0,0 +1,84 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+
+declare const __analyzeSnapshotInTests: (currentTest: string, classes: readonly string[]) => Promise<({ done: Promise<number[]>; file: string })>;
+
+let currentTest: Mocha.Test | undefined;
+
+const snapshotsToAssert: ({ counts: Promise<number[]>; file: string; test: string; opts: ISnapshotAssertOptions })[] = [];
+
+setup(function () {
+	currentTest = this.currentTest;
+});
+
+suiteTeardown(async () => {
+	await Promise.all(snapshotsToAssert.map(async snap => {
+		const counts = await snap.counts;
+
+		const asserts = Object.entries(snap.opts.classes);
+		if (asserts.length !== counts.length) {
+			throw new Error(`expected class counts to equal assertions length for ${snap.test}`);
+		}
+
+		for (const [i, [name, doAssert]] of asserts.entries()) {
+			try {
+				doAssert(counts[i]);
+			} catch (e) {
+				throw new Error(`Unexpected number of ${name} instances (${counts[i]}) after "${snap.test}":\n\n${e.message}\n\nSnapshot saved at: ${snap.file}`);
+			}
+		}
+	}));
+
+	snapshotsToAssert.length = 0;
+});
+
+export interface ISnapshotAssertOptions {
+	classes: Record<string, (count: number) => void>;
+}
+
+const snapshotMinTime = 20_000;
+
+/**
+ * Takes a heap snapshot, and asserts the state of classes in memory. This
+ * works in Node and the Electron sandbox, but is a no-op in the browser.
+ * Snapshots are process asynchronously and will report failures at the end of
+ * the suite.
+ *
+ * This method should be used sparingly (e.g. once at the end of a suite to
+ * ensure nothing leaked before), as gathering a heap snapshot is fairly
+ * slow, at least until V8 11.5.130 (https://v8.dev/blog/speeding-up-v8-heap-snapshots).
+ *
+ * Takes options containing a mapping of class names, and assertion functions
+ * to run on the number of retained instances of that class. For example:
+ *
+ * ```ts
+ * assertSnapshot({
+ *	classes: {
+ *		ShouldNeverLeak: count => assert.strictEqual(count, 0),
+ *		SomeSingleton: count => assert(count <= 1),
+ *	}
+ *});
+ * ```
+ */
+export async function assertHeap(opts: ISnapshotAssertOptions) {
+	if (!currentTest) {
+		throw new Error('assertSnapshot can only be used when a test is running');
+	}
+
+	// snapshotting can take a moment, ensure the test timeout is decently long
+	// so it doesn't immediately fail.
+	if (currentTest.timeout() < snapshotMinTime) {
+		currentTest.timeout(snapshotMinTime);
+	}
+
+	if (typeof __analyzeSnapshotInTests === 'undefined') {
+		return; // running in browser, no-op
+	}
+
+	const { done, file } = await __analyzeSnapshotInTests(currentTest.fullTitle(), Object.keys(opts.classes));
+	snapshotsToAssert.push({ counts: done, file, test: currentTest.fullTitle(), opts });
+}
+

test/unit/analyzeSnapshot.js

@@ -0,0 +1,85 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+//@ts-check
+
+// note: we use a fork here since we can't make a worker from the renderer process
+
+const { fork } = require('child_process');
+const workerData = process.env.SNAPSHOT_WORKER_DATA;
+const fs = require('fs');
+const { pathToFileURL } = require('url');
+
+if (!workerData) {
+	const { join } = require('path');
+	const { tmpdir } = require('os');
+
+	exports.takeSnapshotAndCountClasses = async (/** @type string */currentTest, /** @type string[] */ classes) => {
+		const cleanTitle = currentTest.replace(/[^\w]+/g, '-');
+		const file = join(tmpdir(), `vscode-test-snap-${cleanTitle}.heapsnapshot`);
+
+		if (typeof process.takeHeapSnapshot !== 'function') {
+			// node.js:
+			const inspector = require('inspector');
+			const session = new inspector.Session();
+			session.connect();
+
+			const fd = fs.openSync(file, 'w');
+			await new Promise((resolve, reject) => {
+				session.on('HeapProfiler.addHeapSnapshotChunk', (m) => {
+					fs.writeSync(fd, m.params.chunk);
+				});
+
+				session.post('HeapProfiler.takeHeapSnapshot', null, (err) => {
+					session.disconnect();
+					fs.closeSync(fd);
+					if (err) {
+						reject(err);
+					} else {
+						resolve();
+					}
+				});
+			});
+		} else {
+			// electron exposes this nice method for us:
+			process.takeHeapSnapshot(file);
+		}
+
+		const worker = fork(__filename, {
+			env: {
+				...process.env,
+				SNAPSHOT_WORKER_DATA: JSON.stringify({
+					path: file,
+					classes,
+				})
+			}
+		});
+
+		const promise = new Promise((resolve, reject) => {
+			worker.on('message', (/** @type any */msg) => {
+				if ('err' in msg) {
+					reject(new Error(msg.err));
+				} else {
+					resolve(msg.counts);
+				}
+				worker.kill();
+			});
+		});
+
+		return { done: promise, file: pathToFileURL(file) };
+	};
+} else {
+	const { path, classes } = JSON.parse(workerData);
+	const { decode_bytes } = require('@vscode/v8-heap-parser');
+
+	fs.promises.readFile(path)
+		.then(buf => decode_bytes(buf))
+		.then(graph => graph.get_class_counts(classes))
+		.then(
+			counts => process.send({ counts: Array.from(counts) }),
+			err => process.send({ err: String(err.stack || err) })
+		);
+
+}

test/unit/electron/renderer.js

@@ -67,6 +67,7 @@ const glob = require('glob');
 const util = require('util');
 const bootstrap = require('../../../src/bootstrap');
 const coverage = require('../coverage');
+const { takeSnapshotAndCountClasses } = require('../analyzeSnapshot');
 
 // Disabled custom inspect. See #38847
 if (util.inspect && util.inspect['defaultOptions']) {
@@ -82,6 +83,7 @@ globalThis._VSCODE_PACKAGE_JSON = (require.__$__nodeRequire ?? require)('../../.
 
 // Test file operations that are common across platforms. Used for test infra, namely snapshot tests
 Object.assign(globalThis, {
+	__analyzeSnapshotInTests: takeSnapshotAndCountClasses,
 	__readFileInTests: path => fs.promises.readFile(path, 'utf-8'),
 	__writeFileInTests: (path, contents) => fs.promises.writeFile(path, contents),
 	__readDirInTests: path => fs.promises.readdir(path),

test/unit/node/index.js

@@ -16,6 +16,7 @@ const glob = require('glob');
 const minimatch = require('minimatch');
 const coverage = require('../coverage');
 const minimist = require('minimist');
+const { takeSnapshotAndCountClasses } = require('../analyzeSnapshot');
 
 /**
  * @type {{ build: boolean; run: string; runGlob: string; coverage: boolean; help: boolean; }}
@@ -83,6 +84,7 @@ function main() {
 
 	// Test file operations that are common across platforms. Used for test infra, namely snapshot tests
 	Object.assign(globalThis, {
+		__analyzeSnapshotInTests: takeSnapshotAndCountClasses,
 		__readFileInTests: (/** @type {string} */ path) => fs.promises.readFile(path, 'utf-8'),
 		__writeFileInTests: (/** @type {string} */ path, /** @type {BufferEncoding} */ contents) => fs.promises.writeFile(path, contents),
 		__readDirInTests: (/** @type {string} */ path) => fs.promises.readdir(path),

yarn.lock

@@ -1488,6 +1488,11 @@
     tar-fs "^2.1.1"
     vscode-uri "^3.0.7"
 
+"@vscode/v8-heap-parser@^0.1.0":
+  version "0.1.0"
+  resolved "https://registry.yarnpkg.com/@vscode/v8-heap-parser/-/v8-heap-parser-0.1.0.tgz#f3fe61ce954cc3dd78ed42e09f865450685e351f"
+  integrity sha512-3EvQak7EIOLyIGz+IP9qSwRmP08ZRWgTeoRgAXPVkkDXZ8riqJ7LDtkgx++uHBiJ3MUaSdlUYPZcLFFw7E6zGg==
+
 "@vscode/vscode-languagedetection@1.0.21":
   version "1.0.21"
   resolved "https://registry.yarnpkg.com/@vscode/vscode-languagedetection/-/vscode-languagedetection-1.0.21.tgz#89b48f293f6aa3341bb888c1118d16ff13b032d3"