test support: await fork descendants in formula runs

Kandelo/formula_support/kandelo_formula_support.rb

@@ -314,12 +314,13 @@ def kandelo_install_bin(out_dir, wasm_name, bin_name)
   # `exec_programs:` stages explicit guest exec targets, `guest_files:` stages
   # ordinary files in the guest VFS, `writable_host_directories:` exposes
   # explicit host directories as writable guest mounts for output validation,
-  # and `expected_status:` permits tests for specified nonzero results such as
-  # a grep no-match status.
+  # `expected_fork_descendants:` keeps the host alive until at least that many
+  # fork descendants have exited successfully, and `expected_status:`
+  # permits tests for specified nonzero results such as a grep no-match status.
   def kandelo_run_wasm(
     bin_path, argv, env: {}, stdin: nil, merge_stderr: false, network: false,
     preserve_argv0: false, argv0: nil, exec_programs: {}, guest_files: {},
-    writable_host_directories: {}, expected_status: 0
+    writable_host_directories: {}, expected_fork_descendants: 0, expected_status: 0
   )
     root = kandelo_require_root!
     if !argv0.nil? && (
@@ -328,6 +329,8 @@ def kandelo_run_wasm(
     )
       odie "guest argv0 must be a nonempty normalized absolute path: #{argv0.inspect}"
     end
+    valid_descendant_count = expected_fork_descendants.is_a?(Integer) && expected_fork_descendants >= 0
+    odie "expected fork descendant count must be a nonnegative integer" unless valid_descendant_count
     if (node = ENV.fetch("HOMEBREW_KANDELO_NODE", nil)).to_s != ""
       ENV.prepend_path "PATH", File.dirname(node)
     end
@@ -347,7 +350,8 @@ def kandelo_run_wasm(
     command = +"cd "
     command << Shellwords.escape(root) << " && "
     isolated_runner = network || preserve_argv0 || !argv0.nil? || exec_programs.any? ||
-                      guest_files.any? || writable_host_directories.any?
+                      guest_files.any? || writable_host_directories.any? ||
+                      expected_fork_descendants.positive?
     if isolated_runner
       guest_env = JSON.generate(env.transform_values(&:to_s))
       guest_exec_programs = JSON.generate(exec_programs.transform_values(&:to_s))
@@ -359,6 +363,9 @@ def kandelo_run_wasm(
       command << "KANDELO_FORMULA_WRITABLE_HOST_DIRS_JSON=#{Shellwords.escape(writable_mounts)} "
       command << "KANDELO_FORMULA_ARGV0=#{Shellwords.escape(argv0.to_s)} " if argv0
       command << "KANDELO_FORMULA_ENABLE_NETWORK=#{network ? 1 : 0} "
+      if expected_fork_descendants.positive?
+        command << "KANDELO_FORMULA_EXPECTED_FORK_DESCENDANTS=#{expected_fork_descendants} "
+      end
     else
       env.each { |key, value| command << "#{key}=#{Shellwords.escape(value.to_s)} " }
     end

Kandelo/formula_support/run-network-wasm.ts

@@ -23,6 +23,38 @@ interface WritableRootfs {
   close(fd: number): void;
 }
 
+interface ProcessEvent {
+  kind: "spawn" | "exec" | "exit";
+  pid: number;
+  ppid?: number;
+  exitStatus?: number;
+}
+
+async function waitForForkDescendants(
+  expectedCount: number,
+  activePids: Set<number>,
+  descendantPids: Set<number>,
+  descendantExitStatuses: Map<number, number>,
+  deadline: number,
+): Promise<void> {
+  while (descendantPids.size < expectedCount || activePids.size > 0) {
+    if (Date.now() >= deadline) {
+      throw new Error(
+        `timed out waiting for ${expectedCount} fork descendant(s); ` +
+          `observed ${descendantPids.size}, active ${[...activePids].join(",") || "none"}`,
+      );
+    }
+    await new Promise((resolve) => setTimeout(resolve, 10));
+  }
+
+  const failures = [...descendantExitStatuses]
+    .filter(([, status]) => status !== 0)
+    .map(([pid, status]) => `${pid}:${status}`);
+  if (failures.length > 0) {
+    throw new Error(`descendant process failed: ${failures.join(", ")}`);
+  }
+}
+
 function writeGuestFile(
   rootfs: WritableRootfs,
   guestPath: string,
@@ -77,6 +109,18 @@ async function main(): Promise<void> {
   const writableHostDirectories = JSON.parse(
     process.env.KANDELO_FORMULA_WRITABLE_HOST_DIRS_JSON ?? "{}",
   ) as Record<string, string>;
+  const expectedForkDescendantsValue =
+    process.env.KANDELO_FORMULA_EXPECTED_FORK_DESCENDANTS ?? "0";
+  const expectedForkDescendants = Number(expectedForkDescendantsValue);
+  if (
+    !/^(0|[1-9]\d*)$/.test(expectedForkDescendantsValue) ||
+    !Number.isSafeInteger(expectedForkDescendants) ||
+    expectedForkDescendants < 0
+  ) {
+    throw new Error(
+      `invalid expected fork descendant count: ${expectedForkDescendantsValue}`,
+    );
+  }
   const configuredArgv0 = process.env.KANDELO_FORMULA_ARGV0;
   const argv0 = configuredArgv0 ?? programPath;
   const guestPaths = [...Object.keys(guestFiles), ...Object.keys(execPrograms)];
@@ -175,6 +219,9 @@ async function main(): Promise<void> {
     writeGuestFile(rootfs, entry.guestPath, entry.bytes, entry.mode);
   }
   const rootfsImage = await rootfs.saveImage();
+  const activePids = new Set<number>();
+  const descendantPids = new Set<number>();
+  const descendantExitStatuses = new Map<number, number>();
   const host = new NodeKernelHost({
     maxWorkers: 8,
     execPrograms,
@@ -185,6 +232,19 @@ async function main(): Promise<void> {
     rootfsImage,
     onStdout: (_pid: number, data: Uint8Array) => process.stdout.write(data),
     onStderr: (_pid: number, data: Uint8Array) => process.stderr.write(data),
+    onProcessEvent: (event: ProcessEvent) => {
+      // Fork events carry a parent PID and are posted before fork() returns.
+      // Ignore the synthetic root spawn so a fast root exit cannot be re-added.
+      if (event.kind === "spawn" && event.ppid !== undefined) {
+        activePids.add(event.pid);
+        descendantPids.add(event.pid);
+      } else if (event.kind === "exit") {
+        if (descendantPids.has(event.pid)) {
+          activePids.delete(event.pid);
+          descendantExitStatuses.set(event.pid, event.exitStatus ?? -1);
+        }
+      }
+    },
   });
 
   try {
@@ -208,6 +268,7 @@ async function main(): Promise<void> {
       guestEnv.TIMEOUT ?? process.env.TIMEOUT ?? "30000",
       10,
     );
+    const deadline = Date.now() + timeoutMs;
     const exit = host.spawn(program, [argv0, ...args], {
       cwd: guestEnv.KERNEL_CWD ?? "/tmp",
       env,
@@ -221,7 +282,20 @@ async function main(): Promise<void> {
       );
     });
     try {
-      process.exitCode = await Promise.race([exit, timeout]);
+      const status = await Promise.race([exit, timeout]);
+      process.exitCode = status;
+      if (status === 0 && expectedForkDescendants > 0) {
+        await Promise.race([
+          waitForForkDescendants(
+            expectedForkDescendants,
+            activePids,
+            descendantPids,
+            descendantExitStatuses,
+            deadline,
+          ),
+          timeout,
+        ]);
+      }
     } finally {
       if (timer) clearTimeout(timer);
     }

Kandelo/formula_support/test/kandelo_formula_support_test.rb

@@ -269,20 +269,32 @@ def test_host_tool_executes_from_the_caller_directory
   def test_network_execution_uses_tap_owned_runner
     harness = Harness.new
     output = harness.kandelo_run_wasm(
-      "program.wasm", ["a b"], env: { "TOKEN" => "x y" }, network: true
+      "program.wasm", ["a b"], env: { "TOKEN" => "x y" }, network: true,
+      expected_fork_descendants: 1
     )
 
     assert_equal "runtime-ok\n", output
     assert_includes harness.command, "run-network-wasm.ts"
     assert_includes harness.command, "/tmp/kandelo\\ root"
     assert_includes harness.command, "KANDELO_FORMULA_GUEST_ENV_JSON="
     assert_includes harness.command, "KANDELO_FORMULA_ENABLE_NETWORK=1"
+    assert_includes harness.command, "KANDELO_FORMULA_EXPECTED_FORK_DESCENDANTS=1"
     assert_includes harness.command, "TOKEN"
     refute_includes harness.command, "TOKEN=x\\ y"
     assert_includes harness.command, "program.wasm a\\ b"
     refute_includes harness.command, "examples/run-example.ts"
   end
 
+  def test_execution_rejects_invalid_expected_fork_descendant_count
+    error = assert_raises(RuntimeError) do
+      Harness.new.kandelo_run_wasm(
+        "program.wasm", [], expected_fork_descendants: -1
+      )
+    end
+
+    assert_includes error.message, "expected fork descendant count must be a nonnegative integer"
+  end
+
   def test_default_execution_keeps_standard_runner_and_removes_stale_host_dist
     Dir.mktmpdir("kandelo-formula-support") do |dir|
       root = Pathname(dir)/"kandelo root"