CI: Stacktraces of pytest and xdist-workers on timeout

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')]

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:

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:

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:

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 "<comm>"
+            )
+        except FileNotFoundError:
+            # Process might have exited between reading /proc/<pid>/children and now.
+            comm = "<comm>"
+        try:
+            cmdline = (
+                Path(f"/proc/{pid}/cmdline")
+                .read_text(encoding="utf-8")
+                .replace(
+                    "\x00", " "
+                )  # Replace null bytes with spaces before stripping.
+                .strip()
+                or "<cmdline>"
+            )
+        except FileNotFoundError:
+            # Process might have exited between reading /proc/<pid>/children and now.
+            cmdline = "<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-<pid>.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