cmd/sketch: add a process reaper

josharian · josharian · commit c36d1da8758d · 2025-07-25T21:59:27.000-07:00
When we're PID 1 in a container, reaping zombies
is our responsibility, which we were shirking.

No longer! Now we 🥷 all the 🧟 into ☠️, and not 🐢 either.
diff --git a/cmd/sketch/main.go b/cmd/sketch/main.go
@@ -615,6 +615,14 @@ func setupAndRunAgent(ctx context.Context, flags CLIFlags, modelURL, apiKey, pub
 		os.Setenv("SKETCH_PUB_KEY", pubKey)
 	}
 
+	if inInsideSketch {
+		// Start a process reaper, because we are PID 1.
+		err := startReaper(context.WithoutCancel(ctx))
+		if err != nil {
+			slog.WarnContext(ctx, "failed to start process reaper", "error", err)
+		}
+	}
+
 	wd, err := os.Getwd()
 	if err != nil {
 		return err
diff --git a/cmd/sketch/reaper_linux.go b/cmd/sketch/reaper_linux.go
@@ -0,0 +1,45 @@
+//go:build linux
+
+package main
+
+import (
+	"context"
+	"log/slog"
+	"os"
+	"os/signal"
+
+	"golang.org/x/sys/unix"
+)
+
+func startReaper(ctx context.Context) error {
+	if err := unix.Prctl(unix.PR_SET_CHILD_SUBREAPER, 1, 0, 0, 0); err != nil {
+		return err
+	}
+	go reapZombies(ctx)
+	return nil
+}
+
+// reapZombies runs until ctx is cancelled.
+func reapZombies(ctx context.Context) {
+	sig := make(chan os.Signal, 16)
+	signal.Notify(sig, unix.SIGCHLD)
+
+	for range sig {
+	Reap:
+		for {
+			var status unix.WaitStatus
+			pid, err := unix.Wait4(-1, &status, unix.WNOHANG, nil)
+			switch {
+			case pid > 0:
+				slog.DebugContext(ctx, "reaper ran", "pid", pid, "exit", status.ExitStatus())
+			case err == unix.EINTR:
+				// interrupted: fall through to retry
+			case err == unix.ECHILD || pid == 0:
+				break Reap // no more ready children; wait for next SIGCHLD
+			default:
+				slog.WarnContext(ctx, "wait4 error", "error", err)
+				break Reap
+			}
+		}
+	}
+}
diff --git a/cmd/sketch/reaper_other.go b/cmd/sketch/reaper_other.go
@@ -0,0 +1,9 @@
+//go:build !linux
+
+package main
+
+import "context"
+
+func startReaper(_ context.Context) error {
+	return nil
+}
