Clean up containers started by local dev during process exit hook instead (#10099)

* make cleanup container calls sync and call in process.on("exit")

* fix tests

* changeset

* PR feedback

* cleanup

Author: emily-shen

.changeset/great-terms-roll.md

@@ -0,0 +1,6 @@
+---
+"@cloudflare/vite-plugin": patch
+"wrangler": patch
+---
+
+fix: move local dev container cleanup to process exit hook. This should ensure containers are cleaned up even when Wrangler is shut down programatically.

fixtures/interactive-dev-tests/tests/index.test.ts

@@ -396,7 +396,6 @@ baseDescribe.skipIf(process.platform !== "linux" && process.env.CI === "true")(
 
 			wrangler.pty.write("r");
 
-			// wait for build to finish
 			await vi.waitFor(async () => {
 				const status = await fetch(wrangler.url + "/status");
 				expect(await status.json()).toBe(false);
@@ -435,7 +434,7 @@ baseDescribe.skipIf(process.platform !== "linux" && process.env.CI === "true")(
 			await new Promise<void>((resolve) => {
 				wrangler.pty.onExit(() => resolve());
 			});
-			vi.waitFor(() => {
+			await vi.waitFor(() => {
 				const remainingIds = getContainerIds();
 				expect(remainingIds.length).toBe(0);
 			});
@@ -719,7 +718,7 @@ baseDescribe.skipIf(process.platform !== "linux" && process.env.CI === "true")(
 			await new Promise<void>((resolve) => {
 				wrangler.pty.onExit(() => resolve());
 			});
-			vi.waitFor(() => {
+			await vi.waitFor(() => {
 				const remainingIds = getContainerIds();
 				expect(remainingIds.length).toBe(0);
 			});
@@ -779,7 +778,7 @@ baseDescribe.skipIf(process.platform !== "linux" && process.env.CI === "true")(
 			}
 		});
 
-		it("should be interact with both workers, rebuild the containers with the hotkey and all containers should be cleaned up at the end", async () => {
+		it("should be able to interact with both workers, rebuild the containers with the hotkey and all containers should be cleaned up at the end", async () => {
 			const wrangler = await startWranglerDev([
 				"dev",
 				"-c",

packages/containers-shared/src/utils.ts

@@ -141,33 +141,21 @@ export const isDockerfile = (
 
 /**
  * Kills and removes any containers which come from the given image tag
- *
- * Please note that this function has an almost identical counterpart
- * in the `vite-plugin-cloudflare` package (see `removeContainersByIds`).
- * If you make any changes to this fn, please make sure you persist those
- * changes in `removeContainersByIds` if necessary.
  */
-export const cleanupContainers = async (
+export const cleanupContainers = (
 	dockerPath: string,
 	imageTags: Set<string>
 ) => {
 	try {
 		// Find all containers (stopped and running) for each built image
-		const containerIds = await getContainerIdsByImageTags(
-			dockerPath,
-			imageTags
-		);
+		const containerIds = getContainerIdsByImageTags(dockerPath, imageTags);
 
 		if (containerIds.length === 0) {
 			return true;
 		}
 
 		// Workerd should have stopped all containers, but clean up any in case. Sends a sigkill.
-		await runDockerCmd(
-			dockerPath,
-			["rm", "--force", ...containerIds],
-			["inherit", "pipe", "pipe"]
-		);
+		runDockerCmdWithOutput(dockerPath, ["rm", "--force", ...containerIds]);
 		return true;
 	} catch {
 		return false;
@@ -181,14 +169,14 @@ export const cleanupContainers = async (
  * @param imageTags A set of ancestor image tags
  * @returns The ids of all containers that share the given image tags as ancestors.
  */
-export async function getContainerIdsByImageTags(
+export function getContainerIdsByImageTags(
 	dockerPath: string,
 	imageTags: Set<string>
-): Promise<Array<string>> {
+): string[] {
 	const ids = new Set<string>();
 
 	for (const imageTag of imageTags) {
-		const containerIdsFromImage = await getContainerIdsFromImage(
+		const containerIdsFromImage = getContainerIdsFromImage(
 			dockerPath,
 			imageTag
 		);
@@ -198,7 +186,7 @@ export async function getContainerIdsByImageTags(
 	return Array.from(ids);
 }
 
-export const getContainerIdsFromImage = async (
+export const getContainerIdsFromImage = (
 	dockerPath: string,
 	ancestorImage: string
 ) => {

packages/vite-plugin-cloudflare/src/containers.ts

@@ -2,10 +2,7 @@ import assert from "node:assert";
 import path from "node:path";
 import { prepareContainerImagesForDev } from "@cloudflare/containers-shared/src/images";
 import { getDevContainerImageName } from "@cloudflare/containers-shared/src/knobs";
-import {
-	isDockerfile,
-	runDockerCmd,
-} from "@cloudflare/containers-shared/src/utils";
+import { isDockerfile } from "@cloudflare/containers-shared/src/utils";
 import type { WorkerConfig } from "./plugin-config";
 
 /**
@@ -20,40 +17,6 @@ export function getDockerPath(): string {
 	return process.env[dockerPathEnvVar] || defaultDockerPath;
 }
 
-/**
- * Performs the forced removal of running Docker containers, given their ids.
- *
- * Please note that this function is almost identical to `cleanupContainers`
- * defined in `containers-shared`, with the exception that the current fn
- * expects the list of ids of all containers that are to be removed, as a
- * parameter. This is because extracting these ids is an async operation, and
- * we are currently restricted from performing async cleanup work upon
- * closing/exiting the vite dev process (see vite-plugin-cloudflare/src/index.ts
- * for more details).
- *
- * @param dockerPath The path to the Docker executable
- * @param containerIds The ids of the containers that should be removed
- */
-export async function removeContainersByIds(
-	dockerPath: string,
-	containerIds: string[]
-) {
-	try {
-		if (containerIds.length === 0) {
-			return;
-		}
-
-		await runDockerCmd(
-			dockerPath,
-			["rm", "--force", ...containerIds],
-			["inherit", "pipe", "pipe"]
-		);
-	} catch (error) {
-		// fail silently
-		return;
-	}
-}
-
 /**
  * @returns Container options suitable for building or pulling images,
  * with image tag set to well-known dev format, or undefined if

packages/vite-plugin-cloudflare/src/index.ts

@@ -3,8 +3,8 @@ import * as fsp from "node:fs/promises";
 import * as path from "node:path";
 import * as util from "node:util";
 import {
+	cleanupContainers,
 	generateContainerBuildId,
-	getContainerIdsByImageTags,
 	resolveDockerHost,
 } from "@cloudflare/containers-shared/src/utils";
 import { generateStaticRoutingRuleMatcher } from "@cloudflare/workers-shared/asset-worker/src/utils/rules-engine";
@@ -29,11 +29,7 @@ import {
 	kRequestType,
 	ROUTER_WORKER_NAME,
 } from "./constants";
-import {
-	getDockerPath,
-	prepareContainerImages,
-	removeContainersByIds,
-} from "./containers";
+import { getDockerPath, prepareContainerImages } from "./containers";
 import {
 	addDebugToVitePrintUrls,
 	getDebugPathHtml,
@@ -103,8 +99,7 @@ export function cloudflare(pluginConfig: PluginConfig = {}): vite.Plugin[] {
 
 	const additionalModulePaths = new Set<string>();
 	const nodeJsCompatWarningsMap = new Map<WorkerConfig, NodeJsCompatWarnings>();
-	let containerImageTagsSeen: Set<string> | undefined;
-	let runningContainerIds: Array<string>;
+	let containerImageTagsSeen = new Set<string>();
 
 	/** Used to track whether hooks are being called because of a server restart or a server close event. */
 	let restartingServer = false;
@@ -493,17 +488,6 @@ if (import.meta.hot) {
 							)
 						);
 
-						// poll Docker every two seconds and update the list of ids of all
-						// running containers
-						const dockerPollIntervalId = setInterval(async () => {
-							if (containerImageTagsSeen?.size) {
-								runningContainerIds = await getContainerIdsByImageTags(
-									dockerPath,
-									containerImageTagsSeen
-								);
-							}
-						}, 2000);
-
 						/*
 						 * Upon exiting the dev process we should ensure we perform any
 						 * containers-specific cleanup work. Vite recommends using the
@@ -517,21 +501,11 @@ if (import.meta.hot) {
 						 * `process.exit()` imperatively, and therefore causes `beforeExit`
 						 * not to be emitted).
 						 *
-						 * Furthermore, since the `exit` event handler cannot perform async
-						 * ops as per spec (https://nodejs.org/api/process.html#event-exit),
-						 * we also need a mechanism to ensure that list of containers to be
-						 * cleaned up on exit is up to date. This is what the interval with id
-						 * `dockerPollIntervalId` is for.
-						 *
-						 * It is possible, though very unlikely, that in some rare cases,
-						 * we might be left with some orphaned containers, due to the fact
-						 * that at the point of exiting the dev process, our internal list
-						 * of container ids is out of date. We accept this caveat for now.
-						 *
 						 */
-						process.on("exit", () => {
-							clearInterval(dockerPollIntervalId);
-							removeContainersByIds(dockerPath, runningContainerIds);
+						process.on("exit", async () => {
+							if (containerImageTagsSeen.size) {
+								cleanupContainers(dockerPath, containerImageTagsSeen);
+							}
 						});
 					}
 				}
@@ -629,18 +603,10 @@ if (import.meta.hot) {
 						colors.dim(colors.yellow("\n⚡️ Containers successfully built.\n"))
 					);
 
-					const dockerPollIntervalId = setInterval(async () => {
-						if (containerImageTagsSeen?.size) {
-							runningContainerIds = await getContainerIdsByImageTags(
-								dockerPath,
-								containerImageTagsSeen
-							);
-						}
-					}, 2000);
-
 					process.on("exit", () => {
-						clearInterval(dockerPollIntervalId);
-						removeContainersByIds(dockerPath, runningContainerIds);
+						if (containerImageTagsSeen.size) {
+							cleanupContainers(dockerPath, containerImageTagsSeen);
+						}
 					});
 				}
 
@@ -665,14 +631,7 @@ if (import.meta.hot) {
 					containerImageTagsSeen?.size
 				) {
 					const dockerPath = getDockerPath();
-					runningContainerIds = await getContainerIdsByImageTags(
-						dockerPath,
-						containerImageTagsSeen
-					);
-
-					await removeContainersByIds(dockerPath, runningContainerIds);
-					containerImageTagsSeen.clear();
-					runningContainerIds = [];
+					cleanupContainers(dockerPath, containerImageTagsSeen);
 				}
 
 				debuglog("buildEnd:", restartingServer ? "restarted" : "disposing");

packages/wrangler/e2e/containers.dev.test.ts

@@ -26,7 +26,7 @@ for (const source of imageSource) {
 	const isPullWithoutAccountId = source === "pull" && !CLOUDFLARE_ACCOUNT_ID;
 
 	describe.skipIf(isCINonLinux || isPullWithoutAccountId)(
-		"containers local dev tests: %s",
+		`containers local dev tests: ${source}`,
 		{ timeout: 90_000 },
 		() => {
 			let helper: WranglerE2ETestHelper;
@@ -158,16 +158,23 @@ for (const source of imageSource) {
 				await helper.seed({
 					"wrangler.json": JSON.stringify(wranglerConfig),
 				});
-				// cleanup any running containers
+				/// wait a bit in case the expected cleanup from shutting down wrangler dev is already happening
+				await new Promise((resolve) => setTimeout(resolve, 500));
+				// cleanup any running containers. this does happen automatically when we shut down wrangler,
+				// but treekill is being uncooperative. this is also tested in interactive-dev-fixture
+				// where it is working as expected
 				const ids = getContainerIds("e2econtainer");
 				if (ids.length > 0) {
-					console.log(ids);
 					execSync(`${getDockerPath()} rm -f ${ids.join(" ")}`, {
 						encoding: "utf8",
 					});
 				}
 			});
 			afterAll(async () => {
+				// wait a bit in case the expected cleanup from shutting down wrangler dev is already happening
+				await new Promise((resolve) => setTimeout(resolve, 500));
+				// again this should happen automatically when we shut down wrangler, but treekill is being uncooperative.
+				// this is tested in interactive-dev-fixture where it is working as expected.
 				const ids = getContainerIds("e2econtainer");
 				if (ids.length > 0) {
 					execSync(`${getDockerPath()} rm -f ${ids.join(" ")}`, {
@@ -209,21 +216,29 @@ for (const source of imageSource) {
 				expect(response.status).toBe(200);
 				expect(text).toBe("Container create request sent...");
 
-				// Wait a bit for container to start
-				await new Promise((resolve) => setTimeout(resolve, 2_000));
+				await vi.waitFor(async () => {
+					response = await fetch(`${ready.url}/status`);
+					expect(response.status).toBe(200);
+					const status = await response.json();
+					expect(status).toBe(true);
+				});
 
-				response = await fetch(`${ready.url}/status`);
-				const status = await response.json();
-				expect(response.status).toBe(200);
-				expect(status).toBe(true);
+				await vi.waitFor(
+					async () => {
+						response = await fetch(`${ready.url}/fetch`, {
+							signal: AbortSignal.timeout(3_000),
+							headers: { "MF-Disable-Pretty-Error": "true" },
+						});
+						text = await response.text();
+						expect(text).toBe("Hello World! Have an env var! I'm an env var!");
+					},
+					{ timeout: 5_000 }
+				);
 
-				response = await fetch(`${ready.url}/fetch`);
-				expect(response.status).toBe(200);
-				text = await response.text();
-				expect(text).toBe("Hello World! Have an env var! I'm an env var!");
 				// Check that a container is running using `docker ps`
 				const ids = getContainerIds("e2econtainer");
 				expect(ids.length).toBe(1);
+				await worker.stop();
 			});
 
 			it("won't start the container service if no containers are present", async () => {

packages/wrangler/src/api/startDevWorker/LocalRuntimeController.ts

@@ -191,7 +191,9 @@ export class LocalRuntimeController extends RuntimeController {
 	};
 
 	onBundleStart(_: BundleStartEvent) {
-		// Ignored in local runtime
+		process.on("exit", () => {
+			this.cleanupContainers();
+		});
 	}
 
 	async #onBundleComplete(data: BundleCompleteEvent, id: number) {
@@ -381,12 +383,16 @@ export class LocalRuntimeController extends RuntimeController {
 		// Ignored in local runtime
 	}
 
-	cleanupContainers = async () => {
+	cleanupContainers = () => {
+		if (!this.containerImageTagsSeen.size) {
+			return;
+		}
+
 		assert(
 			this.dockerPath,
 			"Docker path should have been set if containers are enabled"
 		);
-		await cleanupContainers(this.dockerPath, this.containerImageTagsSeen);
+		cleanupContainers(this.dockerPath, this.containerImageTagsSeen);
 	};
 
 	#teardown = async (): Promise<void> => {
@@ -399,15 +405,6 @@ export class LocalRuntimeController extends RuntimeController {
 		await this.#mf?.dispose();
 		this.#mf = undefined;
 
-		if (this.containerImageTagsSeen.size > 0) {
-			try {
-				await this.cleanupContainers();
-			} catch {
-				logger.warn(
-					`Failed to clean up containers. You may have to stop and remove them up manually.`
-				);
-			}
-		}
 		if (this.#remoteProxySessionData) {
 			logger.log(chalk.dim("⎔ Shutting down remote connection..."));
 		}