fix(cli): correct --hardkill-count off-by-one in worker signal handler

taskiq/cli/worker/run.py

@@ -110,7 +110,7 @@ def interrupt_handler(signum: int, _frame: Any) -> None:
         shutdown_event.set()
         # Hard kill is a signal that we should stop
         # everything immediately.
-        if hardkill_counter > args.hardkill_count:
+        if hardkill_counter >= args.hardkill_count:
             logger.warning("Hard kill. Exiting.")
             raise KeyboardInterrupt
         hardkill_counter += 1

tests/cli/worker/hanging_broker.py

@@ -0,0 +1,35 @@
+"""Broker module used by test_hardkill_count subprocess tests.
+
+The broker never yields messages and its graceful shutdown blocks on an
+event that is never set, so a running worker can only exit via the
+hard-kill path (KeyboardInterrupt raised from the signal handler).
+"""
+
+import asyncio
+from collections.abc import AsyncGenerator
+
+from taskiq.abc.broker import AsyncBroker
+from taskiq.message import BrokerMessage
+
+
+class HangingBroker(AsyncBroker):
+    """Broker that idles forever and never finishes graceful shutdown."""
+
+    async def kick(self, message: BrokerMessage) -> None:
+        """No-op kick."""
+
+    def listen(self) -> AsyncGenerator[bytes, None]:
+        """Return an async generator that never yields."""
+
+        async def _gen() -> AsyncGenerator[bytes, None]:
+            await asyncio.Event().wait()
+            yield b""  # pragma: no cover
+
+        return _gen()
+
+    async def shutdown(self) -> None:
+        """Block forever, simulating a graceful shutdown that hangs."""
+        await asyncio.Event().wait()
+
+
+broker = HangingBroker()

tests/cli/worker/test_hardkill_count.py

@@ -0,0 +1,119 @@
+"""Integration test for --hardkill-count signal semantics.
+
+Spawns a real worker subprocess whose broker never finishes graceful
+shutdown, sends SIGINTs directly to a worker child process (the process
+that installs the hard-kill counting signal handler), and checks how
+many signals it takes to hard-kill it.
+
+Documented behavior (``--hardkill-count N``): N termination signals are
+allowed before performing a hardkill, i.e. signals 1..N only set the
+shutdown event and signal N+1 raises KeyboardInterrupt in the worker.
+"""
+
+import os
+import signal
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import psutil
+import pytest
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+SIGNAL_SETTLE_SECONDS = 0.6
+WAIT_TIMEOUT = 15.0
+
+
+def _pid_alive(pid: int) -> bool:
+    """Check whether a process exists."""
+    try:
+        os.kill(pid, 0)
+    except OSError:
+        return False
+    return True
+
+
+def _worker_child_pids(main_pid: int) -> list[int]:
+    """Return pids of worker child processes of the main process."""
+    try:
+        parent = psutil.Process(main_pid)
+    except psutil.NoSuchProcess:
+        return []
+    return [child.pid for child in parent.children(recursive=False)]
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32",
+    reason="SIGINT hard-kill semantics differ on Windows",
+)
+def test_hardkill_count_signals_before_hardkill() -> None:
+    """N signals are graceful; signal N+1 hard-kills the worker process."""
+    hardkill_count = 3
+    env = os.environ.copy()
+    env["PYTHONPATH"] = str(REPO_ROOT) + os.pathsep + env.get("PYTHONPATH", "")
+    proc = subprocess.Popen(  # noqa: S603
+        [
+            sys.executable,
+            "-m",
+            "taskiq",
+            "worker",
+            "tests.cli.worker.hanging_broker:broker",
+            "--hardkill-count",
+            str(hardkill_count),
+            "--workers",
+            "1",
+            "--no-configure-logging",
+            # Keep graceful shutdown pending so only the hard-kill
+            # path can terminate the worker during the test.
+            "--shutdown-timeout",
+            "3600",
+        ],
+        cwd=REPO_ROOT,
+        env=env,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+    try:
+        # Wait for the worker child to spawn and install its handlers.
+        deadline = time.monotonic() + WAIT_TIMEOUT
+        worker_pid: int | None = None
+        while time.monotonic() < deadline:
+            if proc.poll() is not None:
+                pytest.fail(
+                    f"Worker exited during startup: returncode={proc.returncode}",
+                )
+            children = _worker_child_pids(proc.pid)
+            if children:
+                worker_pid = children[0]
+                break
+            time.sleep(0.25)
+        if worker_pid is None:
+            pytest.fail("No worker child process spawned in time")
+        # Extra settle time so the child's signal handlers are installed.
+        time.sleep(2.0)
+
+        # Signals 1..N: only set the shutdown event; the broker's graceful
+        # shutdown hangs, so the worker process must still be alive.
+        for signal_number in range(1, hardkill_count + 1):
+            os.kill(worker_pid, signal.SIGINT)
+            time.sleep(SIGNAL_SETTLE_SECONDS)
+            assert _pid_alive(worker_pid), (
+                f"Worker exited on signal {signal_number} of "
+                f"{hardkill_count} allowed graceful signals"
+            )
+
+        # Signal N+1: hard kill raises KeyboardInterrupt in the worker.
+        os.kill(worker_pid, signal.SIGINT)
+        deadline = time.monotonic() + WAIT_TIMEOUT
+        while time.monotonic() < deadline and _pid_alive(worker_pid):
+            time.sleep(0.1)
+        assert not _pid_alive(worker_pid), (
+            f"Worker still alive after {hardkill_count + 1} signals; "
+            f"--hardkill-count {hardkill_count} should hard-kill on "
+            f"signal {hardkill_count + 1}"
+        )
+    finally:
+        if proc.poll() is None:
+            proc.kill()
+            proc.wait(timeout=10)