experiment separate Broker + Router

docs/examples/router/multiple_brokers.py

@@ -0,0 +1,69 @@
+"""Route one task through several brokers with a shared router."""
+
+import asyncio
+
+from taskiq import Flow, InMemoryBroker, TaskiqRouter
+
+router = TaskiqRouter()
+
+default_email_flow = Flow("emails.default")
+priority_email_flow = Flow("emails.priority")
+bulk_email_flow = Flow("emails.bulk")
+
+default_broker = InMemoryBroker(
+    router=router,
+    broker_name="default",
+    default_flow=default_email_flow,
+    await_inplace=True,
+)
+priority_broker = InMemoryBroker(
+    router=router,
+    broker_name="priority",
+    default_flow=priority_email_flow,
+    await_inplace=True,
+)
+
+
+@default_broker.task(task_name="examples.send_email", domain="notifications")
+async def send_email(user_id: int, template: str) -> str:
+    """Pretend to render and send an email."""
+    return f"{template} email sent to user {user_id}"
+
+
+priority_route = router.route_task(
+    send_email,
+    broker=priority_broker,
+    flow=priority_email_flow,
+)
+
+
+async def _main() -> None:
+    await default_broker.startup()
+    await priority_broker.startup()
+    try:
+        direct_result = await send_email(7, "welcome")
+
+        routed_task = await send_email.kiq(7, "welcome")
+        routed_result = await routed_task.wait_result(timeout=2)
+
+        bulk_task = (
+            await send_email.kicker()
+            .with_route(
+                default_broker,
+                bulk_email_flow,
+            )
+            .kiq(8, "digest")
+        )
+        bulk_result = await bulk_task.wait_result(timeout=2)
+
+        print(f"Direct call: {direct_result}")
+        print(f"Declared route: {priority_route.broker_name}")
+        print(f"Routed call: {routed_result.return_value}")
+        print(f"Route override: {bulk_result.return_value}")
+    finally:
+        await priority_broker.shutdown()
+        await default_broker.shutdown()
+
+
+if __name__ == "__main__":
+    asyncio.run(_main())

docs/examples/router/shared_task_package.py

@@ -0,0 +1,74 @@
+"""Declare shared task definitions and bind them in the final application."""
+
+import asyncio
+from collections.abc import Mapping
+from dataclasses import dataclass
+
+from taskiq import Flow, InMemoryBroker, TaskiqRouter, task_builder
+
+
+@dataclass(frozen=True, slots=True)
+class BillingQueue:
+    """Broker-specific flow that follows the shared flow protocol."""
+
+    name: str
+    priority: int
+
+    def broker_options(self, broker_name: str) -> Mapping[str, object]:
+        """Return options that a billing broker adapter can understand."""
+        return {
+            "broker": broker_name,
+            "priority": self.priority,
+        }
+
+
+@task_builder("billing.calculate_total", domain="billing")
+async def calculate_total(price: int, quantity: int) -> int:
+    """Package-level task definition that is not bound to any broker."""
+    return price * quantity
+
+
+router = TaskiqRouter()
+billing_flow = Flow("billing.tasks")
+priority_billing_flow = BillingQueue(name="billing.priority", priority=10)
+
+billing_broker = InMemoryBroker(
+    router=router,
+    broker_name="billing",
+    default_flow=billing_flow,
+    await_inplace=True,
+)
+
+registered_calculate_total = router.register_task(
+    calculate_total,
+    broker=billing_broker,
+    flow=billing_flow,
+)
+
+
+async def _main() -> None:
+    await billing_broker.startup()
+    try:
+        direct_result = await calculate_total.call(19, 3)
+
+        prepared_task = (
+            registered_calculate_total.kicker()
+            .with_route(
+                billing_broker,
+                priority_billing_flow,
+            )
+            .prepare(19, 3)
+        )
+
+        queued_task = await prepared_task.kiq()
+        queued_result = await queued_task.wait_result(timeout=2)
+
+        print(f"Shared task direct call: {direct_result}")
+        print(f"Prepared message: {prepared_task.message.task_name}")
+        print(f"Registered queued call: {queued_result.return_value}")
+    finally:
+        await billing_broker.shutdown()
+
+
+if __name__ == "__main__":
+    asyncio.run(_main())

