From bc2b50a8705deee0ddb02bcd837ee59611fb2165 Mon Sep 17 00:00:00 2001 From: Sergei Firsov <75792801+Wayfarer545@users.noreply.github.com> Date: Wed, 4 Nov 2026 12:11:18 +0100 Subject: [PATCH] fix(worker): use selector event loop for ZeroMQ workers on Windows --- taskiq/cli/worker/run.py | 15 ++- tests/cli/worker/test_worker_event_loop.py | 122 +++++++++++++++++++++ 2 files changed, 136 insertions(+), 1 deletion(-) create mode 100644 tests/cli/worker/test_worker_event_loop.py diff --git a/taskiq/cli/worker/run.py b/taskiq/cli/worker/run.py index 1443163d..397e2cc6 100644 --- a/taskiq/cli/worker/run.py +++ b/taskiq/cli/worker/run.py @@ -10,6 +10,7 @@ from typing import Any from taskiq.abc.broker import AsyncBroker +from taskiq.brokers.zmq_broker import ZeroMQBroker from taskiq.cli.utils import import_object, import_tasks from taskiq.cli.worker.args import WorkerArgs from taskiq.cli.worker.process_manager import ProcessManager @@ -70,7 +71,7 @@ def get_receiver_type(args: WorkerArgs) -> type[Receiver]: return receiver_type -def start_listen(args: WorkerArgs) -> None: +def start_listen(args: WorkerArgs) -> None: # noqa: C901, PLR0915 """ This function starts actual listening process. @@ -145,6 +146,18 @@ def interrupt_handler(signum: int, _frame: Any) -> None: "or pass broker factory function that returns an AsyncBroker instance.", ) + if ( + sys.platform == "win32" + and isinstance(loop, asyncio.ProactorEventLoop) + and isinstance(broker, ZeroMQBroker) + ): + # Without tornado, pyzmq can't wait for sockets on ProactorEventLoop. + # The broker type is known only after the import, so the loop is replaced + # here to keep subprocess support on ProactorEventLoop for other brokers. + loop.close() + loop = asyncio.SelectorEventLoop() + asyncio.set_event_loop(loop) + broker.is_worker_process = True import_tasks(args.modules, args.tasks_pattern, args.fs_discover) diff --git a/tests/cli/worker/test_worker_event_loop.py b/tests/cli/worker/test_worker_event_loop.py new file mode 100644 index 00000000..f288a9e9 --- /dev/null +++ b/tests/cli/worker/test_worker_event_loop.py @@ -0,0 +1,122 @@ +import subprocess +import sys +import threading +import time +from pathlib import Path + +import pytest +import zmq + +from taskiq import InMemoryBroker +from taskiq.message import TaskiqMessage + +ZMQ_APP = """ +from taskiq import ZeroMQBroker + +broker = ZeroMQBroker(zmq_sub_host="tcp://127.0.0.1:{port}") + + +@broker.task(task_name="ping") +async def ping() -> None: + print("task executed", flush=True) +""" + +OTHER_BROKER_APP = """ +import asyncio + +from taskiq import AsyncBroker + + +class ReportingBroker(AsyncBroker): + async def kick(self, message): + pass + + async def listen(self): + print(type(asyncio.get_running_loop()).__name__, flush=True) + return + yield + + +broker = ReportingBroker() +""" + +WORKER_RUNNER = """ +from taskiq.cli.worker.args import WorkerArgs +from taskiq.cli.worker.run import start_listen + +start_listen(WorkerArgs(broker="app:broker", modules=[])) +""" + +WORKER_TIMEOUT = 30 + + +def test_zmq_worker_executes_published_task(tmp_path: Path) -> None: + """ + A worker with ZeroMQBroker receives and executes tasks. + + The default event loop on Windows can't wait for ZeroMQ sockets. + """ + message = InMemoryBroker().formatter.dumps( + TaskiqMessage( + task_id="1", + task_name="ping", + labels={}, + args=[], + kwargs={}, + ), + ) + with zmq.Context() as context, context.socket(zmq.PUB) as publisher: + publisher.setsockopt(zmq.LINGER, 0) + port = publisher.bind_to_random_port("tcp://127.0.0.1") + (tmp_path / "app.py").write_text(ZMQ_APP.format(port=port)) + (tmp_path / "run_worker.py").write_text(WORKER_RUNNER) + process = subprocess.Popen( + [sys.executable, "run_worker.py"], + cwd=tmp_path, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + output: list[str] = [] + + def read_output() -> None: + assert process.stdout is not None + for line in process.stdout: + output.append(line) + + reader = threading.Thread(target=read_output, daemon=True) + reader.start() + try: + deadline = time.monotonic() + WORKER_TIMEOUT + # A subscriber misses everything published before it connects. + while ( + time.monotonic() < deadline + and process.poll() is None + and not any("task executed" in line for line in output) + ): + publisher.send(message.message) + time.sleep(0.1) + finally: + process.kill() + process.wait() + reader.join() + + assert any("task executed" in line for line in output), "".join(output) + + +@pytest.mark.skipif(sys.platform != "win32", reason="Only Windows has Proactor loop.") +def test_other_brokers_keep_default_loop_on_windows(tmp_path: Path) -> None: + """Only ZeroMQBroker replaces the loop, others may need subprocess support.""" + (tmp_path / "app.py").write_text(OTHER_BROKER_APP) + (tmp_path / "run_worker.py").write_text(WORKER_RUNNER) + + result = subprocess.run( + [sys.executable, "run_worker.py"], + cwd=tmp_path, + capture_output=True, + text=True, + timeout=WORKER_TIMEOUT, + check=False, + ) + + assert "ProactorEventLoop" in result.stdout, result.stdout + result.stderr