Skip to content

Commit b727f04

Browse files
committed
feat(routing): harden broker separation and multi-broker runtime
- separate routing policy, task registration, flows, and broker transport - support explicit multi-broker workers with shared process limits - preserve legacy decorators, scheduling, requeue, and send semantics - harden shutdown, cancellation, delivery ownership, and shared task binding - add compatibility tests, examples, and migration documentation
1 parent cc6d3a3 commit b727f04

37 files changed

Lines changed: 4155 additions & 403 deletions

docs/available-components/brokers.md

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,12 @@ This is a special broker for local development. It uses the same functions to ex
1212
but all tasks are executed locally in the current thread.
1313
By default it uses `InMemoryResultBackend` but this can be overridden.
1414

15+
Because one `InMemoryBroker` instance acts as both client and worker, its
16+
startup and shutdown run both event-handler phases. Middleware and result
17+
backend lifecycle hooks still run once per broker lifecycle. Shutdown waits for
18+
background in-memory tasks before closing those resources and the sync-task
19+
executor.
20+
1521
## ZeroMQBroker
1622

1723
This broker uses [ZMQ](https://zeromq.org/) to communicate between worker and client processes.
Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
"""Expose an explicit multi-broker listener tuple for the worker CLI."""
2+
3+
import asyncio
4+
from collections.abc import AsyncGenerator
5+
6+
from taskiq import AsyncBroker, BrokerMessage, Flow, TaskiqRouter
7+
8+
9+
class DemoBroker(AsyncBroker):
10+
"""Small in-process transport used to keep this configuration executable."""
11+
12+
def __init__(self, router: TaskiqRouter, broker_name: str) -> None:
13+
super().__init__(router=router, broker_name=broker_name)
14+
self.messages: asyncio.Queue[bytes] = asyncio.Queue()
15+
16+
async def kick(self, message: BrokerMessage) -> None:
17+
await self.messages.put(message.message)
18+
19+
async def listen(self) -> AsyncGenerator[bytes, None]:
20+
while True:
21+
yield await self.messages.get()
22+
23+
24+
router = TaskiqRouter()
25+
26+
commands_broker = DemoBroker(router, "commands")
27+
events_broker = DemoBroker(router, "events")
28+
outbound_broker = DemoBroker(router, "outbound")
29+
30+
31+
@commands_broker.task(task_name="orders.notify")
32+
async def notify_order(order_id: str) -> str:
33+
return f"notified:{order_id}"
34+
35+
36+
@commands_broker.task(task_name="orders.process")
37+
async def process_order(order_id: str) -> str:
38+
await notify_order.kiq(order_id)
39+
return f"processed:{order_id}"
40+
41+
42+
orders_flow = Flow("orders.events")
43+
notifications_flow = Flow("orders.notifications")
44+
router.route_task(process_order, broker=events_broker, flow=orders_flow)
45+
router.route_task(
46+
notify_order,
47+
broker=outbound_broker,
48+
flow=notifications_flow,
49+
)
50+
router.subscribe(events_broker, orders_flow, process_order)
51+
52+
# Only these brokers receive listeners. The outbound broker belongs to the
53+
# same Router, so this explicit multi-broker worker starts its client lifecycle
54+
# for routed child sends without consuming from it.
55+
worker_brokers = (commands_broker, events_broker)

docs/guide/cli.md

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,27 @@ Like this:
1818
taskiq worker mybroker:broker_var my_project.module1 my_project.module2
1919
```
2020

21+
### Multiple brokers
22+
23+
The broker import target may expose a module-level sequence of broker instances
24+
that use one shared Router:
25+
26+
```python
27+
# my_project/worker.py
28+
worker_brokers = (rabbit_broker, kafka_broker)
29+
```
30+
31+
```bash
32+
taskiq worker my_project.worker:worker_brokers my_project.tasks
33+
```
34+
35+
The sequence is explicit listener configuration. Router brokers outside the
36+
sequence are started in client mode for routed sends and are not consumed by
37+
this worker. A single broker target keeps its existing lifecycle. Execution,
38+
prefetch and tasks-per-child limits apply to the whole worker process, not once
39+
per listener. See [Routing and flows](./routing-and-flows.md) for validation,
40+
lifecycle and adapter compatibility details.
41+
2142
### Sync function
2243

2344
Taskiq can run synchronous functions. However, since it operates asynchronously, it executes them in a separate thread or process. By default, **ThreadPoolExecutor** is used. But if you're planning to use Taskiq for heavy computations, such as neural network model training or other CPU-intensive tasks, you may want to use **ProcessPoolExecutor** instead.
@@ -146,7 +167,14 @@ kill -HUP <main pid>
146167
If you send `SIGINT` or `SIGKILL` to the main process by pressing <kbd>Ctrl</kbd>+<kbd>C</kbd> or using the `kill` command, it will initiate the shutdown process.
147168
By default, it will stop fetching new messages immediately after receiving the signal but will wait for the completion of all currently executing tasks.
148169

149-
If you don't want to wait too long for tasks to complete each time you shut down the worker, you can either send termination signals three times to the main process to perform a hard kill or configure the `--wait-tasks-timeout` to set a hard time limit for shutting down.
170+
If you don't want to wait indefinitely for tasks to complete during worker shutdown, you can either send termination signals three times to the main process to perform a hard kill or configure `--wait-tasks-timeout` as the graceful callback boundary.
171+
172+
When `--wait-tasks-timeout` is omitted, Taskiq keeps the legacy behavior and
173+
waits for the complete listener shutdown, including running callbacks, without
174+
a time limit. When the task timeout is configured, `--shutdown-timeout` is also
175+
the post-drain listener cleanup allowance. Every managed broker then receives
176+
its own `--shutdown-timeout` deadline, so total multi-broker cleanup can take
177+
longer than that value.
150178

151179
::: tip Cool tip
152180
The number of signals before a hard kill can be configured with the `--hardkill-count` CLI argument.
@@ -170,8 +198,8 @@ The number of signals before a hard kill can be configured with the `--hardkill-
170198
* `--ack-type` - Type of acknowledgement. This parameter is used to set when to acknowledge the task. Possible values are `when_received`, `when_executed`, `when_saved`, `manual`. Default is `when_saved`.
171199
* `--max-tasks-per-child` - maximum number of tasks to be executed by a single worker process before restart.
172200
* `--max-fails` - Maximum number of child process exits.
173-
* `--shutdown-timeout` - maximum amount of time for graceful broker's shutdown in seconds (default 5).
174-
* `--wait-tasks-timeout` - if cannot read new messages from the broker or maximum number of tasks is reached, worker will wait for all current tasks to finish. This parameter sets the maximum amount of time to wait until shutdown.
201+
* `--shutdown-timeout` - maximum graceful shutdown time for each managed broker in seconds (default 5). With `--wait-tasks-timeout`, the same value is also the listener's post-drain cleanup allowance.
202+
* `--wait-tasks-timeout` - maximum graceful wait for active task callbacks before Taskiq cancels and awaits them. When omitted, callback and listener shutdown remains unbounded for backward compatibility.
175203
* `--hardkill-count` - Number of termination signals to the main process before performing a hardkill.
176204

177205
## Scheduler

0 commit comments

Comments
 (0)