diff --git a/taskiq/cli/worker/run.py b/taskiq/cli/worker/run.py index 1443163d..c8f285bc 100644 --- a/taskiq/cli/worker/run.py +++ b/taskiq/cli/worker/run.py @@ -92,6 +92,14 @@ def start_listen(args: WorkerArgs) -> None: format=args.log_format, ) + if uvloop is not None: + logger.debug("UVLOOP found. Using it as async runner") + loop = uvloop.new_event_loop() # type: ignore + else: + loop = asyncio.new_event_loop() + + asyncio.set_event_loop(loop) + def interrupt_handler(signum: int, _frame: Any) -> None: """ Signal handler. @@ -107,7 +115,7 @@ def interrupt_handler(signum: int, _frame: Any) -> None: nonlocal shutdown_event nonlocal hardkill_counter # Soft kill is a signal to start shutdown. - shutdown_event.set() + loop.call_soon_threadsafe(shutdown_event.set) # Hard kill is a signal that we should stop # everything immediately. if hardkill_counter > args.hardkill_count: @@ -120,14 +128,6 @@ def interrupt_handler(signum: int, _frame: Any) -> None: if sys.platform != "win32": signal.signal(signal.SIGHUP, interrupt_handler) - if uvloop is not None: - logger.debug("UVLOOP found. Using it as async runner") - loop = uvloop.new_event_loop() # type: ignore - else: - loop = asyncio.new_event_loop() - - asyncio.set_event_loop(loop) - # This option signals that current # broker is running as a worker. # We must set this field before importing tasks, diff --git a/tests/cli/worker/test_worker_logging.py b/tests/cli/worker/test_worker_logging.py index d970876e..6d5f8e39 100644 --- a/tests/cli/worker/test_worker_logging.py +++ b/tests/cli/worker/test_worker_logging.py @@ -120,3 +120,43 @@ def test_worker_process_logs_listening_started( assert wait_for_output(process, "Listening started.") finally: stop(process) + + +@pytest.mark.skipif( + sys.platform == "win32", + reason="Worker processes can't be stopped with SIGTERM on Windows.", +) +def test_idle_worker_stops_on_sigterm(tmp_path: Path) -> None: + """ + An idle worker shuts down when the main process gets SIGTERM. + + Without uvloop, the event loop waiting for a message doesn't wake up + on its own, so the worker's signal handler has to wake it up. + """ + (tmp_path / "idle_broker.py").write_text(BROKER_MODULE) + (tmp_path / "run_cli.py").write_text(CLI_RUNNER) + # Shadows uvloop, so the worker uses the default asyncio event loop. + (tmp_path / "uvloop.py").write_text("raise ImportError") + process = subprocess.Popen( + [ + sys.executable, + "run_cli.py", + "default", + "worker", + "idle_broker:broker", + "--workers", + "1", + "--log-level", + "INFO", + ], + cwd=tmp_path, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + try: + assert wait_for_output(process, "Listening started.") + process.send_signal(signal.SIGTERM) + assert process.wait(timeout=10) == 0 + finally: + stop(process)