From 190e3e5891cef1d3a94352f16711ef7c09a1d60b Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Tue, 30 Sep 2025 11:52:13 +0530
Subject: [PATCH] Add atexit cleanup via worker process to kittens
 infrastructure

---
 tools/cmd/main.go     |  1 +
 tools/utils/atexit.go | 71 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+)
 create mode 100644 tools/utils/atexit.go

diff --git a/tools/cmd/main.go b/tools/cmd/main.go
index bcbf2efa5..d6cd39cae 100644
--- a/tools/cmd/main.go
+++ b/tools/cmd/main.go
@@ -17,6 +17,7 @@ import (
 )
 
 func KittenMain(args ...string) int {
+	defer utils.WaitForAtexitWorkerToFinish()
 	krm := os.Getenv("KITTY_KITTEN_RUN_MODULE")
 	os.Unsetenv("KITTY_KITTEN_RUN_MODULE")
 	switch krm {
diff --git a/tools/utils/atexit.go b/tools/utils/atexit.go
new file mode 100644
index 000000000..14de67315
--- /dev/null
+++ b/tools/utils/atexit.go
@@ -0,0 +1,71 @@
+package utils
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"os/exec"
+	"sync"
+	"sync/atomic"
+)
+
+var _ = fmt.Print
+
+type worker struct {
+	cmd        *exec.Cmd
+	stdin_pipe io.WriteCloser
+}
+
+var worker_started atomic.Bool
+
+var get_worker = sync.OnceValues(func() (*worker, error) {
+	exe, err := os.Executable()
+	if err != nil {
+		return nil, err
+	}
+	cmd := exec.Command(exe, "__atexit__")
+	cmd.Stdout = nil
+	cmd.Stderr = os.Stderr
+	ans := worker{cmd: cmd}
+	si, err := cmd.StdinPipe()
+	if err != nil {
+		return nil, err
+	}
+	ans.stdin_pipe = si
+	if err = cmd.Run(); err != nil {
+		return nil, err
+	}
+	worker_started.Store(true)
+	return &ans, nil
+})
+
+func WaitForAtexitWorkerToFinish() {
+	if worker_started.Load() {
+		if w, err := get_worker(); err == nil {
+			w.stdin_pipe.Close()
+			_ = w.cmd.Wait()
+		}
+	}
+}
+
+func register(prefix, path string) error {
+	if w, err := get_worker(); err == nil {
+		_, err = fmt.Fprintln(w.stdin_pipe, prefix+" "+path)
+		return err
+	} else {
+		return err
+	}
+
+}
+
+func AtExitUnlink(path string) error {
+	return register("unlink", path)
+}
+
+func AtExitShmUnlink(path string) error {
+	return register("shm_unlink", path)
+}
+
+func AtExitRmtree(path string) error {
+	return register("rmtree", path)
+}