[nkhitrov]

Problem

Core taskiq binds one worker process to a single broker with a single
listen() generator. Running multiple queues today means either multiple
worker deployments or relying on per-broker multi-queue support (which
only some broker packages have). Multi-queue support in core is still an
open discussion (TRFC-2, org discussion #448), with draft PRs #486/#624.

Solution

Additive MultiBroker: a composite broker wrapping N real brokers
(any AsyncBroker implementations — taskiq-aio-pika, taskiq-redis,
taskiq-aio-kafka, etc.) so a single taskiq worker consumes all of
them. No changes to the CLI, the Receiver, or external packages.

fast = AioPikaBroker(url, task_queues=[Queue(name="fast")])
slow = AioPikaBroker(url, task_queues=[Queue(name="slow")])
broker = MultiBroker(fast, slow)
# taskiq worker proj.broker:broker proj.tasks  ->  one worker, two queues

What it does

• listen() merges listen-generators of all inner brokers into a single
  stream (new merge_async_iterables util in taskiq.utils);
• kick() routes by the queue_name label when a routing mapping is
  configured, otherwise uses the default broker;
• startup()/shutdown() delegate to all inner brokers; one failing
  shutdown doesn't prevent the rest from stopping;
• is_worker_process / is_scheduler_process propagate, so inner
  brokers fire their WORKER_* lifecycle events;
• find_task() / get_all_tasks() aggregate task registries, so tasks
  may be registered on any inner broker;
• with_result_backend() propagates the backend to inner brokers to keep
  client-side polling and worker-side saving in sync.

Notes

• For queues all living in a single RabbitMQ broker, native
  task_queues of taskiq-aio-pika >= 0.6.0 keeps a single connection
  and remains preferable; MultiBroker targets the multi-broker /
  multi-cluster case (documented, including the one-connection-per-broker
  trade-off).
• This is intentionally narrower than TRFC-2 (no Queue dataclass / ABC
  contract changes) — it solves the "one worker, many queues/brokers"
  use case additively and can serve as a building block for the full
  TRFC-2 design.

[codecov]

Codecov Report

❌ Patch coverage is 98.44961% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 86.02%. Comparing base (9214379) to head (69d56b1).

Files with missing lines	Patch %	Lines
taskiq/brokers/multi_broker.py	97.87%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #723      +/-   ##
==========================================
+ Coverage   85.49%   86.02%   +0.52%     
==========================================
  Files          68       69       +1     
  Lines        2662     2790     +128     
==========================================
+ Hits         2276     2400     +124     
- Misses        386      390       +4     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[s3rius]

I don't think we need to override those task search methods.

The reason is that if you use this multibroker, you use it with the task decorator, so all tasks will appear in the multibroker's task registry.

With this implementation, the code after line 242 is never expected to be reached. Also, this implementation depends quite a lot on the order of the brokers for instance, if you define a task with the same name on different brokers, only the first broker will be used.

[s3rius]

Suggested change

	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__}.",
	)
	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__}.",
	)
	inner_broker.task = lamba: raise ValueError("Use multibroker instead for registering tasks on indifvidual brokers")

This guardrail might help with remove find_task and get_all_tasks.

[s3rius]

There might be additional validation that this broker is one of brokers. Or, we can remove this parameter in favor or default_routing. This will guarantee that broker is present.

[s3rius]

Since you have routing already, you can remove brokers and default_broker arguments and make signature of something like that:

Suggested change

	*brokers: AsyncBroker,
	default_broker: AsyncBroker | None = None,
	routing: "dict[str, AsyncBroker] | None" = None,
	routing_label: str = "queue_name",
	routing: "dict[str, AsyncBroker] | None" = None,
	default_routing_label: str | None = None,
	routing_label: str = "queue_name",

[s3rius]

We can simplify this pump function using anyio, since it's in our dependencies.

[s3rius]

Such things better define at the top.

[s3rius]

I think that there is a simpler solution:

Suggested change

	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)
	import asyncio
	from collections.abc import AsyncGenerator
	from typing import TypeVar
	import anyio
	from anyio.streams.memory import MemoryObjectSendStream
	_T = TypeVar("_T")
	async def _start_producer(
	stream: AsyncGenerator[_T, None],
	send_stream: MemoryObjectSendStream[_T],
	) -> None:
	async with send_stream:
	async for val in stream:
	await send_stream.send(val)
	async def merge_async_iterables(
	*streams: AsyncGenerator[_T, None],
	) -> AsyncGenerator[_T, None]:
	send_stream, receive_stream = anyio.create_memory_object_stream[_T]()
	async with anyio.create_task_group() as tg, send_stream, receive_stream:
	for stream in streams:
	tg.start_soon(_start_producer, stream, send_stream.clone())
	# We close the original send stream,
	# to only have copies. When all of them are closed,
	# the stream will end.
	send_stream.close()
	async with receive_stream:
	async for message in receive_stream:
	yield message

It uses anyio primitives and should be a bit more robust. Also, we don't have to track manually any sentinel objects and stuff like that.