Reuse Docker container between reloads, closes #369 (#376)

* Reuse Docker container between reloads, closes #369

On my machine, this change reduces reload time from ~1.5s to ~400ms

* Remove redundant `await`s

* Add return type annotations to `Runtime#dispose()` implementations

Author: mrbbot

package-lock.json

@@ -1,2 +1,3 @@
 *
 !.gitignore
+!restart.sh

packages/tre/lib/.gitignore

@@ -0,0 +1,40 @@
+#!/usr/bin/env bash
+# Restarts a process on receiving SIGUSR1.
+# Usage: ./restart.sh <command> <...args>
+
+# Start process and record its PID
+"$@" &
+PID=$!
+echo "[*] Started $PID"
+
+# Trap SIGUSR1 to set $RECEIVED_USR1 to 1, then terminate $PID.
+# Setting $RECEIVED_USR1 will cause the process to be restarted.
+RECEIVED_USR1=0
+trap 'RECEIVED_USR1=1 && kill -TERM $PID' USR1
+
+# Trap SIGINT and SIGTERM to also terminate $PID for cleanup.
+# By not setting $RECEIVED_USR1, we ensure this script exits
+# when $PID exits.
+trap 'kill -TERM $PID' INT TERM
+
+while true
+do
+  # Wait for the started process to exit
+  wait $PID
+  EXIT_CODE=$?
+
+  # If the process exited for any reason other than this script
+  # receiving SIGUSR1, exit the script with the same exit code.
+  if [ $RECEIVED_USR1 -eq 0 ]
+  then
+    echo "[*] Exited with status $EXIT_CODE"
+    exit $EXIT_CODE
+  fi
+
+  # Otherwise, if this script received SIGUSR1, reset the flag,
+  # restart the process, and record its new PID.
+  RECEIVED_USR1=0
+  "$@" &
+  PID=$!
+  echo "[*] Restarted $PID"
+done

packages/tre/lib/restart.sh

@@ -39,6 +39,7 @@
     "acorn": "^8.8.0",
     "acorn-walk": "^8.2.0",
     "capnp-ts": "^0.7.0",
+    "exit-hook": "^2.2.1",
     "get-port": "^5.1.1",
     "kleur": "^4.1.5",
     "stoppable": "^1.1.0",

packages/tre/package.json

@@ -1,6 +1,7 @@
 import assert from "assert";
 import http from "http";
 import path from "path";
+import exitHook from "exit-hook";
 import getPort from "get-port";
 import { bold, green, grey } from "kleur/colors";
 import stoppable from "stoppable";
@@ -99,6 +100,7 @@ type PluginRouters = {
   [Key in keyof Plugins]: OptionalInstanceType<Plugins[Key]["router"]>;
 };
 
+// `__dirname` relative to bundled output `dist/src/index.js`
 const RUNTIME_PATH = path.resolve(__dirname, "..", "..", "lib", "cfwrkr");
 
 type StoppableServer = http.Server & stoppable.WithStop;
