diff --git a/justfile b/justfile index 0cba2f34..4daf7a91 100644 --- a/justfile +++ b/justfile @@ -53,7 +53,7 @@ test-integration *TEST_ARGS: build-bin build-python cd libs/opsqueue_python source "./.setup_local_venv.sh" - timeout {{pytest_timeout_seconds}} pytest --color=yes {{TEST_ARGS}} + timeout --signal term --kill-after 2s {{pytest_timeout_seconds}} pytest --color=yes {{TEST_ARGS}} # Python integration test suite, using artefacts built through Nix. Args are forwarded to pytest [group('nix')] @@ -68,7 +68,7 @@ nix-test-integration *TEST_ARGS: nix-build export OPSQUEUE_BIN="${nix_build_bin_dir}/bin/opsqueue" export RUST_LOG="opsqueue=debug" - timeout {{pytest_timeout_seconds}} pytest --color=yes {{TEST_ARGS}} + timeout --signal term --kill-after 2s {{pytest_timeout_seconds}} pytest --color=yes {{TEST_ARGS}} # Run all linters, fast and slow [group('lint')] diff --git a/libs/opsqueue_python/examples/multiprocessing/multiprocessing_consumer.py b/libs/opsqueue_python/examples/multiprocessing/multiprocessing_consumer.py index 792ea8ac..d99b6472 100644 --- a/libs/opsqueue_python/examples/multiprocessing/multiprocessing_consumer.py +++ b/libs/opsqueue_python/examples/multiprocessing/multiprocessing_consumer.py @@ -19,7 +19,7 @@ def my_operation(data: int) -> int: def main() -> None: n_consumers = 16 processes = [ - multiprocessing.Process(target=run_a_consumer, args=(id,)) + multiprocessing.Process(target=run_a_consumer, args=(id,), daemon=True) for id in range(n_consumers) ] for p in processes: diff --git a/libs/opsqueue_python/examples/multiprocessing/multiprocessing_consumer_preferdistinct.py b/libs/opsqueue_python/examples/multiprocessing/multiprocessing_consumer_preferdistinct.py index bc35f7c3..94aa4794 100644 --- a/libs/opsqueue_python/examples/multiprocessing/multiprocessing_consumer_preferdistinct.py +++ b/libs/opsqueue_python/examples/multiprocessing/multiprocessing_consumer_preferdistinct.py @@ -25,7 +25,7 @@ def my_operation(data: int) -> int: def main() -> None: n_consumers = 4 processes = [ - multiprocessing.Process(target=run_a_consumer, args=(id,)) + multiprocessing.Process(target=run_a_consumer, args=(id,), daemon=True) for id in range(n_consumers) ] for p in processes: diff --git a/libs/opsqueue_python/examples/multiprocessing/multiprocessing_producer_many.py b/libs/opsqueue_python/examples/multiprocessing/multiprocessing_producer_many.py index a91698f7..17ae4777 100644 --- a/libs/opsqueue_python/examples/multiprocessing/multiprocessing_producer_many.py +++ b/libs/opsqueue_python/examples/multiprocessing/multiprocessing_producer_many.py @@ -22,7 +22,7 @@ def main() -> None: f"Starting {n_producers} producers... When `multiprocessing_consumer_preferdistinct.py` is used, expecting all submissions to finish roughly at the same time" ) processes = [ - multiprocessing.Process(target=run_a_producer, args=(id,)) + multiprocessing.Process(target=run_a_producer, args=(id,), daemon=True) for id in range(n_producers) ] for p in processes: diff --git a/libs/opsqueue_python/tests/conftest.py b/libs/opsqueue_python/tests/conftest.py index 0052bca1..4a867a00 100644 --- a/libs/opsqueue_python/tests/conftest.py +++ b/libs/opsqueue_python/tests/conftest.py @@ -1,5 +1,8 @@ import cbor2 import pickle +import faulthandler +import signal +import sys from contextlib import contextmanager, ExitStack from typing import Generator, Callable, Any, Iterable import multiprocess # type: ignore[import-untyped] @@ -14,9 +17,142 @@ from opsqueue.consumer import Strategy +def _stack_dump_path(pid: int) -> Path: + return Path(f"/tmp/opsqueue-pytest-stack-{pid}.log") + + +def _linux_descendant_pids(parent_pid: int = os.getpid()) -> dict[int, tuple[str, str]]: + """ + Returns a dictionary mapping descendant PIDs to their command name and command line. This is a + recursive function that will find all descendants of the given parent PID. The command name and + command line are read from the `/proc` filesystem, which is Linux-specific. + """ + children = ( + Path(f"/proc/{parent_pid}/task/{parent_pid}/children") + .read_text(encoding="utf-8") + .strip() + ) + + def _child_command(pid: int) -> tuple[str, str]: + """ + Returns the command name and command line of the given process id (pid). + """ + try: + comm = ( + Path(f"/proc/{pid}/comm").read_text(encoding="utf-8").strip() + or "" + ) + except FileNotFoundError: + # Process might have exited between reading /proc//children and now. + comm = "" + try: + cmdline = ( + Path(f"/proc/{pid}/cmdline") + .read_text(encoding="utf-8") + .replace( + "\x00", " " + ) # Replace null bytes with spaces before stripping. + .strip() + or "" + ) + except FileNotFoundError: + # Process might have exited between reading /proc//children and now. + cmdline = "" + + return (comm, cmdline) + + pids = [int(pid_text) for pid_text in children.split()] + result = {pid: _child_command(pid) for pid in pids} + for pid in pids: + result.update(_linux_descendant_pids(pid)) + return result + + +def _handle_sigterm(signum: int, frame: object) -> None: + all_descendant_pids = _linux_descendant_pids(os.getpid()) + + # Only signal processes that registered the SIGUSR1 handler (i.e. have a dump file). + # Internal multiprocess infrastructure processes (resource tracker, forkserver) do not + # register the handler and have no dump file. Sending SIGUSR1 to them with the default + # disposition would kill them, which is undesirable. + worker_pids = { + pid: info + for pid, info in all_descendant_pids.items() + if _stack_dump_path(pid).exists() + } + + if worker_pids: + print( + f"[opsqueue pytest] Requesting stack dump from child workers: {list(worker_pids.keys())}", + file=sys.stderr, + flush=True, + ) + for pid in worker_pids.keys(): + try: + os.kill(pid, signal.SIGUSR1) + except ProcessLookupError: + # Worker already exited. + continue + except Exception as exc: # pragma: no cover - defensive diagnostics + print( + f"[opsqueue pytest] Failed sending SIGUSR1 to child pid={pid}: {exc}", + file=sys.stderr, + flush=True, + ) + else: + print( + "[opsqueue pytest] No child worker PIDs found at SIGTERM time", + file=sys.stderr, + flush=True, + ) + + # Dump the stack of the controller process itself, so we can see what it was doing at the time + # of SIGTERM. This also provides the background processes enough time to dump their stacks + # before the controller accesses those. + faulthandler.dump_traceback(file=sys.stderr, all_threads=True) + + for pid in sorted(worker_pids.keys()): + dump_path = _stack_dump_path(pid) + print( + f"[opsqueue pytest] ===== BEGIN CHILD STACK DUMP pid={pid} {worker_pids[pid]} ({dump_path}) =====", + file=sys.stderr, + flush=True, + ) + try: + print(dump_path.read_text(encoding="utf-8"), file=sys.stderr, flush=True) + except Exception as exc: # pragma: no cover - defensive diagnostics + print( + f"[opsqueue pytest] Failed reading dump file for child pid={pid} {worker_pids[pid]}: {exc}", + file=sys.stderr, + flush=True, + ) + print( + f"[opsqueue pytest] ===== END CHILD STACK DUMP pid={pid} {worker_pids[pid]} =====", + file=sys.stderr, + flush=True, + ) + + print( + "[opsqueue pytest] ===== SIGTERM: End of diagnostic output =====", + file=sys.stderr, + flush=True, + ) + os._exit(1) + + +# Register signal handlers at module import time (before pytest_configure hook runs). +# Make SIGUSR1 print a traceback in any process (controller or worker). +faulthandler.register( + signal.SIGUSR1, + file=_stack_dump_path(os.getpid()).open("w", encoding="utf-8"), + all_threads=True, + chain=False, +) +signal.signal(signal.SIGTERM, _handle_sigterm) + + @pytest.hookimpl(tryfirst=True) def pytest_configure(config: pytest.Config) -> None: - print("A") multiprocess.set_start_method("forkserver") @@ -142,18 +278,41 @@ def read_exact_fd(fd: int, num_bytes: int) -> bytes: return bytes(data) +def _background_main(function: Callable[..., None], args: Iterable[Any]) -> None: + """Entry point for background_processes. + + Registers the SIGUSR1 stack-dump handler before delegating to the real + target function, so that the controller can request a traceback from this + process when a timeout fires. + """ + faulthandler.register( + signal.SIGUSR1, + file=_stack_dump_path(os.getpid()).open("w", encoding="utf-8"), + all_threads=True, + chain=False, + ) + function(*args) + + @contextmanager def background_process( function: Callable[..., None], args: Iterable[Any] = (), ) -> Generator[multiprocess.Process, None, None]: - proc = multiprocess.Process(target=function, args=args) + proc = multiprocess.Process( + target=_background_main, + args=(function, args), + daemon=True, + ) try: - proc.daemon = True proc.start() yield proc finally: proc.terminate() + try: + proc.join(timeout=1.0) + except multiprocess.TimeoutError: + proc.kill() @contextmanager