extension/src: add rr test code lens for record and replay debugging

docs/settings.md

@@ -213,11 +213,13 @@ Default:
 Feature level setting to enable/disable code lens for references and run/debug tests
 | Properties | Description |
 | --- | --- |
+| `rrtest` | If true, enables 'rr test' code lens for recording and replaying tests with Mozilla rr. Requires Linux and rr to be installed. <br/> Default: `false` |
 | `runtest` | If true, enables code lens for running and debugging tests <br/> Default: `true` |
 
 Default:
 ```
 {
+	"rrtest" :	false,
 	"runtest" :	true,
 }
 ```

extension/package.json

@@ -1638,11 +1638,17 @@
               "type": "boolean",
               "default": true,
               "description": "If true, enables code lens for running and debugging tests"
+            },
+            "rrtest": {
+              "type": "boolean",
+              "default": false,
+              "description": "If true, enables 'rr test' code lens for recording and replaying tests with Mozilla rr. Requires Linux and rr to be installed."
             }
           },
           "additionalProperties": false,
           "default": {
-            "runtest": true
+            "runtest": true,
+            "rrtest": false
           },
           "description": "Feature level setting to enable/disable code lens for references and run/debug tests",
           "scope": "resource"

extension/src/goMain.ts

@@ -174,6 +174,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<ExtensionA
 	registerCommand('go.subtest.cursor', commands.subTestAtCursor('test'));
 	registerCommand('go.debug.cursor', commands.testAtCursor('debug'));
 	registerCommand('go.debug.subtest.cursor', commands.subTestAtCursor('debug'));
+	registerCommand('go.rr.cursor', commands.rrAtCursor);
 	registerCommand('go.benchmark.cursor', commands.testAtCursor('benchmark'));
 	registerCommand('go.test.package', commands.testCurrentPackage(false));
 	registerCommand('go.benchmark.package', commands.testCurrentPackage(true));

extension/src/goRunTestCodelens.ts

@@ -126,6 +126,15 @@ export class GoRunTestCodeLensProvider extends GoBaseCodeLensProvider {
 						arguments: [{ functionName }]
 					})
 				);
