[CLI] Fix temp dir cleanup when Playground CLI server is killed

packages/playground/cli/src/run-cli.ts

@@ -394,7 +394,33 @@ export async function parseOptionsAndRunCLI() {
 			],
 		} as RunCLIArgs;
 
-		await runCLI(cliArgs);
+		const cliServer = await runCLI(cliArgs);
+		if (cliServer === undefined) {
+			// No server was started, so we are done with our work.
+			process.exit(0);
+		}
+
+		const cleanUpCliAndExit = (() => {
+			// Remember we are already cleaning up to preclude the possibility
+			// of multiple, conflicting cleanup attempts.
+			let promiseToCleanup: Promise<void>;
+
+			return async () => {
+				if (promiseToCleanup !== undefined) {
+					promiseToCleanup = cliServer[Symbol.asyncDispose]();
+				}
+				await promiseToCleanup;
+				process.exit(0);
+			};
+		})();
+
+		// Playground CLI server must be killed to exit. From the terminal,
+		// this may occur via Ctrl+C which sends SIGINT. Let's handle both
+		// SIGINT and SIGTERM (the default kill signal) to make sure we
+		// clean up after ourselves even if this process is being killed.
+		// NOTE: Windows does not support SIGTERM, but Node.js provides some emulation.
+		process.on('SIGINT', cleanUpCliAndExit);
+		process.on('SIGTERM', cleanUpCliAndExit);
 	} catch (e) {
 		if (!(e instanceof Error)) {
 			throw e;
@@ -437,7 +463,6 @@ export interface RunCLIArgs {
 	autoMount?: string;
 	experimentalMultiWorker?: number;
 	experimentalTrace?: boolean;
-	exitOnPrimaryWorkerCrash?: boolean;
 	internalCookieStore?: boolean;
 	'additional-blueprint-steps'?: any[];
 	xdebug?: boolean | { ideKey?: string };
@@ -492,7 +517,17 @@ const italic = (text: string) =>
 const highlight = (text: string) =>
 	process.stdout.isTTY ? `\x1b[33m${text}\x1b[0m` : text;
 
-export async function runCLI(args: RunCLIArgs): Promise<RunCLIServer> {
+// These overloads are declared for convenience so runCLI() can return
+// different things depending on the CLI command without forcing the
+// callers (mostly automated tests) to check return values.
+export async function runCLI(
+	args: RunCLIArgs & { command: 'build-snapshot' | 'run-blueprint' }
+): Promise<void>;
+export async function runCLI(
+	args: RunCLIArgs & { command: 'server' }
+): Promise<RunCLIServer>;
+export async function runCLI(args: RunCLIArgs): Promise<RunCLIServer | void>;
+export async function runCLI(args: RunCLIArgs): Promise<RunCLIServer | void> {
 	let loadBalancer: LoadBalancer;
 	let playground: RemoteAPI<PlaygroundCliWorker>;
 
@@ -562,7 +597,7 @@ export async function runCLI(args: RunCLIArgs): Promise<RunCLIServer> {
 
 	return startServer({
 		port: args['port'] as number,
-		onBind: async (server: Server, port: number): Promise<RunCLIServer> => {
+		onBind: async (server: Server, port: number) => {
 			const host = '127.0.0.1';
 			const serverUrl = `http://${host}:${port}`;
 			const siteUrl = args['site-url'] || serverUrl;
@@ -584,10 +619,10 @@ export async function runCLI(args: RunCLIArgs): Promise<RunCLIServer> {
 			 * because we don't have to create or maintain multiple copies of the same files.
 			 */
 			const tempDirNameDelimiter = '-playground-cli-site-';
-			const nativeDirPath = await createPlaygroundCliTempDir(
+			const nativeDir = await createPlaygroundCliTempDir(
 				tempDirNameDelimiter
 			);
-			logger.debug(`Native temp dir for VFS root: ${nativeDirPath}`);
+			logger.debug(`Native temp dir for VFS root: ${nativeDir.path}`);
 
 			const IDEConfigName = 'WP Playground CLI - Listen for Xdebug';
 
@@ -602,7 +637,7 @@ export async function runCLI(args: RunCLIArgs): Promise<RunCLIServer> {
 			// directory and add the new IDE config.
 			if (args.xdebug && args.experimentalUnsafeIdeIntegration) {
 				await createPlaygroundCliTempDirSymlink(
-					nativeDirPath,
+					nativeDir.path,
 					symlinkPath,
 					process.platform
 				);
@@ -696,17 +731,15 @@ export async function runCLI(args: RunCLIArgs): Promise<RunCLIServer> {
 
 					console.log('');
 				} catch (error) {
-					logger.error(
-						'Could not configure Xdebug:',
-						(error as Error)?.message
-					);
-					process.exit(1);
+					throw new Error('Could not configure Xdebug', {
+						cause: error,
+					});
 				}
 			}
 
 			// We do not know the system temp dir,
 			// but we can try to infer from the location of the current temp dir.
-			const tempDirRoot = path.dirname(nativeDirPath);
+			const tempDirRoot = path.dirname(nativeDir.path);
 
 			const twoDaysInMillis = 2 * 24 * 60 * 60 * 1000;
 			const tempDirStaleAgeInMillis = twoDaysInMillis;
@@ -721,7 +754,7 @@ export async function runCLI(args: RunCLIArgs): Promise<RunCLIServer> {
 
 			// NOTE: We do not add mount declarations for /internal here
 			// because it will be mounted as part of php-wasm init.
-			const nativeInternalDirPath = path.join(nativeDirPath, 'internal');
+			const nativeInternalDirPath = path.join(nativeDir.path, 'internal');
 			mkdirSync(nativeInternalDirPath);
 
 			const userProvidableNativeSubdirs = [
@@ -746,7 +779,7 @@ export async function runCLI(args: RunCLIArgs): Promise<RunCLIServer> {
 					// The user hasn't requested mounting a different native dir for this path,
 					// so let's create a mount from within our native temp dir.
 					const nativeSubdirPath = path.join(
-						nativeDirPath,
+						nativeDir.path,
 						subdirName
 					);
 					mkdirSync(nativeSubdirPath);
@@ -799,26 +832,48 @@ export async function runCLI(args: RunCLIArgs): Promise<RunCLIServer> {
 				}
 			}
 
+			// Remember whether we are already disposing so we can avoid:
+			// - we can avoid multiple, conflicting dispose attempts
+			// - logging that a worker exited while the CLI itself is exiting
+			let disposing = false;
+			const disposeCLI = async function disposeCLI() {
+				if (disposing) {
+					return;
+				}
+
+				disposing = true;
+				await Promise.all(
+					playgroundsToCleanUp.map(async ({ playground, worker }) => {
+						await playground.dispose();
+						await worker.terminate();
+					})
+				);
+				if (server) {
+					await new Promise((resolve) => server.close(resolve));
+				}
+				await nativeDir.cleanup();
+			};
+
 			// Kick off worker threads now to save time later.
 			// There is no need to wait for other async processes to complete.
 			const promisedWorkers = spawnWorkerThreads(
 				totalWorkerCount,
 				handler.getWorkerType(),
-				({ exitCode, isMain, workerIndex }) => {
-					if (exitCode === 0) {
+				({ exitCode, workerIndex }) => {
+					// We are already disposing, so worker exit is expected
+					// and does not need to be logged.
+					if (disposing) {
+						return;
+					}
+
+					if (exitCode !== 0) {
 						return;
 					}
+
 					logger.error(
 						`Worker ${workerIndex} exited with code ${exitCode}\n`
 					);
-					// If the primary worker crashes, exit the entire process.
-					if (!isMain) {
-						return;
-					}
-					if (!args.exitOnPrimaryWorkerCrash) {
-						return;
-					}
-					process.exit(1);
+					// @TODO: Should we respawn the worker if it exited with an error and the CLI is not shutting down?
 				}
 			);
 
@@ -869,10 +924,12 @@ export async function runCLI(args: RunCLIArgs): Promise<RunCLIServer> {
 				if (args.command === 'build-snapshot') {
 					await zipSite(playground, args.outfile as string);
 					logger.log(`WordPress exported to ${args.outfile}`);
-					process.exit(0);
+					await disposeCLI();
+					return;
 				} else if (args.command === 'run-blueprint') {
 					logger.log(`Blueprint executed`);
-					process.exit(0);
+					await disposeCLI();
+					return;
 				}
 
 				if (
@@ -927,17 +984,7 @@ export async function runCLI(args: RunCLIArgs): Promise<RunCLIServer> {
 					playground,
 					server,
 					serverUrl,
-					[Symbol.asyncDispose]: async function disposeCLI() {
-						await Promise.all(
-							playgroundsToCleanUp.map(
-								async ({ playground, worker }) => {
-									await playground.dispose();
-									await worker.terminate();
-								}
-							)
-						);
-						await new Promise((resolve) => server.close(resolve));
-					},
+					[Symbol.asyncDispose]: disposeCLI,
 					workerThreadCount: totalWorkerCount,
 				};
 			} catch (error) {
@@ -993,19 +1040,14 @@ export type SpawnedWorker = {
 async function spawnWorkerThreads(
 	count: number,
 	workerType: WorkerType,
-	onWorkerExit: (options: {
-		exitCode: number;
-		isMain: boolean;
-		workerIndex: number;
-	}) => void
+	onWorkerExit: (options: { exitCode: number; workerIndex: number }) => void
 ): Promise<SpawnedWorker[]> {
 	const promises = [];
 	for (let i = 0; i < count; i++) {
 		const worker = await spawnWorkerThread(workerType);
 		const onExit: (code: number) => void = (code: number) => {
 			onWorkerExit({
 				exitCode: code,
-				isMain: i === 0,
 				workerIndex: i,
 			});
 		};

packages/playground/cli/src/start-server.ts

@@ -8,13 +8,13 @@ import { logger } from '@php-wasm/logger';
 
 export interface ServerOptions {
 	port: number;
-	onBind: (server: Server, port: number) => Promise<RunCLIServer>;
+	onBind: (server: Server, port: number) => Promise<RunCLIServer | void>;
 	handleRequest: (request: PHPRequest) => Promise<PHPResponse>;
 }
 
 export async function startServer(
 	options: ServerOptions
-): Promise<RunCLIServer> {
+): Promise<RunCLIServer | void> {
 	const app = express();
 
 	const server = await new Promise<

packages/playground/cli/src/temp-dir.ts

@@ -32,27 +32,25 @@ export async function createPlaygroundCliTempDir(
 	// Otherwise, we would have to parse the binary name from the full path.
 	const tempDirPrefix = `${nodeBinaryName}${substrToIdentifyTempDirs}${process.pid}-`;
 
-	const nativeDirPath = (
-		await tmpDir({
-			prefix: tempDirPrefix,
-			/*
-			 * Allow recursive cleanup on process exit.
-			 *
-			 * NOTE: I worried about whether this cleanup would follow symlinks
-			 * and delete target files instead of unlinking the symlink,
-			 * but this feature uses rimraf under the hood which respects symlinks:
-			 * https://github.com/raszi/node-tmp/blob/3d2fe387f3f91b13830b9182faa02c3231ea8258/lib/tmp.js#L318
-			 */
-			unsafeCleanup: true,
-		})
-	).path;
+	const nativeDir = await tmpDir({
+		prefix: tempDirPrefix,
+		/*
+		 * Allow recursive cleanup on process exit.
+		 *
+		 * NOTE: I worried about whether this cleanup would follow symlinks
+		 * and delete target files instead of unlinking the symlink,
+		 * but this feature uses rimraf under the hood which respects symlinks:
+		 * https://github.com/raszi/node-tmp/blob/3d2fe387f3f91b13830b9182faa02c3231ea8258/lib/tmp.js#L318
+		 */
+		unsafeCleanup: true,
+	});
 
 	if (autoCleanup) {
 		// Request graceful cleanup on process exit.
 		tmpSetGracefulCleanup();
 	}
 
-	return nativeDirPath;
+	return nativeDir;
 }
 
 /**

packages/playground/cli/tests/temp-dir-test-process.ts

@@ -16,10 +16,10 @@ process.on('message', async (message: any) => {
 		);
 		// Add a file to the temp dir to test that cleanup works
 		// on non-empty dirs.
-		fs.writeFileSync(path.join(tempDir, 'test.txt'), 'test');
+		fs.writeFileSync(path.join(tempDir.path, 'test.txt'), 'test');
 		process.send!({
 			type: 'temp-dir',
-			tempDir,
+			tempDirPath: tempDir.path,
 		});
 	} else if (message.type === 'exit') {
 		process.exit(0);