accurate heap leak checks

test/automation/src/playwrightDriver.ts

@@ -94,6 +94,14 @@ export class PlaywrightDriver {
 		this._cdpSession = await this.page.context().newCDPSession(this.page);
 	}
 
+	async collectGarbage() {
+		if (!this._cdpSession) {
+			throw new Error('CDP not started');
+		}
+
+		await this._cdpSession.send('HeapProfiler.collectGarbage');
+	}
+
 	async evaluate(options: Protocol.Runtime.evaluateParameters): Promise<Protocol.Runtime.evaluateReturnValue> {
 		if (!this._cdpSession) {
 			throw new Error('CDP not started');
@@ -118,6 +126,20 @@ export class PlaywrightDriver {
 		return await this._cdpSession.send('Runtime.callFunctionOn', parameters);
 	}
 
+	async takeHeapSnapshot(): Promise<string> {
+		if (!this._cdpSession) {
+			throw new Error('CDP not started');
+		}
+
+		let snapshot = '';
+		this._cdpSession.addListener('HeapProfiler.addHeapSnapshotChunk', ({ chunk }) => {
+			snapshot += chunk;
+		});
+
+		await this._cdpSession.send('HeapProfiler.takeHeapSnapshot');
+		return snapshot;
+	}
+
 	async getProperties(parameters: Protocol.Runtime.getPropertiesParameters): Promise<Protocol.Runtime.getPropertiesReturnValue> {
 		if (!this._cdpSession) {
 			throw new Error('CDP not started');

test/automation/src/profiler.ts

@@ -4,6 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import type { Protocol } from 'playwright-core/types/protocol';
+const { decode_bytes } = require('@vscode/v8-heap-parser');
 import { Code } from './code';
 import { PlaywrightDriver } from './playwrightDriver';
 
@@ -45,6 +46,26 @@ export class Profiler {
 			throw new Error(leaks.join('\n'));
 		}
 	}
+
+	async checkHeapLeaks(classNames: string | string[], fn: () => Promise<void>): Promise<void> {
+		await this.code.driver.startCDP();
+		await fn();
+
+		const heapSnapshotAfter = await this.code.driver.takeHeapSnapshot();
+		const buff = Buffer.from(heapSnapshotAfter);
+		const graph = await decode_bytes(buff);
+		const counts: number[] = Array.from(graph.get_class_counts(classNames));
+		const leaks: string[] = [];
+		for (let i = 0; i < classNames.length; i++) {
+			if (counts[i] > 0) {
+				leaks.push(`Leaked ${counts[i]} ${classNames[i]}`);
+			}
+		}
+
+		if (leaks.length > 0) {
+			throw new Error(leaks.join('\n'));
+		}
+	}
 }
 
 function generateUuid() {
@@ -106,6 +127,7 @@ function generateUuid() {
  *--------------------------------------------------------------------------------------------*/
 
 const getInstances = async (driver: PlaywrightDriver): Promise<Array<{ name: string; count: number }>> => {
+	await driver.collectGarbage();
 	const objectGroup = `og:${generateUuid()}`;
 	const prototypeDescriptor = await driver.evaluate({
 		expression: 'Object.prototype',

test/smoke/src/areas/notebook/notebook.test.ts

@@ -25,7 +25,16 @@ export function setup(logger: Logger) {
 			cp.execSync('git reset --hard HEAD --quiet', { cwd: app.workspacePathOrFolder });
 		});
 
-		it('leaks check', async function () {
+		it('check heap leaks', async function () {
+			const app = this.app as Application;
+			await app.profiler.checkHeapLeaks(['NotebookTextModel', 'NotebookCellTextModel', 'NotebookEventDispatcher'], async () => {
+				await app.workbench.notebook.openNotebook();
+				await app.workbench.quickaccess.runCommand('workbench.action.files.save');
+				await app.workbench.quickaccess.runCommand('workbench.action.closeActiveEditor');
+			});
+		});
+
+		it('check object leaks', async function () {
 			const app = this.app as Application;
 			await app.profiler.checkLeaks(['NotebookTextModel'], async () => {
 				await app.workbench.notebook.openNotebook();