fix(zipapp): exec the interpreter so signals reach the Python process

news/zipapp-exec-signals.fixed.md

@@ -0,0 +1,6 @@
+(zipapp) The self-executable {obj}`py_zipapp_binary` launcher now `exec`s the
+Python interpreter, so signals sent to the PID the caller started (e.g. by a
+process supervisor) reach the Python program instead of only the launcher
+shell. Its temporary extraction directory is now also removed when the program
+is killed with `SIGKILL`.
+([#2043](https://github.com/bazel-contrib/rules_python/issues/2043))

python/private/zipapp/zip_shell_template.sh

@@ -35,6 +35,7 @@ if [[ -n "${RULES_PYTHON_ADDITIONAL_INTERPRETER_ARGS}" ]]; then
 fi
 
 
+cleanup_zip_dir=""
 if [[ -n "$RULES_PYTHON_EXTRACT_ROOT" ]]; then
   zip_dir="$RULES_PYTHON_EXTRACT_ROOT/$EXTRACT_DIR/$ZIP_HASH"
   if [[ ! -e "$zip_dir/__main__.py" ]]; then
@@ -49,6 +50,7 @@ else
   # Unzip emits a warning and exits 1 with the prelude
   ( unzip -q -d "$zip_dir" "$0" 2>/dev/null || true )
   if [[ -n "$zip_dir" && -z "${RULES_PYTHON_BOOTSTRAP_VERBOSE:-}" ]]; then
+    cleanup_zip_dir=1
     trap 'rm -fr "$zip_dir"' EXIT
   fi
 fi
@@ -77,12 +79,56 @@ command=(
   "$@"
 )
 
-# NOTE: because exec isn't used, signals don't propagate to the child
-# TODO: Use exec and let the program handle cleanup. Without exec,
-# signals don't propagate to the child nicely.
+# Prints an identifier for a running process that changes if the PID is reused,
+# or nothing once the process has exited. A zombie counts as exited, so it
+# doesn't matter whether the caller has reaped it yet.
+process_identity() {
+  local stat fields
+  if [[ -r "/proc/$1/stat" ]]; then
+    stat=$(< "/proc/$1/stat") || return 0
+    # The command name is in parentheses and may contain spaces, so split the
+    # fields after the last ")": field 3 (state) through field 22 (start time).
+    read -r -a fields <<< "${stat##*) }"
+    if [[ "${fields[0]}" != "Z" ]]; then
+      echo "${fields[19]}"
+    fi
+  elif command -v ps >/dev/null 2>&1; then
+    stat=$(ps -o stat=,lstart= -p "$1" 2>/dev/null)
+    stat="${stat#"${stat%%[![:space:]]*}"}"
+    if [[ -n "$stat" && "$stat" != Z* ]]; then
+      echo "${stat#* }"
+    fi
+  elif kill -0 "$1" 2>/dev/null; then
+    echo "running"
+  fi
+}
+
+if [[ -n "$cleanup_zip_dir" ]]; then
+  # exec replaces this shell, so the EXIT trap can't remove the extracted files.
+  # Instead, start a watcher that removes them once this PID, which becomes the
+  # Python process, exits. Unlike the trap, this also works when the process is
+  # killed with SIGKILL. The watcher is double-forked so it isn't a child of the
+  # Python process, and it doesn't hold this process's stdio open. It stays in
+  # the caller's process group, so it ignores the signals sent to a whole group
+  # (e.g. Ctrl-C in a terminal) and exits on its own once the process is gone.
+  # Unlike the implicit SIGINT ignore for background jobs, `trap ''` also
+  # applies to the commands it runs.
+  trap - EXIT
+  launcher_pid=$$
+  launcher_identity=$(process_identity "$launcher_pid")
+  ( (
+    trap '' INT QUIT HUP TERM
+    while [[ "$(process_identity "$launcher_pid")" == "$launcher_identity" ]]; do
+      sleep 1
+    done
+    rm -fr "$zip_dir"
+  ) </dev/null >/dev/null 2>&1 & )
+fi
+
+# We use `exec` instead of a child process so that signals sent directly (e.g.
+# using `kill`) to this process (the PID seen by the calling process) are
+# received by the Python process. Otherwise, this process receives the signal
+# and the Python process keeps running.
 # See https://github.com/bazel-contrib/rules_python/issues/2043#issuecomment-2215469971
 # for more information.
-"${command[@]}"
-# Explicit exit is needed because the implicit next line the zip file this
-# template is prepended to.
-exit 0
+exec "${command[@]}"

tests/py_zipapp/BUILD.bazel

@@ -4,6 +4,7 @@ load("//python:py_library.bzl", "py_library")
 load("//python:py_test.bzl", "py_test")
 load("//python/private:bzlmod_enabled.bzl", "BZLMOD_ENABLED")  # buildifier: disable=bzl-visibility
 load("//python/zipapp:py_zipapp_binary.bzl", "py_zipapp_binary")
+load("//tests/support:support.bzl", "SUPPORTS_BOOTSTRAP_SCRIPT")
 
 py_binary(
     name = "venv_bin",
@@ -120,3 +121,29 @@ py_library(
     experimental_venvs_site_packages = "//python/config_settings:venvs_site_packages",
     imports = ["site-packages"],
 )
+
+py_binary(
+    name = "signals_bin",
+    srcs = ["signals_main.py"],
+    config_settings = {
+        "//python/config_settings:bootstrap_impl": "script",
+    },
+    main = "signals_main.py",
+    target_compatible_with = SUPPORTS_BOOTSTRAP_SCRIPT,
+)
+
+py_zipapp_binary(
+    name = "signals_zipapp",
+    binary = ":signals_bin",
+    target_compatible_with = SUPPORTS_BOOTSTRAP_SCRIPT,
+)
+
+py_test(
+    name = "zipapp_signals_test",
+    srcs = ["zipapp_signals_test.py"],
+    data = [":signals_zipapp"],
+    env = {
+        "TEST_ZIPAPP": "$(location :signals_zipapp)",
+    },
+    target_compatible_with = SUPPORTS_BOOTSTRAP_SCRIPT,
+)

tests/py_zipapp/signals_main.py

@@ -0,0 +1,26 @@
+"A zipapp that reports its PID and extraction directory, then waits for SIGTERM."
+
+import os
+import signal
+import sys
+import time
+
+
+def main():
+    if len(sys.argv) > 1 and sys.argv[1].startswith("--exit-code="):
+        sys.exit(int(sys.argv[1].split("=", 1)[1]))
+
+    def on_sigterm(signum, frame):
+        # print() isn't reentrant; the signal may arrive while main() prints.
+        os.write(sys.stdout.fileno(), b"got SIGTERM\n")
+        sys.exit(0)
+
+    signal.signal(signal.SIGTERM, on_sigterm)
+    print(f"pid={os.getpid()}", flush=True)
+    print(f"zip_dir={sys._xoptions.get('RULES_PYTHON_ZIP_DIR', '')}", flush=True)
+    while True:
+        time.sleep(60)
+
+
+if __name__ == "__main__":
+    main()

tests/py_zipapp/zipapp_signals_test.py

@@ -0,0 +1,108 @@
+import os
+import signal
+import subprocess
+import time
+import unittest
+
+# How long to wait for the extracted files to be removed after the zipapp exits.
+_CLEANUP_TIMEOUT_SECONDS = 15
+
+
+class ZipAppSignalsTest(unittest.TestCase):
+    def setUp(self):
+        self.zipapp_path = os.environ["TEST_ZIPAPP"]
+
+    def _start(self, new_session=False):
+        proc = subprocess.Popen(
+            [self.zipapp_path],
+            stdout=subprocess.PIPE,
+            text=True,
+            start_new_session=new_session,
+        )
+        self.addCleanup(self._kill, proc)
+        assert proc.stdout is not None
+        pid_line = proc.stdout.readline().strip()
+        zip_dir_line = proc.stdout.readline().strip()
+        if not pid_line.startswith("pid=") or not zip_dir_line.startswith("zip_dir="):
+            self.fail(f"unexpected zipapp output: {pid_line!r} {zip_dir_line!r}")
+        pid = int(pid_line.split("=", 1)[1])
+        if pid != proc.pid:
+            # The launcher didn't exec, so the Python process can outlive it and
+            # keep stdout open. Make sure it's killed too.
+            self.addCleanup(self._kill_pid, pid)
+        zip_dir = zip_dir_line.split("=", 1)[1]
+        self.assertTrue(os.path.isdir(zip_dir), f"{zip_dir} does not exist")
+        return proc, pid, zip_dir
+
+    def _kill(self, proc):
+        if proc.poll() is None:
+            proc.kill()
+        try:
+            proc.communicate(timeout=10)
+        except subprocess.TimeoutExpired:
+            pass
+
+    def _kill_pid(self, pid):
+        try:
+            os.kill(pid, signal.SIGKILL)
+        except ProcessLookupError:
+            pass
+
+    def assertRemovedEventually(self, path):
+        deadline = time.monotonic() + _CLEANUP_TIMEOUT_SECONDS
+        while os.path.exists(path):
+            if time.monotonic() > deadline:
+                self.fail(f"{path} was not removed after the zipapp exited")
+            time.sleep(0.2)
+
+    def test_signal_reaches_python_process(self):
+        proc, pid, zip_dir = self._start()
+        # The launcher must exec the interpreter, so the PID the caller started
+        # is the Python process and signals sent to it are delivered there.
+        self.assertEqual(pid, proc.pid)
+
+        proc.send_signal(signal.SIGTERM)
+        output, _ = proc.communicate(timeout=30)
+
+        self.assertIn("got SIGTERM", output)
+        self.assertEqual(proc.returncode, 0)
+        self.assertRemovedEventually(zip_dir)
+
+    def test_extracted_files_removed_after_sigkill(self):
+        proc, _, zip_dir = self._start()
+
+        proc.kill()
+        proc.communicate(timeout=30)
+
+        self.assertRemovedEventually(zip_dir)
+
+    def test_extracted_files_removed_before_process_is_reaped(self):
+        proc, _, zip_dir = self._start()
+
+        # Leave the exited process as a zombie: cleanup must not wait for the
+        # caller to reap it.
+        os.kill(proc.pid, signal.SIGKILL)
+
+        self.assertRemovedEventually(zip_dir)
+
+    def test_extracted_files_removed_after_process_group_signal(self):
+        # A terminal's Ctrl-C (SIGINT), hangup (SIGHUP) or a supervisor stopping
+        # the whole process group (SIGTERM) signals every process in the group,
+        # including the cleanup watcher, which must survive to do its job.
+        for sig in (signal.SIGINT, signal.SIGHUP, signal.SIGTERM):
+            with self.subTest(signal=sig.name):
+                proc, _, zip_dir = self._start(new_session=True)
+
+                os.killpg(proc.pid, sig)
+                proc.communicate(timeout=30)
+
+                self.assertRemovedEventually(zip_dir)
+
+    def test_exit_code_is_propagated(self):
+        result = subprocess.run([self.zipapp_path, "--exit-code=3"], timeout=60)
+
+        self.assertEqual(result.returncode, 3)
+
+
+if __name__ == "__main__":
+    unittest.main()