@@ -114,6 +116,7 @@ export class Miniflare {
 
   readonly #runtimeConstructor: RuntimeConstructor;
   #runtime?: Runtime;
+  #removeRuntimeExitHook?: () => void;
   #runtimeEntryURL?: URL;
 
   readonly #disposeController: AbortController;
@@ -181,6 +184,8 @@ export class Miniflare {
       entryPort,
       loopbackPort
     );
+    this.#removeRuntimeExitHook = exitHook(() => void this.#runtime?.dispose());
+
     this.#runtimeEntryURL = new URL(`http://127.0.0.1:${entryPort}`);
 
     const config = await this.#initialConfigPromise;
@@ -410,6 +415,7 @@ export class Miniflare {
     this.#disposeController.abort();
     await this.#initPromise;
     await this.#updatePromise;
+    this.#removeRuntimeExitHook?.();
     this.#runtime?.dispose();
     await this.#stopLoopbackServer();
   }

packages/tre/src/index.ts

@@ -1,10 +1,20 @@
 import childProcess from "child_process";
-import { Awaitable, MiniflareError } from "../helpers";
+import crypto from "crypto";
+import fs from "fs";
+import os from "os";
+import path from "path";
+import { Awaitable, MiniflareCoreError } from "../helpers";
 import { SERVICE_LOOPBACK, SOCKET_ENTRY } from "../plugins";
 
-export interface Runtime {
-  updateConfig(configBuffer: Buffer): Awaitable<void>;
-  dispose(): Awaitable<void>;
+export abstract class Runtime {
+  constructor(
+    protected readonly runtimeBinaryPath: string,
+    protected readonly entryPort: number,
+    protected readonly loopbackPort: number
+  ) {}
+
+  abstract updateConfig(configBuffer: Buffer): Awaitable<void>;
+  abstract dispose(): Awaitable<void>;
 }
 
 export interface RuntimeConstructor {
@@ -21,14 +31,16 @@ export interface RuntimeConstructor {
 }
 
 const COMMON_RUNTIME_ARGS = ["serve", "--binary", "--verbose"];
+// `__dirname` relative to bundled output `dist/src/index.js`
+const RESTART_PATH = path.resolve(__dirname, "..", "..", "lib", "restart.sh");
 
-function waitForExit(process: childProcess.ChildProcess): Promise<number> {
+function waitForExit(process: childProcess.ChildProcess): Promise<void> {
   return new Promise((resolve) => {
-    process.once("exit", (code) => resolve(code ?? -1));
+    process.once("exit", () => resolve());
   });
 }
 
-class NativeRuntime implements Runtime {
+class NativeRuntime extends Runtime {
   static isSupported() {
     return process.platform === "linux"; // TODO: and "darwin"?
   }
@@ -40,13 +52,14 @@ class NativeRuntime implements Runtime {
   readonly #args: string[];
 
   #process?: childProcess.ChildProcess;
-  #processExitPromise?: Promise<number>;
+  #processExitPromise?: Promise<void>;
 
   constructor(
-    protected readonly runtimeBinaryPath: string,
-    protected readonly entryPort: number,
-    protected readonly loopbackPort: number
+    runtimeBinaryPath: string,
+    entryPort: number,
+    loopbackPort: number
   ) {
+    super(runtimeBinaryPath, entryPort, loopbackPort);
     const [command, ...args] = this.getCommand();
     this.#command = command;
     this.#args = args;
@@ -71,7 +84,7 @@ class NativeRuntime implements Runtime {
     // TODO: what happens if runtime crashes?
 
     // 2. Start new process
-    const runtimeProcess = await childProcess.spawn(this.#command, this.#args, {
+    const runtimeProcess = childProcess.spawn(this.#command, this.#args, {
       stdio: "pipe",
       shell: true,
     });
@@ -86,11 +99,12 @@ class NativeRuntime implements Runtime {
 
     // 3. Write config
     runtimeProcess.stdin.write(configBuffer);
+    runtimeProcess.stdin.end();
   }
 
-  async dispose() {
+  dispose(): Awaitable<void> {
     this.#process?.kill();
-    await this.#processExitPromise;
+    return this.#processExitPromise;
   }
 }
 
@@ -112,7 +126,7 @@ class WSLRuntime extends NativeRuntime {
   }
 }
 
-class DockerRuntime extends NativeRuntime {
+class DockerRuntime extends Runtime {
   static isSupported() {
     const result = childProcess.spawnSync("docker", ["--version"]); // TODO: check daemon running too?
     return result.error === undefined;
@@ -123,23 +137,68 @@ class DockerRuntime extends NativeRuntime {
   static description = "using Docker 🐳";
   static distribution = `linux-${process.arch}`;
 
-  getCommand(): string[] {
-    // TODO: consider reusing container, but just restarting process within
-    return [
+  #configPath = path.join(
+    os.tmpdir(),
+    `miniflare-config-${crypto.randomBytes(16).toString("hex")}.bin`
+  );
+
+  #process?: childProcess.ChildProcess;
+  #processExitPromise?: Promise<void>;
+
+  async updateConfig(configBuffer: Buffer) {
+    // 1. Write config to file (this is much easier than trying to buffer STDIN
+    //    in the restart script)
+    fs.writeFileSync(this.#configPath, configBuffer);
+
+    // 2. If process running, send SIGUSR1 to restart runtime with new config
+    //    (see `lib/restart.sh`)
+    if (this.#process) {
+      this.#process.kill("SIGUSR1");
+      return;
+    }
+
+    // 3. Otherwise, start new process
+    const runtimeProcess = childProcess.spawn(
       "docker",
-      "run",
-      "--platform=linux/amd64",
-      "--interactive",
-      "--rm",
-      `--volume=${this.runtimeBinaryPath}:/runtime`,
-      `--publish=127.0.0.1:${this.entryPort}:8787`,
-      "debian:bullseye-slim",
-      "/runtime",
-      ...COMMON_RUNTIME_ARGS,
-      `--socket-addr=${SOCKET_ENTRY}=*:8787`,
-      `--external-addr=${SERVICE_LOOPBACK}=host.docker.internal:${this.loopbackPort}`,
-      "-",
-    ];
+      [
+        "run",
+        "--platform=linux/amd64",
+        "--interactive",
+        "--rm",
+        `--volume=${RESTART_PATH}:/restart.sh`,
+        `--volume=${this.runtimeBinaryPath}:/runtime`,
+        `--volume=${this.#configPath}:/miniflare-config.bin`,
+        `--publish=127.0.0.1:${this.entryPort}:8787`,
+        "debian:bullseye-slim",
+        "/restart.sh",
+        "/runtime",
+        ...COMMON_RUNTIME_ARGS,
+        `--socket-addr=${SOCKET_ENTRY}=*:8787`,
+        `--external-addr=${SERVICE_LOOPBACK}=host.docker.internal:${this.loopbackPort}`,
+        "/miniflare-config.bin",
+      ],
+      {
+        stdio: "pipe",
+        shell: true,
+      }
+    );
+    this.#process = runtimeProcess;
+    this.#processExitPromise = waitForExit(runtimeProcess);
+
+    // TODO: may want to proxy these and prettify ✨
+    runtimeProcess.stdout.pipe(process.stdout);
+    runtimeProcess.stderr.pipe(process.stderr);
+  }
+
+  dispose(): Awaitable<void> {
+    this.#process?.kill();
+    try {
+      fs.unlinkSync(this.#configPath);
+    } catch (e: any) {
+      // Ignore not found errors if we called dispose() without updateConfig()
+      if (e.code !== "ENOENT") throw e;
+    }
+    return this.#processExitPromise;
   }
 }
 
@@ -160,7 +219,7 @@ export function getSupportedRuntime(): RuntimeConstructor {
   const suggestions = RUNTIMES.map(
     ({ supportSuggestion }) => `- ${supportSuggestion}`
   );
-  throw new MiniflareError(
+  throw new MiniflareCoreError(
     "ERR_RUNTIME_UNSUPPORTED",
     `The 🦄 Cloudflare Workers Runtime 🦄 does not support your system (${
       process.platform