+				if (getGoConfig(document.uri).get<{ [key: string]: boolean }>('enableCodeLens')?.rrtest) {
+					codelens.push(
+						new CodeLens(f.range, {
+							title: 'rr test',
+							command: 'go.rr.cursor',
+							arguments: [{ functionName }]
+						})
+					);
+				}
 
 				for (let i = f.range.start.line; i < f.range.end.line; i++) {
 					const line = document.lineAt(i);

extension/src/goTest.ts

@@ -5,11 +5,15 @@
  *--------------------------------------------------------*/
 'use strict';
 
+import cp = require('child_process');
+import os = require('os');
 import path = require('path');
+import util = require('util');
 import vscode = require('vscode');
 import { CommandFactory } from './commands';
 import { getGoConfig } from './config';
 import { GoExtensionContext } from './context';
+import { toolExecutionEnvironment } from './goEnv';
 import { isModSupported } from './goModules';
 import { escapeSubTestName } from './subTestUtils';
 import {
@@ -25,6 +29,7 @@ import {
 	SuiteToTestMap,
 	getTestFunctions
 } from './testUtils';
+import { getBinPath } from './util';
 
 // lastTestConfig holds a reference to the last executed TestConfig which allows
 // the last test to be easily re-executed.
@@ -346,6 +351,115 @@ export async function debugTestAtCursor(
 	return await vscode.debug.startDebugging(workspaceFolder, debugConfig);
 }
 
+/**
+ * Records and replays the test at cursor using Mozilla rr via `dlv replay`.
+ * Only supported on Linux with rr installed.
+ */
+export const rrAtCursor: CommandFactory = (_, goCtx) => async (args?: { functionName: string }) => {
+	const editor = vscode.window.activeTextEditor;
+	if (!editor) {
+		vscode.window.showInformationMessage('No editor is active.');
+		return;
+	}
+	if (!editor.document.fileName.endsWith('_test.go')) {
+		vscode.window.showInformationMessage('No tests found. Current file is not a test file.');
+		return;
+	}
+	if (os.platform() !== 'linux') {
+		vscode.window.showErrorMessage("'rr test' is only supported on Linux.");
+		return;
+	}
+
+	const rrPath = getBinPath('rr');
+	if (!rrPath || rrPath === 'rr') {
+		// getBinPath returns the input unchanged when not found
+		try {
+			cp.execFileSync('rr', ['--version'], { stdio: 'ignore' });
+		} catch {
+			vscode.window.showErrorMessage("'rr' not found on PATH. Install Mozilla rr to use this feature.");
+			return;
+		}
+	}
+
+	const { testFunctions } = await getTestFunctionsAndTestSuite(false, goCtx, editor.document);
+	const testFunctionName =
+		args && args.functionName
+			? args.functionName
+			: testFunctions?.filter((func) => func.range.contains(editor.selection.start)).map((el) => el.name)[0];
+	if (!testFunctionName) {
+		vscode.window.showInformationMessage('No test function found at cursor.');
+		return;
+	}
+
+	await editor.document.save();
+
+	const goConfig = getGoConfig(editor.document.uri);
+	const pkgDir = path.dirname(editor.document.fileName);
+	const traceDir = path.join(os.tmpdir(), `vscode-go-rr-${Date.now()}`);
+	const testBin = path.join(os.tmpdir(), `vscode-go-rr-bin-${Date.now()}`);
+
+	const tags = getTestTags(goConfig);
+	const buildFlags = tags ? ['-tags', tags] : [];
+	const flagsFromConfig = getTestFlags(goConfig);
+	flagsFromConfig.forEach((x) => {
+		if (x !== '-args') buildFlags.push(x);
+	});
+
+	const goRuntime = getBinPath('go');
+	const execFile = util.promisify(cp.execFile);
+	const env = toolExecutionEnvironment();
+
+	try {
+		await execFile(goRuntime, ['test', '-c', '-o', testBin, ...buildFlags, pkgDir], { env, cwd: pkgDir });
+	} catch (e) {
+		vscode.window.showErrorMessage(`Failed to build test binary: ${e}`);
+		return;
+	}
+
+	const workspaceFolder = vscode.workspace.getWorkspaceFolder(editor.document.uri);
+
+	const exitCode = await new Promise<number>((resolve) => {
+		const writeEmitter = new vscode.EventEmitter<string>();
+		const pty: vscode.Pseudoterminal = {
+			onDidWrite: writeEmitter.event,
+			open() {
+				const proc = cp.spawn(
+					'rr',
+					['record', `--output-trace-dir=${traceDir}`, testBin, '-test.run', `^${testFunctionName}$`],
+					{ cwd: pkgDir, env: { ...process.env, ...env } }
+				);
+				const onData = (chunk: Buffer | string) => {
+					writeEmitter.fire(chunk.toString().replace(/\n/g, '\r\n'));
+				};
+				proc.stdout?.on('data', onData);
+				proc.stderr?.on('data', onData);
+				proc.on('close', (code) => {
+					writeEmitter.fire(`\r\nrr exited with code ${code}.\r\n`);
+					resolve(code ?? 1);
+				});
+			},
+			close() {}
+		};
+		vscode.window.createTerminal({ name: `rr: ${testFunctionName}`, pty }).show();
+	});
+
+	if (exitCode !== 0) {
+		vscode.window.showErrorMessage(`rr record failed (exit code ${exitCode}). Check the terminal for details.`);
+		return;
+	}
+
+	const debugConfig: vscode.DebugConfiguration = {
+		name: `rr replay: ${testFunctionName}`,
+		type: 'go',
+		request: 'launch',
+		mode: 'replay',
+		traceDirPath: traceDir,
+		env: goConfig.get('testEnvVars', {}),
+		envFile: goConfig.get('testEnvFile')
+	};
+	await vscode.debug.startDebugging(workspaceFolder, debugConfig);
+};
+
 /**
  * Runs all tests in the package of the source of the active editor.
  *

extension/src/language/goLanguageServer.ts

@@ -1012,27 +1012,38 @@ async function passLinkifyShowMessageToGopls(cfg: LanguageServerConfig, goplsCon
 	return goplsConfig;
 }
 
-// createTestCodeLens adds the go.test.cursor and go.debug.cursor code lens
+// createTestCodeLens adds the go.test.cursor, go.debug.cursor, and optionally go.rr.cursor code lenses
 function createTestCodeLens(lens: vscode.CodeLens): vscode.CodeLens[] {
 	// CodeLens argument signature in gopls is [fileName: string, testFunctions: string[], benchFunctions: string[]],
 	// so this needs to be deconstructured here
 	// Note that there will always only be one test function name in this context
 	if ((lens.command?.arguments?.length ?? 0) < 2 || (lens.command?.arguments?.[1].length ?? 0) < 1) {
 		return [lens];
 	}
-	return [
+	const functionName = lens.command?.arguments?.[1][0];
+	const lenses = [
 		new vscode.CodeLens(lens.range, {
 			title: '',
 			...lens.command,
 			command: 'go.test.cursor',
-			arguments: [{ functionName: lens.command?.arguments?.[1][0] }]
+			arguments: [{ functionName }]
 		}),
 		new vscode.CodeLens(lens.range, {
 			title: 'debug test',
 			command: 'go.debug.cursor',
-			arguments: [{ functionName: lens.command?.arguments?.[1][0] }]
+			arguments: [{ functionName }]
 		})
 	];
+	if (getGoConfig().get<{ [key: string]: boolean }>('enableCodeLens')?.rrtest) {
+		lenses.push(
+			new vscode.CodeLens(lens.range, {
+				title: 'rr test',
+				command: 'go.rr.cursor',
+				arguments: [{ functionName }]
+			})
+		);
+	}
+	return lenses;
 }
 
 function createBenchmarkCodeLens(lens: vscode.CodeLens): vscode.CodeLens[] {