diff --git a/news/zipapp-exec-signals.fixed.md b/news/zipapp-exec-signals.fixed.md new file mode 100644 index 0000000000..d4eafb2704 --- /dev/null +++ b/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)) diff --git a/python/private/zipapp/zip_shell_template.sh b/python/private/zipapp/zip_shell_template.sh index d79331444a..47f85d75d4 100644 --- a/python/private/zipapp/zip_shell_template.sh +++ b/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 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[@]}" diff --git a/tests/py_zipapp/BUILD.bazel b/tests/py_zipapp/BUILD.bazel index e0f878cbe1..ede19baff8 100644 --- a/tests/py_zipapp/BUILD.bazel +++ b/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, +) diff --git a/tests/py_zipapp/signals_main.py b/tests/py_zipapp/signals_main.py new file mode 100644 index 0000000000..1c5be788c2 --- /dev/null +++ b/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() diff --git a/tests/py_zipapp/zipapp_signals_test.py b/tests/py_zipapp/zipapp_signals_test.py new file mode 100644 index 0000000000..4831c09150 --- /dev/null +++ b/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()