diff --git a/docs/available-components/brokers.md b/docs/available-components/brokers.md index 5ba91dbb..3fbff8b4 100644 --- a/docs/available-components/brokers.md +++ b/docs/available-components/brokers.md @@ -80,6 +80,163 @@ await my_task.kicker().with_broker(broker).kiq() ::: +## MultiBroker + +This broker combines multiple real brokers into one, so a single worker +can consume messages from all of them at the same time. +Every inner broker keeps its own connection and queue. The `MultiBroker` +only merges the listening and routes kicks. + +This is useful when queues live in different clusters or even different +technologies (two RabbitMQ clusters, RabbitMQ and Kafka, etc.), but you +want a single worker deployment processing all of them. + +### Usage + +Tasks can be registered on inner brokers directly: + +```python +from taskiq import MultiBroker +from taskiq_aio_pika import AioPikaBroker, Queue + +fast_broker = AioPikaBroker( + "amqp://guest:guest@localhost:5672/", + task_queues=[Queue(name="fast_tasks")], +) +slow_broker = AioPikaBroker( + "amqp://guest:guest@localhost:5672/", + task_queues=[Queue(name="slow_tasks")], +) +broker = MultiBroker(fast_broker, slow_broker) +``` + +```python +# tasks.py +from my_project.broker import fast_broker, slow_broker + +@fast_broker.task +async def fast_task() -> None: + ... + +@slow_broker.task +async def slow_task() -> None: + ... +``` + +With this setup a single worker command consumes messages from both queues: + +```bash +taskiq worker my_project.broker:broker my_project.tasks +``` + +It's also possible to register tasks on the `MultiBroker` itself. +In this case pass a `routing` mapping and mark tasks with the `queue_name` label: + +```python +broker = MultiBroker( + fast_broker, + slow_broker, + routing={"fast": fast_broker, "slow": slow_broker}, +) + +@broker.task(queue_name="slow") +async def slow_task() -> None: + ... +``` + +Messages without the routing label are kicked using the default broker +(the first one, unless `default_broker` is passed). + +### How it works + +- Every inner broker is a full broker with its own connection, queue + declarations and settings (for example, its own `qos` when using + taskiq-aio-pika). The composite opens one connection per inner broker + in every worker process. +- `listen()` of all inner brokers run concurrently and are merged into + a single stream for the worker. There are no ordering or fairness + guarantees between queues: messages are processed as they arrive. +- Tasks registered on inner brokers are kicked directly by their broker; + the composite is not involved. Only tasks registered on the + `MultiBroker` itself are routed by the `queue_name` label. +- Startup and shutdown are delegated to all inner brokers. A failing + shutdown of one broker is logged and doesn't prevent the others + from stopping. The `is_worker_process` and `is_scheduler_process` + flags propagate, so inner brokers fire their WORKER_* events. +- Task lookup aggregates registries of all inner brokers, so tasks can + be registered on any of them and the worker finds them all. +- On shutdown the merged listening and all inner generators are closed + deterministically: every inner broker runs its cleanup logic, and + running tasks are awaited (bounded by `--wait-tasks-timeout`). + +### Failure handling + +The MultiBroker fails soft: an error of one broker doesn't stop the others. + +- **Connection flaps are handled by brokers themselves.** For example, + taskiq-aio-pika uses aio-pika's RobustConnection and reconnects + automatically. Such errors never reach the MultiBroker and change + nothing in the worker. +- **If an inner broker's listening fails anyway** (an unrecoverable + error, or a dead cluster for brokers without built-in reconnection), + the error is logged with a traceback, the failed broker is excluded + from listening, and the worker keeps consuming the remaining brokers. +- **Exclusion is permanent for the lifetime of the worker process.** + The MultiBroker doesn't re-subscribe failed brokers in place. Restart + the worker (or let the supervisor do it) to re-enable all brokers; + messages accumulated in the recovered queues are picked up after that. +- **If every inner broker is down**, the listening ends and the worker + process stops. The supervisor restarts it: with the default + `--max-fails -1` it restarts indefinitely until the brokers come back. +- **In-flight messages keep processing.** Messages already handed over + to the worker before the failure complete normally. Messages that + stayed in the failed broker follow that broker's delivery semantics + (see the guarantees table below). +- **Kicks are not affected by worker-side degradation.** Exclusion only + stops consuming. If the target cluster is down, a kick fails on the + caller side with the inner broker's own error. +- **Task-level errors don't degrade anything.** As with any taskiq + broker, an exception inside a task is logged and saved to the result + backend, the message is acknowledged, and the worker keeps running. + +### Guarantees + +The composite only merges brokers. All delivery guarantees come from +the inner brokers: + +| Aspect | MultiBroker | Inner brokers | +| --- | --- | --- | +| One worker consuming all queues | yes | | +| Kick routing by `queue_name` label | yes | | +| Partial degradation on broker failure | yes | | +| Graceful shutdown of all brokers | yes | | +| Reconnection of a failed broker | no | broker's own (aio-pika reconnects; Redis list queues don't) | +| Delivery guarantee (at-least/at-most once) | no | broker's own: RabbitMQ requeues unacked messages; Redis list queues lose popped messages on crash | +| Message redelivery, dead-lettering | no | broker's own | +| Per-queue concurrency limits | no | configured per inner broker (e.g. `qos` of taskiq-aio-pika) | +| Cross-queue ordering or fairness | no (messages arrive interleaved) | per-queue ordering per broker | + +::: caution Be careful! +The worker side uses formatters and serializers of the `MultiBroker`, +while kicks go through inner brokers. Default settings are compatible, +but if you customize them, make sure they match on both sides. + +Also, middlewares of inner brokers work only on the kick side, +and middlewares of the `MultiBroker` — only on the execution side. + +Use `MultiBroker.with_result_backend()` to set a result backend: +it propagates the backend to all inner brokers, keeping client-side +polling and worker-side saving in sync. +::: + +If all of your queues live in a single RabbitMQ broker, prefer the native +multi-queue support of [taskiq-aio-pika](https://github.com/taskiq-python/taskiq-aio-pika) +(it listens to all `task_queues` over a single connection). +The `MultiBroker` is useful when queues are served by different brokers, +or by brokers without native multi-queue support. Note that it opens +one connection per inner broker in every worker process. + + ## Custom brokers These brokers are not part of the core Taskiq lib. But they are maintained by Taskiq developers. diff --git a/taskiq/__init__.py b/taskiq/__init__.py index 1a764bc4..0c2d1244 100644 --- a/taskiq/__init__.py +++ b/taskiq/__init__.py @@ -11,6 +11,7 @@ from taskiq.abc.schedule_source import ScheduleSource from taskiq.acks import AckableMessage, AcknowledgeType from taskiq.brokers.inmemory_broker import InMemoryBroker +from taskiq.brokers.multi_broker import MultiBroker from taskiq.brokers.shared_broker import async_shared_broker from taskiq.brokers.zmq_broker import ZeroMQBroker from taskiq.context import Context @@ -49,6 +50,7 @@ "BrokerMessage", "Context", "InMemoryBroker", + "MultiBroker", "NoResultError", "PrometheusMiddleware", "ResultGetError", diff --git a/taskiq/brokers/multi_broker.py b/taskiq/brokers/multi_broker.py new file mode 100644 index 00000000..bbd35c50 --- /dev/null +++ b/taskiq/brokers/multi_broker.py @@ -0,0 +1,282 @@ +import contextlib +import sys +from collections.abc import AsyncGenerator +from logging import getLogger +from typing import TYPE_CHECKING, Any + +from taskiq.abc.broker import AckableMessage, AsyncBroker +from taskiq.decor import AsyncTaskiqDecoratedTask +from taskiq.message import BrokerMessage +from taskiq.utils import merge_async_iterables + +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing_extensions import Self + +if TYPE_CHECKING: # pragma: no cover + from taskiq.abc.result_backend import AsyncResultBackend + +logger = getLogger(__name__) + + +class MultiBroker(AsyncBroker): + """ + Broker that combines multiple brokers into one. + + It wraps several real brokers and allows a single worker + to consume messages from all of them at the same time. + Every inner broker keeps its own connection and queue, + this broker only merges the listening and routes kicks. + + If an inner broker fails, it's excluded from listening + and the worker keeps consuming the remaining brokers. + The exclusion lasts until the worker process restarts. + When all inner brokers are down, the worker stops + and is restarted by the supervisor. + + Example of running one worker for multiple queues: + + >>> fast_broker = AioPikaBroker(url, queue_name="fast") + >>> slow_broker = AioPikaBroker(url, queue_name="slow") + >>> broker = MultiBroker(fast_broker, slow_broker) + + Tasks must be registered on inner brokers in this case, + and a single `taskiq worker` command will consume + messages from both queues. + + It's also possible to register tasks on the MultiBroker + itself. To route such tasks, pass a `routing` mapping and + mark tasks with the `queue_name` label: + + >>> broker = MultiBroker( + ... fast_broker, + ... slow_broker, + ... routing={"fast": fast_broker, "slow": slow_broker}, + ... ) + + :param brokers: brokers to combine, at least one is required. + :param default_broker: broker used to kick messages + without a routing label, defaults to the first broker. + :param routing: mapping of routing label values to brokers. + :param routing_label: message label used for routing, defaults to `queue_name`. + :raises ValueError: if no brokers are provided, + or any of them has a wrong type. + """ + + def __init__( + self, + *brokers: AsyncBroker, + default_broker: AsyncBroker | None = None, + routing: "dict[str, AsyncBroker] | None" = None, + routing_label: str = "queue_name", + ) -> None: + if not brokers: + raise ValueError("MultiBroker requires at least one broker.") + for inner_broker in brokers: + if not isinstance(inner_broker, AsyncBroker): + raise ValueError( + "All brokers must be AsyncBroker instances, " + f"got {type(inner_broker).__name__}.", + ) + if default_broker is not None and not isinstance( + default_broker, + AsyncBroker, + ): + raise ValueError( + "default_broker must be an AsyncBroker instance, " + f"got {type(default_broker).__name__}.", + ) + # Inner brokers must be saved before calling super().__init__(), + # since it sets the is_worker_process flag, + # which propagates to all inner brokers. + self.brokers: list[AsyncBroker] = list(brokers) + self.default_broker = default_broker or self.brokers[0] + self.routing: dict[str, AsyncBroker] = routing or {} + self.routing_label = routing_label + super().__init__() + + @property + def is_worker_process(self) -> bool: + """True only if broker runs in worker process.""" + return self._is_worker_process + + @is_worker_process.setter + def is_worker_process(self, value: bool) -> None: + self._is_worker_process = value + for inner_broker in self.brokers: + inner_broker.is_worker_process = value + + @property + def is_scheduler_process(self) -> bool: + """True only if broker runs in scheduler process.""" + return self._is_scheduler_process + + @is_scheduler_process.setter + def is_scheduler_process(self, value: bool) -> None: + self._is_scheduler_process = value + for inner_broker in self.brokers: + inner_broker.is_scheduler_process = value + + async def kick(self, message: BrokerMessage) -> None: + """ + Kick a message using one of the inner brokers. + + If the message has a routing label and a routing mapping + is configured, then the broker mapped to the label value + is used. In any other case the default broker is used. + + :param message: message to send. + :raises ValueError: if the routing label value + has no mapped broker. + """ + label_value = message.labels.get(self.routing_label) + if label_value is not None and self.routing: + target_broker = self.routing.get(label_value) + if target_broker is None: + raise ValueError( + f"There is no broker mapped for queue '{label_value}'. " + f"Known queues: {sorted(self.routing)}.", + ) + else: + target_broker = self.default_broker + await target_broker.kick(message) + + async def listen(self) -> AsyncGenerator["bytes | AckableMessage", None]: + """ + Listen to all inner brokers simultaneously. + + This method merges listen generators of all inner brokers + into a single generator, so the worker receives messages + from every broker at the same time. + + If one of the inner brokers fails, it's excluded from + listening, but all remaining brokers keep working. Only when + every inner broker is down, this generator ends and the worker + restarts. + + :yield: incoming messages from all inner brokers. + :return: nothing. + """ + merged = merge_async_iterables( + *(self._listen_one(inner_broker) for inner_broker in self.brokers), + ) + async for message in merged: + yield message + + async def _listen_one( + self, + inner_broker: AsyncBroker, + ) -> AsyncGenerator["bytes | AckableMessage", None]: + """ + Listen to a single inner broker, degrading on its failure. + + If the inner broker's listen fails, the error is logged + and this generator ends, so the failed broker is excluded + from the merged listening, but doesn't stop the whole worker. + + :param inner_broker: broker to listen to. + :yield: incoming messages from the inner broker. + :return: nothing. + """ + iterator = inner_broker.listen() + try: + async for message in iterator: + yield message + except Exception as exc: + logger.exception( + "Inner broker %s failed: %r. " + "It's excluded from listening, " + "continuing with remaining brokers.", + inner_broker, + exc, + ) + finally: + # Catches both the failure case and the outer shutdown: + # a suspended generator must run its cleanup logic. + with contextlib.suppress(RuntimeError): + await iterator.aclose() + + async def startup(self) -> None: + """ + Startup the composite broker and all inner brokers. + + Event handlers and middlewares of this broker are called + first, then every inner broker is started. + """ + await super().startup() + for inner_broker in self.brokers: + await inner_broker.startup() + + async def shutdown(self) -> None: + """ + Shut down all inner brokers and the composite broker. + + Exceptions raised by inner brokers are logged + and don't prevent other brokers from shutting down. + """ + for inner_broker in self.brokers: + try: + await inner_broker.shutdown() + except Exception as exc: + logger.warning( + "Cannot shutdown broker %s: %s", + inner_broker, + exc, + exc_info=True, + ) + await super().shutdown() + + def find_task( + self, + task_name: str, + ) -> "AsyncTaskiqDecoratedTask[Any, Any] | None": + """ + Search for a task in the composite broker and inner brokers. + + :param task_name: name of a task. + :returns: found task or None. + """ + task = self.local_task_registry.get(task_name) + if task is not None: + return task + for inner_broker in self.brokers: + task = inner_broker.find_task(task_name) + if task is not None: + return task + return self.global_task_registry.get(task_name) + + def get_all_tasks(self) -> "dict[str, AsyncTaskiqDecoratedTask[Any, Any]]": + """ + Fetch all tasks available in the composite and inner brokers. + + :return: dict of all tasks. Keys are task names, + values are tasks. + """ + all_tasks = { + **self.global_task_registry, + **self.local_task_registry, + } + for inner_broker in self.brokers: + all_tasks.update(inner_broker.get_all_tasks()) + return all_tasks + + def with_result_backend( + self, + result_backend: "AsyncResultBackend[Any]", + ) -> Self: + """ + Set a result backend for the composite and all inner brokers. + + Inner brokers use their result backend to serve + `wait_result` calls on the client side, while the worker + saves results using the result backend of this broker. + Setting the same backend everywhere keeps them in sync. + + :param result_backend: new result backend. + :return: self + """ + self.result_backend = result_backend + for inner_broker in self.brokers: + inner_broker.result_backend = result_backend + return self diff --git a/taskiq/utils.py b/taskiq/utils.py index 725adc1f..1604c3ea 100644 --- a/taskiq/utils.py +++ b/taskiq/utils.py @@ -1,5 +1,7 @@ +import asyncio +import contextlib import inspect -from collections.abc import Awaitable, Coroutine +from collections.abc import AsyncGenerator, Awaitable, Coroutine from types import CoroutineType from typing import Any, TypeVar, Union @@ -29,3 +31,69 @@ async def maybe_awaitable( if inspect.isawaitable(possible_coroutine): return await possible_coroutine return possible_coroutine # type: ignore + + +async def merge_async_iterables( + *iterables: AsyncGenerator[_T, None], +) -> AsyncGenerator[_T, None]: + """ + Merge multiple async generators into a single one. + + All values from every generator are yielded + as soon as they arrive, in no particular order. + + If any of the inner generators raises an exception, + this exception is re-raised for the consumer. + + When the merged generator is closed or cancelled, + all inner generators are closed as well. + + :param iterables: async generators to merge. + :return: single generator with values from all generators. + """ + if not iterables: + return + + queue: asyncio.Queue[Any] = asyncio.Queue() + pumps = [ + asyncio.create_task(_pump_iterable(iterable, queue)) for iterable in iterables + ] + stopped_count = 0 + try: + while stopped_count < len(iterables): + value = await queue.get() + if value is _MERGE_STOP: + stopped_count += 1 + continue + if isinstance(value, BaseException): + raise value + yield value + finally: + for pump_task in pumps: + pump_task.cancel() + await asyncio.gather(*pumps, return_exceptions=True) + for iterable in iterables: + # A generator is already closed if its pump was + # cancelled in the middle of iteration. For others + # this triggers their cleanup logic. + with contextlib.suppress(RuntimeError): + await iterable.aclose() + + +_MERGE_STOP = object() + + +async def _pump_iterable( + iterable: AsyncGenerator[_T, None], + queue: "asyncio.Queue[Any]", +) -> None: + """Forward all values from a single iterable to the shared queue.""" + try: + async for value in iterable: + await queue.put(value) + except asyncio.CancelledError: + raise + except BaseException as exc: + await queue.put(exc) + else: + await queue.put(_MERGE_STOP) diff --git a/tests/brokers/test_multi_broker.py b/tests/brokers/test_multi_broker.py new file mode 100644 index 00000000..3a973824 --- /dev/null +++ b/tests/brokers/test_multi_broker.py @@ -0,0 +1,364 @@ +import asyncio +from collections.abc import AsyncGenerator + +import pytest + +from taskiq import MultiBroker +from taskiq.acks import AckableMessage +from taskiq.brokers.inmemory_broker import InmemoryResultBackend +from taskiq.events import TaskiqEvents +from taskiq.message import BrokerMessage +from taskiq.receiver import Receiver +from tests.utils import AsyncQueueBroker + + +def build_message(labels: dict[str, str] | None = None) -> BrokerMessage: + """Build a minimal broker message for kick tests.""" + return BrokerMessage( + task_id="task-id", + task_name="tests:test_task", + message=b"payload", + labels=labels or {}, + ) + + +def test_multi_broker_requires_at_least_one_broker() -> None: + with pytest.raises(ValueError, match="at least one broker"): + MultiBroker() + + +def test_multi_broker_validates_broker_types() -> None: + with pytest.raises(ValueError, match="AsyncBroker"): + MultiBroker(42) # type: ignore[arg-type] + + +def test_multi_broker_validates_default_broker_type() -> None: + with pytest.raises(ValueError, match="AsyncBroker"): + MultiBroker(AsyncQueueBroker(), default_broker=42) # type: ignore[arg-type] + + +async def test_kick_uses_first_broker_as_default() -> None: + first, second = AsyncQueueBroker(), AsyncQueueBroker() + broker = MultiBroker(first, second) + + await broker.kick(build_message()) + + assert first.queue.qsize() == 1 + assert second.queue.qsize() == 0 + + +async def test_kick_uses_explicit_default_broker() -> None: + first, second = AsyncQueueBroker(), AsyncQueueBroker() + broker = MultiBroker(first, second, default_broker=second) + + await broker.kick(build_message()) + + assert first.queue.qsize() == 0 + assert second.queue.qsize() == 1 + + +async def test_kick_routes_by_queue_name_label() -> None: + first, second = AsyncQueueBroker(), AsyncQueueBroker() + broker = MultiBroker(first, second, routing={"slow": second}) + + await broker.kick(build_message(labels={"queue_name": "slow"})) + + assert first.queue.qsize() == 0 + assert second.queue.qsize() == 1 + + +async def test_kick_without_label_goes_to_default_broker() -> None: + first, second = AsyncQueueBroker(), AsyncQueueBroker() + broker = MultiBroker(first, second, routing={"slow": second}) + + await broker.kick(build_message()) + + assert first.queue.qsize() == 1 + assert second.queue.qsize() == 0 + + +async def test_kick_ignores_label_when_routing_not_configured() -> None: + first, second = AsyncQueueBroker(), AsyncQueueBroker() + broker = MultiBroker(first, second) + + await broker.kick(build_message(labels={"queue_name": "slow"})) + + assert first.queue.qsize() == 1 + assert second.queue.qsize() == 0 + + +async def test_kick_with_unknown_queue_name_raises() -> None: + first, second = AsyncQueueBroker(), AsyncQueueBroker() + broker = MultiBroker(first, second, routing={"slow": second}) + + with pytest.raises(ValueError, match="no broker mapped for queue 'fast'"): + await broker.kick(build_message(labels={"queue_name": "fast"})) + + +async def test_kick_uses_custom_routing_label() -> None: + first, second = AsyncQueueBroker(), AsyncQueueBroker() + broker = MultiBroker( + first, + second, + routing={"emails": second}, + routing_label="topic", + ) + + await broker.kick(build_message(labels={"topic": "emails"})) + + assert first.queue.qsize() == 0 + assert second.queue.qsize() == 1 + + +async def test_listen_yields_messages_from_all_brokers() -> None: + first, second = AsyncQueueBroker(), AsyncQueueBroker() + broker = MultiBroker(first, second) + await first.kick(build_message()) + await second.kick(build_message()) + await second.kick(build_message()) + + received = 0 + iterator = broker.listen() + for _ in range(3): + await asyncio.wait_for(iterator.__anext__(), timeout=1) + received += 1 + await iterator.aclose() + + assert received == 3 + assert first.queue.qsize() == 0 + assert second.queue.qsize() == 0 + + +async def test_listen_with_single_broker() -> None: + first = AsyncQueueBroker() + broker = MultiBroker(first) + await first.kick(build_message()) + + iterator = broker.listen() + await asyncio.wait_for(iterator.__anext__(), timeout=1) + await iterator.aclose() + + assert first.queue.qsize() == 0 + + +def test_process_flags_propagate_to_inner_brokers() -> None: + first, second = AsyncQueueBroker(), AsyncQueueBroker() + broker = MultiBroker(first, second) + + broker.is_worker_process = True + assert first.is_worker_process is True + assert second.is_worker_process is True + + broker.is_worker_process = False + assert first.is_worker_process is False + + broker.is_scheduler_process = True + assert first.is_scheduler_process is True + assert second.is_scheduler_process is True + + +async def test_startup_delegates_to_inner_brokers() -> None: + started: list[str] = [] + first, second = AsyncQueueBroker(), AsyncQueueBroker() + first.add_event_handler( + TaskiqEvents.WORKER_STARTUP, + lambda _: started.append("first"), + ) + second.add_event_handler( + TaskiqEvents.WORKER_STARTUP, + lambda _: started.append("second"), + ) + broker = MultiBroker(first, second) + broker.add_event_handler( + TaskiqEvents.WORKER_STARTUP, + lambda _: started.append("composite"), + ) + broker.is_worker_process = True + + await broker.startup() + + assert sorted(started) == ["composite", "first", "second"] + + +async def test_shutdown_delegates_to_inner_brokers() -> None: + shut_down: list[str] = [] + first, second = AsyncQueueBroker(), AsyncQueueBroker() + first.add_event_handler( + TaskiqEvents.WORKER_SHUTDOWN, + lambda _: shut_down.append("first"), + ) + second.add_event_handler( + TaskiqEvents.WORKER_SHUTDOWN, + lambda _: shut_down.append("second"), + ) + broker = MultiBroker(first, second) + broker.add_event_handler( + TaskiqEvents.WORKER_SHUTDOWN, + lambda _: shut_down.append("composite"), + ) + broker.is_worker_process = True + + await broker.shutdown() + + assert sorted(shut_down) == ["composite", "first", "second"] + + +async def test_shutdown_tolerates_failing_inner_brokers() -> None: + shut_down: list[str] = [] + + class FailingBroker(AsyncQueueBroker): + async def shutdown(self) -> None: + raise RuntimeError("shutdown failed") + + failing = FailingBroker() + ok_broker = AsyncQueueBroker() + ok_broker.add_event_handler( + TaskiqEvents.CLIENT_SHUTDOWN, + lambda _: shut_down.append("ok"), + ) + broker = MultiBroker(failing, ok_broker) + broker.add_event_handler( + TaskiqEvents.CLIENT_SHUTDOWN, + lambda _: shut_down.append("composite"), + ) + + await broker.shutdown() + + assert "ok" in shut_down + assert "composite" in shut_down + + +async def test_find_task_searches_inner_brokers() -> None: + first, second = AsyncQueueBroker(), AsyncQueueBroker() + + @first.task + async def task_one() -> None: ... + + @second.task + async def task_two() -> None: ... + + broker = MultiBroker(first, second) + + assert broker.find_task(task_one.task_name) is task_one + assert broker.find_task(task_two.task_name) is task_two + assert broker.find_task("missing:task") is None + + +async def test_find_task_finds_tasks_registered_on_composite() -> None: + broker = MultiBroker(AsyncQueueBroker()) + + @broker.task + async def composite_task() -> None: ... + + assert broker.find_task(composite_task.task_name) is composite_task + + +async def test_get_all_tasks_aggregates_inner_brokers() -> None: + first, second = AsyncQueueBroker(), AsyncQueueBroker() + + @first.task + async def task_one() -> None: ... + + @second.task + async def task_two() -> None: ... + + broker = MultiBroker(first, second) + + all_tasks = broker.get_all_tasks() + assert all_tasks[task_one.task_name] is task_one + assert all_tasks[task_two.task_name] is task_two + + +async def test_with_result_backend_propagates_to_inner_brokers() -> None: + first, second = AsyncQueueBroker(), AsyncQueueBroker() + broker = MultiBroker(first, second) + backend = InmemoryResultBackend[None]() + + broker.with_result_backend(backend) + + assert broker.result_backend is backend + assert first.result_backend is backend + assert second.result_backend is backend + + +async def test_receiver_processes_tasks_from_all_brokers() -> None: + first, second = AsyncQueueBroker(), AsyncQueueBroker() + broker = MultiBroker(first, second) + broker.with_result_backend(InmemoryResultBackend()) + executed: list[str] = [] + + @first.task + async def task_one() -> str: + executed.append("one") + return "one" + + @second.task + async def task_two() -> str: + executed.append("two") + return "two" + + first_handle = await task_one.kiq() + second_handle = await task_two.kiq() + + receiver = Receiver(broker=broker) + finish_event = asyncio.Event() + listener = asyncio.create_task(receiver.listen(finish_event)) + + first_result = await first_handle.wait_result(timeout=5) + second_result = await second_handle.wait_result(timeout=5) + + assert first_result.return_value == "one" + assert second_result.return_value == "two" + assert sorted(executed) == ["one", "two"] + + finish_event.set() + await asyncio.wait_for(listener, timeout=5) + await broker.shutdown() + + +class DyingBroker(AsyncQueueBroker): + """Broker that yields N messages and then fails like a dead cluster.""" + + def __init__(self, messages_before_death: int) -> None: + super().__init__() + self.messages_before_death = messages_before_death + + async def listen(self) -> AsyncGenerator[AckableMessage, None]: + while self.messages_before_death > 0: + self.messages_before_death -= 1 + data = await self.queue.get() + yield AckableMessage(data=data, ack=self.queue.task_done) + raise RuntimeError("cluster is down") + + +async def test_failed_broker_does_not_stop_healthy_brokers( + caplog: pytest.LogCaptureFixture, +) -> None: + dying = DyingBroker(messages_before_death=2) + healthy = AsyncQueueBroker() + broker = MultiBroker(dying, healthy) + + for _ in range(2): + await dying.kick(build_message()) + for _ in range(3): + await healthy.kick(build_message()) + + received = 0 + iterator = broker.listen() + for _ in range(5): + await asyncio.wait_for(iterator.__anext__(), timeout=1) + received += 1 + await iterator.aclose() + + assert received == 5 + assert "cluster is down" in caplog.text + + +async def test_listen_ends_when_all_inner_brokers_fail() -> None: + broker = MultiBroker( + DyingBroker(messages_before_death=0), + DyingBroker(messages_before_death=0), + ) + + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(broker.listen().__anext__(), timeout=1) diff --git a/tests/test_utils.py b/tests/test_utils.py index 485ab319..cb89e609 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,4 +1,10 @@ -from taskiq.utils import maybe_awaitable +import asyncio +from collections.abc import AsyncGenerator +from typing import cast + +import pytest + +from taskiq.utils import maybe_awaitable, merge_async_iterables async def test_maybe_awaitable_coroutine() -> None: @@ -14,3 +20,58 @@ def meme() -> int: return 1 assert await maybe_awaitable(meme()) == 1 + + +async def test_merge_async_iterables_merges_all_iterables() -> None: + async def gen(*items: int) -> AsyncGenerator[int, None]: + for item in items: + await asyncio.sleep(0) + yield item + + merged = [ + item + async for item in merge_async_iterables( + gen(1, 2, 3), + gen(4, 5, 6), + ) + ] + + assert sorted(merged) == [1, 2, 3, 4, 5, 6] + + +async def test_merge_async_iterables_without_iterables() -> None: + # The TypeVar cannot be inferred from an empty argument list, + # so the return type is provided explicitly. + merged = cast("AsyncGenerator[int, None]", merge_async_iterables()) + assert [item async for item in merged] == [] + + +async def test_merge_async_iterables_propagates_inner_errors() -> None: + async def failing() -> AsyncGenerator[int, None]: + yield 1 + raise ValueError("inner generator failed") + + with pytest.raises(ValueError, match="inner generator failed"): + async for _ in merge_async_iterables(failing()): + pass + + +async def test_merge_async_iterables_closes_inner_iterables() -> None: + closed: list[str] = [] + + def tracking(name: str) -> AsyncGenerator[int, None]: + async def gen() -> AsyncGenerator[int, None]: + try: + while True: + yield 1 + await asyncio.sleep(0) + finally: + closed.append(name) + + return gen() + + merged = merge_async_iterables(tracking("first"), tracking("second")) + await merged.__anext__() + await merged.aclose() + + assert sorted(closed) == ["first", "second"]