taskiq/__init__.py

@@ -24,18 +24,22 @@
     TaskiqError,
     TaskiqResultTimeoutError,
 )
+from taskiq.flow import Flow, FlowProtocol
 from taskiq.funcs import gather
+from taskiq.kicker import PreparedKiq
 from taskiq.message import BrokerMessage, TaskiqMessage
 from taskiq.middlewares import (
     PrometheusMiddleware,
     SimpleRetryMiddleware,
     SmartRetryMiddleware,
 )
 from taskiq.result import TaskiqResult
+from taskiq.router import TaskiqRoute, TaskiqRouter
 from taskiq.scheduler.scheduled_task import ScheduledTask
 from taskiq.scheduler.scheduler import TaskiqScheduler
 from taskiq.state import TaskiqState
 from taskiq.task import AsyncTaskiqTask
+from taskiq.task_builder import TaskDefinition, task_builder
 
 __version__ = version("taskiq")
 
@@ -47,8 +51,11 @@
     "AsyncTaskiqTask",
     "BrokerMessage",
     "Context",
+    "Flow",
+    "FlowProtocol",
     "InMemoryBroker",
     "NoResultError",
+    "PreparedKiq",
     "PrometheusMiddleware",
     "ResultGetError",
     "ResultIsReadyError",
@@ -58,6 +65,7 @@
     "SendTaskError",
     "SimpleRetryMiddleware",
     "SmartRetryMiddleware",
+    "TaskDefinition",
     "TaskiqDepends",
     "TaskiqError",
     "TaskiqEvents",
@@ -66,10 +74,13 @@
     "TaskiqMiddleware",
     "TaskiqResult",
     "TaskiqResultTimeoutError",
+    "TaskiqRoute",
+    "TaskiqRouter",
     "TaskiqScheduler",
     "TaskiqState",
     "ZeroMQBroker",
     "__version__",
     "async_shared_broker",
     "gather",
+    "task_builder",
 ]

taskiq/abc/broker.py

@@ -24,11 +24,14 @@
 from taskiq.decor import AsyncTaskiqDecoratedTask
 from taskiq.events import TaskiqEvents
 from taskiq.exceptions import TaskBrokerMismatchError
+from taskiq.flow import FlowProtocol
 from taskiq.formatters.proxy_formatter import ProxyFormatter
 from taskiq.message import BrokerMessage
 from taskiq.result_backends.dummy import DummyResultBackend
+from taskiq.router import TaskiqRouter
 from taskiq.serializers.json_serializer import JSONSerializer
 from taskiq.state import TaskiqState
+from taskiq.task_builder import TaskDefinition
 from taskiq.utils import maybe_awaitable
 from taskiq.warnings import TaskiqDeprecationWarning
 
@@ -78,6 +81,10 @@ def __init__(
         self,
         result_backend: "AsyncResultBackend[_T] | None" = None,
         task_id_generator: Callable[[], str] | None = None,
+        *,
+        router: TaskiqRouter | None = None,
+        broker_name: str | None = None,
+        default_flow: FlowProtocol | None = None,
     ) -> None:
         if result_backend is None:
             result_backend = DummyResultBackend()
