fix(zipapp): forward termination signals

news/3809.fixed.md

@@ -0,0 +1,2 @@
+(zipapp) Forwarded termination signals to the application and preserved its exit status.
+([#3809](https://github.com/bazel-contrib/rules_python/issues/3809))

python/private/python_bootstrap_template.txt

@@ -28,6 +28,7 @@ if (
 
 from os.path import abspath, dirname, join, basename, normpath
 import os
+import signal
 import shutil
 import subprocess
 
@@ -488,6 +489,43 @@ def runfiles_envvar(runfiles_root):
 
   return (None, None)
 
+def _run_subprocess(argv, env, cwd):
+  if IS_WINDOWS:
+    return subprocess.call(argv, env=env, cwd=cwd)
+
+  child = None
+  pending_signals = []
+
+  def forward_signal(signum, _frame):
+    if child is None:
+      pending_signals.append(signum)
+    else:
+      try:
+        # Keep wait() as the sole child reaper. Popen.send_signal() may call
+        # poll(), which can race wait() and lose the child's exit status.
+        os.kill(child.pid, signum)
+      except ProcessLookupError:
+        pass
+
+  previous_handlers = {}
+  for name in ("SIGHUP", "SIGINT", "SIGQUIT", "SIGTERM"):
+    signum = getattr(signal, name, None)
+    if signum is not None:
+      previous_handlers[signum] = signal.signal(signum, forward_signal)
+
+  try:
+    child = subprocess.Popen(argv, env=env, cwd=cwd)
+    for signum in pending_signals:
+      forward_signal(signum, None)
+    ret_code = child.wait()
+  finally:
+    for signum, handler in previous_handlers.items():
+      signal.signal(signum, handler)
+
+  if ret_code < 0:
+    return 128 - ret_code
+  return ret_code
+
 def execute_file(python_program, main_filename, args, env, runfiles_root,
                 workspace, delete_dirs):
   # type: (str, str, list[str], dict[str, str], str, str|None, str|None) -> ...
@@ -535,8 +573,7 @@ def execute_file(python_program, main_filename, args, env, runfiles_root,
   print_verbose("run: subproc: environ:", mapping=os.environ)
   print_verbose("run: subproc: cwd:", workspace)
   print_verbose("run: subproc: argv:", values=argv)
-  ret_code = subprocess.call(
-      argv, env=env, cwd=workspace)
+  ret_code = _run_subprocess(argv, env=env, cwd=workspace)
   print_verbose("run: subproc: exit code:", ret_code)
 
   if delete_dirs:

python/private/stage1_bootstrap_template.sh

@@ -153,7 +153,9 @@ python_exe=$(find_python_interpreter $RUNFILES_DIR $PYTHON_BINARY)
 # Zip files have to re-create the venv bin/python3 symlink because they
 # don't contain it already.
 if [[ "$IS_ZIPFILE" == "1" ]]; then
-  use_exec=0
+  # Stage 2 removes the extracted runfiles through RULES_PYTHON_ZIP_DIR, so
+  # this bootstrap does not need to remain alive for cleanup.
+  use_exec=1
   # It should always be under runfiles, but double check this. We don't
   # want to accidentally create symlinks elsewhere.
   if [[ "$python_exe" != $RUNFILES_DIR/* ]]; then
@@ -333,8 +335,8 @@ command=(
 # for more information.
 #
 # However, we can't use exec when there is cleanup to do afterwards. Control
-# must return to this process so it can run the trap handlers. Such cases
-# occur when zip mode or recreate_venv_at_runtime creates temporary files.
+# must return to this process so it can run the trap handlers. This case
+# occurs when recreate_venv_at_runtime creates a temporary venv.
 if [[ "$use_exec" == "0" ]]; then
   "${command[@]}"
   exit $?

python/private/zipapp/zip_main_template.py

@@ -25,6 +25,7 @@
 
 import os  # noqa: E402
 import shutil  # noqa: E402
+import signal  # noqa: E402
 import stat  # noqa: E402
 import subprocess  # noqa: E402
 import tempfile  # noqa: E402
@@ -234,6 +235,44 @@ def create_runfiles_root():
     return join(extract_root, "runfiles")
 
 
+def run_subprocess(subprocess_argv, env, cwd):
+    if IS_WINDOWS:
+        return subprocess.call(subprocess_argv, env=env, cwd=cwd)
+
+    child = None
+    pending_signals = []
+
+    def forward_signal(signum, _frame):
+        if child is None:
+            pending_signals.append(signum)
+        else:
+            try:
+                # Keep wait() as the sole child reaper. Popen.send_signal() may call
+                # poll(), which can race wait() and lose the child's exit status.
+                os.kill(child.pid, signum)
+            except ProcessLookupError:
+                pass
+
+    previous_handlers = {}
+    for name in ("SIGHUP", "SIGINT", "SIGQUIT", "SIGTERM"):
+        signum = getattr(signal, name, None)
+        if signum is not None:
+            previous_handlers[signum] = signal.signal(signum, forward_signal)
+
+    try:
+        child = subprocess.Popen(subprocess_argv, env=env, cwd=cwd)
+        for signum in pending_signals:
+            forward_signal(signum, None)
+        ret_code = child.wait()
+    finally:
+        for signum, handler in previous_handlers.items():
+            signal.signal(signum, handler)
+
+    if ret_code < 0:
+        return 128 - ret_code
+    return ret_code
+
+
 def execute_file(
     python_program,
     main_filename,
@@ -276,7 +315,7 @@ def execute_file(
         print_verbose("subprocess env:", mapping=env)
         print_verbose("subprocess cwd:", workspace)
         print_verbose("subprocess argv:", values=subprocess_argv)
-        ret_code = subprocess.call(subprocess_argv, env=env, cwd=workspace)
+        ret_code = run_subprocess(subprocess_argv, env=env, cwd=workspace)
         print_verbose("subprocess exit code:", ret_code)
         sys.exit(ret_code)
     finally:

python/private/zipapp/zip_shell_template.sh

@@ -77,12 +77,5 @@ 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.
-# 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
+# The stage 2 bootstrap removes the extracted runfiles when the program exits.
+exec "${command[@]}"

tests/bootstrap_impls/bin.py

@@ -13,8 +13,21 @@
 # limitations under the License.
 
 import os
+import signal
 import sys
 
+if sys.argv[1:] in (["handled"], ["unhandled"]):
+    if sys.argv[1:] == ["handled"]:
+
+        def handle(signum, _frame):
+            print(f"received:{signum}", flush=True)
+
+        signal.signal(signal.SIGTERM, handle)
+
+    print(f"ready:{os.getpid()}", flush=True)
+    signal.pause()
+    raise SystemExit(0)
+
 print("Hello")
 print(
     "RULES_PYTHON_ZIP_DIR:{}".format(sys._xoptions.get("RULES_PYTHON_ZIP_DIR", "UNSET"))

tests/bootstrap_impls/run_binary_zip_yes_test.sh

@@ -42,3 +42,68 @@ if ! (echo "$actual" | grep "$expected_pattern" ) >/dev/null; then
   exit 1
 fi
 
+case "$(uname -s)" in
+  CYGWIN*|MINGW*|MSYS*) exit 0 ;;
+esac
+
+test_dir=$(mktemp -d)
+launcher_pid=""
+application_pid=""
+watchdog_pid=""
+
+cleanup() {
+  if [[ -n "${watchdog_pid}" ]]; then
+    kill "${watchdog_pid}" 2>/dev/null || true
+  fi
+  if [[ -n "${launcher_pid}" ]]; then
+    kill -KILL "${launcher_pid}" 2>/dev/null || true
+  fi
+  if [[ -n "${application_pid}" ]]; then
+    kill -KILL "${application_pid}" 2>/dev/null || true
+  fi
+  rm -rf "${test_dir}"
+}
+trap cleanup EXIT
+
+run_signal_case() {
+  local mode="$1"
+  local expected_exit="$2"
+  local expected_output="$3"
+  local log="${test_dir}/${mode}.log"
+
+  "$bin" "${mode}" >"${log}" 2>&1 &
+  launcher_pid=$!
+  for _ in {1..100}; do
+    if grep -F "ready:" "${log}" >/dev/null; then
+      break
+    fi
+    sleep 0.1
+  done
+  grep -F "ready:" "${log}" >/dev/null || return 1
+  application_pid=$(sed -n 's/^ready://p' "${log}")
+
+  kill -TERM "${launcher_pid}"
+  (
+    sleep 10
+    kill -KILL "${launcher_pid}" "${application_pid}" 2>/dev/null || true
+  ) &
+  watchdog_pid=$!
+  wait "${launcher_pid}"
+  exit_code=$?
+  kill "${watchdog_pid}" 2>/dev/null || true
+  watchdog_pid=""
+
+  if [[ "${exit_code}" != "${expected_exit}" ]]; then
+    echo "expected exit ${expected_exit}, got ${exit_code}" >&2
+    cat "${log}" >&2
+    return 1
+  fi
+  if [[ -n "${expected_output}" ]]; then
+    grep -F "${expected_output}" "${log}" >/dev/null || return 1
+  fi
+  launcher_pid=""
+  application_pid=""
+}
+
+run_signal_case handled 0 "received:15" || exit 1
+run_signal_case unhandled 143 "" || exit 1

tests/py_zipapp/main.py

@@ -1,7 +1,31 @@
-"A trivial zipapp that prints a message"
+"A trivial zipapp that prints a message or waits for a signal."
+
+import os
+import signal
+import sys
+
+
+def wait_for_signal(handle_signal):
+    if handle_signal:
+
+        def handle(signum, _frame):
+            print(f"received:{signum}", flush=True)
+
+        signal.signal(signal.SIGTERM, handle)
+
+    print(f"ready:{os.getpid()}", flush=True)
+    signal.pause()
+    return 0
 
 
 def main():
+    if sys.argv[1:] == ["wait-for-sigterm"]:
+        return wait_for_signal(handle_signal=True)
+    if sys.argv[1:] == ["wait-for-unhandled-sigterm"]:
+        return wait_for_signal(handle_signal=False)
+    if len(sys.argv) == 3 and sys.argv[1] == "exit":
+        return int(sys.argv[2])
+
     print("Hello from zipapp")
     try:
         import some_dep
@@ -12,13 +36,12 @@ def main():
 
         print(f"dep: {pkgdep.pkgmod}")
     except ImportError as e:
-        import sys
-
         e.add_note(
             "Failed to import a dependency.\n" + "sys.path:\n" + "\n".join(sys.path)
         )
         raise
+    return 0
 
 
 if __name__ == "__main__":
-    main()
+    raise SystemExit(main())

tests/py_zipapp/system_python_zipapp_test.py

@@ -1,5 +1,7 @@
 import os
+import signal
 import subprocess
+import sys
 import unittest
 
 
@@ -28,5 +30,79 @@ def test_zipapp_runnable(self):
         self.assertIn("dep:", output)
 
 
+@unittest.skipUnless(os.name == "posix", "POSIX signals are required")
+class PosixSignalZipAppTest(unittest.TestCase):
+    def zipapp_command(self, *args, invoke_with_python=False):
+        zipapp_path = os.environ["TEST_ZIPAPP"]
+        command = [zipapp_path]
+        if invoke_with_python:
+            command.insert(0, sys.executable)
+        return [*command, *args]
+
+    def start_signal_app(self, mode, invoke_with_python):
+        process = subprocess.Popen(
+            self.zipapp_command(mode, invoke_with_python=invoke_with_python),
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+        )
+        self.assertIsNotNone(process.stdout)
+        ready = process.stdout.readline().strip()
+        self.assertTrue(ready.startswith("ready:"), ready)
+        return process, int(ready.removeprefix("ready:"))
+
+    def stop_process(self, process, application_pid):
+        if process.poll() is None:
+            process.kill()
+            process.wait()
+        try:
+            os.kill(application_pid, signal.SIGTERM)
+        except ProcessLookupError:
+            pass
+        if process.stdout is not None:
+            process.stdout.close()
+
+    def test_zipapp_forwards_sigterm(self):
+        for invoke_with_python in (False, True):
+            with self.subTest(invoke_with_python=invoke_with_python):
+                process, application_pid = self.start_signal_app(
+                    "wait-for-sigterm", invoke_with_python
+                )
+                try:
+                    process.terminate()
+                    output, _ = process.communicate(timeout=10)
+                    self.assertEqual(0, process.returncode, output)
+                    self.assertIn(f"received:{signal.SIGTERM}", output)
+                finally:
+                    self.stop_process(process, application_pid)
+
+    def test_zipapp_preserves_signal_termination(self):
+        for invoke_with_python in (False, True):
+            with self.subTest(invoke_with_python=invoke_with_python):
+                process, application_pid = self.start_signal_app(
+                    "wait-for-unhandled-sigterm", invoke_with_python
+                )
+                try:
+                    process.terminate()
+                    output, _ = process.communicate(timeout=10)
+                    expected = (
+                        128 + signal.SIGTERM if invoke_with_python else -signal.SIGTERM
+                    )
+                    self.assertEqual(expected, process.returncode, output)
+                finally:
+                    self.stop_process(process, application_pid)
+
+    def test_zipapp_preserves_nonzero_exit_status(self):
+        for invoke_with_python in (False, True):
+            with self.subTest(invoke_with_python=invoke_with_python):
+                process = subprocess.run(
+                    self.zipapp_command(
+                        "exit", "17", invoke_with_python=invoke_with_python
+                    ),
+                    check=False,
+                )
+                self.assertEqual(17, process.returncode)
+
+
 if __name__ == "__main__":
     unittest.main()