diff --git a/justfile b/justfile index 7c8ad88e..9803a5c3 100644 --- a/justfile +++ b/justfile @@ -5,6 +5,8 @@ _default: # Timeout (seconds) for Python integration test recipes. # Override with OPSQUEUE_PYTEST_TIMEOUT, e.g. OPSQUEUE_PYTEST_TIMEOUT=600 just nix-test-integration pytest_timeout_seconds := env_var_or_default("OPSQUEUE_PYTEST_TIMEOUT", "60") +# The kill window is chosen to be above the `1s` join window on `multiprocess.Process`. +pytest_timeout_kill_seconds := "2s" # Build-and-run the opsqueue binary (development profile) [group('run')] @@ -53,7 +55,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 {{pytest_timeout_kill_seconds}} {{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 +70,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 {{pytest_timeout_kill_seconds}} {{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..e8c0f7ed 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,171 @@ from opsqueue.consumer import Strategy +def _stack_dump_path(pid: int = os.getpid()) -> 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. + """ + # The `/proc/{pid}/task/{tid}/children` file contains a space-separated list of child PIDs for + # the given task. This interface is also not completely stable while there are running + # processes. It can omit child PIDs if any of the child processes exits while reading the file. + # + # The canonical way is to read all `/proc/{pid}/stat` and reconstruct the tree in memory. As + # this is currently only used for debug logging on SIGTERM, we can accept the risk of missing + # some child PIDs. + 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: + """ + Handle `SIGTERM` by dumping the stack traces of all child processes that have registered the + `SIGUSR1` handler (i.e. have a stack dump file). This is used for debugging tests that hang, as + `pytest` will receive `SIGTERM` on the controller process when `timeout` fires. Do note that we + will receive `SIGKILL` after a few seconds, so dump the stack traces quickly. + """ + 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} {worker_pids[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]} ({dump_path}) =====", + file=sys.stderr, + flush=True, + ) + + print( + "[opsqueue pytest] ===== SIGTERM: End of diagnostic output =====", + file=sys.stderr, + flush=True, + ) + os._exit(1) + + +def register_sigusr1_stack_dump_handler() -> None: + """ + Make SIGUSR1 print a traceback in any process (controller, worker or background process). + + Registering this handler in subprocesses spawned by `multiprocess.Process` is a bit tricky, as + the `multiprocess` module does not provide a hook for this and doesn't run any module imports in + the child processes. We therefore register the handler at `_background_main` call instead. + + This is used for debugging tests that hang, as `timeout` will send SIGTERM to the controller + and the controller will request stack dumps via `SIGUSR1` from all child processes that have + created the stack dump file. The stack dumps will be written to already opened file when the + process receives `SIGUSR1`, and the controller will read those files and print them to `stderr`. + The file path is `/tmp/opsqueue-pytest-stack-.log`. + """ + faulthandler.register( + signal.SIGUSR1, + file=_stack_dump_path().open("w", encoding="utf-8"), + all_threads=True, + chain=False, + ) + + +# Register signal handlers at module import time (before pytest_configure hook runs). +register_sigusr1_stack_dump_handler() +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 +307,37 @@ 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. + """ + register_sigusr1_stack_dump_handler() + 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() # Process can lock-up or ignore the termination signal, kill it. + proc.join() # Process should exit immediate after sending the kill signal. @contextmanager