@@ -103,6 +110,13 @@ def __init__(
         self.serializer: TaskiqSerializer = JSONSerializer()
         self.formatter: TaskiqFormatter = ProxyFormatter(self)
         self.id_generator = task_id_generator
+        self.router = router or TaskiqRouter()
+        self.default_flow = default_flow
+        self.broker_name = self.router.set_broker(
+            self,
+            name=broker_name,
+            default_flow=default_flow,
+        )
         self.local_task_registry: dict[str, AsyncTaskiqDecoratedTask[Any, Any]] = {}
         # Every event has a list of handlers.
         # Every handler is a function which takes state as a first argument.
@@ -133,10 +147,14 @@ def find_task(self, task_name: str) -> AsyncTaskiqDecoratedTask[Any, Any] | None
         :param task_name: name of a task.
         :returns: found task or None.
         """
-        return self.local_task_registry.get(
-            task_name,
-        ) or self.global_task_registry.get(
-            task_name,
+        return (
+            self.local_task_registry.get(
+                task_name,
+            )
+            or self.router.find_task(task_name)
+            or self.global_task_registry.get(
+                task_name,
+            )
         )
 
     def get_all_tasks(self) -> dict[str, AsyncTaskiqDecoratedTask[Any, Any]]:
@@ -152,7 +170,11 @@ def get_all_tasks(self) -> dict[str, AsyncTaskiqDecoratedTask[Any, Any]]:
 
         :return: dict of all tasks. Keys are task names, values are tasks.
         """
-        return {**self.global_task_registry, **self.local_task_registry}
+        return {
+            **self.global_task_registry,
+            **self.router.get_all_tasks(),
+            **self.local_task_registry,
+        }
 
     def add_dependency_context(self, new_ctx: dict[Any, Any]) -> None:
         """
@@ -237,6 +259,23 @@ async def kick(
         :param message: name of a task.
         """
 
+    async def kick_to_flow(
+        self,
+        message: BrokerMessage,
+        flow: FlowProtocol | None = None,
+    ) -> None:
+        """
+        Send message to a flow-aware broker.
+
+        Existing brokers can keep implementing only `kick`. New brokers may
+        override this method and use `flow` to route to a concrete queue, topic,
+        stream or any other transport address.
+
+        :param message: message to send.
+        :param flow: optional transport-neutral flow.
+        """
+        await self.kick(message)
+
     @abstractmethod
     def listen(self) -> AsyncGenerator[bytes | AckableMessage, None]:
         """
@@ -362,7 +401,10 @@ def inner(
 
     def register_task(
         self,
-        func: Callable[_FuncParams, _ReturnType],
+        func: (
+            Callable[_FuncParams, _ReturnType]
+            | TaskDefinition[_FuncParams, _ReturnType]
+        ),
         task_name: str | None = None,
         **labels: Any,
     ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]:
@@ -380,6 +422,12 @@ def register_task(
 
         :returns: registered task.
         """
+        if isinstance(func, TaskDefinition):
+            return self.router.register_task(
+                func,
+                broker=self,
+                flow=self.default_flow,
+            )
         return self.task(task_name=task_name, **labels)(func)
 
     def on_event(self, *events: TaskiqEvents) -> Callable[[EventHandler], EventHandler]:
@@ -533,6 +581,11 @@ def _register_task(
         if task.broker != self:
             raise TaskBrokerMismatchError(broker=task.broker)
         self.local_task_registry[task_name] = task
+        self.router.register_task(
+            task,
+            broker=self,
+            flow=self.default_flow,
+        )
 
     async def __aenter__(self) -> None:
         """Starts the broker as ctx manager."""

taskiq/brokers/inmemory_broker.py

@@ -9,8 +9,10 @@
 from taskiq.depends.progress_tracker import TaskProgress
 from taskiq.events import TaskiqEvents
 from taskiq.exceptions import UnknownTaskError
+from taskiq.flow import FlowProtocol
 from taskiq.message import BrokerMessage
 from taskiq.receiver import Receiver
+from taskiq.router import TaskiqRouter
 from taskiq.utils import maybe_awaitable
 
 _ReturnType = TypeVar("_ReturnType")
@@ -130,8 +132,16 @@ def __init__(
         max_async_tasks_jitter: int = 0,
         propagate_exceptions: bool = True,
         await_inplace: bool = False,
+        *,
+        router: TaskiqRouter | None = None,
+        broker_name: str | None = None,
+        default_flow: FlowProtocol | None = None,
     ) -> None:
-        super().__init__()
+        super().__init__(
+            router=router,
+            broker_name=broker_name,
+            default_flow=default_flow,
+        )
         self.result_backend: InmemoryResultBackend[Any] = InmemoryResultBackend(
             max_stored_results=max_stored_results,
         )

taskiq/context.py

@@ -30,7 +30,7 @@ async def requeue(self) -> None:
         requeue_count = int(self.message.labels.get("X-Taskiq-requeue", 0))
         requeue_count += 1
         self.message.labels["X-Taskiq-requeue"] = str(requeue_count)
-        await self.broker.kick(self.broker.formatter.dumps(self.message))
+        await self.broker.router.requeue(self.message, broker=self.broker)
         raise NoResultError
 
     def reject(self) -> None: