From 3c6854bcf2e74591dc5ad0896c6bce596f152d4b Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Sun, 17 May 2026 00:12:57 +0200 Subject: [PATCH 01/10] experiment --- docs/examples/router/multiple_brokers.py | 65 +++++ docs/examples/router/shared_task_package.py | 48 ++++ taskiq/__init__.py | 11 + taskiq/abc/broker.py | 63 ++++- taskiq/brokers/inmemory_broker.py | 12 +- taskiq/context.py | 2 +- taskiq/flow.py | 58 +++++ taskiq/kicker.py | 142 +++++++++-- taskiq/router.py | 268 ++++++++++++++++++++ taskiq/task_builder.py | 107 ++++++++ tests/test_router.py | 164 ++++++++++++ 11 files changed, 910 insertions(+), 30 deletions(-) create mode 100644 docs/examples/router/multiple_brokers.py create mode 100644 docs/examples/router/shared_task_package.py create mode 100644 taskiq/flow.py create mode 100644 taskiq/router.py create mode 100644 taskiq/task_builder.py create mode 100644 tests/test_router.py diff --git a/docs/examples/router/multiple_brokers.py b/docs/examples/router/multiple_brokers.py new file mode 100644 index 00000000..c1d5c1aa --- /dev/null +++ b/docs/examples/router/multiple_brokers.py @@ -0,0 +1,65 @@ +"""Route one task through several brokers with a shared router.""" + +import asyncio + +from taskiq import Flow, InMemoryBroker, TaskiqRouter + +router = TaskiqRouter() + +default_email_flow = Flow.queue("emails.default") +priority_email_flow = Flow.queue("emails.priority") +bulk_email_flow = Flow.queue("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}" + + +router.route_task( + send_email.task_name, + broker="priority", + 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", + bulk_email_flow, + ).kiq(8, "digest") + bulk_result = await bulk_task.wait_result(timeout=2) + + print(f"Direct call: {direct_result}") + print(f"Default route: {router.resolve_route(send_email.task_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()) diff --git a/docs/examples/router/shared_task_package.py b/docs/examples/router/shared_task_package.py new file mode 100644 index 00000000..0614679e --- /dev/null +++ b/docs/examples/router/shared_task_package.py @@ -0,0 +1,48 @@ +"""Declare shared task definitions and bind them in the final application.""" + +import asyncio + +from taskiq import Flow, InMemoryBroker, TaskiqRouter, task_builder + + +@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.queue("billing.tasks") +priority_billing_flow = Flow.queue("billing.priority") + +billing_broker = InMemoryBroker( + router=router, + broker_name="billing", + default_flow=billing_flow, + await_inplace=True, +) + +registered_calculate_total = billing_broker.register_task(calculate_total) + + +async def _main() -> None: + await billing_broker.startup() + try: + direct_result = await calculate_total.call(19, 3) + + prepared_task = registered_calculate_total.kicker().with_flow( + 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()) diff --git a/taskiq/__init__.py b/taskiq/__init__.py index 1a764bc4..8efaf3f6 100644 --- a/taskiq/__init__.py +++ b/taskiq/__init__.py @@ -24,7 +24,9 @@ TaskiqError, TaskiqResultTimeoutError, ) +from taskiq.flow import Flow, FlowKind from taskiq.funcs import gather +from taskiq.kicker import PreparedKiq from taskiq.message import BrokerMessage, TaskiqMessage from taskiq.middlewares import ( PrometheusMiddleware, @@ -32,10 +34,12 @@ 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") @@ -48,8 +52,11 @@ "AsyncTaskiqTask", "BrokerMessage", "Context", + "Flow", + "FlowKind", "InMemoryBroker", "NoResultError", + "PreparedKiq", "PrometheusMiddleware", "ResultGetError", "ResultIsReadyError", @@ -59,6 +66,7 @@ "SendTaskError", "SimpleRetryMiddleware", "SmartRetryMiddleware", + "TaskDefinition", "TaskiqDepends", "TaskiqError", "TaskiqEvents", @@ -67,10 +75,13 @@ "TaskiqMiddleware", "TaskiqResult", "TaskiqResultTimeoutError", + "TaskiqRoute", + "TaskiqRouter", "TaskiqScheduler", "TaskiqState", "ZeroMQBroker", "__version__", "async_shared_broker", "gather", + "task_builder", ] diff --git a/taskiq/abc/broker.py b/taskiq/abc/broker.py index ea2e86c0..71cd1c79 100644 --- a/taskiq/abc/broker.py +++ b/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 Flow 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: Flow | 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: Flow | 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,8 @@ 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 +420,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 +579,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.""" diff --git a/taskiq/brokers/inmemory_broker.py b/taskiq/brokers/inmemory_broker.py index 0a7cc98e..f341d4dc 100644 --- a/taskiq/brokers/inmemory_broker.py +++ b/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 Flow 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: Flow | 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, ) diff --git a/taskiq/context.py b/taskiq/context.py index 48058ebb..1e380cf0 100644 --- a/taskiq/context.py +++ b/taskiq/context.py @@ -57,7 +57,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: diff --git a/taskiq/flow.py b/taskiq/flow.py new file mode 100644 index 00000000..5fde2b46 --- /dev/null +++ b/taskiq/flow.py @@ -0,0 +1,58 @@ +import enum +from dataclasses import dataclass, field, replace +from typing import Any + +__all__ = ("Flow", "FlowKind") + + +@enum.unique +class FlowKind(str, enum.Enum): + """Transport-neutral flow shape.""" + + QUEUE = "queue" + TOPIC = "topic" + STREAM = "stream" + + +@dataclass(frozen=True, slots=True) +class Flow: + """Transport-neutral publish or subscribe address. + + Plain flows are intentionally generic. Every broker may interpret a flow + using its own defaults: queue name, topic, stream, channel, list key, or any + other transport address. + + Broker packages can subclass this value object to expose transport-specific + details while still accepting plain Flow instances. + """ + + name: str + kind: FlowKind = FlowKind.QUEUE + options: dict[str, Any] = field( + default_factory=dict, + compare=False, + hash=False, + ) + + @classmethod + def queue(cls, name: str, **options: Any) -> "Flow": + """Create a queue-like flow.""" + return cls(name=name, kind=FlowKind.QUEUE, options=options) + + @classmethod + def topic(cls, name: str, **options: Any) -> "Flow": + """Create a topic-like flow.""" + return cls(name=name, kind=FlowKind.TOPIC, options=options) + + @classmethod + def stream(cls, name: str, **options: Any) -> "Flow": + """Create a stream-like flow.""" + return cls(name=name, kind=FlowKind.STREAM, options=options) + + def with_options(self, **options: Any) -> "Flow": + """Return the same flow with additional generic options.""" + return replace(self, options={**self.options, **options}) + + def broker_options(self, broker_name: str) -> dict[str, Any]: + """Return transport options for broker-specific implementations.""" + return dict(self.options) diff --git a/taskiq/kicker.py b/taskiq/kicker.py index dc113a7e..96ee9f88 100644 --- a/taskiq/kicker.py +++ b/taskiq/kicker.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from collections.abc import Coroutine from dataclasses import asdict, is_dataclass from datetime import datetime, timedelta @@ -9,17 +11,20 @@ Generic, ParamSpec, TypeVar, - Union, + cast, overload, ) from pydantic import BaseModel from taskiq.abc.middleware import TaskiqMiddleware +from taskiq.abc.result_backend import AsyncResultBackend from taskiq.compat import model_dump from taskiq.exceptions import SendTaskError +from taskiq.flow import Flow from taskiq.labels import prepare_label from taskiq.message import TaskiqMessage +from taskiq.router import TaskiqRouter from taskiq.scheduler.created_schedule import CreatedSchedule from taskiq.scheduler.scheduled_task import CronSpec, ScheduledTask from taskiq.task import AsyncTaskiqTask @@ -36,13 +41,29 @@ logger = getLogger("taskiq") +class PreparedKiq(Generic[_ReturnType]): + """Prepared task invocation that can be sent later.""" + + def __init__( + self, + kicker: AsyncKicker[..., _ReturnType], + message: TaskiqMessage, + ) -> None: + self.kicker = kicker + self.message = message + + async def kiq(self) -> AsyncTaskiqTask[_ReturnType]: + """Send prepared invocation.""" + return await self.kicker.kiq_message(self.message) + + class AsyncKicker(Generic[_FuncParams, _ReturnType]): """Class that used to modify data before sending it to broker.""" def __init__( self, task_name: str, - broker: "AsyncBroker", + broker: AsyncBroker, labels: dict[str, Any], return_type: type[_ReturnType] | None = None, ) -> None: @@ -52,11 +73,13 @@ def __init__( self.custom_task_id: str | None = None self.custom_schedule_id: str | None = None self.return_type = return_type + self.route_broker: AsyncBroker | str | None = None + self.route_flow: Flow | None = None def with_labels( self, **labels: str | float, - ) -> "AsyncKicker[_FuncParams, _ReturnType]": + ) -> AsyncKicker[_FuncParams, _ReturnType]: """ Update function's labels before sending. @@ -69,7 +92,7 @@ def with_labels( def with_task_id( self, task_id: str | None, - ) -> "AsyncKicker[_FuncParams, _ReturnType]": + ) -> AsyncKicker[_FuncParams, _ReturnType]: """ Set task_id for current execution. @@ -85,7 +108,7 @@ def with_task_id( def with_schedule_id( self, schedule_id: str, - ) -> "AsyncKicker[_FuncParams, _ReturnType]": + ) -> AsyncKicker[_FuncParams, _ReturnType]: """ Set schedule_id for current execution. @@ -97,8 +120,8 @@ def with_schedule_id( def with_broker( self, - broker: "AsyncBroker", - ) -> "AsyncKicker[_FuncParams, _ReturnType]": + broker: AsyncBroker, + ) -> AsyncKicker[_FuncParams, _ReturnType]: """ Replace broker for the function. @@ -109,11 +132,55 @@ def with_broker( :return: Kicker with new broker. """ self.broker = broker + self.route_broker = broker + return self + + def with_flow( + self, + flow: Flow | None, + ) -> AsyncKicker[_FuncParams, _ReturnType]: + """ + Replace flow for the current invocation. + + :param flow: flow to send message to. + :return: Kicker with a route flow override. + """ + self.route_flow = flow return self + def with_route( + self, + broker: AsyncBroker | str, + flow: Flow | None, + ) -> AsyncKicker[_FuncParams, _ReturnType]: + """ + Replace broker and flow for the current invocation. + + :param broker: broker instance or broker name. + :param flow: flow to send message to. + :return: Kicker with a route override. + """ + self.route_broker = broker + self.route_flow = flow + return self + + def prepare( + self, + *args: _FuncParams.args, + **kwargs: _FuncParams.kwargs, + ) -> PreparedKiq[_ReturnType]: + """ + Prepare a task invocation without sending it. + + :param args: function's arguments. + :param kwargs: function's key word arguments. + :return: prepared task invocation. + """ + return PreparedKiq(self, self._prepare_message(*args, **kwargs)) + @overload async def kiq( - self: "AsyncKicker[_FuncParams, CoroutineType[Any, Any, _T]]", + self: AsyncKicker[_FuncParams, CoroutineType[Any, Any, _T]], *args: _FuncParams.args, **kwargs: _FuncParams.kwargs, ) -> AsyncTaskiqTask[_T]: # pragma: no cover @@ -121,7 +188,7 @@ async def kiq( @overload async def kiq( - self: "AsyncKicker[_FuncParams, Coroutine[Any, Any, _T]]", + self: AsyncKicker[_FuncParams, Coroutine[Any, Any, _T]], *args: _FuncParams.args, **kwargs: _FuncParams.kwargs, ) -> AsyncTaskiqTask[_T]: # pragma: no cover @@ -129,7 +196,7 @@ async def kiq( @overload async def kiq( - self: "AsyncKicker[_FuncParams, _ReturnType]", + self: AsyncKicker[_FuncParams, _ReturnType], *args: _FuncParams.args, **kwargs: _FuncParams.kwargs, ) -> AsyncTaskiqTask[_ReturnType]: # pragma: no cover @@ -156,29 +223,60 @@ async def kiq( logger.debug( f"Kicking {self.task_name} with args={args} and kwargs={kwargs}.", ) - message = self._prepare_message(*args, **kwargs) - for middleware in self.broker.middlewares: - if middleware.__class__.pre_send != TaskiqMiddleware.pre_send: - message = await maybe_awaitable(middleware.pre_send(message)) + return await self.kiq_message(self._prepare_message(*args, **kwargs)) + + async def kiq_message( + self, + message: TaskiqMessage, + ) -> AsyncTaskiqTask[_ReturnType]: + """Send a prepared message.""" try: - await self.broker.kick(self.broker.formatter.dumps(message)) + router = getattr(self.broker, "router", None) + if isinstance(router, TaskiqRouter): + return await router.kiq( + message, + broker=self.route_broker, + flow=self.route_flow, + return_type=self.return_type, + ) + return await self._legacy_kiq(message) except Exception as exc: raise SendTaskError from exc - for middleware in reversed(self.broker.middlewares): + async def _legacy_kiq(self, message: TaskiqMessage) -> AsyncTaskiqTask[_ReturnType]: + """ + Send message through the pre-router broker path. + + This keeps middleware tests and external broker-like mocks compatible + while real AsyncBroker instances use TaskiqRouter. + """ + middlewares = getattr(self.broker, "middlewares", []) + if not isinstance(middlewares, list): + middlewares = [] + + for middleware in middlewares: + if middleware.__class__.pre_send != TaskiqMiddleware.pre_send: + message = await maybe_awaitable(middleware.pre_send(message)) + + await self.broker.kick(self.broker.formatter.dumps(message)) + + for middleware in reversed(middlewares): if middleware.__class__.post_send != TaskiqMiddleware.post_send: await maybe_awaitable(middleware.post_send(message)) return AsyncTaskiqTask( task_id=message.task_id, - result_backend=self.broker.result_backend, - return_type=self.return_type, # type: ignore # (pyright issue) + result_backend=cast( + AsyncResultBackend[_ReturnType], + self.broker.result_backend, + ), + return_type=self.return_type, ) async def schedule_by_cron( self, - source: "ScheduleSource", - cron: Union[str, "CronSpec"], + source: ScheduleSource, + cron: str | CronSpec, *args: _FuncParams.args, **kwargs: _FuncParams.kwargs, ) -> CreatedSchedule[_ReturnType]: @@ -217,7 +315,7 @@ async def schedule_by_cron( async def schedule_by_interval( self, - source: "ScheduleSource", + source: ScheduleSource, interval: int | timedelta, *args: _FuncParams.args, **kwargs: _FuncParams.kwargs, @@ -249,7 +347,7 @@ async def schedule_by_interval( async def schedule_by_time( self, - source: "ScheduleSource", + source: ScheduleSource, time: datetime, *args: _FuncParams.args, **kwargs: _FuncParams.kwargs, diff --git a/taskiq/router.py b/taskiq/router.py new file mode 100644 index 00000000..07129022 --- /dev/null +++ b/taskiq/router.py @@ -0,0 +1,268 @@ +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from logging import getLogger +from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, cast, overload + +from taskiq.abc.middleware import TaskiqMiddleware +from taskiq.abc.result_backend import AsyncResultBackend +from taskiq.flow import Flow +from taskiq.message import TaskiqMessage +from taskiq.task import AsyncTaskiqTask +from taskiq.task_builder import TaskDefinition +from taskiq.utils import maybe_awaitable + +if TYPE_CHECKING: # pragma: no cover + from taskiq.abc.broker import AsyncBroker + from taskiq.decor import AsyncTaskiqDecoratedTask + +__all__ = ("TaskiqRoute", "TaskiqRouter") + +_FuncParams = ParamSpec("_FuncParams") +_ReturnType = TypeVar("_ReturnType") + +logger = getLogger("taskiq.router") + + +@dataclass(frozen=True, slots=True) +class TaskiqRoute: + """Resolved outbound route for a task invocation.""" + + broker_name: str + flow: Flow | None = None + + +class TaskiqRouter: + """Registry and routing layer shared by one or more brokers.""" + + def __init__(self) -> None: + self.brokers: dict[str, AsyncBroker] = {} + self.default_broker_name: str | None = None + self.task_registry: dict[str, AsyncTaskiqDecoratedTask[Any, Any]] = {} + self.routes: dict[str, TaskiqRoute] = {} + + def set_broker( + self, + broker: AsyncBroker, + name: str | None = None, + default_flow: Flow | None = None, + ) -> str: + """Register broker as a transport in this router.""" + broker_name = name or broker.__class__.__name__ + registered = self.brokers.get(broker_name) + if registered is not None and registered is not broker: + raise ValueError( + f"Broker name {broker_name!r} is already registered. " + "Please provide an explicit unique broker_name.", + ) + self.brokers[broker_name] = broker + if self.default_broker_name is None: + self.default_broker_name = broker_name + return broker_name + + def find_task( + self, + task_name: str, + ) -> AsyncTaskiqDecoratedTask[Any, Any] | None: + """Find a task by name.""" + return self.task_registry.get(task_name) + + def get_all_tasks(self) -> dict[str, AsyncTaskiqDecoratedTask[Any, Any]]: + """Return all tasks registered in this router.""" + return dict(self.task_registry) + + def register_task( + self, + task: ( + AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType] + | TaskDefinition[_FuncParams, _ReturnType] + ), + broker: AsyncBroker | str | None = None, + flow: Flow | None = None, + ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: + """Register a bound task or bind a task definition to a broker.""" + if isinstance(task, TaskDefinition): + target_broker = self._resolve_broker(broker) + registered_task = target_broker.register_task( + task.original_func, + task_name=task.task_name, + **task.labels, + ) + if flow is not None: + self.route_task(task.task_name, broker=target_broker, flow=flow) + return registered_task + + self.task_registry[task.task_name] = task + route_broker: AsyncBroker | str | None = broker + if route_broker is None: + route_broker = getattr(task, "broker", None) + if route_broker is not None or flow is not None: + self.route_task(task.task_name, broker=route_broker, flow=flow) + return task + + @overload + def task( + self, + task_name: Callable[_FuncParams, _ReturnType], + *, + broker: AsyncBroker | str | None = None, + flow: Flow | None = None, + **labels: Any, + ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: ... + + @overload + def task( + self, + task_name: str | None = None, + *, + broker: AsyncBroker | str | None = None, + flow: Flow | None = None, + **labels: Any, + ) -> Callable[ + [Callable[_FuncParams, _ReturnType]], + AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType], + ]: ... + + def task( + self, + task_name: str | Callable[_FuncParams, _ReturnType] | None = None, + *, + broker: AsyncBroker | str | None = None, + flow: Flow | None = None, + **labels: Any, + ) -> Any: + """Decorate and register a task through this router.""" + + def register( + func: Callable[_FuncParams, _ReturnType], + ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: + target_broker = self._resolve_broker(broker) + real_task_name = task_name if not callable(task_name) else None + task = target_broker.task(task_name=real_task_name, **labels)(func) + if flow is not None: + self.route_task(task.task_name, broker=target_broker, flow=flow) + return task + + if callable(task_name): + function = task_name + task_name = None + return register(function) + + return register + + def route_task( + self, + task_name: str, + broker: AsyncBroker | str | None = None, + flow: Flow | None = None, + ) -> TaskiqRoute: + """Set default outbound route for a task.""" + broker_name = self._resolve_broker_name(broker) + route = TaskiqRoute(broker_name=broker_name, flow=flow) + self.routes[task_name] = route + return route + + def resolve_route( + self, + task_name: str, + broker: AsyncBroker | str | None = None, + flow: Flow | None = None, + ) -> TaskiqRoute: + """Resolve outbound route for a task invocation.""" + if broker is not None: + broker_name = self._resolve_broker_name(broker) + route_flow = flow + if route_flow is None: + route_flow = self._broker_default_flow(broker_name) + return TaskiqRoute( + broker_name=broker_name, + flow=route_flow, + ) + + route = self.routes.get(task_name) + if route is not None: + if flow is None: + return route + return TaskiqRoute(broker_name=route.broker_name, flow=flow) + + broker_name = self._resolve_broker_name(None) + route_flow = flow + if route_flow is None: + route_flow = self._broker_default_flow(broker_name) + return TaskiqRoute( + broker_name=broker_name, + flow=route_flow, + ) + + async def kiq( + self, + message: TaskiqMessage, + *, + broker: AsyncBroker | str | None = None, + flow: Flow | None = None, + return_type: type[_ReturnType] | None = None, + ) -> AsyncTaskiqTask[_ReturnType]: + """Send message through the resolved broker and flow.""" + route = self.resolve_route(message.task_name, broker=broker, flow=flow) + target_broker = self.brokers[route.broker_name] + + for middleware in target_broker.middlewares: + if middleware.__class__.pre_send != TaskiqMiddleware.pre_send: + message = await maybe_awaitable(middleware.pre_send(message)) + broker_message = target_broker.formatter.dumps(message) + await target_broker.kick_to_flow(broker_message, route.flow) + + for middleware in reversed(target_broker.middlewares): + if middleware.__class__.post_send != TaskiqMiddleware.post_send: + await maybe_awaitable(middleware.post_send(message)) + + return AsyncTaskiqTask( + task_id=message.task_id, + result_backend=cast( + AsyncResultBackend[_ReturnType], + target_broker.result_backend, + ), + return_type=return_type, + ) + + async def requeue( + self, + message: TaskiqMessage, + *, + broker: AsyncBroker | str | None = None, + flow: Flow | None = None, + ) -> None: + """Send an existing message again through the resolved route.""" + route = self.resolve_route(message.task_name, broker=broker, flow=flow) + target_broker = self.brokers[route.broker_name] + await target_broker.kick_to_flow( + target_broker.formatter.dumps(message), + route.flow, + ) + + def _resolve_broker(self, broker: AsyncBroker | str | None) -> AsyncBroker: + broker_name = self._resolve_broker_name(broker) + return self.brokers[broker_name] + + def _resolve_broker_name(self, broker: AsyncBroker | str | None) -> str: + if isinstance(broker, str): + if broker not in self.brokers: + raise ValueError(f"Unknown broker {broker!r}.") + return broker + + if broker is not None: + broker_name = getattr(broker, "broker_name", None) + if broker_name is not None and broker_name in self.brokers: + return broker_name + for registered_name, registered_broker in self.brokers.items(): + if registered_broker is broker: + return registered_name + raise ValueError("Broker is not registered in this router.") + + if self.default_broker_name is None: + raise ValueError("Router doesn't have registered brokers.") + return self.default_broker_name + + def _broker_default_flow(self, broker_name: str) -> Flow | None: + return getattr(self.brokers[broker_name], "default_flow", None) diff --git a/taskiq/task_builder.py b/taskiq/task_builder.py new file mode 100644 index 00000000..d2d3b0d1 --- /dev/null +++ b/taskiq/task_builder.py @@ -0,0 +1,107 @@ +import inspect +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any, Generic, ParamSpec, TypeVar, overload + +from taskiq.message import TaskiqMessage + +__all__ = ("TaskDefinition", "task_builder") + +_FuncParams = ParamSpec("_FuncParams") +_ReturnType = TypeVar("_ReturnType") + + +@dataclass(frozen=True, slots=True) +class TaskDefinition(Generic[_FuncParams, _ReturnType]): + """Unbound task declaration that can be registered later.""" + + task_name: str + original_func: Callable[_FuncParams, _ReturnType] + labels: dict[str, Any] = field(default_factory=dict) + return_type: type[_ReturnType] | None = None + + def __call__( + self, + *args: _FuncParams.args, + **kwargs: _FuncParams.kwargs, + ) -> _ReturnType: + """Call original function directly.""" + return self.original_func(*args, **kwargs) + + async def call( + self, + *args: _FuncParams.args, + **kwargs: _FuncParams.kwargs, + ) -> _ReturnType: + """Execute original function in the current process.""" + result = self.original_func(*args, **kwargs) + if inspect.isawaitable(result): + return await result + return result + + def message( + self, + task_id: str, + *args: _FuncParams.args, + **kwargs: _FuncParams.kwargs, + ) -> TaskiqMessage: + """Build a TaskiqMessage without binding this definition to a router.""" + return TaskiqMessage( + task_id=task_id, + task_name=self.task_name, + labels=dict(self.labels), + args=list(args), + kwargs=dict(kwargs), + ) + + +@overload +def task_builder( + task_name: Callable[_FuncParams, _ReturnType], + **labels: Any, +) -> TaskDefinition[_FuncParams, _ReturnType]: ... + + +@overload +def task_builder( + task_name: str | None = None, + **labels: Any, +) -> Callable[ + [Callable[_FuncParams, _ReturnType]], + TaskDefinition[_FuncParams, _ReturnType], +]: ... + + +def task_builder( + task_name: str | Callable[_FuncParams, _ReturnType] | None = None, + **labels: Any, +) -> Any: + """Build an unbound task definition. + + This decorator is intended for library/package tasks that should be + registered by the final application. + """ + + def build( + func: Callable[_FuncParams, _ReturnType], + ) -> TaskDefinition[_FuncParams, _ReturnType]: + real_task_name = task_name + if real_task_name is None or callable(real_task_name): + real_task_name = f"{func.__module__}:{func.__name__}" + return_type = None + signature = inspect.signature(func) + if signature.return_annotation is not inspect.Signature.empty: + return_type = signature.return_annotation + return TaskDefinition( + task_name=real_task_name, + original_func=func, + labels=dict(labels), + return_type=return_type, + ) + + if callable(task_name): + function = task_name + task_name = None + return build(function) + + return build diff --git a/tests/test_router.py b/tests/test_router.py new file mode 100644 index 00000000..00b1883d --- /dev/null +++ b/tests/test_router.py @@ -0,0 +1,164 @@ +from collections.abc import AsyncGenerator + +import pytest + +from taskiq import Flow, TaskiqRouter, task_builder +from taskiq.abc.broker import AsyncBroker +from taskiq.message import BrokerMessage + + +class RecordingBroker(AsyncBroker): + """Broker that records sent messages and flows.""" + + def __init__( + self, + *, + router: TaskiqRouter | None = None, + broker_name: str | None = None, + default_flow: Flow | None = None, + ) -> None: + self.sent: list[tuple[BrokerMessage, Flow | None]] = [] + super().__init__( + router=router, + broker_name=broker_name, + default_flow=default_flow, + ) + + async def kick(self, message: BrokerMessage) -> None: + """Record old-style send.""" + self.sent.append((message, None)) + + async def kick_to_flow( + self, + message: BrokerMessage, + flow: Flow | None = None, + ) -> None: + """Record flow-aware send.""" + self.sent.append((message, flow)) + + async def listen(self) -> AsyncGenerator[bytes, None]: + """Recording broker doesn't listen in these tests.""" + if False: + yield b"" + + +def test_broker_creates_default_router() -> None: + broker = RecordingBroker() + + assert broker.router.brokers[broker.broker_name] is broker + assert broker.router.default_broker_name == broker.broker_name + + +async def test_old_broker_task_api_registers_task_in_router() -> None: + broker = RecordingBroker() + + @broker.task(task_name="demo.task") + async def demo_task() -> None: + return None + + assert broker.find_task("demo.task") is demo_task + assert broker.router.find_task("demo.task") is demo_task + + await demo_task.kiq() + + assert broker.sent[0][0].task_name == "demo.task" + assert broker.sent[0][1] is None + + +async def test_router_can_route_task_to_another_broker_flow() -> None: + router = TaskiqRouter() + source = RecordingBroker(router=router, broker_name="source") + target = RecordingBroker(router=router, broker_name="target") + flow = Flow("events") + + @source.task(task_name="demo.task") + async def demo_task() -> None: + return None + + router.route_task("demo.task", broker="target", flow=flow) + + await demo_task.kiq() + + assert source.sent == [] + assert target.sent[0][0].task_name == "demo.task" + assert target.sent[0][1] == flow + + +async def test_kicker_route_override_wins_over_registered_route() -> None: + router = TaskiqRouter() + first = RecordingBroker(router=router, broker_name="first") + second = RecordingBroker(router=router, broker_name="second") + first_flow = Flow("first") + second_flow = Flow("second") + + @first.task(task_name="demo.task") + async def demo_task() -> None: + return None + + router.route_task("demo.task", broker="first", flow=first_flow) + + await demo_task.kicker().with_route("second", second_flow).kiq() + + assert first.sent == [] + assert second.sent[0][1] == second_flow + + +async def test_kicker_can_prepare_invocation_for_later() -> None: + broker = RecordingBroker() + + @broker.task(task_name="demo.task") + async def demo_task(value: int) -> None: + return None + + prepared = demo_task.kicker().with_labels(trace_id="abc").prepare(1) + + assert prepared.message.task_name == "demo.task" + assert prepared.message.args == [1] + assert prepared.message.labels["trace_id"] == "abc" + + await prepared.kiq() + + assert broker.sent[0][0].task_id == prepared.message.task_id + + +async def test_task_builder_can_be_registered_later() -> None: + broker = RecordingBroker() + + @task_builder("shared.add", queue="shared") + def add(left: int, right: int) -> int: + return left + right + + assert await add.call(1, 2) == 3 + + registered = broker.register_task(add) + + assert registered.task_name == "shared.add" + assert registered.labels == {"queue": "shared"} + assert broker.router.find_task("shared.add") is registered + + await registered.kiq(1, 2) + + assert broker.sent[0][0].task_name == "shared.add" + + +async def test_router_task_decorator_can_choose_broker_and_flow() -> None: + router = TaskiqRouter() + target = RecordingBroker(router=router, broker_name="target") + flow = Flow("target-flow") + + @router.task("demo.task", broker="target", flow=flow) + async def demo_task() -> None: + return None + + await demo_task.kiq() + + assert target.sent[0][0].task_name == "demo.task" + assert target.sent[0][1] == flow + + +def test_router_rejects_duplicate_broker_names() -> None: + router = TaskiqRouter() + RecordingBroker(router=router, broker_name="broker") + + with pytest.raises(ValueError, match="already registered"): + RecordingBroker(router=router, broker_name="broker") From 3de13d0ffe91d26b2ada0244bf6058b7ccdfced1 Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Sat, 23 May 2026 22:37:37 +0200 Subject: [PATCH 02/10] experiment --- docs/examples/router/multiple_brokers.py | 26 +++--- docs/examples/router/shared_task_package.py | 38 +++++++-- taskiq/__init__.py | 4 +- taskiq/abc/broker.py | 12 +-- taskiq/brokers/inmemory_broker.py | 4 +- taskiq/flow.py | 54 +++++------- taskiq/kicker.py | 10 ++- taskiq/router.py | 93 +++++++++++++-------- tests/test_router.py | 72 ++++++++++++++-- 9 files changed, 202 insertions(+), 111 deletions(-) diff --git a/docs/examples/router/multiple_brokers.py b/docs/examples/router/multiple_brokers.py index c1d5c1aa..7d6977ae 100644 --- a/docs/examples/router/multiple_brokers.py +++ b/docs/examples/router/multiple_brokers.py @@ -6,9 +6,9 @@ router = TaskiqRouter() -default_email_flow = Flow.queue("emails.default") -priority_email_flow = Flow.queue("emails.priority") -bulk_email_flow = Flow.queue("emails.bulk") +default_email_flow = Flow("emails.default") +priority_email_flow = Flow("emails.priority") +bulk_email_flow = Flow("emails.bulk") default_broker = InMemoryBroker( router=router, @@ -30,9 +30,9 @@ async def send_email(user_id: int, template: str) -> str: return f"{template} email sent to user {user_id}" -router.route_task( - send_email.task_name, - broker="priority", +priority_route = router.route_task( + send_email, + broker=priority_broker, flow=priority_email_flow, ) @@ -46,14 +46,18 @@ async def _main() -> None: 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", - bulk_email_flow, - ).kiq(8, "digest") + 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"Default route: {router.resolve_route(send_email.task_name)}") + 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: diff --git a/docs/examples/router/shared_task_package.py b/docs/examples/router/shared_task_package.py index 0614679e..3a9e4902 100644 --- a/docs/examples/router/shared_task_package.py +++ b/docs/examples/router/shared_task_package.py @@ -1,10 +1,27 @@ """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.""" @@ -12,8 +29,8 @@ async def calculate_total(price: int, quantity: int) -> int: router = TaskiqRouter() -billing_flow = Flow.queue("billing.tasks") -priority_billing_flow = Flow.queue("billing.priority") +billing_flow = Flow("billing.tasks") +priority_billing_flow = BillingQueue(name="billing.priority", priority=10) billing_broker = InMemoryBroker( router=router, @@ -22,7 +39,11 @@ async def calculate_total(price: int, quantity: int) -> int: await_inplace=True, ) -registered_calculate_total = billing_broker.register_task(calculate_total) +registered_calculate_total = router.register_task( + calculate_total, + broker=billing_broker, + flow=billing_flow, +) async def _main() -> None: @@ -30,9 +51,14 @@ async def _main() -> None: try: direct_result = await calculate_total.call(19, 3) - prepared_task = registered_calculate_total.kicker().with_flow( - priority_billing_flow, - ).prepare(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) diff --git a/taskiq/__init__.py b/taskiq/__init__.py index 8efaf3f6..8aeb203c 100644 --- a/taskiq/__init__.py +++ b/taskiq/__init__.py @@ -24,7 +24,7 @@ TaskiqError, TaskiqResultTimeoutError, ) -from taskiq.flow import Flow, FlowKind +from taskiq.flow import Flow, FlowProtocol from taskiq.funcs import gather from taskiq.kicker import PreparedKiq from taskiq.message import BrokerMessage, TaskiqMessage @@ -53,7 +53,7 @@ "BrokerMessage", "Context", "Flow", - "FlowKind", + "FlowProtocol", "InMemoryBroker", "NoResultError", "PreparedKiq", diff --git a/taskiq/abc/broker.py b/taskiq/abc/broker.py index 71cd1c79..42a1817b 100644 --- a/taskiq/abc/broker.py +++ b/taskiq/abc/broker.py @@ -24,7 +24,7 @@ from taskiq.decor import AsyncTaskiqDecoratedTask from taskiq.events import TaskiqEvents from taskiq.exceptions import TaskBrokerMismatchError -from taskiq.flow import Flow +from taskiq.flow import FlowProtocol from taskiq.formatters.proxy_formatter import ProxyFormatter from taskiq.message import BrokerMessage from taskiq.result_backends.dummy import DummyResultBackend @@ -84,7 +84,7 @@ def __init__( *, router: TaskiqRouter | None = None, broker_name: str | None = None, - default_flow: Flow | None = None, + default_flow: FlowProtocol | None = None, ) -> None: if result_backend is None: result_backend = DummyResultBackend() @@ -262,7 +262,7 @@ async def kick( async def kick_to_flow( self, message: BrokerMessage, - flow: Flow | None = None, + flow: FlowProtocol | None = None, ) -> None: """ Send message to a flow-aware broker. @@ -401,8 +401,10 @@ def inner( def register_task( self, - func: Callable[_FuncParams, _ReturnType] - | TaskDefinition[_FuncParams, _ReturnType], + func: ( + Callable[_FuncParams, _ReturnType] + | TaskDefinition[_FuncParams, _ReturnType] + ), task_name: str | None = None, **labels: Any, ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: diff --git a/taskiq/brokers/inmemory_broker.py b/taskiq/brokers/inmemory_broker.py index f341d4dc..e63a97f6 100644 --- a/taskiq/brokers/inmemory_broker.py +++ b/taskiq/brokers/inmemory_broker.py @@ -9,7 +9,7 @@ from taskiq.depends.progress_tracker import TaskProgress from taskiq.events import TaskiqEvents from taskiq.exceptions import UnknownTaskError -from taskiq.flow import Flow +from taskiq.flow import FlowProtocol from taskiq.message import BrokerMessage from taskiq.receiver import Receiver from taskiq.router import TaskiqRouter @@ -135,7 +135,7 @@ def __init__( *, router: TaskiqRouter | None = None, broker_name: str | None = None, - default_flow: Flow | None = None, + default_flow: FlowProtocol | None = None, ) -> None: super().__init__( router=router, diff --git a/taskiq/flow.py b/taskiq/flow.py index 5fde2b46..8754492e 100644 --- a/taskiq/flow.py +++ b/taskiq/flow.py @@ -1,58 +1,42 @@ -import enum +from collections.abc import Mapping from dataclasses import dataclass, field, replace -from typing import Any +from typing import Protocol, runtime_checkable -__all__ = ("Flow", "FlowKind") +__all__ = ("Flow", "FlowProtocol") -@enum.unique -class FlowKind(str, enum.Enum): - """Transport-neutral flow shape.""" +@runtime_checkable +class FlowProtocol(Protocol): + """Transport-neutral flow contract accepted by routers and brokers.""" - QUEUE = "queue" - TOPIC = "topic" - STREAM = "stream" + name: str + + def broker_options(self, broker_name: str) -> Mapping[str, object]: + """Return options relevant for a concrete broker implementation.""" + ... @dataclass(frozen=True, slots=True) class Flow: - """Transport-neutral publish or subscribe address. - - Plain flows are intentionally generic. Every broker may interpret a flow - using its own defaults: queue name, topic, stream, channel, list key, or any - other transport address. + """Generic transport-neutral flow address. - Broker packages can subclass this value object to expose transport-specific - details while still accepting plain Flow instances. + Broker packages can provide their own flow objects that implement + FlowProtocol, for example RabbitQueue, KafkaTopic, NatsSubject or RedisQueue. + This generic value object is intentionally small and works as a common + fallback for brokers that only need a named address with optional metadata. """ name: str - kind: FlowKind = FlowKind.QUEUE - options: dict[str, Any] = field( + options: Mapping[str, object] = field( default_factory=dict, compare=False, hash=False, ) - @classmethod - def queue(cls, name: str, **options: Any) -> "Flow": - """Create a queue-like flow.""" - return cls(name=name, kind=FlowKind.QUEUE, options=options) - - @classmethod - def topic(cls, name: str, **options: Any) -> "Flow": - """Create a topic-like flow.""" - return cls(name=name, kind=FlowKind.TOPIC, options=options) - - @classmethod - def stream(cls, name: str, **options: Any) -> "Flow": - """Create a stream-like flow.""" - return cls(name=name, kind=FlowKind.STREAM, options=options) - - def with_options(self, **options: Any) -> "Flow": + def with_options(self, **options: object) -> "Flow": """Return the same flow with additional generic options.""" return replace(self, options={**self.options, **options}) - def broker_options(self, broker_name: str) -> dict[str, Any]: + def broker_options(self, broker_name: str) -> Mapping[str, object]: """Return transport options for broker-specific implementations.""" return dict(self.options) diff --git a/taskiq/kicker.py b/taskiq/kicker.py index 96ee9f88..d0182e4c 100644 --- a/taskiq/kicker.py +++ b/taskiq/kicker.py @@ -21,7 +21,7 @@ from taskiq.abc.result_backend import AsyncResultBackend from taskiq.compat import model_dump from taskiq.exceptions import SendTaskError -from taskiq.flow import Flow +from taskiq.flow import FlowProtocol from taskiq.labels import prepare_label from taskiq.message import TaskiqMessage from taskiq.router import TaskiqRouter @@ -74,7 +74,7 @@ def __init__( self.custom_schedule_id: str | None = None self.return_type = return_type self.route_broker: AsyncBroker | str | None = None - self.route_flow: Flow | None = None + self.route_flow: FlowProtocol | None = None def with_labels( self, @@ -137,7 +137,7 @@ def with_broker( def with_flow( self, - flow: Flow | None, + flow: FlowProtocol | None, ) -> AsyncKicker[_FuncParams, _ReturnType]: """ Replace flow for the current invocation. @@ -151,12 +151,14 @@ def with_flow( def with_route( self, broker: AsyncBroker | str, - flow: Flow | None, + flow: FlowProtocol | None, ) -> AsyncKicker[_FuncParams, _ReturnType]: """ Replace broker and flow for the current invocation. :param broker: broker instance or broker name. + Broker instances are preferred; names are kept for configuration + and backward-compatible lookup. :param flow: flow to send message to. :return: Kicker with a route override. """ diff --git a/taskiq/router.py b/taskiq/router.py index 07129022..9f770f90 100644 --- a/taskiq/router.py +++ b/taskiq/router.py @@ -7,7 +7,7 @@ from taskiq.abc.middleware import TaskiqMiddleware from taskiq.abc.result_backend import AsyncResultBackend -from taskiq.flow import Flow +from taskiq.flow import FlowProtocol from taskiq.message import TaskiqMessage from taskiq.task import AsyncTaskiqTask from taskiq.task_builder import TaskDefinition @@ -29,8 +29,13 @@ class TaskiqRoute: """Resolved outbound route for a task invocation.""" - broker_name: str - flow: Flow | None = None + broker: AsyncBroker + flow: FlowProtocol | None = None + + @property + def broker_name(self) -> str: + """Return registered broker name for compatibility and diagnostics.""" + return self.broker.broker_name class TaskiqRouter: @@ -46,7 +51,7 @@ def set_broker( self, broker: AsyncBroker, name: str | None = None, - default_flow: Flow | None = None, + default_flow: FlowProtocol | None = None, ) -> str: """Register broker as a transport in this router.""" broker_name = name or broker.__class__.__name__ @@ -79,7 +84,7 @@ def register_task( | TaskDefinition[_FuncParams, _ReturnType] ), broker: AsyncBroker | str | None = None, - flow: Flow | None = None, + flow: FlowProtocol | None = None, ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: """Register a bound task or bind a task definition to a broker.""" if isinstance(task, TaskDefinition): @@ -107,7 +112,7 @@ def task( task_name: Callable[_FuncParams, _ReturnType], *, broker: AsyncBroker | str | None = None, - flow: Flow | None = None, + flow: FlowProtocol | None = None, **labels: Any, ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: ... @@ -117,7 +122,7 @@ def task( task_name: str | None = None, *, broker: AsyncBroker | str | None = None, - flow: Flow | None = None, + flow: FlowProtocol | None = None, **labels: Any, ) -> Callable[ [Callable[_FuncParams, _ReturnType]], @@ -129,7 +134,7 @@ def task( task_name: str | Callable[_FuncParams, _ReturnType] | None = None, *, broker: AsyncBroker | str | None = None, - flow: Flow | None = None, + flow: FlowProtocol | None = None, **labels: Any, ) -> Any: """Decorate and register a task through this router.""" @@ -153,30 +158,32 @@ def register( def route_task( self, - task_name: str, + task: str | AsyncTaskiqDecoratedTask[Any, Any], broker: AsyncBroker | str | None = None, - flow: Flow | None = None, + flow: FlowProtocol | None = None, ) -> TaskiqRoute: """Set default outbound route for a task.""" - broker_name = self._resolve_broker_name(broker) - route = TaskiqRoute(broker_name=broker_name, flow=flow) + task_name = self._resolve_task_name(task) + target_broker = self._resolve_broker(broker) + route = TaskiqRoute(broker=target_broker, flow=flow) self.routes[task_name] = route return route def resolve_route( self, - task_name: str, + task: str | AsyncTaskiqDecoratedTask[Any, Any], broker: AsyncBroker | str | None = None, - flow: Flow | None = None, + flow: FlowProtocol | None = None, ) -> TaskiqRoute: """Resolve outbound route for a task invocation.""" + task_name = self._resolve_task_name(task) if broker is not None: - broker_name = self._resolve_broker_name(broker) + target_broker = self._resolve_broker(broker) route_flow = flow if route_flow is None: - route_flow = self._broker_default_flow(broker_name) + route_flow = self._broker_default_flow(target_broker) return TaskiqRoute( - broker_name=broker_name, + broker=target_broker, flow=route_flow, ) @@ -184,14 +191,14 @@ def resolve_route( if route is not None: if flow is None: return route - return TaskiqRoute(broker_name=route.broker_name, flow=flow) + return TaskiqRoute(broker=route.broker, flow=flow) - broker_name = self._resolve_broker_name(None) + target_broker = self._resolve_broker(None) route_flow = flow if route_flow is None: - route_flow = self._broker_default_flow(broker_name) + route_flow = self._broker_default_flow(target_broker) return TaskiqRoute( - broker_name=broker_name, + broker=target_broker, flow=route_flow, ) @@ -200,12 +207,12 @@ async def kiq( message: TaskiqMessage, *, broker: AsyncBroker | str | None = None, - flow: Flow | None = None, + flow: FlowProtocol | None = None, return_type: type[_ReturnType] | None = None, ) -> AsyncTaskiqTask[_ReturnType]: """Send message through the resolved broker and flow.""" route = self.resolve_route(message.task_name, broker=broker, flow=flow) - target_broker = self.brokers[route.broker_name] + target_broker = route.broker for middleware in target_broker.middlewares: if middleware.__class__.pre_send != TaskiqMiddleware.pre_send: @@ -231,38 +238,50 @@ async def requeue( message: TaskiqMessage, *, broker: AsyncBroker | str | None = None, - flow: Flow | None = None, + flow: FlowProtocol | None = None, ) -> None: """Send an existing message again through the resolved route.""" route = self.resolve_route(message.task_name, broker=broker, flow=flow) - target_broker = self.brokers[route.broker_name] + target_broker = route.broker await target_broker.kick_to_flow( target_broker.formatter.dumps(message), route.flow, ) def _resolve_broker(self, broker: AsyncBroker | str | None) -> AsyncBroker: - broker_name = self._resolve_broker_name(broker) - return self.brokers[broker_name] - - def _resolve_broker_name(self, broker: AsyncBroker | str | None) -> str: if isinstance(broker, str): if broker not in self.brokers: raise ValueError(f"Unknown broker {broker!r}.") - return broker + return self.brokers[broker] if broker is not None: broker_name = getattr(broker, "broker_name", None) - if broker_name is not None and broker_name in self.brokers: - return broker_name - for registered_name, registered_broker in self.brokers.items(): + if isinstance(broker_name, str): + registered_broker = self.brokers.get(broker_name) + if registered_broker is broker: + return broker + for registered_broker in self.brokers.values(): if registered_broker is broker: - return registered_name + return registered_broker raise ValueError("Broker is not registered in this router.") if self.default_broker_name is None: raise ValueError("Router doesn't have registered brokers.") - return self.default_broker_name + return self.brokers[self.default_broker_name] + + def _resolve_broker_name(self, broker: AsyncBroker | str | None) -> str: + return self._resolve_broker(broker).broker_name + + def _resolve_task_name( + self, + task: str | AsyncTaskiqDecoratedTask[Any, Any], + ) -> str: + if isinstance(task, str): + return task + task_name = getattr(task, "task_name", None) + if isinstance(task_name, str): + return task_name + raise TypeError("Route task must be a task name or decorated task.") - def _broker_default_flow(self, broker_name: str) -> Flow | None: - return getattr(self.brokers[broker_name], "default_flow", None) + def _broker_default_flow(self, broker: AsyncBroker) -> FlowProtocol | None: + return getattr(broker, "default_flow", None) diff --git a/tests/test_router.py b/tests/test_router.py index 00b1883d..f20d7223 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -1,12 +1,28 @@ -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Mapping +from dataclasses import dataclass import pytest -from taskiq import Flow, TaskiqRouter, task_builder +from taskiq import Flow, FlowProtocol, TaskiqRouter, task_builder from taskiq.abc.broker import AsyncBroker from taskiq.message import BrokerMessage +@dataclass(frozen=True, slots=True) +class BrokerQueue: + """Broker-specific flow used to prove protocol-based routing.""" + + name: str + durable: bool = True + + def broker_options(self, broker_name: str) -> Mapping[str, object]: + """Return options for the target broker.""" + return { + "broker": broker_name, + "durable": self.durable, + } + + class RecordingBroker(AsyncBroker): """Broker that records sent messages and flows.""" @@ -15,9 +31,9 @@ def __init__( *, router: TaskiqRouter | None = None, broker_name: str | None = None, - default_flow: Flow | None = None, + default_flow: FlowProtocol | None = None, ) -> None: - self.sent: list[tuple[BrokerMessage, Flow | None]] = [] + self.sent: list[tuple[BrokerMessage, FlowProtocol | None]] = [] super().__init__( router=router, broker_name=broker_name, @@ -31,7 +47,7 @@ async def kick(self, message: BrokerMessage) -> None: async def kick_to_flow( self, message: BrokerMessage, - flow: Flow | None = None, + flow: FlowProtocol | None = None, ) -> None: """Record flow-aware send.""" self.sent.append((message, flow)) @@ -75,10 +91,12 @@ async def test_router_can_route_task_to_another_broker_flow() -> None: async def demo_task() -> None: return None - router.route_task("demo.task", broker="target", flow=flow) + route = router.route_task(demo_task, broker=target, flow=flow) await demo_task.kiq() + assert route.broker is target + assert route.broker_name == "target" assert source.sent == [] assert target.sent[0][0].task_name == "demo.task" assert target.sent[0][1] == flow @@ -95,14 +113,50 @@ async def test_kicker_route_override_wins_over_registered_route() -> None: async def demo_task() -> None: return None - router.route_task("demo.task", broker="first", flow=first_flow) + router.route_task(demo_task, broker=first, flow=first_flow) - await demo_task.kicker().with_route("second", second_flow).kiq() + await demo_task.kicker().with_route(second, second_flow).kiq() assert first.sent == [] assert second.sent[0][1] == second_flow +async def test_router_keeps_string_broker_lookup_for_compatibility() -> None: + router = TaskiqRouter() + source = RecordingBroker(router=router, broker_name="source") + target = RecordingBroker(router=router, broker_name="target") + flow = Flow("compat") + + @source.task(task_name="demo.task") + async def demo_task() -> None: + return None + + router.route_task("demo.task", broker="target", flow=flow) + + await demo_task.kicker().with_route("target", flow).kiq() + + assert target.sent[0][0].task_name == "demo.task" + assert target.sent[0][1] == flow + + +async def test_router_accepts_broker_specific_flow_protocol() -> None: + broker = RecordingBroker(broker_name="transport") + flow = BrokerQueue(name="critical", durable=False) + + @broker.task(task_name="demo.task") + async def demo_task() -> None: + return None + + await demo_task.kicker().with_flow(flow).kiq() + + assert isinstance(flow, FlowProtocol) + assert broker.sent[0][1] is flow + assert flow.broker_options("transport") == { + "broker": "transport", + "durable": False, + } + + async def test_kicker_can_prepare_invocation_for_later() -> None: broker = RecordingBroker() @@ -146,7 +200,7 @@ async def test_router_task_decorator_can_choose_broker_and_flow() -> None: target = RecordingBroker(router=router, broker_name="target") flow = Flow("target-flow") - @router.task("demo.task", broker="target", flow=flow) + @router.task("demo.task", broker=target, flow=flow) async def demo_task() -> None: return None From d1f1d5d155742618a9f063b38a42e1a4f7f235c9 Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Sat, 6 Jun 2026 16:11:15 +0200 Subject: [PATCH 03/10] feat(routing): separate task routing from broker transport Introduce router-owned task routing, flow subscriptions, and shared task declarations while preserving the existing broker-first task API. Add transport-neutral Flow contracts, explicit route/subscription semantics, multi-broker router ownership, task_builder base_cls support, prepared invocation route snapshots, and scheduler/requeue compatibility behavior. Document migration guidance and add executable routing examples for multiple brokers and shared task packages. --- docs/examples/router/multiple_brokers.py | 47 ++- docs/examples/router/shared_task_package.py | 40 ++- docs/guide/architecture-overview.md | 22 ++ docs/guide/routing-and-flows.md | 345 ++++++++++++++++++++ docs/guide/scheduling-tasks.md | 17 +- taskiq/__init__.py | 7 +- taskiq/abc/broker.py | 141 ++++++-- taskiq/flow.py | 39 ++- taskiq/kicker.py | 160 +++++---- taskiq/message.py | 60 +++- taskiq/router.py | 299 +++++++++-------- taskiq/routing/__init__.py | 17 + taskiq/routing/dispatcher.py | 84 +++++ taskiq/routing/models.py | 38 +++ taskiq/routing/references.py | 18 + taskiq/routing/registries.py | 106 ++++++ taskiq/routing/routes.py | 106 ++++++ taskiq/routing/subscriptions.py | 136 ++++++++ taskiq/scheduler/scheduler.py | 36 +- taskiq/task_builder.py | 58 +++- tests/routing/__init__.py | 0 tests/routing/models.py | 117 +++++++ tests/routing/test_dispatch.py | 324 ++++++++++++++++++ tests/routing/test_router_core.py | 74 +++++ tests/routing/test_shared_tasks.py | 94 ++++++ tests/routing/test_subscriptions.py | 186 +++++++++++ tests/scheduler/test_scheduler.py | 114 +++++++ tests/test_requeue.py | 90 ++++- tests/test_router.py | 218 ------------- tests/test_task_builder.py | 130 ++++++++ tests/utils.py | 37 ++- 31 files changed, 2652 insertions(+), 508 deletions(-) create mode 100644 docs/guide/routing-and-flows.md create mode 100644 taskiq/routing/__init__.py create mode 100644 taskiq/routing/dispatcher.py create mode 100644 taskiq/routing/models.py create mode 100644 taskiq/routing/references.py create mode 100644 taskiq/routing/registries.py create mode 100644 taskiq/routing/routes.py create mode 100644 taskiq/routing/subscriptions.py create mode 100644 tests/routing/__init__.py create mode 100644 tests/routing/models.py create mode 100644 tests/routing/test_dispatch.py create mode 100644 tests/routing/test_router_core.py create mode 100644 tests/routing/test_shared_tasks.py create mode 100644 tests/routing/test_subscriptions.py delete mode 100644 tests/test_router.py create mode 100644 tests/test_task_builder.py diff --git a/docs/examples/router/multiple_brokers.py b/docs/examples/router/multiple_brokers.py index 7d6977ae..c845bbf0 100644 --- a/docs/examples/router/multiple_brokers.py +++ b/docs/examples/router/multiple_brokers.py @@ -2,7 +2,7 @@ import asyncio -from taskiq import Flow, InMemoryBroker, TaskiqRouter +from taskiq import Flow, InMemoryBroker, TaskiqRoute, TaskiqRouter router = TaskiqRouter() @@ -35,6 +35,23 @@ async def send_email(user_id: int, template: str) -> str: broker=priority_broker, flow=priority_email_flow, ) +priority_subscription = router.subscribe( + priority_broker, + priority_email_flow, + send_email, +) + + +def _format_route(task_name: str, route: TaskiqRoute) -> str: + """Return a readable route diagnostic for the example output.""" + flow_name = route.flow.name if route.flow is not None else "" + return f"{task_name} -> broker={route.broker_name}, flow={flow_name}" + + +def _format_listen_plan() -> str: + """Return flows that the priority broker should subscribe to.""" + flow_names = ", ".join(flow.name for flow in priority_broker.get_subscribed_flows()) + return f"priority listens to: {flow_names}" async def _main() -> None: @@ -43,22 +60,34 @@ async def _main() -> None: 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) + declared_route = router.resolve_route(send_email) + assert declared_route == priority_route - bulk_task = ( + routed_task = ( await send_email.kicker() - .with_route( - default_broker, - bulk_email_flow, + .with_route(declared_route) + .kiq( + 7, + "welcome", ) - .kiq(8, "digest") ) + routed_result = await routed_task.wait_result(timeout=2) + + bulk_route = router.resolve_route( + send_email, + broker=default_broker, + flow=bulk_email_flow, + ) + bulk_task = await send_email.kicker().with_route(bulk_route).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"Router rule: {_format_route(send_email.task_name, priority_route)}") + print(f"Subscription tasks: {sorted(priority_subscription.task_names)}") + print(_format_listen_plan()) + print(f"Resolved route: {_format_route(send_email.task_name, declared_route)}") print(f"Routed call: {routed_result.return_value}") + print(f"Override route: {_format_route(send_email.task_name, bulk_route)}") print(f"Route override: {bulk_result.return_value}") finally: await priority_broker.shutdown() diff --git a/docs/examples/router/shared_task_package.py b/docs/examples/router/shared_task_package.py index 3a9e4902..0bbd6616 100644 --- a/docs/examples/router/shared_task_package.py +++ b/docs/examples/router/shared_task_package.py @@ -3,8 +3,15 @@ import asyncio from collections.abc import Mapping from dataclasses import dataclass - -from taskiq import Flow, InMemoryBroker, TaskiqRouter, task_builder +from typing import Any + +from taskiq import ( + AsyncTaskiqDecoratedTask, + Flow, + InMemoryBroker, + TaskiqRouter, + task_builder, +) @dataclass(frozen=True, slots=True) @@ -14,15 +21,22 @@ class BillingQueue: name: str priority: int - def broker_options(self, broker_name: str) -> Mapping[str, object]: + def broker_options(self) -> 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") +class BillingTask(AsyncTaskiqDecoratedTask[Any, Any]): + """Custom task class shared by billing package tasks.""" + + def billing_name(self) -> str: + """Return a billing-specific task name.""" + return self.task_name + + +@task_builder("billing.calculate_total", base_cls=BillingTask, 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 @@ -44,6 +58,11 @@ async def calculate_total(price: int, quantity: int) -> int: broker=billing_broker, flow=billing_flow, ) +router.subscribe( + billing_broker, + billing_flow, + registered_calculate_total, +) async def _main() -> None: @@ -51,11 +70,15 @@ async def _main() -> None: try: direct_result = await calculate_total.call(19, 3) + priority_route = router.resolve_route( + registered_calculate_total, + broker=billing_broker, + flow=priority_billing_flow, + ) prepared_task = ( registered_calculate_total.kicker() .with_route( - billing_broker, - priority_billing_flow, + priority_route, ) .prepare(19, 3) ) @@ -64,6 +87,9 @@ async def _main() -> None: queued_result = await queued_task.wait_result(timeout=2) print(f"Shared task direct call: {direct_result}") + print(f"Registered task class: {registered_calculate_total.billing_name()}") + listen_flow = billing_broker.get_subscribed_flows()[0] + print(f"Registered listen flow: {listen_flow.name}") print(f"Prepared message: {prepared_task.message.task_name}") print(f"Registered queued call: {queued_result.return_value}") finally: diff --git a/docs/guide/architecture-overview.md b/docs/guide/architecture-overview.md index d4b825f1..f84437d2 100644 --- a/docs/guide/architecture-overview.md +++ b/docs/guide/architecture-overview.md @@ -82,6 +82,28 @@ asyncio.run(main()) ``` +## Router and flows + +Taskiq can use a `TaskiqRouter` to keep routing rules outside of broker +implementations. Brokers remain transport adapters, while the router owns task +registration, route resolution and flow subscriptions. +This section describes the `experiment/separate_broker` branch contract. The +old `@broker.task(...)`, `.kiq()`, labels, scheduler and result backend behavior +remain compatible, while router/flow APIs are additive review material for the +branch. + +`Flow` is a transport-neutral delivery address. Broker packages may provide +their own flow classes for queue, topic, subject or stream options, as long as +they implement the same flow protocol. The router deduplicates subscriptions by +flow name and rejects same-name flows with incompatible broker options. + +Routing and subscribing are separate responsibilities. `route_task(...)` +chooses the outbound broker and flow for task invocations. `subscribe(...)` +adds flows to a broker listen plan for flow-aware broker adapters. Worker task +lookup still uses `task_name`; flow does not select the Python task. + +Read more in the [Routing and flows](./routing-and-flows.md) section. + ## Messages Every message has labels. You can define labels diff --git a/docs/guide/routing-and-flows.md b/docs/guide/routing-and-flows.md new file mode 100644 index 00000000..208d2502 --- /dev/null +++ b/docs/guide/routing-and-flows.md @@ -0,0 +1,345 @@ +--- +title: Routing and flows +order: 10 +--- + +# Routing and flows + +::: warning Experimental branch +This page describes the routing and flow contract on the +`experiment/separate_broker` branch. The old task declaration API remains +compatible. Broker-specific packages may keep using their existing send/listen +implementation until they opt into flow-aware dispatch and listen-plan support. +::: + +Taskiq keeps the old task declaration API: + +```python +@broker.task(task_name="billing.charge") +async def charge(user_id: int) -> None: + ... +``` + +The broker is still the transport adapter, but routing decisions can be moved to +a `TaskiqRouter`. This is useful when one application owns several brokers or +when a task package declares tasks before the final application knows which +broker must handle them. + +The important responsibility split is: + +- Task declaration describes a Python callable, task name, labels and direct + call behavior. +- Kicker prepares a task invocation: task id, labels, args, kwargs and optional + route overrides. +- Router owns outbound routing policy: task invocation to broker and flow. +- Router also owns inbound listen-plan records for flow-aware brokers. +- Broker owns transport lifecycle: startup, shutdown, send, listen, ack/nack, + serializer/formatter, middleware and result backend integration. +- Worker still executes by `task_name`; flow does not select the Python task. + +## Router + +A router owns broker registration, task registration and routing rules. +One router can own several brokers. A broker has a single router owner; pass the +router when creating the broker instead of registering the same broker in +several routers. + +```python +from taskiq import Flow, InMemoryBroker, TaskiqRouter + +router = TaskiqRouter() + +default_broker = InMemoryBroker(router=router, broker_name="default") +priority_broker = InMemoryBroker(router=router, broker_name="priority") + +priority_route = router.route_task( + "billing.charge", + broker=priority_broker, + flow=Flow("billing.priority"), +) +``` + +If a broker is created without an explicit router, Taskiq creates a default +router for that broker. Use an explicit shared router when one worker or app +needs several brokers to follow the same routing policy. + +Pass broker objects to routing APIs. If broker names come from configuration, +resolve them explicitly first: + +```python +broker = router.get_broker("priority") +route = router.route_task("billing.charge", broker=broker) +``` + +This keeps routing errors local and avoids hidden string references. + +## Flow + +`Flow` is a transport-neutral delivery address. It is intentionally small. +The routing identity of a flow is its logical name: + +```python +from taskiq import Flow, FlowIdentity, get_flow_identity + +flow = Flow("billing.priority").with_options(priority=10) + +assert flow.identity == FlowIdentity("billing.priority") +assert get_flow_identity(flow) == flow.identity +``` + +Broker packages can expose their own flow objects when they need +transport-specific options: + +```python +from collections.abc import Mapping +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class RabbitQueue: + name: str + durable: bool = True + + def broker_options(self) -> Mapping[str, object]: + return {"durable": self.durable} +``` + +The core protocol requires `name` and `broker_options()`. Broker-specific +settings stay in the flow object and are returned through `broker_options()`; +they are not encoded into router task names or message labels. + +Flow identity and declaration options are separate: + +- Same broker, same flow name and same `broker_options()` means the listen plan + can deduplicate the flow and merge task names. +- Same broker and same flow name with different `broker_options()` is a + conflict, so Taskiq raises `ValueError` instead of silently choosing one + queue/topic declaration. +- Different brokers may use the same flow name independently. + +## Sending with a route + +`router.route_task()` returns a `TaskiqRoute`. You can pass that object to a +kicker when you want the route to be explicit at the call site. + +```python +route = router.resolve_route("billing.charge") +task = await charge.kicker().with_route(route).kiq(42) +``` + +For one-off overrides, resolve a route with the target broker and flow: + +```python +bulk_route = router.resolve_route( + "billing.charge", + broker=default_broker, + flow=Flow("billing.bulk"), +) + +task = await charge.kicker().with_route(bulk_route).kiq(42) +``` + +`with_broker(...)` remains available for shared-task style invocation. When a +kicker has both a route and a later broker override, the broker override clears +the explicit route. Taskiq then resolves the task for the override broker: it +uses that broker's task route flow when one exists, otherwise the broker +`default_flow`. + +`prepare(...)` creates a prepared invocation for later use. Prepared invocations +keep a snapshot of the route that was resolved at prepare time, so later changes +to the mutable kicker object or router route table do not change that prepared +send. + +## Subscriptions + +Routers also keep inbound flow subscriptions. Routing and subscribing are +separate operations: `route_task()` selects the outbound broker/flow for a task +invocation, while `subscribe()` adds a flow to the broker listen plan. + +```python +billing_flow = Flow("billing.priority") + +route = router.route_task( + "billing.charge", + broker=priority_broker, + flow=billing_flow, +) + +router.subscribe(priority_broker, billing_flow, "billing.charge") + +subscriptions = router.get_subscriptions(priority_broker) +flows = priority_broker.get_subscribed_flows() +``` + +`route_task(...)` does not subscribe by default. This is intentional: outbound +routing answers "where should this invocation be sent?", while subscribing +answers "which flows should this broker listen to?". + +The deprecated `route_task(..., subscribe=True)` shim still performs this +subscription when a flow is resolved, but new code should call `subscribe()` +directly. + +Existing brokers can keep implementing `listen()` as before. New flow-aware +brokers may use `get_subscribed_flows()` to subscribe to queues, topics, +subjects or streams while the routing rules stay in the router. + +`get_subscribed_flows()` returns the broker `default_flow` plus explicit +subscriptions. It deduplicates by flow identity and checks broker options for +conflicts. + +## Scheduler and requeue + +`ScheduledTask` remains a transport-neutral invocation payload. It stores task +name, labels, args, kwargs, optional task id and schedule timing fields; it does +not store broker objects, broker names, route objects, flows or +transport-specific flow options. + +When a scheduled task is ready, `TaskiqScheduler` resolves the task route in the +scheduler process through the scheduler broker's router. Route changes made +before `on_ready(...)` runs affect that scheduled dispatch. If no router route +exists for the scheduled task, Taskiq keeps the old behavior and sends through +the scheduler broker. + +`with_broker(...)`, `with_route(...)` and `with_flow(...)` on a kicker are not +persisted schedule route metadata. They can affect `CreatedSchedule.kiq()`, +which is an immediate queued invocation helper, but they do not change the +stored schedule payload. + +`Context.requeue()` is current-broker sticky. It sends the same task message +through the broker currently executing the task, even if the task's default +route now points to another broker. Requeue flow selection is explicit +low-level override, same-broker task route flow, current broker `default_flow`, +then no flow. Taskiq does not currently preserve the inbound source flow because +flow provenance is not stored in `TaskiqMessage` or `Context`. + +## Shared task declarations + +Libraries can declare tasks without importing an application broker: + +```python +from taskiq import task_builder + + +@task_builder("billing.calculate_total", domain="billing") +async def calculate_total(price: int, quantity: int) -> int: + return price * quantity +``` + +The final application binds the task definition to its router and broker: + +```python +billing_flow = Flow("billing.tasks") + +registered = router.register_task( + calculate_total, + broker=billing_broker, + flow=billing_flow, +) + +router.subscribe( + billing_broker, + billing_flow, + registered, +) + +result = await registered.kiq(19, 3) +``` + +Task name and labels belong to the shared declaration. Register a +`TaskDefinition` without `task_name` or label overrides; Taskiq rejects those +overrides instead of silently ignoring them. + +The unbound task can still run locally: + +```python +total = await calculate_total.call(19, 3) +``` + +Shared task declarations can also provide a custom task class for one task: + +```python +from typing import Any + +from taskiq import AsyncTaskiqDecoratedTask, task_builder + + +class TracingTask(AsyncTaskiqDecoratedTask[Any, Any]): + def tracing_name(self) -> str: + return self.task_name + + +@task_builder("billing.traced_charge", base_cls=TracingTask) +async def traced_charge(user_id: int) -> None: + ... +``` + +When the final application registers this definition, Taskiq creates the bound +task using `TracingTask`. If `base_cls` is not provided, Taskiq uses the native +decorated task class. + +For low-level integrations, `TaskDefinition.message(...)` builds a +`TaskiqMessage` without binding the task. It uses the same argument and label +preparation contract as a normal `.kicker().prepare(...)` invocation. + +## Compatibility and migration notes + +The routing/flow branch keeps the old broker-first task API working: + +```python +@broker.task(task_name="billing.charge") +async def charge(user_id: int) -> None: + ... + +await charge.kiq(42) +await charge.kicker().with_labels(source="api").kiq(42) +``` + +Migration guidance for application code: + +- Keep using `@broker.task(...)` when one broker owns task declaration and send + path. No router code is required for this case. +- Use one shared `TaskiqRouter` when several brokers must share task registry + and routing policy. +- Pass broker objects to routing APIs. If configuration stores broker names, + resolve them with `router.get_broker(name)` before calling `route_task(...)`. +- Replace `route_task(..., subscribe=True)` usage with explicit + `router.subscribe(...)`. +- Do not use labels as a transport routing schema. Labels remain task metadata. +- Register a `TaskDefinition` without `task_name` or label overrides; plain + callable registration still supports those overrides. + +Migration guidance for broker packages: + +- Existing brokers that implement only `kick(message)` and `listen()` remain + valid. The default `kick_to_flow(message, flow=None)` calls `kick(message)`. +- Flow-aware brokers can override `kick_to_flow(...)` and map `flow.name` to a + queue, topic, subject or stream. +- Flow-aware listeners can call `get_subscribed_flows()` during listen setup to + add router-owned subscriptions while keeping old default listen behavior when + there are no explicit subscriptions. +- Broker-specific declaration settings belong in broker-specific flow objects + and `broker_options()`, not in task names, labels or scheduler payloads. +- Ack/nack behavior, serializer/formatter behavior, retry, middleware hooks and + result backend ownership remain broker responsibilities. + +Scheduler and requeue compatibility: + +- `ScheduledTask` does not store routes, broker names, flow objects or + broker-specific flow options in this iteration. +- Scheduled dispatch resolves the route late in the scheduler process through + the scheduler broker's router. +- `Context.requeue()` sends through the broker currently executing the task. It + does not re-resolve a task route across brokers. +- Taskiq does not preserve inbound source-flow provenance yet because that data + is not stored in `TaskiqMessage` or `Context`. + +## Examples + +See executable examples for the two main usage shapes: + +- [Multiple brokers and explicit routes](../examples/router/multiple_brokers.py) + shows one task routed through several brokers with explicit subscription + ownership. +- [Shared task package](../examples/router/shared_task_package.py) shows + `task_builder(...)`, a custom `base_cls`, later registration and a prepared + queued invocation. diff --git a/docs/guide/scheduling-tasks.md b/docs/guide/scheduling-tasks.md index 44e646f2..ab6ce848 100644 --- a/docs/guide/scheduling-tasks.md +++ b/docs/guide/scheduling-tasks.md @@ -159,13 +159,26 @@ If you want to pass additional labels, you can call these methods on the `Kicker ) ``` -::: warning Cool warning! +::: warning Routing warning -The `with_broker` method won't do anything in this case, since we have a broker assigned to each scheduler. +`with_broker`, `with_route` and `with_flow` are immediate invocation settings. +They are not stored in the `ScheduledTask` payload. Scheduled dispatch resolves +the route later in the scheduler process through the scheduler broker's router. +Use router configuration available to the scheduler process when a scheduled +task must go to a specific broker or flow. ::: +The scheduler keeps route data out of schedule sources on purpose. A stored +`ScheduledTask` remains a transport-neutral payload with task name, labels, +args, kwargs and schedule timing fields. If you change a router rule before a +schedule becomes ready, the scheduler uses the updated route at send time. If no +route exists, the scheduler sends through its own broker as before. + Each of these methods return you an instance of the `CreatedSchedule` class. This object has unique schedule ID and some helper methods. For example, you can use the `unschedule` method to remove the schedule from the source. +Calling `kiq()` on a `CreatedSchedule` is an immediate queued invocation helper; +it uses the original kicker settings for that immediate send and does not change +the stored schedule payload. ```python schedule = await my_task.schedule_by_time( diff --git a/taskiq/__init__.py b/taskiq/__init__.py index 8aeb203c..6ebb34ab 100644 --- a/taskiq/__init__.py +++ b/taskiq/__init__.py @@ -24,7 +24,7 @@ TaskiqError, TaskiqResultTimeoutError, ) -from taskiq.flow import Flow, FlowProtocol +from taskiq.flow import Flow, FlowIdentity, FlowProtocol, get_flow_identity from taskiq.funcs import gather from taskiq.kicker import PreparedKiq from taskiq.message import BrokerMessage, TaskiqMessage @@ -34,7 +34,7 @@ SmartRetryMiddleware, ) from taskiq.result import TaskiqResult -from taskiq.router import TaskiqRoute, TaskiqRouter +from taskiq.router import TaskiqRoute, TaskiqRouter, TaskiqSubscription from taskiq.scheduler.scheduled_task import ScheduledTask from taskiq.scheduler.scheduler import TaskiqScheduler from taskiq.state import TaskiqState @@ -53,6 +53,7 @@ "BrokerMessage", "Context", "Flow", + "FlowIdentity", "FlowProtocol", "InMemoryBroker", "NoResultError", @@ -79,9 +80,11 @@ "TaskiqRouter", "TaskiqScheduler", "TaskiqState", + "TaskiqSubscription", "ZeroMQBroker", "__version__", "async_shared_broker", "gather", + "get_flow_identity", "task_builder", ] diff --git a/taskiq/abc/broker.py b/taskiq/abc/broker.py index 42a1817b..a25a98a6 100644 --- a/taskiq/abc/broker.py +++ b/taskiq/abc/broker.py @@ -176,6 +176,16 @@ def get_all_tasks(self) -> dict[str, AsyncTaskiqDecoratedTask[Any, Any]]: **self.local_task_registry, } + def get_subscribed_flows(self) -> tuple[FlowProtocol, ...]: + """ + Return flows this broker should subscribe to. + + Existing brokers can keep their current `listen` implementation. New + flow-aware brokers may use this method to configure queue, topic, + stream or subject subscriptions from router-owned rules. + """ + return self.router.get_broker_flows(self) + def add_dependency_context(self, new_ctx: dict[Any, Any]) -> None: """ Add first-level dependencies. @@ -351,40 +361,12 @@ def make_decorated_task( def inner( func: Callable[_FuncParams, _ReturnType], ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: - nonlocal inner_task_name - if inner_task_name is None: - fmodule = func.__module__ - if fmodule == "__main__": # pragma: no cover - fmodule = ".".join( - os.path.normpath(sys.argv[0]) - .removesuffix(".py") - .split(os.path.sep), - ) - fname = func.__name__ - if fname == "": - fname = f"lambda_{uuid4().hex}" - inner_task_name = f"{fmodule}:{fname}" - wrapper = wraps(func) - - sign = get_type_hints(func) - return_type = None - if "return" in sign: - return_type = sign["return"] - - decorated_task = wrapper( - self.decorator_class( - broker=self, - original_func=func, - labels=inner_labels, - task_name=inner_task_name, - return_type=return_type, # type: ignore - ), + return self._decorate_task( + func, + task_name=inner_task_name, + labels=inner_labels, ) - self._register_task(decorated_task.task_name, decorated_task) # type: ignore - - return decorated_task # type: ignore - return inner if callable(task_name): @@ -423,12 +405,87 @@ def register_task( :returns: registered task. """ if isinstance(func, TaskDefinition): + if task_name is not None or labels: + raise ValueError( + "TaskDefinition already defines task_name and labels. " + "Register it without task_name or label overrides.", + ) return self.router.register_task( func, broker=self, flow=self.default_flow, ) - return self.task(task_name=task_name, **labels)(func) + return self._decorate_task( + func, + task_name=task_name, + labels=labels or {}, + ) + + def bind_task_definition( + self, + task: TaskDefinition[_FuncParams, _ReturnType], + *, + register: bool = True, + ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: + """Bind an unbound task definition to this broker.""" + task_cls = AsyncTaskiqDecoratedTask if task.base_cls is None else task.base_cls + return self._decorate_task( + task.original_func, + task_name=task.task_name, + labels=dict(task.labels), + base_cls=task_cls, + register=register, + ) + + def _decorate_task( + self, + func: Callable[_FuncParams, _ReturnType], + task_name: str | None, + labels: dict[str, Any], + base_cls: type[AsyncTaskiqDecoratedTask[Any, Any]] | None = None, + *, + register: bool = True, + ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: + """Build and register a decorated task object.""" + real_task_name = task_name + if real_task_name is None: + fmodule = func.__module__ + if fmodule == "__main__": # pragma: no cover + fmodule = ".".join( + os.path.normpath(sys.argv[0]) + .removesuffix(".py") + .split(os.path.sep), + ) + fname = func.__name__ + if fname == "": + fname = f"lambda_{uuid4().hex}" + real_task_name = f"{fmodule}:{fname}" + + task_cls = self.decorator_class if base_cls is None else base_cls + if not isinstance(task_cls, type) or not issubclass( + task_cls, + AsyncTaskiqDecoratedTask, + ): + raise TypeError("base_cls must be a subclass of AsyncTaskiqDecoratedTask.") + + sign = get_type_hints(func) + return_type = None + if "return" in sign: + return_type = sign["return"] + + decorated_task = wraps(func)( + task_cls( + broker=self, + original_func=func, + labels=labels, + task_name=real_task_name, + return_type=return_type, # type: ignore + ), + ) + + if register: + self._register_task(decorated_task.task_name, decorated_task) # type: ignore + return decorated_task # type: ignore def on_event(self, *events: TaskiqEvents) -> Callable[[EventHandler], EventHandler]: """ @@ -580,12 +637,28 @@ 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, ) + self._store_task(task_name, task) + + def _store_task( + self, + task_name: str, + task: AsyncTaskiqDecoratedTask[Any, Any], + ) -> None: + """ + Store a decorated task in this broker. + + Router-managed binding uses this method after router state has accepted + the task, so shared task registration does not re-enter router + registration while it is already in progress. + """ + if task.broker != self: + raise TaskBrokerMismatchError(broker=task.broker) + self.local_task_registry[task_name] = task async def __aenter__(self) -> None: """Starts the broker as ctx manager.""" diff --git a/taskiq/flow.py b/taskiq/flow.py index 8754492e..1e5c06b1 100644 --- a/taskiq/flow.py +++ b/taskiq/flow.py @@ -1,17 +1,28 @@ from collections.abc import Mapping from dataclasses import dataclass, field, replace +from types import MappingProxyType from typing import Protocol, runtime_checkable -__all__ = ("Flow", "FlowProtocol") +__all__ = ("Flow", "FlowIdentity", "FlowProtocol", "get_flow_identity") + + +@dataclass(frozen=True, slots=True) +class FlowIdentity: + """Stable identity for one logical flow in routing and listen plans.""" + + name: str @runtime_checkable class FlowProtocol(Protocol): """Transport-neutral flow contract accepted by routers and brokers.""" - name: str + @property + def name(self) -> str: + """Return transport-neutral flow name.""" + ... - def broker_options(self, broker_name: str) -> Mapping[str, object]: + def broker_options(self) -> Mapping[str, object]: """Return options relevant for a concrete broker implementation.""" ... @@ -33,10 +44,30 @@ class Flow: hash=False, ) + def __post_init__(self) -> None: + """Freeze options to keep Flow a stable value object.""" + object.__setattr__(self, "options", MappingProxyType(dict(self.options))) + def with_options(self, **options: object) -> "Flow": """Return the same flow with additional generic options.""" return replace(self, options={**self.options, **options}) - def broker_options(self, broker_name: str) -> Mapping[str, object]: + @property + def identity(self) -> FlowIdentity: + """Return routing and subscription identity for this flow.""" + return get_flow_identity(self) + + def broker_options(self) -> Mapping[str, object]: """Return transport options for broker-specific implementations.""" return dict(self.options) + + +def get_flow_identity(flow: FlowProtocol) -> FlowIdentity: + """ + Return stable identity for a flow. + + Flow identity is intentionally transport-neutral and based on the logical + flow name. Broker-specific declaration options are validated separately by + the subscription plan. + """ + return FlowIdentity(name=flow.name) diff --git a/taskiq/kicker.py b/taskiq/kicker.py index d0182e4c..8debca3e 100644 --- a/taskiq/kicker.py +++ b/taskiq/kicker.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Coroutine -from dataclasses import asdict, is_dataclass +from dataclasses import replace from datetime import datetime, timedelta from logging import getLogger from types import CoroutineType @@ -15,16 +15,12 @@ overload, ) -from pydantic import BaseModel - from taskiq.abc.middleware import TaskiqMiddleware from taskiq.abc.result_backend import AsyncResultBackend -from taskiq.compat import model_dump from taskiq.exceptions import SendTaskError from taskiq.flow import FlowProtocol -from taskiq.labels import prepare_label -from taskiq.message import TaskiqMessage -from taskiq.router import TaskiqRouter +from taskiq.message import TaskiqMessage, _build_taskiq_message +from taskiq.router import TaskiqRoute, TaskiqRouter from taskiq.scheduler.created_schedule import CreatedSchedule from taskiq.scheduler.scheduled_task import CronSpec, ScheduledTask from taskiq.task import AsyncTaskiqTask @@ -48,13 +44,25 @@ def __init__( self, kicker: AsyncKicker[..., _ReturnType], message: TaskiqMessage, + broker: AsyncBroker, + route: TaskiqRoute | None, + flow: FlowProtocol | None, ) -> None: self.kicker = kicker self.message = message + self.broker = broker + self.route = route + self.flow = flow async def kiq(self) -> AsyncTaskiqTask[_ReturnType]: """Send prepared invocation.""" - return await self.kicker.kiq_message(self.message) + return await self.kicker.kiq_message( + self.message, + broker=self.broker, + route=self.route, + flow=self.flow, + use_current_route=False, + ) class AsyncKicker(Generic[_FuncParams, _ReturnType]): @@ -73,8 +81,9 @@ def __init__( self.custom_task_id: str | None = None self.custom_schedule_id: str | None = None self.return_type = return_type - self.route_broker: AsyncBroker | str | None = None + self.route: TaskiqRoute | None = None self.route_flow: FlowProtocol | None = None + self._broker_overridden = False def with_labels( self, @@ -132,7 +141,9 @@ def with_broker( :return: Kicker with new broker. """ self.broker = broker - self.route_broker = broker + self.route = None + self.route_flow = None + self._broker_overridden = True return self def with_flow( @@ -145,25 +156,23 @@ def with_flow( :param flow: flow to send message to. :return: Kicker with a route flow override. """ + if self.route is not None: + self.route = replace(self.route, flow=flow) self.route_flow = flow return self def with_route( self, - broker: AsyncBroker | str, - flow: FlowProtocol | None, + route: TaskiqRoute, ) -> AsyncKicker[_FuncParams, _ReturnType]: """ - Replace broker and flow for the current invocation. + Replace route for the current invocation. - :param broker: broker instance or broker name. - Broker instances are preferred; names are kept for configuration - and backward-compatible lookup. - :param flow: flow to send message to. + :param route: route to send message through. :return: Kicker with a route override. """ - self.route_broker = broker - self.route_flow = flow + self.route = route + self.route_flow = route.flow return self def prepare( @@ -178,7 +187,14 @@ def prepare( :param kwargs: function's key word arguments. :return: prepared task invocation. """ - return PreparedKiq(self, self._prepare_message(*args, **kwargs)) + broker, route, flow = self._prepare_route_snapshot() + return PreparedKiq( + self, + self._prepare_message(*args, **kwargs), + broker=broker, + route=route, + flow=flow, + ) @overload async def kiq( @@ -230,29 +246,56 @@ async def kiq( async def kiq_message( self, message: TaskiqMessage, + *, + broker: AsyncBroker | None = None, + route: TaskiqRoute | None = None, + flow: FlowProtocol | None = None, + use_current_route: bool = True, ) -> AsyncTaskiqTask[_ReturnType]: """Send a prepared message.""" try: - router = getattr(self.broker, "router", None) + target_broker = broker or self.broker + if use_current_route: + target_route = self.route if route is None else route + target_flow = self.route_flow if flow is None else flow + else: + target_route = route + target_flow = flow + if target_route is not None: + if broker is not None and broker is not target_route.broker: + raise ValueError("Pass either route or broker override.") + target_broker = target_route.broker + target_flow = None + router = getattr(target_broker, "router", None) if isinstance(router, TaskiqRouter): + broker_override = None + if target_route is None and ( + broker is not None or self._broker_overridden + ): + broker_override = target_broker return await router.kiq( message, - broker=self.route_broker, - flow=self.route_flow, + route=target_route, + broker=broker_override, + flow=target_flow, return_type=self.return_type, ) - return await self._legacy_kiq(message) + return await self._legacy_kiq(target_broker, message) except Exception as exc: raise SendTaskError from exc - async def _legacy_kiq(self, message: TaskiqMessage) -> AsyncTaskiqTask[_ReturnType]: + async def _legacy_kiq( + self, + broker: AsyncBroker, + message: TaskiqMessage, + ) -> AsyncTaskiqTask[_ReturnType]: """ Send message through the pre-router broker path. This keeps middleware tests and external broker-like mocks compatible while real AsyncBroker instances use TaskiqRouter. """ - middlewares = getattr(self.broker, "middlewares", []) + middlewares = getattr(broker, "middlewares", []) if not isinstance(middlewares, list): middlewares = [] @@ -260,7 +303,7 @@ async def _legacy_kiq(self, message: TaskiqMessage) -> AsyncTaskiqTask[_ReturnTy if middleware.__class__.pre_send != TaskiqMiddleware.pre_send: message = await maybe_awaitable(middleware.pre_send(message)) - await self.broker.kick(self.broker.formatter.dumps(message)) + await broker.kick(broker.formatter.dumps(message)) for middleware in reversed(middlewares): if middleware.__class__.post_send != TaskiqMiddleware.post_send: @@ -270,11 +313,30 @@ async def _legacy_kiq(self, message: TaskiqMessage) -> AsyncTaskiqTask[_ReturnTy task_id=message.task_id, result_backend=cast( AsyncResultBackend[_ReturnType], - self.broker.result_backend, + broker.result_backend, ), return_type=self.return_type, ) + def _prepare_route_snapshot( + self, + ) -> tuple[AsyncBroker, TaskiqRoute | None, FlowProtocol | None]: + """Resolve the route that a prepared invocation must keep.""" + router = getattr(self.broker, "router", None) + if not isinstance(router, TaskiqRouter): + return self.broker, None, self.route_flow + + if self.route is not None: + return self.route.broker, self.route, None + + broker_override = self.broker if self._broker_overridden else None + route = router.resolve_route( + self.task_name, + broker=broker_override, + flow=self.route_flow, + ) + return route.broker, route, None + async def schedule_by_cron( self, source: ScheduleSource, @@ -378,27 +440,6 @@ async def schedule_by_time( await source.add_schedule(scheduled) return CreatedSchedule(self, source, scheduled) - @classmethod - def _prepare_arg(cls, arg: Any) -> Any: - """ - Parses argument if possible. - - This function is used to construct dicts - from pydantic models or dataclasses. - - :param arg: argument to format. - :return: Formatted argument. - """ - if isinstance(arg, BaseModel): - arg = model_dump(arg) - if is_dataclass(arg): - if isinstance(arg, type): - raise ValueError( - f"Cannot serialize types. The {arg} is not serializable.", - ) - arg = asdict(arg) - return arg - def _prepare_message( self, *args: Any, @@ -411,27 +452,14 @@ def _prepare_message( :param kwargs: function's kwargs. :return: constructed message. """ - formatted_args = [] - formatted_kwargs = {} - labels = {} - labels_types = {} - for arg in args: - formatted_args.append(self._prepare_arg(arg)) - for kwarg_name, kwarg_val in kwargs.items(): - formatted_kwargs[kwarg_name] = self._prepare_arg(kwarg_val) - - for label, label_val in self.labels.items(): - labels[label], labels_types[label] = prepare_label(label_val) - task_id = self.custom_task_id if task_id is None: task_id = self.broker.id_generator() - return TaskiqMessage( + return _build_taskiq_message( task_id=task_id, task_name=self.task_name, - labels=labels, - labels_types=labels_types, - args=formatted_args, - kwargs=formatted_kwargs, + labels=self.labels, + args=args, + kwargs=kwargs, ) diff --git a/taskiq/message.py b/taskiq/message.py index 39008eb8..00ed4791 100644 --- a/taskiq/message.py +++ b/taskiq/message.py @@ -1,8 +1,11 @@ +from collections.abc import Iterable, Mapping +from dataclasses import asdict, is_dataclass from typing import Any from pydantic import BaseModel -from taskiq.labels import parse_label +from taskiq.compat import model_dump +from taskiq.labels import parse_label, prepare_label class TaskiqMessage(BaseModel): @@ -35,6 +38,61 @@ def parse_labels(self) -> None: self.labels[label] = parse_label(self.labels[label], label_type) +def _prepare_message_arg(arg: Any) -> Any: + """ + Prepare invocation argument for message payloads. + + :param arg: argument to prepare. + :return: serializable argument representation. + """ + if isinstance(arg, BaseModel): + arg = model_dump(arg) + if is_dataclass(arg): + if isinstance(arg, type): + raise ValueError( + f"Cannot serialize types. The {arg} is not serializable.", + ) + arg = asdict(arg) + return arg + + +def _build_taskiq_message( + task_id: str, + task_name: str, + labels: Mapping[str, Any], + args: Iterable[Any], + kwargs: Mapping[str, Any], +) -> TaskiqMessage: + """ + Build a taskiq message using the common invocation contract. + + :param task_id: task id. + :param task_name: task name. + :param labels: task invocation labels. + :param args: task positional arguments. + :param kwargs: task keyword arguments. + :return: prepared taskiq message. + """ + prepared_labels: dict[str, Any] = {} + labels_types: dict[str, int] = {} + for label_name, label_value in labels.items(): + prepared_labels[label_name], labels_types[label_name] = prepare_label( + label_value, + ) + + return TaskiqMessage( + task_id=task_id, + task_name=task_name, + labels=prepared_labels, + labels_types=labels_types, + args=[_prepare_message_arg(arg) for arg in args], + kwargs={ + kwarg_name: _prepare_message_arg(kwarg_value) + for kwarg_name, kwarg_value in kwargs.items() + }, + ) + + class BrokerMessage(BaseModel): """Format of messages for brokers.""" diff --git a/taskiq/router.py b/taskiq/router.py index 9f770f90..b62d16d3 100644 --- a/taskiq/router.py +++ b/taskiq/router.py @@ -1,51 +1,81 @@ from __future__ import annotations +import warnings from collections.abc import Callable -from dataclasses import dataclass -from logging import getLogger -from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, cast, overload +from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, overload -from taskiq.abc.middleware import TaskiqMiddleware -from taskiq.abc.result_backend import AsyncResultBackend from taskiq.flow import FlowProtocol from taskiq.message import TaskiqMessage +from taskiq.routing import ( + BrokerRegistry, + RouterDispatcher, + RouteRegistry, + SubscriptionPlan, + TaskiqRoute, + TaskiqSubscription, + TaskRegistry, +) +from taskiq.routing.references import resolve_task_name from taskiq.task import AsyncTaskiqTask from taskiq.task_builder import TaskDefinition -from taskiq.utils import maybe_awaitable +from taskiq.warnings import TaskiqDeprecationWarning if TYPE_CHECKING: # pragma: no cover from taskiq.abc.broker import AsyncBroker from taskiq.decor import AsyncTaskiqDecoratedTask -__all__ = ("TaskiqRoute", "TaskiqRouter") +__all__ = ("TaskiqRoute", "TaskiqRouter", "TaskiqSubscription") + +TaskiqRoute.__module__ = __name__ +TaskiqSubscription.__module__ = __name__ _FuncParams = ParamSpec("_FuncParams") _ReturnType = TypeVar("_ReturnType") -logger = getLogger("taskiq.router") +class TaskiqRouter: + """Facade for task registry, routing policy, subscriptions and dispatch.""" -@dataclass(frozen=True, slots=True) -class TaskiqRoute: - """Resolved outbound route for a task invocation.""" + def __init__(self) -> None: + self._brokers = BrokerRegistry() + self._tasks = TaskRegistry() + self._routes = RouteRegistry(self._brokers) + self._subscriptions = SubscriptionPlan(self._brokers) + self._dispatcher = RouterDispatcher(self._routes) - broker: AsyncBroker - flow: FlowProtocol | None = None + @property + def brokers(self) -> dict[str, AsyncBroker]: + """Return mutable broker registry for compatibility.""" + return self._brokers.brokers @property - def broker_name(self) -> str: - """Return registered broker name for compatibility and diagnostics.""" - return self.broker.broker_name + def default_broker(self) -> AsyncBroker | None: + """Return default broker for compatibility.""" + return self._brokers.default_broker + @default_broker.setter + def default_broker(self, broker: AsyncBroker | None) -> None: + self._brokers.default_broker = broker -class TaskiqRouter: - """Registry and routing layer shared by one or more brokers.""" + @property + def task_registry(self) -> dict[str, AsyncTaskiqDecoratedTask[Any, Any]]: + """Return mutable task registry for compatibility.""" + return self._tasks.tasks - def __init__(self) -> None: - self.brokers: dict[str, AsyncBroker] = {} - self.default_broker_name: str | None = None - self.task_registry: dict[str, AsyncTaskiqDecoratedTask[Any, Any]] = {} - self.routes: dict[str, TaskiqRoute] = {} + @property + def routes(self) -> dict[str, TaskiqRoute]: + """Return mutable route registry for compatibility.""" + return self._routes.routes + + @property + def subscriptions(self) -> list[TaskiqSubscription]: + """Return mutable subscription registry for compatibility.""" + return self._subscriptions.subscriptions + + @property + def default_broker_name(self) -> str | None: + """Return default broker name for compatibility and diagnostics.""" + return self._brokers.default_broker_name def set_broker( self, @@ -54,28 +84,30 @@ def set_broker( default_flow: FlowProtocol | None = None, ) -> str: """Register broker as a transport in this router.""" - broker_name = name or broker.__class__.__name__ - registered = self.brokers.get(broker_name) - if registered is not None and registered is not broker: + if getattr(broker, "router", self) is not self: raise ValueError( - f"Broker name {broker_name!r} is already registered. " - "Please provide an explicit unique broker_name.", + "Broker is attached to another router. " + "Pass router=... when creating the broker.", ) - self.brokers[broker_name] = broker - if self.default_broker_name is None: - self.default_broker_name = broker_name + broker_name = self._brokers.register(broker, name=name) + if default_flow is not None: + broker.default_flow = default_flow return broker_name + def get_broker(self, name: str) -> AsyncBroker: + """Return a broker by registered name.""" + return self._brokers.get(name) + def find_task( self, task_name: str, ) -> AsyncTaskiqDecoratedTask[Any, Any] | None: """Find a task by name.""" - return self.task_registry.get(task_name) + return self._tasks.find(task_name) def get_all_tasks(self) -> dict[str, AsyncTaskiqDecoratedTask[Any, Any]]: """Return all tasks registered in this router.""" - return dict(self.task_registry) + return self._tasks.get_all() def register_task( self, @@ -83,35 +115,36 @@ def register_task( AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType] | TaskDefinition[_FuncParams, _ReturnType] ), - broker: AsyncBroker | str | None = None, + broker: AsyncBroker | None = None, flow: FlowProtocol | None = None, ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: """Register a bound task or bind a task definition to a broker.""" if isinstance(task, TaskDefinition): - target_broker = self._resolve_broker(broker) - registered_task = target_broker.register_task( - task.original_func, - task_name=task.task_name, - **task.labels, + target_broker = self._brokers.resolve(broker) + registered_task = target_broker.bind_task_definition( + task, + register=False, + ) + self._register_bound_task( + registered_task, + broker=target_broker, + flow=flow, + ) + # Router and broker share this internal registration boundary. + target_broker._store_task( # noqa: SLF001 + registered_task.task_name, + registered_task, ) - if flow is not None: - self.route_task(task.task_name, broker=target_broker, flow=flow) return registered_task - self.task_registry[task.task_name] = task - route_broker: AsyncBroker | str | None = broker - if route_broker is None: - route_broker = getattr(task, "broker", None) - if route_broker is not None or flow is not None: - self.route_task(task.task_name, broker=route_broker, flow=flow) - return task + return self._register_bound_task(task, broker=broker, flow=flow) @overload def task( self, task_name: Callable[_FuncParams, _ReturnType], *, - broker: AsyncBroker | str | None = None, + broker: AsyncBroker | None = None, flow: FlowProtocol | None = None, **labels: Any, ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: ... @@ -121,7 +154,7 @@ def task( self, task_name: str | None = None, *, - broker: AsyncBroker | str | None = None, + broker: AsyncBroker | None = None, flow: FlowProtocol | None = None, **labels: Any, ) -> Callable[ @@ -133,7 +166,7 @@ def task( self, task_name: str | Callable[_FuncParams, _ReturnType] | None = None, *, - broker: AsyncBroker | str | None = None, + broker: AsyncBroker | None = None, flow: FlowProtocol | None = None, **labels: Any, ) -> Any: @@ -142,16 +175,15 @@ def task( def register( func: Callable[_FuncParams, _ReturnType], ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: - target_broker = self._resolve_broker(broker) + target_broker = self._brokers.resolve(broker) real_task_name = task_name if not callable(task_name) else None task = target_broker.task(task_name=real_task_name, **labels)(func) if flow is not None: - self.route_task(task.task_name, broker=target_broker, flow=flow) + self.route_task(task, broker=target_broker, flow=flow) return task if callable(task_name): function = task_name - task_name = None return register(function) return register @@ -159,77 +191,74 @@ def register( def route_task( self, task: str | AsyncTaskiqDecoratedTask[Any, Any], - broker: AsyncBroker | str | None = None, + broker: AsyncBroker | None = None, flow: FlowProtocol | None = None, + *, + subscribe: bool = False, ) -> TaskiqRoute: """Set default outbound route for a task.""" - task_name = self._resolve_task_name(task) - target_broker = self._resolve_broker(broker) - route = TaskiqRoute(broker=target_broker, flow=flow) - self.routes[task_name] = route + task_name = resolve_task_name(task) + route = self._routes.set_route(task_name, broker=broker, flow=flow) + if subscribe: + warnings.warn( + "`route_task(..., subscribe=True)` is deprecated. " + "Use `router.subscribe(...)` to add inbound flow subscriptions.", + TaskiqDeprecationWarning, + stacklevel=2, + ) + if subscribe and route.flow is not None: + self.subscribe(route.broker, route.flow, task_name) return route def resolve_route( self, task: str | AsyncTaskiqDecoratedTask[Any, Any], - broker: AsyncBroker | str | None = None, + broker: AsyncBroker | None = None, flow: FlowProtocol | None = None, ) -> TaskiqRoute: """Resolve outbound route for a task invocation.""" - task_name = self._resolve_task_name(task) - if broker is not None: - target_broker = self._resolve_broker(broker) - route_flow = flow - if route_flow is None: - route_flow = self._broker_default_flow(target_broker) - return TaskiqRoute( - broker=target_broker, - flow=route_flow, - ) + return self._routes.resolve_route(task, broker=broker, flow=flow) - route = self.routes.get(task_name) - if route is not None: - if flow is None: - return route - return TaskiqRoute(broker=route.broker, flow=flow) - - target_broker = self._resolve_broker(None) - route_flow = flow - if route_flow is None: - route_flow = self._broker_default_flow(target_broker) - return TaskiqRoute( - broker=target_broker, - flow=route_flow, + def subscribe( + self, + broker: AsyncBroker, + flow: FlowProtocol, + *tasks: str | AsyncTaskiqDecoratedTask[Any, Any], + ) -> TaskiqSubscription: + """Register an inbound flow subscription for a broker.""" + task_names = tuple(resolve_task_name(task) for task in tasks) + return self._subscriptions.subscribe( + broker, + flow, + task_names, ) + def get_subscriptions( + self, + broker: AsyncBroker | None = None, + ) -> tuple[TaskiqSubscription, ...]: + """Return registered inbound subscriptions.""" + return self._subscriptions.get(broker) + + def get_broker_flows(self, broker: AsyncBroker) -> tuple[FlowProtocol, ...]: + """Return flows a broker should subscribe to.""" + return self._subscriptions.get_broker_flows(broker) + async def kiq( self, message: TaskiqMessage, *, - broker: AsyncBroker | str | None = None, + route: TaskiqRoute | None = None, + broker: AsyncBroker | None = None, flow: FlowProtocol | None = None, return_type: type[_ReturnType] | None = None, ) -> AsyncTaskiqTask[_ReturnType]: """Send message through the resolved broker and flow.""" - route = self.resolve_route(message.task_name, broker=broker, flow=flow) - target_broker = route.broker - - for middleware in target_broker.middlewares: - if middleware.__class__.pre_send != TaskiqMiddleware.pre_send: - message = await maybe_awaitable(middleware.pre_send(message)) - broker_message = target_broker.formatter.dumps(message) - await target_broker.kick_to_flow(broker_message, route.flow) - - for middleware in reversed(target_broker.middlewares): - if middleware.__class__.post_send != TaskiqMiddleware.post_send: - await maybe_awaitable(middleware.post_send(message)) - - return AsyncTaskiqTask( - task_id=message.task_id, - result_backend=cast( - AsyncResultBackend[_ReturnType], - target_broker.result_backend, - ), + return await self._dispatcher.kiq( + message, + route=route, + broker=broker, + flow=flow, return_type=return_type, ) @@ -237,51 +266,33 @@ async def requeue( self, message: TaskiqMessage, *, - broker: AsyncBroker | str | None = None, + route: TaskiqRoute | None = None, + broker: AsyncBroker | None = None, flow: FlowProtocol | None = None, ) -> None: """Send an existing message again through the resolved route.""" - route = self.resolve_route(message.task_name, broker=broker, flow=flow) - target_broker = route.broker - await target_broker.kick_to_flow( - target_broker.formatter.dumps(message), - route.flow, + await self._dispatcher.requeue( + message, + route=route, + broker=broker, + flow=flow, ) - def _resolve_broker(self, broker: AsyncBroker | str | None) -> AsyncBroker: - if isinstance(broker, str): - if broker not in self.brokers: - raise ValueError(f"Unknown broker {broker!r}.") - return self.brokers[broker] - - if broker is not None: - broker_name = getattr(broker, "broker_name", None) - if isinstance(broker_name, str): - registered_broker = self.brokers.get(broker_name) - if registered_broker is broker: - return broker - for registered_broker in self.brokers.values(): - if registered_broker is broker: - return registered_broker - raise ValueError("Broker is not registered in this router.") - - if self.default_broker_name is None: - raise ValueError("Router doesn't have registered brokers.") - return self.brokers[self.default_broker_name] - - def _resolve_broker_name(self, broker: AsyncBroker | str | None) -> str: - return self._resolve_broker(broker).broker_name - - def _resolve_task_name( + def _register_bound_task( self, - task: str | AsyncTaskiqDecoratedTask[Any, Any], - ) -> str: - if isinstance(task, str): - return task - task_name = getattr(task, "task_name", None) - if isinstance(task_name, str): - return task_name - raise TypeError("Route task must be a task name or decorated task.") + task: AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType], + broker: AsyncBroker | None = None, + flow: FlowProtocol | None = None, + ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: + route_broker = broker + if route_broker is None: + route_broker = getattr(task, "broker", None) + + route = None + if route_broker is not None or flow is not None: + route = self._routes.build_route(broker=route_broker, flow=flow) - def _broker_default_flow(self, broker: AsyncBroker) -> FlowProtocol | None: - return getattr(broker, "default_flow", None) + self._tasks.register(task) + if route is not None: + self._routes.set_resolved_route(task.task_name, route) + return task diff --git a/taskiq/routing/__init__.py b/taskiq/routing/__init__.py new file mode 100644 index 00000000..09dfe0dc --- /dev/null +++ b/taskiq/routing/__init__.py @@ -0,0 +1,17 @@ +"""Internal routing components used by the public router facade.""" + +from taskiq.routing.dispatcher import RouterDispatcher +from taskiq.routing.models import TaskiqRoute, TaskiqSubscription +from taskiq.routing.registries import BrokerRegistry, TaskRegistry +from taskiq.routing.routes import RouteRegistry +from taskiq.routing.subscriptions import SubscriptionPlan + +__all__ = ( + "BrokerRegistry", + "RouteRegistry", + "RouterDispatcher", + "SubscriptionPlan", + "TaskRegistry", + "TaskiqRoute", + "TaskiqSubscription", +) diff --git a/taskiq/routing/dispatcher.py b/taskiq/routing/dispatcher.py new file mode 100644 index 00000000..d5b1ad43 --- /dev/null +++ b/taskiq/routing/dispatcher.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, TypeVar, cast + +from taskiq.abc.middleware import TaskiqMiddleware +from taskiq.abc.result_backend import AsyncResultBackend +from taskiq.flow import FlowProtocol +from taskiq.message import TaskiqMessage +from taskiq.routing.models import TaskiqRoute +from taskiq.routing.routes import RouteRegistry +from taskiq.task import AsyncTaskiqTask +from taskiq.utils import maybe_awaitable + +if TYPE_CHECKING: # pragma: no cover + from taskiq.abc.broker import AsyncBroker + +__all__ = ("RouterDispatcher",) + +_ReturnType = TypeVar("_ReturnType") + + +class RouterDispatcher: + """Send task messages through resolved router routes.""" + + def __init__(self, routes: RouteRegistry) -> None: + self.routes = routes + + async def kiq( + self, + message: TaskiqMessage, + *, + route: TaskiqRoute | None = None, + broker: AsyncBroker | None = None, + flow: FlowProtocol | None = None, + return_type: type[_ReturnType] | None = None, + ) -> AsyncTaskiqTask[_ReturnType]: + """Send message through the resolved broker and flow.""" + target_route = self.routes.resolve_explicit_or_default_route( + message.task_name, + route=route, + broker=broker, + flow=flow, + ) + target_broker = target_route.broker + + for middleware in target_broker.middlewares: + if middleware.__class__.pre_send != TaskiqMiddleware.pre_send: + message = await maybe_awaitable(middleware.pre_send(message)) + broker_message = target_broker.formatter.dumps(message) + await target_broker.kick_to_flow(broker_message, target_route.flow) + + for middleware in reversed(target_broker.middlewares): + if middleware.__class__.post_send != TaskiqMiddleware.post_send: + await maybe_awaitable(middleware.post_send(message)) + + return AsyncTaskiqTask( + task_id=message.task_id, + result_backend=cast( + AsyncResultBackend[_ReturnType], + target_broker.result_backend, + ), + return_type=return_type, + ) + + async def requeue( + self, + message: TaskiqMessage, + *, + route: TaskiqRoute | None = None, + broker: AsyncBroker | None = None, + flow: FlowProtocol | None = None, + ) -> None: + """Send an existing message again through the resolved route.""" + target_route = self.routes.resolve_explicit_or_default_route( + message.task_name, + route=route, + broker=broker, + flow=flow, + ) + target_broker = target_route.broker + await target_broker.kick_to_flow( + target_broker.formatter.dumps(message), + target_route.flow, + ) diff --git a/taskiq/routing/models.py b/taskiq/routing/models.py new file mode 100644 index 00000000..439ae7da --- /dev/null +++ b/taskiq/routing/models.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from taskiq.flow import FlowProtocol + +if TYPE_CHECKING: # pragma: no cover + from taskiq.abc.broker import AsyncBroker + +__all__ = ("TaskiqRoute", "TaskiqSubscription") + + +@dataclass(frozen=True, slots=True) +class TaskiqRoute: + """Resolved outbound route for a task invocation.""" + + broker: AsyncBroker + flow: FlowProtocol | None = None + + @property + def broker_name(self) -> str: + """Return registered broker name for diagnostics.""" + return self.broker.broker_name + + +@dataclass(frozen=True, slots=True) +class TaskiqSubscription: + """Inbound flow subscription owned by a router.""" + + broker: AsyncBroker + flow: FlowProtocol + task_names: frozenset[str] = field(default_factory=frozenset) + + @property + def broker_name(self) -> str: + """Return registered broker name for diagnostics.""" + return self.broker.broker_name diff --git a/taskiq/routing/references.py b/taskiq/routing/references.py new file mode 100644 index 00000000..f43270bc --- /dev/null +++ b/taskiq/routing/references.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: # pragma: no cover + from taskiq.decor import AsyncTaskiqDecoratedTask + +__all__ = ("resolve_task_name",) + + +def resolve_task_name(task: str | AsyncTaskiqDecoratedTask[Any, Any]) -> str: + """Resolve a task reference to a task name.""" + if isinstance(task, str): + return task + task_name = getattr(task, "task_name", None) + if isinstance(task_name, str): + return task_name + raise TypeError("Route task must be a task name or decorated task.") diff --git a/taskiq/routing/registries.py b/taskiq/routing/registries.py new file mode 100644 index 00000000..9226cdaa --- /dev/null +++ b/taskiq/routing/registries.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from taskiq.flow import FlowProtocol + +if TYPE_CHECKING: # pragma: no cover + from taskiq.abc.broker import AsyncBroker + from taskiq.decor import AsyncTaskiqDecoratedTask + +__all__ = ("BrokerRegistry", "TaskRegistry") + + +class BrokerRegistry: + """Registry of brokers attached to one router.""" + + def __init__(self) -> None: + self.brokers: dict[str, AsyncBroker] = {} + self.default_broker: AsyncBroker | None = None + + @property + def default_broker_name(self) -> str | None: + """Return default broker name for compatibility and diagnostics.""" + if self.default_broker is None: + return None + return self.default_broker.broker_name + + def register(self, broker: AsyncBroker, name: str | None = None) -> str: + """Register broker as a transport in this registry.""" + broker_name = name or broker.__class__.__name__ + registered = self.brokers.get(broker_name) + if registered is not None and registered is not broker: + raise ValueError( + f"Broker name {broker_name!r} is already registered. " + "Please provide an explicit unique broker_name.", + ) + for registered_name, registered_broker in self.brokers.items(): + if registered_broker is broker and registered_name != broker_name: + raise ValueError( + f"Broker is already registered as {registered_name!r}.", + ) + + self.brokers[broker_name] = broker + if self.default_broker is None: + self.default_broker = broker + return broker_name + + def get(self, name: str) -> AsyncBroker: + """Return a broker by registered name.""" + try: + return self.brokers[name] + except KeyError as exc: + raise ValueError(f"Unknown broker {name!r}.") from exc + + def resolve(self, broker: AsyncBroker | None) -> AsyncBroker: + """Resolve an explicit broker or return the default broker.""" + if isinstance(broker, str): + raise TypeError( + "Broker string references are not accepted here. " + "Use router.get_broker(name) and pass the broker object.", + ) + + if broker is not None: + for registered_broker in self.brokers.values(): + if registered_broker is broker: + return registered_broker + raise ValueError("Broker is not registered in this router.") + + if self.default_broker is None: + raise ValueError("Router doesn't have registered brokers.") + return self.default_broker + + def default_flow(self, broker: AsyncBroker) -> FlowProtocol | None: + """Return broker-owned default flow.""" + return getattr(broker, "default_flow", None) + + +class TaskRegistry: + """Registry of tasks attached to one router.""" + + def __init__(self) -> None: + self.tasks: dict[str, AsyncTaskiqDecoratedTask[Any, Any]] = {} + + def register( + self, + task: AsyncTaskiqDecoratedTask[Any, Any], + ) -> None: + """Register a decorated task and reject name conflicts.""" + existing_task = self.tasks.get(task.task_name) + if existing_task is not None and existing_task is not task: + raise ValueError( + f"Task name {task.task_name!r} is already registered " + "in this router.", + ) + self.tasks[task.task_name] = task + + def find( + self, + task_name: str, + ) -> AsyncTaskiqDecoratedTask[Any, Any] | None: + """Find a task by name.""" + return self.tasks.get(task_name) + + def get_all(self) -> dict[str, AsyncTaskiqDecoratedTask[Any, Any]]: + """Return all tasks registered in this registry.""" + return dict(self.tasks) diff --git a/taskiq/routing/routes.py b/taskiq/routing/routes.py new file mode 100644 index 00000000..96a866f8 --- /dev/null +++ b/taskiq/routing/routes.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +from dataclasses import replace +from typing import TYPE_CHECKING, Any + +from taskiq.flow import FlowProtocol +from taskiq.routing.models import TaskiqRoute +from taskiq.routing.references import resolve_task_name +from taskiq.routing.registries import BrokerRegistry + +if TYPE_CHECKING: # pragma: no cover + from taskiq.abc.broker import AsyncBroker + from taskiq.decor import AsyncTaskiqDecoratedTask + +__all__ = ("RouteRegistry",) + + +class RouteRegistry: + """Outbound route policy for task invocations.""" + + def __init__(self, brokers: BrokerRegistry) -> None: + self.brokers = brokers + self.routes: dict[str, TaskiqRoute] = {} + + def build_route( + self, + broker: AsyncBroker | None = None, + flow: FlowProtocol | None = None, + ) -> TaskiqRoute: + """Build a route without mutating route state.""" + target_broker = self.brokers.resolve(broker) + route_flow = flow + if route_flow is None: + route_flow = self.brokers.default_flow(target_broker) + return TaskiqRoute(broker=target_broker, flow=route_flow) + + def set_route( + self, + task: str | AsyncTaskiqDecoratedTask[Any, Any], + broker: AsyncBroker | None = None, + flow: FlowProtocol | None = None, + ) -> TaskiqRoute: + """Set default outbound route for a task.""" + task_name = resolve_task_name(task) + route = self.build_route(broker=broker, flow=flow) + self.routes[task_name] = route + return route + + def set_resolved_route(self, task_name: str, route: TaskiqRoute) -> TaskiqRoute: + """Store an already resolved route for a task.""" + self.routes[task_name] = route + return route + + def resolve_route( + self, + task: str | AsyncTaskiqDecoratedTask[Any, Any], + broker: AsyncBroker | None = None, + flow: FlowProtocol | None = None, + ) -> TaskiqRoute: + """Resolve outbound route for a task invocation.""" + task_name = resolve_task_name(task) + if broker is not None: + target_broker = self.brokers.resolve(broker) + route_flow = self._resolve_flow_for_broker_override( + task_name, + target_broker, + flow, + ) + return TaskiqRoute(broker=target_broker, flow=route_flow) + + route = self.routes.get(task_name) + if route is not None: + if flow is None: + return route + return replace(route, flow=flow) + + return self.build_route(flow=flow) + + def resolve_explicit_or_default_route( + self, + task: str | AsyncTaskiqDecoratedTask[Any, Any], + *, + route: TaskiqRoute | None, + broker: AsyncBroker | None, + flow: FlowProtocol | None, + ) -> TaskiqRoute: + """Resolve a route from explicit route or broker/flow overrides.""" + if route is not None: + if broker is not None or flow is not None: + raise ValueError("Pass either route or broker/flow overrides.") + self.brokers.resolve(route.broker) + return route + return self.resolve_route(task, broker=broker, flow=flow) + + def _resolve_flow_for_broker_override( + self, + task_name: str, + broker: AsyncBroker, + flow: FlowProtocol | None, + ) -> FlowProtocol | None: + if flow is not None: + return flow + registered_route = self.routes.get(task_name) + if registered_route is not None and registered_route.broker is broker: + return registered_route.flow + return self.brokers.default_flow(broker) diff --git a/taskiq/routing/subscriptions.py b/taskiq/routing/subscriptions.py new file mode 100644 index 00000000..2c6c6b23 --- /dev/null +++ b/taskiq/routing/subscriptions.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import replace +from typing import TYPE_CHECKING + +from taskiq.flow import FlowProtocol, get_flow_identity +from taskiq.routing.models import TaskiqSubscription +from taskiq.routing.registries import BrokerRegistry + +if TYPE_CHECKING: # pragma: no cover + from taskiq.abc.broker import AsyncBroker + +__all__ = ("SubscriptionPlan",) + + +class SubscriptionPlan: + """Inbound listen-plan subscriptions for flow-aware brokers.""" + + def __init__(self, brokers: BrokerRegistry) -> None: + self.brokers = brokers + self.subscriptions: list[TaskiqSubscription] = [] + + def subscribe( + self, + broker: AsyncBroker, + flow: FlowProtocol, + task_names: Iterable[str], + ) -> TaskiqSubscription: + """Register an inbound flow subscription for a broker.""" + target_broker = self.brokers.resolve(broker) + resolved_task_names = frozenset(task_names) + default_flow = self.brokers.default_flow(target_broker) + if default_flow is not None: + self._ensure_compatible_flow( + broker=target_broker, + registered_flow=default_flow, + new_flow=flow, + ) + + for subscription_index, subscription in enumerate(self.subscriptions): + if subscription.broker is target_broker and self._is_same_flow( + subscription.flow, + flow, + ): + self._ensure_compatible_flow( + broker=target_broker, + registered_flow=subscription.flow, + new_flow=flow, + ) + updated = replace( + subscription, + task_names=subscription.task_names | resolved_task_names, + ) + self.subscriptions[subscription_index] = updated + return updated + + subscription = TaskiqSubscription( + broker=target_broker, + flow=flow, + task_names=resolved_task_names, + ) + self.subscriptions.append(subscription) + return subscription + + def get( + self, + broker: AsyncBroker | None = None, + ) -> tuple[TaskiqSubscription, ...]: + """Return registered inbound subscriptions.""" + if broker is None: + return tuple(self.subscriptions) + target_broker = self.brokers.resolve(broker) + return tuple( + subscription + for subscription in self.subscriptions + if subscription.broker is target_broker + ) + + def get_broker_flows(self, broker: AsyncBroker) -> tuple[FlowProtocol, ...]: + """Return flows a broker should subscribe to.""" + target_broker = self.brokers.resolve(broker) + flows = [subscription.flow for subscription in self.get(target_broker)] + default_flow = self.brokers.default_flow(target_broker) + if default_flow is not None and not self._contains_flow_identity( + flows, + default_flow, + ): + flows.insert(0, default_flow) + elif default_flow is not None: + for flow in flows: + if self._is_same_flow(flow, default_flow): + self._ensure_compatible_flow( + broker=target_broker, + registered_flow=flow, + new_flow=default_flow, + ) + break + return tuple(flows) + + def _contains_flow_identity( + self, + flows: Iterable[FlowProtocol], + flow: FlowProtocol, + ) -> bool: + return any( + self._is_same_flow(registered_flow, flow) for registered_flow in flows + ) + + def _is_same_flow( + self, + first_flow: FlowProtocol, + second_flow: FlowProtocol, + ) -> bool: + return get_flow_identity(first_flow) == get_flow_identity(second_flow) + + def _ensure_compatible_flow( + self, + *, + broker: AsyncBroker, + registered_flow: FlowProtocol, + new_flow: FlowProtocol, + ) -> None: + if not self._is_same_flow(registered_flow, new_flow): + return + + registered_options = dict(registered_flow.broker_options()) + new_options = dict(new_flow.broker_options()) + if registered_options == new_options: + return + + flow_identity = get_flow_identity(new_flow) + raise ValueError( + f"Flow {flow_identity.name!r} is already registered for broker " + f"{broker.broker_name!r} with different broker options.", + ) diff --git a/taskiq/scheduler/scheduler.py b/taskiq/scheduler/scheduler.py index 16677554..b0a1e00d 100644 --- a/taskiq/scheduler/scheduler.py +++ b/taskiq/scheduler/scheduler.py @@ -1,8 +1,9 @@ from logging import getLogger -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from taskiq.exceptions import ScheduledTaskCancelledError from taskiq.kicker import AsyncKicker +from taskiq.router import TaskiqRouter from taskiq.scheduler.scheduled_task import ScheduledTask from taskiq.utils import maybe_awaitable @@ -47,19 +48,34 @@ async def on_ready(self, source: "ScheduleSource", task: ScheduledTask) -> None: except ScheduledTaskCancelledError: logger.info("Scheduled task %s has been cancelled.", task.task_name) else: - await ( - AsyncKicker(task.task_name, self.broker, task.labels) - .with_labels( - schedule_id=task.schedule_id, - ) + kicker: AsyncKicker[Any, Any] = ( + AsyncKicker(task.task_name, self.broker, dict(task.labels)) + .with_labels(schedule_id=task.schedule_id) .with_task_id(task_id=task.task_id) - .kiq( - *task.args, - **task.kwargs, - ) ) + self._apply_scheduled_route(kicker, task.task_name) + await kicker.kiq(*task.args, **task.kwargs) await maybe_awaitable(source.post_send(task)) async def shutdown(self) -> None: """Shutdown the scheduler process.""" await self.broker.shutdown() + + def _apply_scheduled_route( + self, + kicker: AsyncKicker[Any, Any], + task_name: str, + ) -> None: + """ + Apply scheduler dispatch routing without changing schedule payloads. + + Registered router routes are resolved at send time. If there is no + router route for the scheduled task, keep the old scheduler behavior and + send through the scheduler broker instead of the router default broker. + """ + router = getattr(self.broker, "router", None) + if isinstance(router, TaskiqRouter) and task_name in router.routes: + kicker.with_route(router.resolve_route(task_name)) + return + + kicker.with_broker(self.broker) diff --git a/taskiq/task_builder.py b/taskiq/task_builder.py index d2d3b0d1..eeec6ee2 100644 --- a/taskiq/task_builder.py +++ b/taskiq/task_builder.py @@ -1,13 +1,20 @@ +from __future__ import annotations + import inspect -from collections.abc import Callable +from collections.abc import Callable, Coroutine, Mapping from dataclasses import dataclass, field -from typing import Any, Generic, ParamSpec, TypeVar, overload +from types import CoroutineType, MappingProxyType +from typing import TYPE_CHECKING, Any, Generic, ParamSpec, TypeVar, overload + +from taskiq.message import TaskiqMessage, _build_taskiq_message -from taskiq.message import TaskiqMessage +if TYPE_CHECKING: # pragma: no cover + from taskiq.decor import AsyncTaskiqDecoratedTask __all__ = ("TaskDefinition", "task_builder") _FuncParams = ParamSpec("_FuncParams") +_T = TypeVar("_T") _ReturnType = TypeVar("_ReturnType") @@ -17,8 +24,13 @@ class TaskDefinition(Generic[_FuncParams, _ReturnType]): task_name: str original_func: Callable[_FuncParams, _ReturnType] - labels: dict[str, Any] = field(default_factory=dict) + labels: Mapping[str, Any] = field(default_factory=dict) return_type: type[_ReturnType] | None = None + base_cls: type[AsyncTaskiqDecoratedTask[Any, Any]] | None = None + + def __post_init__(self) -> None: + """Freeze labels to make task definitions stable declarations.""" + object.__setattr__(self, "labels", MappingProxyType(dict(self.labels))) def __call__( self, @@ -28,11 +40,32 @@ def __call__( """Call original function directly.""" return self.original_func(*args, **kwargs) + @overload + async def call( + self: TaskDefinition[_FuncParams, CoroutineType[Any, Any, _T]], + *args: _FuncParams.args, + **kwargs: _FuncParams.kwargs, + ) -> _T: ... + + @overload + async def call( + self: TaskDefinition[_FuncParams, Coroutine[Any, Any, _T]], + *args: _FuncParams.args, + **kwargs: _FuncParams.kwargs, + ) -> _T: ... + + @overload async def call( self, *args: _FuncParams.args, **kwargs: _FuncParams.kwargs, - ) -> _ReturnType: + ) -> _ReturnType: ... + + async def call( + self, + *args: _FuncParams.args, + **kwargs: _FuncParams.kwargs, + ) -> Any: """Execute original function in the current process.""" result = self.original_func(*args, **kwargs) if inspect.isawaitable(result): @@ -46,18 +79,20 @@ def message( **kwargs: _FuncParams.kwargs, ) -> TaskiqMessage: """Build a TaskiqMessage without binding this definition to a router.""" - return TaskiqMessage( + return _build_taskiq_message( task_id=task_id, task_name=self.task_name, - labels=dict(self.labels), - args=list(args), - kwargs=dict(kwargs), + labels=self.labels, + args=args, + kwargs=kwargs, ) @overload def task_builder( task_name: Callable[_FuncParams, _ReturnType], + *, + base_cls: type[AsyncTaskiqDecoratedTask[Any, Any]] | None = None, **labels: Any, ) -> TaskDefinition[_FuncParams, _ReturnType]: ... @@ -65,6 +100,8 @@ def task_builder( @overload def task_builder( task_name: str | None = None, + *, + base_cls: type[AsyncTaskiqDecoratedTask[Any, Any]] | None = None, **labels: Any, ) -> Callable[ [Callable[_FuncParams, _ReturnType]], @@ -74,6 +111,8 @@ def task_builder( def task_builder( task_name: str | Callable[_FuncParams, _ReturnType] | None = None, + *, + base_cls: type[AsyncTaskiqDecoratedTask[Any, Any]] | None = None, **labels: Any, ) -> Any: """Build an unbound task definition. @@ -97,6 +136,7 @@ def build( original_func=func, labels=dict(labels), return_type=return_type, + base_cls=base_cls, ) if callable(task_name): diff --git a/tests/routing/__init__.py b/tests/routing/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/routing/models.py b/tests/routing/models.py new file mode 100644 index 00000000..e6a4c4f5 --- /dev/null +++ b/tests/routing/models.py @@ -0,0 +1,117 @@ +from collections.abc import AsyncGenerator, Mapping +from dataclasses import dataclass +from typing import Any + +from taskiq import AsyncTaskiqDecoratedTask, FlowProtocol, TaskiqRouter +from taskiq.abc.broker import AsyncBroker +from taskiq.abc.middleware import TaskiqMiddleware +from taskiq.abc.result_backend import AsyncResultBackend +from taskiq.message import BrokerMessage, TaskiqMessage +from taskiq.result import TaskiqResult + + +@dataclass(frozen=True, slots=True) +class BrokerQueue: + """Broker-specific flow used to prove protocol-based routing.""" + + name: str + durable: bool = True + + def broker_options(self) -> Mapping[str, object]: + """Return options for the target broker.""" + return { + "durable": self.durable, + } + + +class TracingTask(AsyncTaskiqDecoratedTask[Any, Any]): + """Custom task class used to prove task_builder base_cls binding.""" + + def tracing_name(self) -> str: + """Return a trace-friendly task name.""" + return self.task_name + + +class CountingRouter(TaskiqRouter): + """Router that counts task registration calls.""" + + def __init__(self) -> None: + self.register_task_calls = 0 + super().__init__() + + def register_task( + self, + task: Any, + broker: AsyncBroker | None = None, + flow: FlowProtocol | None = None, + ) -> Any: + """Count registration calls and delegate to the router.""" + self.register_task_calls += 1 + return super().register_task(task, broker=broker, flow=flow) + + +class OldStyleRecordingBroker(AsyncBroker): + """Broker that only implements the old required kick/listen API.""" + + def __init__( + self, + *, + router: TaskiqRouter | None = None, + broker_name: str | None = None, + ) -> None: + self.sent: list[BrokerMessage] = [] + super().__init__(router=router, broker_name=broker_name) + + async def kick(self, message: BrokerMessage) -> None: + """Record old-style send.""" + self.sent.append(message) + + async def listen(self) -> AsyncGenerator[bytes, None]: + """Old-style recording broker does not listen in tests.""" + if False: + yield b"" + + +class RecordingResultBackend(AsyncResultBackend[Any]): + """Result backend marker that records stored results.""" + + def __init__(self) -> None: + self.results: dict[str, TaskiqResult[Any]] = {} + + async def set_result(self, task_id: str, result: TaskiqResult[Any]) -> None: + """Store result by task id.""" + self.results[task_id] = result + + async def is_result_ready(self, task_id: str) -> bool: + """Return whether result exists.""" + return task_id in self.results + + async def get_result( + self, + task_id: str, + with_logs: bool = False, + ) -> TaskiqResult[Any]: + """Return stored result.""" + return self.results[task_id] + + +class RecordingMiddleware(TaskiqMiddleware): + """Middleware that records client-side send hooks.""" + + def __init__( + self, + name: str, + events: list[tuple[str, str, str]], + ) -> None: + super().__init__() + self.name = name + self.events = events + + def pre_send(self, message: TaskiqMessage) -> TaskiqMessage: + """Record pre-send call.""" + self.events.append((self.name, "pre_send", message.task_name)) + return message + + def post_send(self, message: TaskiqMessage) -> None: + """Record post-send call.""" + self.events.append((self.name, "post_send", message.task_name)) diff --git a/tests/routing/test_dispatch.py b/tests/routing/test_dispatch.py new file mode 100644 index 00000000..0317f4e3 --- /dev/null +++ b/tests/routing/test_dispatch.py @@ -0,0 +1,324 @@ +from taskiq import Flow, FlowProtocol, InMemoryBroker, TaskiqRouter +from tests.routing.models import ( + BrokerQueue, + OldStyleRecordingBroker, + RecordingMiddleware, + RecordingResultBackend, +) +from tests.utils import RecordingBroker + + +async def test_old_broker_task_api_registers_task_in_router() -> None: + broker = RecordingBroker() + + @broker.task(task_name="demo.task") + async def demo_task() -> None: + return None + + assert broker.find_task("demo.task") is demo_task + assert broker.router.find_task("demo.task") is demo_task + + await demo_task.kiq() + + assert broker.sent[0][0].task_name == "demo.task" + assert broker.sent[0][1] is None + + +async def test_broker_task_api_uses_default_flow() -> None: + flow = Flow("default") + broker = RecordingBroker(default_flow=flow) + + @broker.task(task_name="demo.task") + async def demo_task() -> None: + return None + + await demo_task.kiq() + + assert broker.sent[0][0].task_name == "demo.task" + assert broker.sent[0][1] == flow + assert broker.get_subscribed_flows() == (flow,) + + +def test_router_set_broker_can_configure_default_flow() -> None: + router = TaskiqRouter() + broker = RecordingBroker(router=router, broker_name="broker") + flow = Flow("manual.default") + + router.set_broker(broker, name="broker", default_flow=flow) + route = router.resolve_route("demo.task", broker=broker) + + assert broker.default_flow == flow + assert route.flow == flow + assert broker.get_subscribed_flows() == (flow,) + + +async def test_router_can_route_task_to_another_broker_flow() -> None: + router = TaskiqRouter() + source = RecordingBroker(router=router, broker_name="source") + target = RecordingBroker(router=router, broker_name="target") + flow = Flow("events") + + @source.task(task_name="demo.task") + async def demo_task() -> None: + return None + + route = router.route_task(demo_task, broker=target, flow=flow) + resolved_route = router.resolve_route(demo_task) + + await demo_task.kiq() + + assert route.broker is target + assert route.broker_name == "target" + assert resolved_route.broker is target + assert resolved_route.flow == flow + assert source.sent == [] + assert target.sent[0][0].task_name == "demo.task" + assert target.sent[0][1] == flow + + +async def test_old_style_broker_uses_kick_to_flow_fallback() -> None: + router = TaskiqRouter() + source = RecordingBroker(router=router, broker_name="source") + target = OldStyleRecordingBroker(router=router, broker_name="target") + flow = Flow("legacy") + + @source.task(task_name="demo.task") + async def demo_task() -> None: + return None + + router.route_task(demo_task, broker=target, flow=flow) + + await demo_task.kiq() + + assert source.sent == [] + assert target.sent[0].task_name == "demo.task" + + +async def test_kicker_route_override_wins_over_registered_route() -> None: + router = TaskiqRouter() + first = RecordingBroker(router=router, broker_name="first") + second = RecordingBroker(router=router, broker_name="second") + first_flow = Flow("first") + second_flow = Flow("second") + + @first.task(task_name="demo.task") + async def demo_task() -> None: + return None + + router.route_task(demo_task, broker=first, flow=first_flow) + + route = router.resolve_route(demo_task, broker=second, flow=second_flow) + + await demo_task.kicker().with_route(route).kiq() + + assert first.sent == [] + assert second.sent[0][1] == second_flow + + +async def test_kicker_broker_override_wins_over_registered_route() -> None: + router = TaskiqRouter() + first = RecordingBroker(router=router, broker_name="first") + second_flow = Flow("second") + second = RecordingBroker( + router=router, + broker_name="second", + default_flow=second_flow, + ) + first_flow = Flow("first") + + @first.task(task_name="demo.task") + async def demo_task() -> None: + return None + + first_route = router.route_task(demo_task, broker=first, flow=first_flow) + + await demo_task.kicker().with_route(first_route).with_broker(second).kiq() + + assert first.sent == [] + assert second.sent[0][1] == second_flow + + +async def test_router_uses_explicit_broker_lookup_for_config_names() -> None: + router = TaskiqRouter() + source = RecordingBroker(router=router, broker_name="source") + target = RecordingBroker(router=router, broker_name="target") + flow = Flow("compat") + + @source.task(task_name="demo.task") + async def demo_task() -> None: + return None + + broker = router.get_broker("target") + route = router.route_task("demo.task", broker=broker, flow=flow) + + await demo_task.kicker().with_route(route).kiq() + + assert target.sent[0][0].task_name == "demo.task" + assert target.sent[0][1] == flow + + +async def test_router_accepts_broker_specific_flow_protocol() -> None: + broker = RecordingBroker(broker_name="transport") + flow = BrokerQueue(name="critical", durable=False) + + @broker.task(task_name="demo.task") + async def demo_task() -> None: + return None + + await demo_task.kicker().with_flow(flow).kiq() + + assert isinstance(flow, FlowProtocol) + assert broker.sent[0][1] is flow + assert flow.broker_options() == { + "durable": False, + } + + +async def test_kicker_preserves_labels_task_id_and_result_backend() -> None: + broker = RecordingBroker() + backend = RecordingResultBackend() + broker.with_id_generator(lambda: "generated-id") + broker.with_result_backend(backend) + + @broker.task(task_name="demo.task", task_label="declared") + async def demo_task() -> None: + return None + + sent_task = await demo_task.kicker().with_labels(call_label=42).kiq() + broker_message = broker.sent[0][0] + + assert sent_task.task_id == "generated-id" + assert sent_task.result_backend is backend + assert broker_message.task_id == "generated-id" + assert broker_message.labels["task_label"] == "declared" + assert broker_message.labels["call_label"] == "42" + + +async def test_kicker_can_prepare_invocation_for_later() -> None: + broker = RecordingBroker() + + @broker.task(task_name="demo.task") + async def demo_task(value: int) -> None: + return None + + prepared = demo_task.kicker().with_labels(trace_id="abc").prepare(1) + + assert prepared.message.task_name == "demo.task" + assert prepared.message.args == [1] + assert prepared.message.labels["trace_id"] == "abc" + + await prepared.kiq() + + assert broker.sent[0][0].task_id == prepared.message.task_id + + +async def test_prepared_invocation_keeps_route_snapshot() -> None: + router = TaskiqRouter() + first = RecordingBroker(router=router, broker_name="first") + second = RecordingBroker(router=router, broker_name="second") + first_flow = Flow("first") + second_flow = Flow("second") + + @first.task(task_name="demo.task") + async def demo_task(value: int) -> None: + return None + + first_route = router.route_task(demo_task, broker=first, flow=first_flow) + second_route = router.resolve_route(demo_task, broker=second, flow=second_flow) + kicker = demo_task.kicker().with_route(first_route) + + prepared = kicker.prepare(1) + kicker.with_route(second_route) + + await prepared.kiq() + + assert first.sent[0][1] == first_flow + assert second.sent == [] + + +async def test_prepared_invocation_resolves_default_route_snapshot() -> None: + router = TaskiqRouter() + source = RecordingBroker(router=router, broker_name="source") + first = RecordingBroker(router=router, broker_name="first") + second = RecordingBroker(router=router, broker_name="second") + first_flow = Flow("first") + second_flow = Flow("second") + + @source.task(task_name="demo.task") + async def demo_task(value: int) -> None: + return None + + router.route_task(demo_task, broker=first, flow=first_flow) + prepared = demo_task.kicker().prepare(1) + + router.route_task(demo_task, broker=second, flow=second_flow) + + await prepared.kiq() + + assert source.sent == [] + assert first.sent[0][1] == first_flow + assert second.sent == [] + + +async def test_router_task_decorator_can_choose_broker_and_flow() -> None: + router = TaskiqRouter() + target = RecordingBroker(router=router, broker_name="target") + flow = Flow("target-flow") + + @router.task("demo.task", broker=target, flow=flow) + async def demo_task() -> None: + return None + + await demo_task.kiq() + + assert target.sent[0][0].task_name == "demo.task" + assert target.sent[0][1] == flow + + +async def test_routed_dispatch_uses_target_middleware_and_result_backend() -> None: + router = TaskiqRouter() + source = RecordingBroker(router=router, broker_name="source") + target = RecordingBroker(router=router, broker_name="target") + source_events: list[tuple[str, str, str]] = [] + target_events: list[tuple[str, str, str]] = [] + target_backend = RecordingResultBackend() + + source.add_middlewares(RecordingMiddleware("source", source_events)) + target.add_middlewares(RecordingMiddleware("target", target_events)) + target.with_result_backend(target_backend) + + @source.task(task_name="demo.task") + async def demo_task() -> None: + return None + + router.route_task(demo_task, broker=target, flow=Flow("target")) + + sent_task = await demo_task.kiq() + + assert sent_task.result_backend is target_backend + assert source_events == [] + assert target_events == [ + ("target", "pre_send", "demo.task"), + ("target", "post_send", "demo.task"), + ] + + +async def test_worker_lookup_uses_task_name_not_flow() -> None: + broker = InMemoryBroker(await_inplace=True) + calls: list[str] = [] + + @broker.task(task_name="demo.first") + async def first_task() -> str: + calls.append("first") + return "first" + + @broker.task(task_name="demo.second") + async def second_task() -> str: + calls.append("second") + return "second" + + task = await first_task.kicker().with_flow(Flow(second_task.task_name)).kiq() + result = await task.wait_result(timeout=2) + + assert result.return_value == "first" + assert calls == ["first"] diff --git a/tests/routing/test_router_core.py b/tests/routing/test_router_core.py new file mode 100644 index 00000000..8f9f359b --- /dev/null +++ b/tests/routing/test_router_core.py @@ -0,0 +1,74 @@ +import pytest + +from taskiq import TaskiqRoute, TaskiqRouter, TaskiqSubscription, task_builder +from tests.routing.models import CountingRouter +from tests.utils import RecordingBroker + + +def test_broker_creates_default_router() -> None: + broker = RecordingBroker() + + assert broker.router.brokers[broker.broker_name] is broker + assert broker.router.default_broker_name == broker.broker_name + + +def test_public_router_models_keep_router_module() -> None: + assert TaskiqRoute.__module__ == "taskiq.router" + assert TaskiqSubscription.__module__ == "taskiq.router" + + +def test_router_rejects_duplicate_broker_names() -> None: + router = TaskiqRouter() + RecordingBroker(router=router, broker_name="broker") + + with pytest.raises(ValueError, match="already registered"): + RecordingBroker(router=router, broker_name="broker") + + +def test_router_rejects_broker_attached_to_another_router() -> None: + first_router = TaskiqRouter() + second_router = TaskiqRouter() + broker = RecordingBroker(router=first_router, broker_name="broker") + + with pytest.raises(ValueError, match="attached to another router"): + second_router.set_broker(broker, name="broker") + + +def test_router_rejects_string_broker_references() -> None: + router = TaskiqRouter() + RecordingBroker(router=router, broker_name="broker") + + with pytest.raises(TypeError, match="Broker string references"): + router.route_task("demo.task", broker="broker") # type: ignore[arg-type] + + +def test_router_rejects_duplicate_task_names() -> None: + router = TaskiqRouter() + first = RecordingBroker(router=router, broker_name="first") + second = RecordingBroker(router=router, broker_name="second") + + @first.task(task_name="demo.task") + async def first_task() -> None: + return None + + with pytest.raises(ValueError, match="already registered"): + + @second.task(task_name="demo.task") + async def second_task() -> None: + return None + + +def test_task_definition_binding_does_not_reenter_router_registration() -> None: + router = CountingRouter() + broker = RecordingBroker(router=router) + + @task_builder("shared.once") + def shared_task() -> None: + return None + + registered = broker.register_task(shared_task) + + assert registered.task_name == "shared.once" + assert router.register_task_calls == 1 + assert broker.router.find_task("shared.once") is registered + assert broker.find_task("shared.once") is registered diff --git a/tests/routing/test_shared_tasks.py b/tests/routing/test_shared_tasks.py new file mode 100644 index 00000000..4efab51d --- /dev/null +++ b/tests/routing/test_shared_tasks.py @@ -0,0 +1,94 @@ +import pytest + +from taskiq import AsyncTaskiqDecoratedTask, Flow, task_builder +from tests.routing.models import TracingTask +from tests.utils import RecordingBroker + + +async def test_task_builder_can_be_registered_later() -> None: + broker = RecordingBroker() + + @task_builder("shared.add", queue="shared") + def add(left: int, right: int) -> int: + return left + right + + assert await add.call(1, 2) == 3 + + registered = broker.register_task(add) + + assert type(registered) is AsyncTaskiqDecoratedTask + assert registered.task_name == "shared.add" + assert registered.labels == {"queue": "shared"} + assert broker.router.find_task("shared.add") is registered + + await registered.kiq(1, 2) + + assert broker.sent[0][0].task_name == "shared.add" + + +async def test_router_register_task_definition_binds_to_selected_broker() -> None: + source = RecordingBroker(broker_name="source") + target = RecordingBroker(router=source.router, broker_name="target") + flow = Flow("target") + + @task_builder("shared.routed") + async def shared_task() -> None: + return None + + registered = source.router.register_task(shared_task, broker=target, flow=flow) + + await registered.kiq() + + assert source.sent == [] + assert target.find_task("shared.routed") is registered + assert target.sent[0][0].task_name == "shared.routed" + assert target.sent[0][1] == flow + + +async def test_task_builder_can_use_custom_base_cls() -> None: + broker = RecordingBroker() + + @task_builder("shared.traced", base_cls=TracingTask, queue="shared") + async def traced(value: int) -> int: + return value + 1 + + registered = broker.register_task(traced) + + assert isinstance(registered, TracingTask) + assert registered.tracing_name() == "shared.traced" + assert registered.labels == {"queue": "shared"} + assert await traced.call(1) == 2 + + await registered.kiq(1) + + assert broker.sent[0][0].task_name == "shared.traced" + + +def test_task_definition_default_flow_does_not_create_subscription() -> None: + flow = Flow("shared.default") + broker = RecordingBroker(default_flow=flow) + + @task_builder("shared.default") + def shared_task() -> None: + return None + + registered = broker.register_task(shared_task) + route = broker.router.resolve_route(registered) + + assert route.flow == flow + assert broker.router.get_subscriptions(broker) == () + assert broker.get_subscribed_flows() == (flow,) + + +def test_register_task_definition_rejects_overrides() -> None: + broker = RecordingBroker() + + @task_builder("shared.add", queue="shared") + def add(left: int, right: int) -> int: + return left + right + + with pytest.raises(ValueError, match="TaskDefinition already defines"): + broker.register_task(add, task_name="other.name") + + with pytest.raises(ValueError, match="TaskDefinition already defines"): + broker.register_task(add, queue="other") diff --git a/tests/routing/test_subscriptions.py b/tests/routing/test_subscriptions.py new file mode 100644 index 00000000..cc82603b --- /dev/null +++ b/tests/routing/test_subscriptions.py @@ -0,0 +1,186 @@ +import pytest + +from taskiq import ( + Flow, + FlowIdentity, + FlowProtocol, + TaskiqRouter, + get_flow_identity, +) +from taskiq.warnings import TaskiqDeprecationWarning +from tests.routing.models import BrokerQueue +from tests.utils import RecordingBroker + + +def test_flow_identity_uses_logical_name() -> None: + generic_flow = Flow("events").with_options(durable=True) + broker_flow = BrokerQueue(name="events", durable=True) + + assert generic_flow.identity == FlowIdentity("events") + assert get_flow_identity(broker_flow) == FlowIdentity("events") + assert generic_flow.broker_options() == {"durable": True} + + +def test_route_task_does_not_register_subscription_by_default() -> None: + router = TaskiqRouter() + broker = RecordingBroker(router=router, broker_name="broker") + flow = Flow("events") + + route = router.route_task("demo.task", broker=broker, flow=flow) + + assert route.flow == flow + assert router.get_subscriptions(broker) == () + assert broker.get_subscribed_flows() == () + + +def test_route_task_subscribe_true_is_deprecated_shim() -> None: + router = TaskiqRouter() + broker = RecordingBroker(router=router, broker_name="broker") + flow = Flow("events") + + with pytest.warns(TaskiqDeprecationWarning, match="router.subscribe"): + route = router.route_task( + "demo.task", + broker=broker, + flow=flow, + subscribe=True, + ) + + subscription = router.get_subscriptions(broker)[0] + + assert route.flow == flow + assert subscription.broker is broker + assert subscription.flow == flow + assert subscription.task_names == frozenset({"demo.task"}) + assert broker.get_subscribed_flows() == (flow,) + + +def test_route_update_does_not_mutate_existing_subscriptions() -> None: + router = TaskiqRouter() + broker = RecordingBroker(router=router, broker_name="broker") + first_flow = Flow("events.first") + second_flow = Flow("events.second") + + router.route_task("demo.task", broker=broker, flow=first_flow) + router.subscribe(broker, first_flow, "demo.task") + router.route_task("demo.task", broker=broker, flow=second_flow) + + subscription = router.get_subscriptions(broker)[0] + route = router.resolve_route("demo.task") + + assert route.flow == second_flow + assert subscription.flow == first_flow + assert subscription.task_names == frozenset({"demo.task"}) + assert broker.get_subscribed_flows() == (first_flow,) + + +def test_subscribe_dedupes_flow_and_merges_task_names() -> None: + router = TaskiqRouter() + broker = RecordingBroker(router=router, broker_name="broker") + flow = Flow("events") + + router.subscribe(broker, flow, "first.task") + subscription = router.subscribe(broker, flow, "second.task") + + assert router.get_subscriptions(broker) == (subscription,) + assert subscription.task_names == frozenset({"first.task", "second.task"}) + assert broker.get_subscribed_flows() == (flow,) + + +def test_subscribe_merges_same_identity_with_compatible_options() -> None: + router = TaskiqRouter() + broker = RecordingBroker(router=router, broker_name="broker") + generic_flow = Flow("events").with_options(durable=True) + broker_flow = BrokerQueue(name="events", durable=True) + + router.subscribe(broker, generic_flow, "first.task") + subscription = router.subscribe(broker, broker_flow, "second.task") + + assert router.get_subscriptions(broker) == (subscription,) + assert subscription.flow is generic_flow + assert subscription.task_names == frozenset({"first.task", "second.task"}) + assert broker.get_subscribed_flows() == (generic_flow,) + + +@pytest.mark.parametrize( + "conflicting_flow", + [ + pytest.param( + Flow("events").with_options(durable=False), + id="generic-flow", + ), + pytest.param( + BrokerQueue(name="events", durable=False), + id="broker-specific-flow", + ), + ], +) +def test_subscribe_rejects_same_identity_with_different_options( + conflicting_flow: FlowProtocol, +) -> None: + router = TaskiqRouter() + broker = RecordingBroker(router=router, broker_name="broker") + + router.subscribe(broker, Flow("events").with_options(durable=True), "first.task") + + with pytest.raises(ValueError, match="different broker options"): + router.subscribe( + broker, + conflicting_flow, + "second.task", + ) + + +def test_subscribe_rejects_default_flow_option_conflict() -> None: + router = TaskiqRouter() + default_flow = Flow("events").with_options(durable=True) + broker = RecordingBroker( + router=router, + broker_name="broker", + default_flow=default_flow, + ) + + with pytest.raises(ValueError, match="different broker options"): + router.subscribe( + broker, + BrokerQueue(name="events", durable=False), + "demo.task", + ) + + +def test_subscribe_allows_same_identity_on_different_brokers() -> None: + router = TaskiqRouter() + first_broker = RecordingBroker(router=router, broker_name="first") + second_broker = RecordingBroker(router=router, broker_name="second") + + first_subscription = router.subscribe( + first_broker, + Flow("events").with_options(durable=True), + "first.task", + ) + second_subscription = router.subscribe( + second_broker, + BrokerQueue(name="events", durable=False), + "second.task", + ) + + assert router.get_subscriptions(first_broker) == (first_subscription,) + assert router.get_subscriptions(second_broker) == (second_subscription,) + + +def test_subscribed_flows_dedupes_default_flow_by_identity() -> None: + router = TaskiqRouter() + default_flow = Flow("events").with_options(durable=True) + broker = RecordingBroker( + router=router, + broker_name="broker", + default_flow=default_flow, + ) + + subscription = router.subscribe( + broker, + BrokerQueue(name="events", durable=True), + "demo.task", + ) + + assert broker.get_subscribed_flows() == (subscription.flow,) diff --git a/tests/scheduler/test_scheduler.py b/tests/scheduler/test_scheduler.py index 05e6358a..eeb9d35c 100644 --- a/tests/scheduler/test_scheduler.py +++ b/tests/scheduler/test_scheduler.py @@ -1,8 +1,28 @@ +from datetime import datetime, timezone + +from taskiq import Flow, TaskiqRouter from taskiq.abc.schedule_source import ScheduleSource from taskiq.brokers.inmemory_broker import InMemoryBroker +from taskiq.compat import model_dump from taskiq.exceptions import ScheduledTaskCancelledError from taskiq.scheduler.scheduled_task import ScheduledTask from taskiq.scheduler.scheduler import TaskiqScheduler +from tests.utils import RecordingBroker + + +class RecordingScheduleSource(ScheduleSource): + """Schedule source that records dynamically added schedules.""" + + def __init__(self) -> None: + self.schedules: list[ScheduledTask] = [] + + async def get_schedules(self) -> list["ScheduledTask"]: + """Return stored schedules.""" + return self.schedules + + async def add_schedule(self, schedule: "ScheduledTask") -> None: + """Store a schedule.""" + self.schedules.append(schedule) class CancellingScheduleSource(ScheduleSource): @@ -31,3 +51,97 @@ async def test_scheduled_task_cancelled() -> None: ) await scheduler.on_ready(source, task) # error is caught + + +async def test_scheduler_resolves_route_at_send_time() -> None: + router = TaskiqRouter() + scheduler_broker = RecordingBroker(router=router, broker_name="scheduler") + first_broker = RecordingBroker(router=router, broker_name="first") + second_broker = RecordingBroker(router=router, broker_name="second") + first_flow = Flow("first") + second_flow = Flow("second") + source = RecordingScheduleSource() + scheduler = TaskiqScheduler(broker=scheduler_broker, sources=[source]) + + @scheduler_broker.task(task_name="demo.task") + async def demo_task() -> None: + return None + + router.route_task(demo_task, broker=first_broker, flow=first_flow) + scheduled_task = ScheduledTask( + task_name="demo.task", + labels={}, + args=[], + kwargs={}, + cron="* * * * *", + ) + router.route_task(demo_task, broker=second_broker, flow=second_flow) + + await scheduler.on_ready(source, scheduled_task) + + assert first_broker.sent == [] + assert second_broker.sent[0][0].task_name == "demo.task" + assert second_broker.sent[0][1] == second_flow + + +async def test_scheduler_without_route_uses_scheduler_broker() -> None: + router = TaskiqRouter() + default_broker = RecordingBroker(router=router, broker_name="default") + scheduler_flow = Flow("scheduler.default") + scheduler_broker = RecordingBroker( + router=router, + broker_name="scheduler", + default_flow=scheduler_flow, + ) + source = RecordingScheduleSource() + scheduler = TaskiqScheduler(broker=scheduler_broker, sources=[source]) + scheduled_task = ScheduledTask( + task_name="external.task", + labels={"source": "database"}, + args=[], + kwargs={}, + cron="* * * * *", + ) + + await scheduler.on_ready(source, scheduled_task) + + sent_message, sent_flow = scheduler_broker.sent[0] + assert default_broker.sent == [] + assert sent_message.task_name == "external.task" + assert sent_message.labels["source"] == "database" + assert sent_message.labels["schedule_id"] == scheduled_task.schedule_id + assert sent_flow == scheduler_flow + assert scheduled_task.labels == {"source": "database"} + + +async def test_created_schedule_kiq_does_not_mutate_schedule_payload() -> None: + router = TaskiqRouter() + source_broker = RecordingBroker(router=router, broker_name="source") + target_broker = RecordingBroker(router=router, broker_name="target") + target_flow = Flow("target") + source = RecordingScheduleSource() + + @source_broker.task(task_name="demo.task") + async def demo_task(value: int) -> None: + return None + + route = router.resolve_route( + demo_task, + broker=target_broker, + flow=target_flow, + ) + schedule = ( + await demo_task.kicker() + .with_route(route) + .with_labels(trace_id="trace") + .schedule_by_time(source, datetime.now(timezone.utc), 1) + ) + schedule_payload = model_dump(schedule.task) + + await schedule.kiq() + + assert source.schedules == [schedule.task] + assert {"route", "flow", "broker", "broker_name"}.isdisjoint(schedule_payload) + assert schedule.task.labels == {"trace_id": "trace"} + assert target_broker.sent[0][0].task_name == "demo.task" + assert target_broker.sent[0][1] == target_flow diff --git a/tests/test_requeue.py b/tests/test_requeue.py index c451b6df..a23f98c3 100644 --- a/tests/test_requeue.py +++ b/tests/test_requeue.py @@ -1,4 +1,20 @@ -from taskiq import Context, InMemoryBroker, TaskiqDepends +import pytest + +from taskiq import Context, Flow, InMemoryBroker, TaskiqDepends, TaskiqRouter +from taskiq.exceptions import NoResultError +from taskiq.message import TaskiqMessage +from tests.utils import RecordingBroker + + +def build_message(task_name: str = "demo.task") -> TaskiqMessage: + """Build a message for context requeue tests.""" + return TaskiqMessage( + task_id="task-id", + task_name=task_name, + labels={}, + args=[], + kwargs={}, + ) async def test_requeue() -> None: @@ -46,3 +62,75 @@ async def task(_: None = TaskiqDepends(dep_func)) -> None: ) assert runs_count == 2 + + +async def test_requeue_preserves_current_broker_when_route_points_elsewhere() -> None: + router = TaskiqRouter() + current_flow = Flow("current.default") + current_broker = RecordingBroker( + router=router, + broker_name="current", + default_flow=current_flow, + ) + routed_broker = RecordingBroker(router=router, broker_name="routed") + routed_flow = Flow("routed") + message = build_message() + context = Context(message, current_broker) + + router.route_task(message.task_name, broker=routed_broker, flow=routed_flow) + + with pytest.raises(NoResultError): + await context.requeue() + + sent_message, sent_flow = current_broker.sent[0] + assert routed_broker.sent == [] + assert sent_message.task_name == message.task_name + assert sent_message.labels["X-Taskiq-requeue"] == "1" + assert sent_flow == current_flow + + +async def test_requeue_prefers_same_broker_route_flow() -> None: + router = TaskiqRouter() + default_flow = Flow("current.default") + route_flow = Flow("current.route") + current_broker = RecordingBroker( + router=router, + broker_name="current", + default_flow=default_flow, + ) + message = build_message() + context = Context(message, current_broker) + + router.route_task(message.task_name, broker=current_broker, flow=route_flow) + + with pytest.raises(NoResultError): + await context.requeue() + + assert current_broker.sent[0][1] == route_flow + + +async def test_requeue_without_route_or_default_flow_sends_without_flow() -> None: + broker = RecordingBroker() + message = build_message() + context = Context(message, broker) + + with pytest.raises(NoResultError): + await context.requeue() + + assert broker.sent[0][1] is None + + +async def test_router_requeue_accepts_explicit_flow_override() -> None: + router = TaskiqRouter() + default_flow = Flow("default") + explicit_flow = Flow("explicit") + broker = RecordingBroker( + router=router, + broker_name="broker", + default_flow=default_flow, + ) + message = build_message() + + await router.requeue(message, broker=broker, flow=explicit_flow) + + assert broker.sent[0][1] == explicit_flow diff --git a/tests/test_router.py b/tests/test_router.py deleted file mode 100644 index f20d7223..00000000 --- a/tests/test_router.py +++ /dev/null @@ -1,218 +0,0 @@ -from collections.abc import AsyncGenerator, Mapping -from dataclasses import dataclass - -import pytest - -from taskiq import Flow, FlowProtocol, TaskiqRouter, task_builder -from taskiq.abc.broker import AsyncBroker -from taskiq.message import BrokerMessage - - -@dataclass(frozen=True, slots=True) -class BrokerQueue: - """Broker-specific flow used to prove protocol-based routing.""" - - name: str - durable: bool = True - - def broker_options(self, broker_name: str) -> Mapping[str, object]: - """Return options for the target broker.""" - return { - "broker": broker_name, - "durable": self.durable, - } - - -class RecordingBroker(AsyncBroker): - """Broker that records sent messages and flows.""" - - def __init__( - self, - *, - router: TaskiqRouter | None = None, - broker_name: str | None = None, - default_flow: FlowProtocol | None = None, - ) -> None: - self.sent: list[tuple[BrokerMessage, FlowProtocol | None]] = [] - super().__init__( - router=router, - broker_name=broker_name, - default_flow=default_flow, - ) - - async def kick(self, message: BrokerMessage) -> None: - """Record old-style send.""" - self.sent.append((message, None)) - - async def kick_to_flow( - self, - message: BrokerMessage, - flow: FlowProtocol | None = None, - ) -> None: - """Record flow-aware send.""" - self.sent.append((message, flow)) - - async def listen(self) -> AsyncGenerator[bytes, None]: - """Recording broker doesn't listen in these tests.""" - if False: - yield b"" - - -def test_broker_creates_default_router() -> None: - broker = RecordingBroker() - - assert broker.router.brokers[broker.broker_name] is broker - assert broker.router.default_broker_name == broker.broker_name - - -async def test_old_broker_task_api_registers_task_in_router() -> None: - broker = RecordingBroker() - - @broker.task(task_name="demo.task") - async def demo_task() -> None: - return None - - assert broker.find_task("demo.task") is demo_task - assert broker.router.find_task("demo.task") is demo_task - - await demo_task.kiq() - - assert broker.sent[0][0].task_name == "demo.task" - assert broker.sent[0][1] is None - - -async def test_router_can_route_task_to_another_broker_flow() -> None: - router = TaskiqRouter() - source = RecordingBroker(router=router, broker_name="source") - target = RecordingBroker(router=router, broker_name="target") - flow = Flow("events") - - @source.task(task_name="demo.task") - async def demo_task() -> None: - return None - - route = router.route_task(demo_task, broker=target, flow=flow) - - await demo_task.kiq() - - assert route.broker is target - assert route.broker_name == "target" - assert source.sent == [] - assert target.sent[0][0].task_name == "demo.task" - assert target.sent[0][1] == flow - - -async def test_kicker_route_override_wins_over_registered_route() -> None: - router = TaskiqRouter() - first = RecordingBroker(router=router, broker_name="first") - second = RecordingBroker(router=router, broker_name="second") - first_flow = Flow("first") - second_flow = Flow("second") - - @first.task(task_name="demo.task") - async def demo_task() -> None: - return None - - router.route_task(demo_task, broker=first, flow=first_flow) - - await demo_task.kicker().with_route(second, second_flow).kiq() - - assert first.sent == [] - assert second.sent[0][1] == second_flow - - -async def test_router_keeps_string_broker_lookup_for_compatibility() -> None: - router = TaskiqRouter() - source = RecordingBroker(router=router, broker_name="source") - target = RecordingBroker(router=router, broker_name="target") - flow = Flow("compat") - - @source.task(task_name="demo.task") - async def demo_task() -> None: - return None - - router.route_task("demo.task", broker="target", flow=flow) - - await demo_task.kicker().with_route("target", flow).kiq() - - assert target.sent[0][0].task_name == "demo.task" - assert target.sent[0][1] == flow - - -async def test_router_accepts_broker_specific_flow_protocol() -> None: - broker = RecordingBroker(broker_name="transport") - flow = BrokerQueue(name="critical", durable=False) - - @broker.task(task_name="demo.task") - async def demo_task() -> None: - return None - - await demo_task.kicker().with_flow(flow).kiq() - - assert isinstance(flow, FlowProtocol) - assert broker.sent[0][1] is flow - assert flow.broker_options("transport") == { - "broker": "transport", - "durable": False, - } - - -async def test_kicker_can_prepare_invocation_for_later() -> None: - broker = RecordingBroker() - - @broker.task(task_name="demo.task") - async def demo_task(value: int) -> None: - return None - - prepared = demo_task.kicker().with_labels(trace_id="abc").prepare(1) - - assert prepared.message.task_name == "demo.task" - assert prepared.message.args == [1] - assert prepared.message.labels["trace_id"] == "abc" - - await prepared.kiq() - - assert broker.sent[0][0].task_id == prepared.message.task_id - - -async def test_task_builder_can_be_registered_later() -> None: - broker = RecordingBroker() - - @task_builder("shared.add", queue="shared") - def add(left: int, right: int) -> int: - return left + right - - assert await add.call(1, 2) == 3 - - registered = broker.register_task(add) - - assert registered.task_name == "shared.add" - assert registered.labels == {"queue": "shared"} - assert broker.router.find_task("shared.add") is registered - - await registered.kiq(1, 2) - - assert broker.sent[0][0].task_name == "shared.add" - - -async def test_router_task_decorator_can_choose_broker_and_flow() -> None: - router = TaskiqRouter() - target = RecordingBroker(router=router, broker_name="target") - flow = Flow("target-flow") - - @router.task("demo.task", broker=target, flow=flow) - async def demo_task() -> None: - return None - - await demo_task.kiq() - - assert target.sent[0][0].task_name == "demo.task" - assert target.sent[0][1] == flow - - -def test_router_rejects_duplicate_broker_names() -> None: - router = TaskiqRouter() - RecordingBroker(router=router, broker_name="broker") - - with pytest.raises(ValueError, match="already registered"): - RecordingBroker(router=router, broker_name="broker") diff --git a/tests/test_task_builder.py b/tests/test_task_builder.py new file mode 100644 index 00000000..0acf3066 --- /dev/null +++ b/tests/test_task_builder.py @@ -0,0 +1,130 @@ +from dataclasses import dataclass +from typing import Any + +import pytest +from pydantic import BaseModel + +from taskiq import AsyncTaskiqDecoratedTask, InMemoryBroker, task_builder +from taskiq.labels import LabelType + + +class CustomBrokerTask(AsyncTaskiqDecoratedTask[Any, Any]): + """Custom broker-level task class used by old decorator tests.""" + + +@dataclass(frozen=True, slots=True) +class Payload: + """Payload dataclass used to check message argument preparation.""" + + value: int + + +class User(BaseModel): + """Payload model used to check pydantic argument preparation.""" + + name: str + + +async def test_task_definition_call_executes_sync_and_async_functions() -> None: + @task_builder("shared.add") + def add(left: int, right: int) -> int: + return left + right + + @task_builder("shared.multiply") + async def multiply(left: int, right: int) -> int: + return left * right + + assert add(1, 2) == 3 + assert await add.call(1, 2) == 3 + assert await multiply(2, 3) == 6 + assert await multiply.call(2, 3) == 6 + + +def test_task_definition_message_matches_registered_kicker_message() -> None: + @task_builder( + "shared.message", + priority=7, + enabled=True, + payload=b"abc", + ) + def process(payload: Payload, user: User) -> None: + return None + + broker = InMemoryBroker() + registered = broker.register_task(process) + + definition_message = process.message( + "task-id", + Payload(value=1), + User(name="Ada"), + ) + prepared_message = ( + registered.kicker() + .with_task_id("task-id") + .prepare( + Payload(value=1), + User(name="Ada"), + ) + .message + ) + + assert definition_message == prepared_message + assert definition_message.args == [{"value": 1}, {"name": "Ada"}] + assert definition_message.labels == { + "enabled": "True", + "payload": "YWJj", + "priority": "7", + } + assert definition_message.labels_types == { + "enabled": LabelType.BOOL.value, + "payload": LabelType.BYTES.value, + "priority": LabelType.INT.value, + } + + +def test_task_definition_message_rejects_dataclass_types() -> None: + @task_builder("shared.message") + def process(payload: Payload) -> None: + return None + + with pytest.raises(ValueError, match="Cannot serialize types"): + process.message("task-id", Payload) # type: ignore[arg-type] + + +def test_shared_task_without_base_cls_uses_native_task_class() -> None: + broker = InMemoryBroker() + broker.decorator_class = CustomBrokerTask + + @broker.task(task_name="old.decorator") + async def old_task() -> None: + return None + + @task_builder("shared.default") + async def shared_task() -> None: + return None + + registered = broker.register_task(shared_task) + + assert isinstance(old_task, CustomBrokerTask) + assert type(registered) is AsyncTaskiqDecoratedTask + + +@pytest.mark.parametrize( + "base_cls", + [ + pytest.param(object, id="non-task-type"), + pytest.param(0, id="falsey-runtime-value"), + ], +) +def test_invalid_task_definition_base_cls_fails_when_bound(base_cls: Any) -> None: + @task_builder("shared.invalid", base_cls=base_cls) + def invalid_task() -> None: + return None + + broker = InMemoryBroker() + + with pytest.raises( + TypeError, + match="base_cls must be a subclass of AsyncTaskiqDecoratedTask", + ): + broker.register_task(invalid_task) diff --git a/tests/utils.py b/tests/utils.py index c6c659d2..a90b59bf 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1,10 +1,45 @@ import asyncio from collections.abc import AsyncGenerator -from taskiq import AsyncBroker, BrokerMessage +from taskiq import AsyncBroker, BrokerMessage, FlowProtocol, TaskiqRouter from taskiq.acks import AckableMessage +class RecordingBroker(AsyncBroker): + """Broker for tests that records messages and selected flows.""" + + def __init__( + self, + *, + router: TaskiqRouter | None = None, + broker_name: str | None = None, + default_flow: FlowProtocol | None = None, + ) -> None: + self.sent: list[tuple[BrokerMessage, FlowProtocol | None]] = [] + super().__init__( + router=router, + broker_name=broker_name, + default_flow=default_flow, + ) + + async def kick(self, message: BrokerMessage) -> None: + """Record old-style send.""" + self.sent.append((message, None)) + + async def kick_to_flow( + self, + message: BrokerMessage, + flow: FlowProtocol | None = None, + ) -> None: + """Record flow-aware send.""" + self.sent.append((message, flow)) + + async def listen(self) -> AsyncGenerator[bytes, None]: + """Recording broker does not listen in tests.""" + if False: + yield b"" + + class AsyncQueueBroker(AsyncBroker): """ Broker for testing. From 99a34cabbce4484b53ed9c8db4355d83d30b2ede Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Sat, 6 Jun 2026 22:12:17 +0200 Subject: [PATCH 04/10] feat(routing): separate task routing from broker transport Introduce Flow, TaskiqRouter, route and subscription contracts to move task dispatch decisions out of brokers while keeping brokers as transport adapters. Preserve backward-compatible --- taskiq/abc/broker.py | 8 +++++--- taskiq/router.py | 2 ++ tests/routing/test_router_core.py | 9 +++++++++ 3 files changed, 16 insertions(+), 3 deletions(-) diff --git a/taskiq/abc/broker.py b/taskiq/abc/broker.py index a25a98a6..9405edbb 100644 --- a/taskiq/abc/broker.py +++ b/taskiq/abc/broker.py @@ -142,7 +142,8 @@ def find_task(self, task_name: str) -> AsyncTaskiqDecoratedTask[Any, Any] | None It searches task by name in dict of tasks that were registered for this broker directly. - If it fails, it checks global dict of all available tasks. + If it fails, it checks tasks registered in the broker's router, + then global dict of all available tasks. :param task_name: name of a task. :returns: found task or None. @@ -161,8 +162,9 @@ def get_all_tasks(self) -> dict[str, AsyncTaskiqDecoratedTask[Any, Any]]: """ Method to fetch all tasks available in broker. - This method returns all tasks, globally and locally - available in broker. With local tasks having higher priority. + This method returns all tasks globally, through the broker's router + and locally available in broker. With local tasks having higher + priority. So, if you have two tasks with the same name, one registered in global registry and one registered diff --git a/taskiq/router.py b/taskiq/router.py index b62d16d3..fa054a4a 100644 --- a/taskiq/router.py +++ b/taskiq/router.py @@ -55,6 +55,8 @@ def default_broker(self) -> AsyncBroker | None: @default_broker.setter def default_broker(self, broker: AsyncBroker | None) -> None: + if broker is not None: + self._brokers.resolve(broker) self._brokers.default_broker = broker @property diff --git a/tests/routing/test_router_core.py b/tests/routing/test_router_core.py index 8f9f359b..d2b92301 100644 --- a/tests/routing/test_router_core.py +++ b/tests/routing/test_router_core.py @@ -34,6 +34,15 @@ def test_router_rejects_broker_attached_to_another_router() -> None: second_router.set_broker(broker, name="broker") +def test_router_default_broker_setter_rejects_foreign_broker() -> None: + first_router = TaskiqRouter() + second_router = TaskiqRouter() + broker = RecordingBroker(router=first_router, broker_name="broker") + + with pytest.raises(ValueError, match="not registered"): + second_router.default_broker = broker + + def test_router_rejects_string_broker_references() -> None: router = TaskiqRouter() RecordingBroker(router=router, broker_name="broker") From cc6d3a38047c64c076519d6994d917cf406174f3 Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Sun, 7 Jun 2026 11:13:00 +0200 Subject: [PATCH 05/10] fix(routing): tighten router contracts after review Remove direct router access to broker private task storage and add an explicit broker-owned storage boundary for router-managed task registration. Clarify prepared route snapshots, subscription task-name semantics and setup-time router mutation contracts. Extend routing docs with flow resolution precedence and custom task base class middleware behavior. Add regression coverage for shared task local registration, custom base_cls middleware hooks and scheduler fallback when a route is removed before dispatch. --- docs/guide/routing-and-flows.md | 22 +++++++++++++++++++ taskiq/abc/broker.py | 13 ++++++++++++ taskiq/kicker.py | 9 +++++++- taskiq/router.py | 8 ++----- taskiq/routing/models.py | 8 ++++++- taskiq/routing/subscriptions.py | 7 +++++- tests/routing/test_shared_tasks.py | 24 ++++++++++++++++++++- tests/scheduler/test_scheduler.py | 34 ++++++++++++++++++++++++++++++ 8 files changed, 115 insertions(+), 10 deletions(-) diff --git a/docs/guide/routing-and-flows.md b/docs/guide/routing-and-flows.md index 208d2502..12db2fc9 100644 --- a/docs/guide/routing-and-flows.md +++ b/docs/guide/routing-and-flows.md @@ -150,6 +150,16 @@ keep a snapshot of the route that was resolved at prepare time, so later changes to the mutable kicker object or router route table do not change that prepared send. +Flow selection is resolved in this order: + +| Source | Effect | +| --- | --- | +| Explicit route | `.with_route(route)` uses that route's broker and flow. A later `.with_flow(flow)` updates the route flow, while a later `.with_broker(broker)` clears the explicit route. | +| Explicit broker | `.with_broker(broker)` sends through that broker. If `.with_flow(flow)` is also set, that flow wins; otherwise Taskiq uses the same-broker task route flow, then the broker `default_flow`. | +| Explicit flow | `.with_flow(flow)` without a broker override replaces the flow of the resolved route or broker default route. | +| Router task route | `router.route_task(task, broker=..., flow=...)` is the default outbound route for that task. | +| Broker default flow | `InMemoryBroker(..., default_flow=...)` and other brokers with `default_flow` provide the fallback flow for that broker. | + ## Subscriptions Routers also keep inbound flow subscriptions. Routing and subscribing are @@ -179,6 +189,11 @@ The deprecated `route_task(..., subscribe=True)` shim still performs this subscription when a flow is resolved, but new code should call `subscribe()` directly. +`TaskiqSubscription.task_names` is diagnostic listen-plan metadata. It records +which task names caused the broker to listen to a flow, but it does not make the +flow an inbound task router. Workers still execute messages by +`TaskiqMessage.task_name`. + Existing brokers can keep implementing `listen()` as before. New flow-aware brokers may use `get_subscribed_flows()` to subscribe to queues, topics, subjects or streams while the routing rules stay in the router. @@ -187,6 +202,10 @@ subjects or streams while the routing rules stay in the router. subscriptions. It deduplicates by flow identity and checks broker options for conflicts. +Router routes and subscriptions are mutable setup-time configuration. They are +not thread-safe runtime coordination primitives; configure them before starting +concurrent worker or client activity. + ## Scheduler and requeue `ScheduledTask` remains a transport-neutral invocation payload. It stores task @@ -277,6 +296,9 @@ When the final application registers this definition, Taskiq creates the bound task using `TracingTask`. If `base_cls` is not provided, Taskiq uses the native decorated task class. +Custom task classes do not bypass broker middleware. Send and execute lifecycle +hooks are still owned by the selected broker and worker path. + For low-level integrations, `TaskDefinition.message(...)` builds a `TaskiqMessage` without binding the task. It uses the same argument and label preparation contract as a normal `.kicker().prepare(...)` invocation. diff --git a/taskiq/abc/broker.py b/taskiq/abc/broker.py index 9405edbb..06926827 100644 --- a/taskiq/abc/broker.py +++ b/taskiq/abc/broker.py @@ -439,6 +439,19 @@ def bind_task_definition( register=register, ) + def store_registered_task( + self, + task: AsyncTaskiqDecoratedTask[Any, Any], + ) -> None: + """ + Store a task after router-owned registration accepted it. + + Application code should use `task` or `register_task`. This method is + the broker/router integration boundary for cases where the router owns + registration ordering and the broker owns local task storage. + """ + self._store_task(task.task_name, task) + def _decorate_task( self, func: Callable[_FuncParams, _ReturnType], diff --git a/taskiq/kicker.py b/taskiq/kicker.py index 8debca3e..8de79f90 100644 --- a/taskiq/kicker.py +++ b/taskiq/kicker.py @@ -321,7 +321,14 @@ async def _legacy_kiq( def _prepare_route_snapshot( self, ) -> tuple[AsyncBroker, TaskiqRoute | None, FlowProtocol | None]: - """Resolve the route that a prepared invocation must keep.""" + """ + Resolve route state that a prepared invocation must keep. + + The returned tuple is `(broker, route, explicit_flow_override)`. + When a route is present, the route already carries its flow, so the + separate flow override is `None`. A separate flow is returned only for + legacy broker paths that cannot snapshot a `TaskiqRoute`. + """ router = getattr(self.broker, "router", None) if not isinstance(router, TaskiqRouter): return self.broker, None, self.route_flow diff --git a/taskiq/router.py b/taskiq/router.py index fa054a4a..8b526066 100644 --- a/taskiq/router.py +++ b/taskiq/router.py @@ -132,11 +132,7 @@ def register_task( broker=target_broker, flow=flow, ) - # Router and broker share this internal registration boundary. - target_broker._store_task( # noqa: SLF001 - registered_task.task_name, - registered_task, - ) + target_broker.store_registered_task(registered_task) return registered_task return self._register_bound_task(task, broker=broker, flow=flow) @@ -178,7 +174,7 @@ def register( func: Callable[_FuncParams, _ReturnType], ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: target_broker = self._brokers.resolve(broker) - real_task_name = task_name if not callable(task_name) else None + real_task_name: str | None = None if callable(task_name) else task_name task = target_broker.task(task_name=real_task_name, **labels)(func) if flow is not None: self.route_task(task, broker=target_broker, flow=flow) diff --git a/taskiq/routing/models.py b/taskiq/routing/models.py index 439ae7da..2446aadd 100644 --- a/taskiq/routing/models.py +++ b/taskiq/routing/models.py @@ -26,7 +26,13 @@ def broker_name(self) -> str: @dataclass(frozen=True, slots=True) class TaskiqSubscription: - """Inbound flow subscription owned by a router.""" + """ + Inbound flow subscription owned by a router. + + `task_names` is diagnostic listen-plan metadata. It records which task + declarations caused a broker to listen to this flow; worker execution still + resolves inbound messages by `TaskiqMessage.task_name`. + """ broker: AsyncBroker flow: FlowProtocol diff --git a/taskiq/routing/subscriptions.py b/taskiq/routing/subscriptions.py index 2c6c6b23..7f38ba2d 100644 --- a/taskiq/routing/subscriptions.py +++ b/taskiq/routing/subscriptions.py @@ -15,7 +15,12 @@ class SubscriptionPlan: - """Inbound listen-plan subscriptions for flow-aware brokers.""" + """ + Inbound listen-plan subscriptions for flow-aware brokers. + + This is mutable setup-time configuration, not a thread-safe runtime + coordination primitive. + """ def __init__(self, brokers: BrokerRegistry) -> None: self.brokers = brokers diff --git a/tests/routing/test_shared_tasks.py b/tests/routing/test_shared_tasks.py index 4efab51d..e7aa596c 100644 --- a/tests/routing/test_shared_tasks.py +++ b/tests/routing/test_shared_tasks.py @@ -1,7 +1,7 @@ import pytest from taskiq import AsyncTaskiqDecoratedTask, Flow, task_builder -from tests.routing.models import TracingTask +from tests.routing.models import RecordingMiddleware, TracingTask from tests.utils import RecordingBroker @@ -41,6 +41,7 @@ async def shared_task() -> None: assert source.sent == [] assert target.find_task("shared.routed") is registered + assert target.local_task_registry["shared.routed"] is registered assert target.sent[0][0].task_name == "shared.routed" assert target.sent[0][1] == flow @@ -64,6 +65,27 @@ async def traced(value: int) -> int: assert broker.sent[0][0].task_name == "shared.traced" +async def test_task_builder_custom_base_cls_uses_broker_middleware() -> None: + broker = RecordingBroker() + events: list[tuple[str, str, str]] = [] + broker.add_middlewares(RecordingMiddleware("broker", events)) + + @task_builder("shared.traced.middleware", base_cls=TracingTask) + async def traced(value: int) -> int: + return value + 1 + + registered = broker.register_task(traced) + + assert isinstance(registered, TracingTask) + + await registered.kiq(1) + + assert events == [ + ("broker", "pre_send", "shared.traced.middleware"), + ("broker", "post_send", "shared.traced.middleware"), + ] + + def test_task_definition_default_flow_does_not_create_subscription() -> None: flow = Flow("shared.default") broker = RecordingBroker(default_flow=flow) diff --git a/tests/scheduler/test_scheduler.py b/tests/scheduler/test_scheduler.py index eeb9d35c..bc37effc 100644 --- a/tests/scheduler/test_scheduler.py +++ b/tests/scheduler/test_scheduler.py @@ -114,6 +114,40 @@ async def test_scheduler_without_route_uses_scheduler_broker() -> None: assert scheduled_task.labels == {"source": "database"} +async def test_scheduler_removed_route_uses_scheduler_broker() -> None: + router = TaskiqRouter() + scheduler_flow = Flow("scheduler.default") + scheduler_broker = RecordingBroker( + router=router, + broker_name="scheduler", + default_flow=scheduler_flow, + ) + target_broker = RecordingBroker(router=router, broker_name="target") + source = RecordingScheduleSource() + scheduler = TaskiqScheduler(broker=scheduler_broker, sources=[source]) + + @scheduler_broker.task(task_name="demo.task") + async def demo_task() -> None: + return None + + router.route_task(demo_task, broker=target_broker, flow=Flow("target")) + scheduled_task = ScheduledTask( + task_name="demo.task", + labels={}, + args=[], + kwargs={}, + cron="* * * * *", + ) + del router.routes["demo.task"] + + await scheduler.on_ready(source, scheduled_task) + + sent_message, sent_flow = scheduler_broker.sent[0] + assert target_broker.sent == [] + assert sent_message.task_name == "demo.task" + assert sent_flow == scheduler_flow + + async def test_created_schedule_kiq_does_not_mutate_schedule_payload() -> None: router = TaskiqRouter() source_broker = RecordingBroker(router=router, broker_name="source") From b727f04e774d0c0721fc24dfe50d73cd27425289 Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Sat, 11 Jul 2026 13:40:16 +0200 Subject: [PATCH 06/10] 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 --- docs/available-components/brokers.md | 6 + docs/examples/router/multi_broker_worker.py | 55 + docs/guide/cli.md | 34 +- docs/guide/routing-and-flows.md | 201 +++- taskiq/_task_function.py | 138 +++ taskiq/abc/broker.py | 88 +- taskiq/acks.py | 30 + taskiq/brokers/inmemory_broker.py | 53 +- taskiq/brokers/shared_broker.py | 7 + taskiq/cli/worker/args.py | 14 +- taskiq/cli/worker/run.py | 92 +- taskiq/context.py | 38 +- taskiq/decor.py | 8 + taskiq/kicker.py | 99 +- taskiq/receiver/receiver.py | 438 +++++-- taskiq/receiver/runtime.py | 567 +++++++++ taskiq/router.py | 73 +- taskiq/routing/dispatcher.py | 35 +- taskiq/routing/registries.py | 74 +- taskiq/routing/routes.py | 93 +- taskiq/routing/subscriptions.py | 44 +- taskiq/scheduler/scheduler.py | 2 +- taskiq/sending.py | 57 + taskiq/task_builder.py | 42 +- tests/brokers/test_inmemory.py | 178 +++ tests/cli/worker/test_broker_selection.py | 108 ++ .../middlewares/admin_middleware/conftest.py | 13 +- tests/receiver/test_runtime.py | 1072 +++++++++++++++++ tests/routing/test_dispatch.py | 128 +- tests/routing/test_router_core.py | 40 + tests/routing/test_send_lifecycle.py | 258 ++++ tests/routing/test_shared_tasks.py | 77 ++ tests/routing/test_subscriptions.py | 31 + tests/scheduler/test_scheduler.py | 5 +- tests/test_exceptions_flow.py | 10 + tests/test_requeue.py | 185 ++- tests/test_task_builder.py | 165 ++- 37 files changed, 4155 insertions(+), 403 deletions(-) create mode 100644 docs/examples/router/multi_broker_worker.py create mode 100644 taskiq/_task_function.py create mode 100644 taskiq/receiver/runtime.py create mode 100644 taskiq/sending.py create mode 100644 tests/cli/worker/test_broker_selection.py create mode 100644 tests/receiver/test_runtime.py create mode 100644 tests/routing/test_send_lifecycle.py diff --git a/docs/available-components/brokers.md b/docs/available-components/brokers.md index c4dad6df..6b558827 100644 --- a/docs/available-components/brokers.md +++ b/docs/available-components/brokers.md @@ -12,6 +12,12 @@ This is a special broker for local development. It uses the same functions to ex but all tasks are executed locally in the current thread. By default it uses `InMemoryResultBackend` but this can be overridden. +Because one `InMemoryBroker` instance acts as both client and worker, its +startup and shutdown run both event-handler phases. Middleware and result +backend lifecycle hooks still run once per broker lifecycle. Shutdown waits for +background in-memory tasks before closing those resources and the sync-task +executor. + ## ZeroMQBroker This broker uses [ZMQ](https://zeromq.org/) to communicate between worker and client processes. diff --git a/docs/examples/router/multi_broker_worker.py b/docs/examples/router/multi_broker_worker.py new file mode 100644 index 00000000..7ed0522e --- /dev/null +++ b/docs/examples/router/multi_broker_worker.py @@ -0,0 +1,55 @@ +"""Expose an explicit multi-broker listener tuple for the worker CLI.""" + +import asyncio +from collections.abc import AsyncGenerator + +from taskiq import AsyncBroker, BrokerMessage, Flow, TaskiqRouter + + +class DemoBroker(AsyncBroker): + """Small in-process transport used to keep this configuration executable.""" + + def __init__(self, router: TaskiqRouter, broker_name: str) -> None: + super().__init__(router=router, broker_name=broker_name) + self.messages: asyncio.Queue[bytes] = asyncio.Queue() + + async def kick(self, message: BrokerMessage) -> None: + await self.messages.put(message.message) + + async def listen(self) -> AsyncGenerator[bytes, None]: + while True: + yield await self.messages.get() + + +router = TaskiqRouter() + +commands_broker = DemoBroker(router, "commands") +events_broker = DemoBroker(router, "events") +outbound_broker = DemoBroker(router, "outbound") + + +@commands_broker.task(task_name="orders.notify") +async def notify_order(order_id: str) -> str: + return f"notified:{order_id}" + + +@commands_broker.task(task_name="orders.process") +async def process_order(order_id: str) -> str: + await notify_order.kiq(order_id) + return f"processed:{order_id}" + + +orders_flow = Flow("orders.events") +notifications_flow = Flow("orders.notifications") +router.route_task(process_order, broker=events_broker, flow=orders_flow) +router.route_task( + notify_order, + broker=outbound_broker, + flow=notifications_flow, +) +router.subscribe(events_broker, orders_flow, process_order) + +# Only these brokers receive listeners. The outbound broker belongs to the +# same Router, so this explicit multi-broker worker starts its client lifecycle +# for routed child sends without consuming from it. +worker_brokers = (commands_broker, events_broker) diff --git a/docs/guide/cli.md b/docs/guide/cli.md index 8b8a7ab4..d6df2339 100644 --- a/docs/guide/cli.md +++ b/docs/guide/cli.md @@ -18,6 +18,27 @@ Like this: taskiq worker mybroker:broker_var my_project.module1 my_project.module2 ``` +### Multiple brokers + +The broker import target may expose a module-level sequence of broker instances +that use one shared Router: + +```python +# my_project/worker.py +worker_brokers = (rabbit_broker, kafka_broker) +``` + +```bash +taskiq worker my_project.worker:worker_brokers my_project.tasks +``` + +The sequence is explicit listener configuration. Router brokers outside the +sequence are started in client mode for routed sends and are not consumed by +this worker. A single broker target keeps its existing lifecycle. Execution, +prefetch and tasks-per-child limits apply to the whole worker process, not once +per listener. See [Routing and flows](./routing-and-flows.md) for validation, +lifecycle and adapter compatibility details. + ### Sync function 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
If you send `SIGINT` or `SIGKILL` to the main process by pressing Ctrl+C or using the `kill` command, it will initiate the shutdown process. By default, it will stop fetching new messages immediately after receiving the signal but will wait for the completion of all currently executing tasks. -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. +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. + +When `--wait-tasks-timeout` is omitted, Taskiq keeps the legacy behavior and +waits for the complete listener shutdown, including running callbacks, without +a time limit. When the task timeout is configured, `--shutdown-timeout` is also +the post-drain listener cleanup allowance. Every managed broker then receives +its own `--shutdown-timeout` deadline, so total multi-broker cleanup can take +longer than that value. ::: tip Cool tip 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- * `--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`. * `--max-tasks-per-child` - maximum number of tasks to be executed by a single worker process before restart. * `--max-fails` - Maximum number of child process exits. -* `--shutdown-timeout` - maximum amount of time for graceful broker's shutdown in seconds (default 5). -* `--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. +* `--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. +* `--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. * `--hardkill-count` - Number of termination signals to the main process before performing a hardkill. ## Scheduler diff --git a/docs/guide/routing-and-flows.md b/docs/guide/routing-and-flows.md index 12db2fc9..72119e93 100644 --- a/docs/guide/routing-and-flows.md +++ b/docs/guide/routing-and-flows.md @@ -60,8 +60,55 @@ priority_route = router.route_task( ``` If a broker is created without an explicit router, Taskiq creates a default -router for that broker. Use an explicit shared router when one worker or app -needs several brokers to follow the same routing policy. +router for that broker. Use an explicit shared router when one application +sends tasks through several brokers under the same routing policy. + +The worker CLI can explicitly listen to several brokers from one process. Export +a non-empty module-level sequence and pass its import path as the existing +broker argument: + +```python +# application/worker.py +worker_brokers = (commands_broker, events_broker) +``` + +```bash +taskiq worker application.worker:worker_brokers application.tasks +``` + +The sequence is the complete listener set. Other brokers registered in the +same Router do not receive listeners; the runtime starts them in client mode so +tasks can publish routed child invocations through them. Passing one broker +keeps the legacy lifecycle and does not start Router peers. + +All listener brokers must be distinct, use the same Router and have unique +broker names. One child `Receiver` owns each transport's formatter, middleware, +result backend, acknowledgement and iterator. The executor, +`--max-async-tasks`, `--max-prefetch` and `--max-tasks-per-child` budgets remain +process-wide. Startup is sequential, listener shutdown overlaps, and broker +cleanup runs in reverse startup order. + +Without `--wait-tasks-timeout`, listener shutdown, including running callbacks, +has no deadline, preserving single-broker CLI behavior. When the option is set, +each Receiver cancels and awaits callbacks at that boundary; the coordinator +allows the configured `--shutdown-timeout` as post-drain cleanup time before +hard-cancelling a listener. Each broker shutdown then receives its own timeout. + +Each listener owns one pending transport read. If several deliveries complete +at the exact `--max-tasks-per-child` boundary, the runtime drains those accepted +messages instead of discarding them; the final count can exceed the threshold +by at most the number of additional listeners. It never becomes a separate +full task budget per broker. + +Router routes and subscriptions must be configured before worker startup. A +custom Receiver may inherit the base runtime-state behavior; an override of +`attach_runtime_state()` must call `super()` and preserve the returned state. +Importing a Router itself is not a worker target. + +The core runtime accepts existing adapters, but transport-specific constructor +attachment and Flow translation are adapter capabilities. Verify those against +the adapter version used by the application before combining RabbitMQ, Kafka, +Redis, NATS or another transport in one worker. Pass broker objects to routing APIs. If broker names come from configuration, resolve them explicitly first: @@ -71,7 +118,17 @@ broker = router.get_broker("priority") route = router.route_task("billing.charge", broker=broker) ``` -This keeps routing errors local and avoids hidden string references. +This keeps routing errors local and avoids hidden string references. A task +route declared without `flow=` follows the selected broker's current +`default_flow` at dispatch time. Changing that default after task registration +therefore affects later normal sends; an explicit route flow remains fixed. + +Router registry properties are immutable point-in-time snapshots for +diagnostics: `brokers`, `task_registry` and `routes` return read-only mappings, +while `subscriptions` returns a tuple. They are not live views. Use +`set_broker()`, `register_task()`, `route_task()`, `remove_route()`, +`subscribe()` and `unsubscribe()` so every mutation passes through the owning +registry's validation. ## Flow @@ -119,8 +176,9 @@ Flow identity and declaration options are separate: ## Sending with a route -`router.route_task()` returns a `TaskiqRoute`. You can pass that object to a -kicker when you want the route to be explicit at the call site. +`router.route_task()` returns a resolved `TaskiqRoute`. You can pass that object +to a kicker when you want the current broker and flow to be an explicit +call-site snapshot. ```python route = router.resolve_route("billing.charge") @@ -139,26 +197,48 @@ bulk_route = router.resolve_route( task = await charge.kicker().with_route(bulk_route).kiq(42) ``` -`with_broker(...)` remains available for shared-task style invocation. When a -kicker has both a route and a later broker override, the broker override clears -the explicit route. Taskiq then resolves the task for the override broker: it -uses that broker's task route flow when one exists, otherwise the broker -`default_flow`. +Kicker routing modifiers are applied in call order. The last modifier wins for +the state it controls: `with_route()` replaces broker and flow selection, +`with_broker()` clears every earlier route and flow override, and `with_flow()` +replaces only the flow selected so far. -`prepare(...)` creates a prepared invocation for later use. Prepared invocations -keep a snapshot of the route that was resolved at prepare time, so later changes -to the mutable kicker object or router route table do not change that prepared -send. +`with_broker(...)` remains available for shared-task style invocation. Taskiq +then resolves the task for the override broker: it uses that broker's explicit +same-broker task route flow when one exists, otherwise the broker's current +`default_flow`. A later `with_flow(flow)` overrides that fallback. -Flow selection is resolved in this order: +`with_flow()` requires a flow object. `None` is intentionally rejected because +it could mean either "remove my override" or "send without a flow". To bypass a +configured broker default explicitly, use a fully specified route: + +```python +from taskiq import TaskiqRoute + +no_flow_route = TaskiqRoute(broker=default_broker, flow=None) +task = await charge.kicker().with_route(no_flow_route).kiq(42) +``` + +The resulting resolution contract is: | Source | Effect | | --- | --- | -| Explicit route | `.with_route(route)` uses that route's broker and flow. A later `.with_flow(flow)` updates the route flow, while a later `.with_broker(broker)` clears the explicit route. | -| Explicit broker | `.with_broker(broker)` sends through that broker. If `.with_flow(flow)` is also set, that flow wins; otherwise Taskiq uses the same-broker task route flow, then the broker `default_flow`. | -| Explicit flow | `.with_flow(flow)` without a broker override replaces the flow of the resolved route or broker default route. | -| Router task route | `router.route_task(task, broker=..., flow=...)` is the default outbound route for that task. | -| Broker default flow | `InMemoryBroker(..., default_flow=...)` and other brokers with `default_flow` provide the fallback flow for that broker. | +| `with_route(route)` | Uses exactly that route's broker and flow, including an explicit `flow=None`. A later broker or flow modifier replaces the corresponding selection. | +| `with_broker(broker)` | Clears earlier route/flow overrides. A later `with_flow(flow)` wins; otherwise Taskiq uses an explicit same-broker task route flow, then the broker's current `default_flow`. | +| `with_flow(flow)` | Replaces the flow selected by an earlier route/broker/default policy. A later route or broker modifier replaces or clears it. | +| Router task route | Selects the default outbound broker. Its explicit flow wins; an omitted flow is resolved from the broker's current `default_flow` for each normal send. | +| Broker default flow | Provides the dispatch-time fallback when no explicit route flow exists. | + +`prepare(...)` resolves this table once. The resulting `PreparedKiq` retains the +resolved broker and flow even if the route table or broker `default_flow` +changes before `prepared.kiq()`. + +Routed and legacy sends use the same client middleware/error contract. Taskiq +runs `pre_send` in registration order and `post_send` in reverse order. +Middleware and route-configuration exceptions keep their original types; +formatter and transport exceptions are raised as `SendTaskError` with the +original exception as their cause. Cancellation is never converted into a send +failure, and `post_send` is not called when formatting or transport dispatch +fails. ## Subscriptions @@ -185,6 +265,14 @@ flows = priority_broker.get_subscribed_flows() routing answers "where should this invocation be sent?", while subscribing answers "which flows should this broker listen to?". +`router.unsubscribe(broker, flow)` removes one compatible explicit +subscription and returns the removed `TaskiqSubscription`, or `None` when no +such subscription exists. A flow declaration with the same identity but +different broker options is rejected instead of removing an ambiguous entry. +Removing an explicit subscription does not disable the broker's +`default_flow`; change broker configuration explicitly when that default should +no longer be listened to. + The deprecated `route_task(..., subscribe=True)` shim still performs this subscription when a flow is resolved, but new code should call `subscribe()` directly. @@ -204,7 +292,8 @@ conflicts. Router routes and subscriptions are mutable setup-time configuration. They are not thread-safe runtime coordination primitives; configure them before starting -concurrent worker or client activity. +concurrent worker or client activity. Immutable snapshots prevent accidental +invariant bypass, but they do not add synchronization for concurrent writers. ## Scheduler and requeue @@ -224,13 +313,34 @@ persisted schedule route metadata. They can affect `CreatedSchedule.kiq()`, which is an immediate queued invocation helper, but they do not change the stored schedule payload. -`Context.requeue()` is current-broker sticky. It sends the same task message +`Context.requeue()` is current-broker sticky. It publishes a replacement copy through the broker currently executing the task, even if the task's default route now points to another broker. Requeue flow selection is explicit low-level override, same-broker task route flow, current broker `default_flow`, then no flow. Taskiq does not currently preserve the inbound source flow because flow provenance is not stored in `TaskiqMessage` or `Context`. +Replacement publication and acknowledgement follow an explicit ownership +contract: + +- The replacement gets an incremented `X-Taskiq-requeue` label. The current + Context message is updated only after publication succeeds. +- After successful publication, normal `when_executed` and `when_saved` timing + remains intact. A `manual` delivery is acknowledged by Receiver after the + terminal requeue path. `AckController` keeps every acknowledgement + idempotent. +- A failed or cancelled publication remains unacknowledged for + `when_executed`, `when_saved` and `manual`, allowing broker redelivery. The + original label is not incremented. +- `when_received` is intentionally different: the original was acknowledged + before task execution, so a later publication failure cannot restore it. +- If acknowledgement itself fails after successful publication, that error is + propagated and the controller remains unacknowledged. The replacement may + already exist, so broker redelivery can produce a duplicate and task handlers + must remain idempotent. +- A non-ackable broker can publish the replacement but cannot settle the + original through Taskiq. + ## Shared task declarations Libraries can declare tasks without importing an application broker: @@ -274,6 +384,23 @@ The unbound task can still run locally: total = await calculate_total.call(19, 3) ``` +`TaskDefinition` preserves the wrapped callable's name, qualified name, +module, documentation and inspectable signature. Binding the same definition +to one or several brokers does not rename or otherwise mutate the package-level +function. Bound task objects expose the same callable metadata, while Taskiq +uses a stable internal module-level callable for worker execution and process +pool serialization. + +Factory-generated declarations from the same Python source location must use +stable, unique task names. Taskiq includes `task_name` in the deterministic +module export identity, so each declaration remains independently importable +and picklable in spawned worker processes. + +As with regular synchronous Taskiq tasks, functions intended for a spawned +process pool must be declared at module scope so the worker process can import +their module. The internal binding callable is an implementation detail and +does not change the task name or message contract. + Shared task declarations can also provide a custom task class for one task: ```python @@ -293,12 +420,18 @@ async def traced_charge(user_id: int) -> None: ``` When the final application registers this definition, Taskiq creates the bound -task using `TracingTask`. If `base_cls` is not provided, Taskiq uses the native -decorated task class. +task using `TracingTask`. An explicit `base_cls` takes precedence over broker +defaults. If `base_cls` is not provided, Taskiq uses the selected broker's +`decorator_class`, preserving broker-specific task and kicker APIs. Custom task classes do not bypass broker middleware. Send and execute lifecycle hooks are still owned by the selected broker and worker path. +Binding is published in two phases: Router validates task-name availability +before constructing a bound task, the broker stores it through +`store_registered_task()`, and Router then publishes the task and route. A +storage failure therefore does not leave a partial Router entry. + For low-level integrations, `TaskDefinition.message(...)` builds a `TaskiqMessage` without binding the task. It uses the same argument and label preparation contract as a normal `.kicker().prepare(...)` invocation. @@ -324,11 +457,20 @@ Migration guidance for application code: and routing policy. - Pass broker objects to routing APIs. If configuration stores broker names, resolve them with `router.get_broker(name)` before calling `route_task(...)`. +- Replace direct mutation of Router registry properties with the validated + registration, route and subscription methods described above. - Replace `route_task(..., subscribe=True)` usage with explicit `router.subscribe(...)`. - Do not use labels as a transport routing schema. Labels remain task metadata. - Register a `TaskDefinition` without `task_name` or label overrides; plain callable registration still supports those overrides. +- To listen on several brokers in one worker process, export a module-level + broker tuple and pass that single import path to `taskiq worker`. Do not pass + several broker paths as positional arguments and do not use Router membership + as implicit listener configuration. +- Brokers in the listener tuple share process-wide CLI limits. Remaining + brokers in that Router receive client lifecycle only; keep them out of the + tuple when this worker must not consume their transport. Migration guidance for broker packages: @@ -336,11 +478,19 @@ Migration guidance for broker packages: valid. The default `kick_to_flow(message, flow=None)` calls `kick(message)`. - Flow-aware brokers can override `kick_to_flow(...)` and map `flow.name` to a queue, topic, subject or stream. +- A successful `kick_to_flow(...)` return is the publication boundary used by + `Context.requeue()` before Receiver acknowledges the original delivery. + Flow-aware adapters must not return success before their normal send contract + considers the replacement accepted. - Flow-aware listeners can call `get_subscribed_flows()` during listen setup to add router-owned subscriptions while keeping old default listen behavior when there are no explicit subscriptions. - Broker-specific declaration settings belong in broker-specific flow objects and `broker_options()`, not in task names, labels or scheduler payloads. +- Brokers that store registered tasks outside `local_task_registry` must + override `store_registered_task()`. `AsyncSharedBroker`, for example, keeps + bound definitions in the global task registry just like its old decorator + path. - Ack/nack behavior, serializer/formatter behavior, retry, middleware hooks and result backend ownership remain broker responsibilities. @@ -362,6 +512,9 @@ See executable examples for the two main usage shapes: - [Multiple brokers and explicit routes](../examples/router/multiple_brokers.py) shows one task routed through several brokers with explicit subscription ownership. +- [Multi-broker worker](../examples/router/multi_broker_worker.py) shows the + explicit listener tuple and an outbound-only Router broker used by the CLI + runtime. - [Shared task package](../examples/router/shared_task_package.py) shows `task_builder(...)`, a custom `base_cls`, later registration and a prepared queued invocation. diff --git a/taskiq/_task_function.py b/taskiq/_task_function.py new file mode 100644 index 00000000..0fecd18e --- /dev/null +++ b/taskiq/_task_function.py @@ -0,0 +1,138 @@ +"""Internal callable identity helpers for reusable task declarations.""" + +from __future__ import annotations + +import hashlib +import sys +from collections.abc import Callable +from types import FunctionType +from typing import Any, ParamSpec, TypeVar, cast + +_FuncParams = ParamSpec("_FuncParams") +_ReturnType = TypeVar("_ReturnType") + +_BINDING_MARKER = "__taskiq_binding_function__" +_BINDING_SOURCE = "__taskiq_binding_source__" +_EMPTY_CELL = object() + + +def _values_match(left: object, right: object) -> bool: + """Compare identity or recursively stable immutable values.""" + if left is right: + return True + if type(left) is not type(right): + return False + if isinstance(left, (str, bytes, int, float, bool, type(None))): + return bool(left == right) + if isinstance(left, tuple): + right_tuple = cast(tuple[object, ...], right) + return len(left) == len(right_tuple) and all( + _values_match(left_item, right_item) + for left_item, right_item in zip(left, right_tuple, strict=True) + ) + if isinstance(left, frozenset): + return left == cast(frozenset[object], right) + return False + + +def _mappings_match( + left: dict[str, Any] | None, + right: dict[str, Any] | None, +) -> bool: + """Compare keyword defaults without trusting mutable value equality.""" + if left is right: + return True + if left is None or right is None or left.keys() != right.keys(): + return False + return all(_values_match(left[key], right[key]) for key in left) + + +def _closure_values(func: FunctionType) -> tuple[object, ...]: + """Return closure cell contents while preserving empty-cell identity.""" + values = [] + for cell in func.__closure__ or (): + try: + values.append(cell.cell_contents) + except ValueError: + values.append(_EMPTY_CELL) + return tuple(values) + + +def _execution_state_matches(left: object, right: object) -> bool: + """Return whether two runtime functions execute equivalent captured state.""" + if not isinstance(left, FunctionType) or not isinstance(right, FunctionType): + return False + if left.__code__ is not right.__code__ or left.__globals__ is not right.__globals__: + return False + if not _values_match(left.__defaults__, right.__defaults__): + return False + if not _mappings_match(left.__kwdefaults__, right.__kwdefaults__): + return False + return _values_match(_closure_values(left), _closure_values(right)) + + +def ensure_task_binding_function( + func: Callable[_FuncParams, _ReturnType], + *, + binding_key: str, +) -> Callable[_FuncParams, _ReturnType]: + """Return a stable module-level callable without mutating ``func``.""" + module_name = getattr(func, "__module__", type(func).__module__) + function_qualname = getattr(func, "__qualname__", type(func).__qualname__) + identity_module = ( + "__main__" if module_name in {"__main__", "__mp_main__"} else module_name + ) + function_code = getattr(func, "__code__", None) + first_line = getattr(function_code, "co_firstlineno", 0) + identity = f"{identity_module}\0{function_qualname}\0{first_line}\0{binding_key}" + digest = hashlib.blake2s(identity.encode(), digest_size=10).hexdigest() + export_name = f"_taskiq_shared_{digest}" + module = sys.modules[module_name] + + existing = getattr(module, export_name, None) + existing_source = getattr(existing, _BINDING_SOURCE, None) + if existing_source is func: + return cast(Callable[_FuncParams, _ReturnType], existing) + if existing is not None and _execution_state_matches(existing_source, func): + return cast(Callable[_FuncParams, _ReturnType], existing) + if existing is not None and ( + existing_source is None + or getattr(existing_source, "__code__", None) is function_code + ): + raise ValueError( + f"Task binding identity collision for {binding_key!r}. " + "Factory-generated TaskDefinitions must use unique task names.", + ) + + if isinstance(func, FunctionType): + binding_function = FunctionType( + func.__code__, + func.__globals__, + export_name, + func.__defaults__, + func.__closure__, + ) + binding_function.__kwdefaults__ = func.__kwdefaults__ + else: + + def forwarding_function(*args: Any, **kwargs: Any) -> Any: + return func(*args, **kwargs) + + binding_function = cast(FunctionType, forwarding_function) + + binding_function.__module__ = module_name + binding_function.__name__ = export_name + binding_function.__qualname__ = export_name + binding_function.__doc__ = getattr(func, "__doc__", None) + binding_function.__annotations__ = dict(getattr(func, "__annotations__", {})) + binding_function.__dict__.update(getattr(func, "__dict__", {})) + binding_function.__dict__["__wrapped__"] = func + setattr(binding_function, _BINDING_SOURCE, func) + setattr(binding_function, _BINDING_MARKER, True) + setattr(module, export_name, binding_function) + return cast(Callable[_FuncParams, _ReturnType], binding_function) + + +def is_task_binding_function(func: Callable[..., Any]) -> bool: + """Return whether ``func`` already has a stable process-pool identity.""" + return bool(getattr(func, _BINDING_MARKER, False)) diff --git a/taskiq/abc/broker.py b/taskiq/abc/broker.py index 06926827..55e4d4c9 100644 --- a/taskiq/abc/broker.py +++ b/taskiq/abc/broker.py @@ -31,7 +31,7 @@ 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.task_builder import TaskDefinition, _get_task_binding_function from taskiq.utils import maybe_awaitable from taskiq.warnings import TaskiqDeprecationWarning @@ -220,12 +220,9 @@ def add_middlewares(self, *middlewares: "TaskiqMiddleware") -> None: async def startup(self) -> None: """Do something when starting broker.""" - event = TaskiqEvents.CLIENT_STARTUP - if self.is_worker_process: - event = TaskiqEvents.WORKER_STARTUP - - for handler in self.event_handlers[event]: - await maybe_awaitable(handler(self.state)) + for event in self._get_startup_events(): + for handler in self.event_handlers[event]: + await maybe_awaitable(handler(self.state)) for middleware in self.middlewares: if middleware.__class__.startup != TaskiqMiddleware.startup: @@ -233,6 +230,12 @@ async def startup(self) -> None: await self.result_backend.startup() + def _get_startup_events(self) -> tuple[TaskiqEvents, ...]: + """Return event phases owned by this broker startup.""" + if self.is_worker_process: + return (TaskiqEvents.WORKER_STARTUP,) + return (TaskiqEvents.CLIENT_STARTUP,) + async def shutdown(self) -> None: """ Close the broker. @@ -240,19 +243,55 @@ async def shutdown(self) -> None: This method is called, when broker is closing. """ - event = TaskiqEvents.CLIENT_SHUTDOWN - if self.is_worker_process: - event = TaskiqEvents.WORKER_SHUTDOWN + shutdown_error: BaseException | None = None - # Call all shutdown events. - for handler in self.event_handlers[event]: - await maybe_awaitable(handler(self.state)) + for event in self._get_shutdown_events(): + for handler in self.event_handlers[event]: + try: + await maybe_awaitable(handler(self.state)) + except BaseException as exc: + shutdown_error = self._remember_shutdown_error( + shutdown_error, + exc, + ) for middleware in self.middlewares: if middleware.__class__.shutdown != TaskiqMiddleware.shutdown: - await maybe_awaitable(middleware.shutdown()) - - await self.result_backend.shutdown() + try: + await maybe_awaitable(middleware.shutdown()) + except BaseException as exc: + shutdown_error = self._remember_shutdown_error( + shutdown_error, + exc, + ) + + try: + await self.result_backend.shutdown() + except BaseException as exc: + shutdown_error = self._remember_shutdown_error(shutdown_error, exc) + + if shutdown_error is not None: + raise shutdown_error + + def _get_shutdown_events(self) -> tuple[TaskiqEvents, ...]: + """Return event phases owned by this broker shutdown.""" + if self.is_worker_process: + return (TaskiqEvents.WORKER_SHUTDOWN,) + return (TaskiqEvents.CLIENT_SHUTDOWN,) + + @staticmethod + def _remember_shutdown_error( + first_error: BaseException | None, + current_error: BaseException, + ) -> BaseException: + """Keep the first shutdown failure while allowing cleanup to continue.""" + if first_error is None: + return current_error + logger.error( + "Additional error while shutting down broker resources.", + exc_info=current_error, + ) + return first_error @abstractmethod async def kick( @@ -415,7 +454,6 @@ def register_task( return self.router.register_task( func, broker=self, - flow=self.default_flow, ) return self._decorate_task( func, @@ -430,13 +468,13 @@ def bind_task_definition( register: bool = True, ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: """Bind an unbound task definition to this broker.""" - task_cls = AsyncTaskiqDecoratedTask if task.base_cls is None else task.base_cls return self._decorate_task( - task.original_func, + _get_task_binding_function(task), task_name=task.task_name, labels=dict(task.labels), - base_cls=task_cls, + base_cls=task.base_cls, register=register, + metadata_func=task.original_func, ) def store_registered_task( @@ -448,7 +486,8 @@ def store_registered_task( Application code should use `task` or `register_task`. This method is the broker/router integration boundary for cases where the router owns - registration ordering and the broker owns local task storage. + registration ordering and the broker owns task storage. Brokers with a + custom registration store should override this method. """ self._store_task(task.task_name, task) @@ -460,8 +499,10 @@ def _decorate_task( base_cls: type[AsyncTaskiqDecoratedTask[Any, Any]] | None = None, *, register: bool = True, + metadata_func: Callable[_FuncParams, _ReturnType] | None = None, ) -> AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType]: """Build and register a decorated task object.""" + metadata_source = func if metadata_func is None else metadata_func real_task_name = task_name if real_task_name is None: fmodule = func.__module__ @@ -483,12 +524,12 @@ def _decorate_task( ): raise TypeError("base_cls must be a subclass of AsyncTaskiqDecoratedTask.") - sign = get_type_hints(func) + sign = get_type_hints(metadata_source) return_type = None if "return" in sign: return_type = sign["return"] - decorated_task = wraps(func)( + decorated_task = wraps(metadata_source)( task_cls( broker=self, original_func=func, @@ -655,7 +696,6 @@ def _register_task( self.router.register_task( task, broker=self, - flow=self.default_flow, ) self._store_task(task_name, task) diff --git a/taskiq/acks.py b/taskiq/acks.py index 894a9160..2bc4097e 100644 --- a/taskiq/acks.py +++ b/taskiq/acks.py @@ -57,11 +57,19 @@ class AckableMessage(BaseModel): ack: Callable[[], None | Awaitable[None]] +class _RequeueState(enum.Enum): + NOT_REQUESTED = enum.auto() + PUBLISHING = enum.auto() + PUBLISHED = enum.auto() + FAILED = enum.auto() + + class AckController: """Controls acknowledgement state for a received message.""" def __init__(self, ack: Callable[[], None | Awaitable[None]] | None) -> None: self._ack = ack + self._requeue_state = _RequeueState.NOT_REQUESTED self.is_acked = False @property @@ -69,6 +77,28 @@ def is_ackable(self) -> bool: """Whether the current message supports acknowledgement.""" return self._ack is not None + @property + def requeue_published(self) -> bool: + """Whether a replacement message was published successfully.""" + return self._requeue_state == _RequeueState.PUBLISHED + + @property + def requeue_failed(self) -> bool: + """Whether replacement publication failed or was cancelled.""" + return self._requeue_state == _RequeueState.FAILED + + def start_requeue(self) -> None: + """Mark replacement publication as in progress.""" + self._requeue_state = _RequeueState.PUBLISHING + + def complete_requeue(self) -> None: + """Mark replacement publication as successful.""" + self._requeue_state = _RequeueState.PUBLISHED + + def fail_requeue(self) -> None: + """Mark replacement publication as failed or cancelled.""" + self._requeue_state = _RequeueState.FAILED + async def ack(self) -> None: """Acknowledge the current message once.""" if self._ack is None: diff --git a/taskiq/brokers/inmemory_broker.py b/taskiq/brokers/inmemory_broker.py index e63a97f6..d37bdd75 100644 --- a/taskiq/brokers/inmemory_broker.py +++ b/taskiq/brokers/inmemory_broker.py @@ -13,7 +13,6 @@ from taskiq.message import BrokerMessage from taskiq.receiver import Receiver from taskiq.router import TaskiqRouter -from taskiq.utils import maybe_awaitable _ReturnType = TypeVar("_ReturnType") @@ -198,19 +197,45 @@ async def wait_all(self) -> None: Useful when used in testing and you need to await all sent tasks before asserting results. """ - to_await = list(self._running_tasks) - for task in to_await: - await task + task_errors: list[BaseException] = [] + while self._running_tasks: + results = await asyncio.gather( + *tuple(self._running_tasks), + return_exceptions=True, + ) + task_errors.extend( + result for result in results if isinstance(result, BaseException) + ) - async def startup(self) -> None: - """Runs startup events for client and worker side.""" - for event in (TaskiqEvents.CLIENT_STARTUP, TaskiqEvents.WORKER_STARTUP): - for handler in self.event_handlers.get(event, []): - await maybe_awaitable(handler(self.state)) + if task_errors: + raise task_errors[0] + + def _get_startup_events(self) -> tuple[TaskiqEvents, ...]: + """Run both sides because this broker executes tasks in the client process.""" + return (TaskiqEvents.CLIENT_STARTUP, TaskiqEvents.WORKER_STARTUP) + + def _get_shutdown_events(self) -> tuple[TaskiqEvents, ...]: + """Shut down both client and worker event phases in their legacy order.""" + return (TaskiqEvents.CLIENT_SHUTDOWN, TaskiqEvents.WORKER_SHUTDOWN) async def shutdown(self) -> None: - """Runs shutdown events for client and worker side.""" - for event in (TaskiqEvents.CLIENT_SHUTDOWN, TaskiqEvents.WORKER_SHUTDOWN): - for handler in self.event_handlers.get(event, []): - await maybe_awaitable(handler(self.state)) - self.executor.shutdown() + """Drain local execution, close broker resources and stop the executor.""" + shutdown_error: BaseException | None = None + + try: + await self.wait_all() + except BaseException as exc: + shutdown_error = exc + + try: + await super().shutdown() + except BaseException as exc: + shutdown_error = self._remember_shutdown_error(shutdown_error, exc) + + try: + self.executor.shutdown() + except BaseException as exc: # pragma: no cover + shutdown_error = self._remember_shutdown_error(shutdown_error, exc) + + if shutdown_error is not None: + raise shutdown_error diff --git a/taskiq/brokers/shared_broker.py b/taskiq/brokers/shared_broker.py index 7d12edf3..d5170a27 100644 --- a/taskiq/brokers/shared_broker.py +++ b/taskiq/brokers/shared_broker.py @@ -75,6 +75,13 @@ def _register_task( ) -> None: self.global_task_registry[task_name] = task + def store_registered_task( + self, + task: AsyncTaskiqDecoratedTask[Any, Any], + ) -> None: + """Store a router-bound shared task in the global task registry.""" + self._register_task(task.task_name, task) + async_shared_broker = AsyncSharedBroker() shared_task = async_shared_broker.task diff --git a/taskiq/cli/worker/args.py b/taskiq/cli/worker/args.py index df3eb957..7cbf72bd 100644 --- a/taskiq/cli/worker/args.py +++ b/taskiq/cli/worker/args.py @@ -1,8 +1,11 @@ +from __future__ import annotations + from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser, Namespace from collections.abc import Sequence from dataclasses import dataclass, field from typing import Any +from taskiq.abc.broker import AsyncBroker from taskiq.acks import AcknowledgeType from taskiq.cli.common_args import LogLevel @@ -25,7 +28,7 @@ def receiver_arg_type(string: str) -> tuple[str, str]: class WorkerArgs: """Taskiq worker CLI arguments.""" - broker: str + broker: str | AsyncBroker | Sequence[AsyncBroker] modules: list[str] app_dir: str | None = None tasks_pattern: Sequence[str] = ("**/tasks.py",) @@ -61,7 +64,7 @@ def from_cli( cls, args: Sequence[str] | None = None, defaults: dict[str, Any] | None = None, - ) -> "WorkerArgs": + ) -> WorkerArgs: """ Construct TaskiqArgs instanc from CLI arguments. @@ -73,7 +76,8 @@ def from_cli( parser.add_argument( "broker", help=( - "Where to search for broker or broker factory function. " + "Where to search for a broker, broker listener sequence, " + "or factory function. " "This string must be specified in " "'module.module:variable' format." ), @@ -179,7 +183,7 @@ def from_cli( "--shutdown-timeout", type=float, default=5, - help="Maximum amount of time for graceful broker's shutdown is seconds.", + help="Maximum time in seconds for each broker's graceful shutdown.", ) parser.add_argument( "--reload", @@ -256,7 +260,7 @@ def from_cli( "--wait-tasks-timeout", type=float, default=None, - help="Maximum time to wait for all current tasks to finish before exiting.", + help="Maximum graceful wait before cancelling active task callbacks.", ) parser.add_argument( "--hardkill-count", diff --git a/taskiq/cli/worker/run.py b/taskiq/cli/worker/run.py index 24d8f8db..a23f20b9 100644 --- a/taskiq/cli/worker/run.py +++ b/taskiq/cli/worker/run.py @@ -4,6 +4,7 @@ import os import signal import sys +import warnings from concurrent.futures import Executor, ProcessPoolExecutor, ThreadPoolExecutor from multiprocessing import get_start_method, set_start_method from sys import platform @@ -14,6 +15,14 @@ from taskiq.cli.worker.args import WorkerArgs from taskiq.cli.worker.process_manager import ProcessManager from taskiq.receiver import Receiver +from taskiq.receiver.runtime import ( + BrokerSelection, + WorkerRuntime, + WorkerRuntimeOptions, + _shutdown_broker, + normalize_broker_source, +) +from taskiq.warnings import TaskiqDeprecationWarning try: import uvloop @@ -30,30 +39,16 @@ async def shutdown_broker(broker: AsyncBroker, timeout: float) -> None: - """ - This function used to shutdown broker. - - Broker can throw errors during shutdown, - or it may return some value. - - We need to handle such situations. - - :param broker: current broker. - :param timeout: maximum amount of time to shutdown the broker. - """ - logger.info("Shutting down the broker.") - try: - ret_val = await asyncio.wait_for(broker.shutdown(), timeout) # type: ignore - if ret_val is not None: - logger.info("Broker has returned value on shutdown: '%s'", str(ret_val)) - except asyncio.TimeoutError: - logger.warning("Broker.shutdown cannot be completed in %s seconds.", timeout) - except Exception as exc: - logger.warning( - "Exception found while shutting down broker: %s", - exc, - exc_info=True, - ) + """Shut down one broker through the legacy CLI helper.""" + warnings.warn( + "taskiq.cli.worker.run.shutdown_broker() is deprecated; worker resource " + "lifecycle is owned by WorkerRuntime.", + TaskiqDeprecationWarning, + stacklevel=2, + ) + cancellation = await _shutdown_broker(broker, timeout) + if cancellation is not None: + raise cancellation def get_receiver_type(args: WorkerArgs) -> type[Receiver]: @@ -70,6 +65,17 @@ def get_receiver_type(args: WorkerArgs) -> type[Receiver]: return receiver_type +def get_broker_selection(args: WorkerArgs) -> BrokerSelection: + """Resolve one broker or an explicit broker-listener sequence.""" + if isinstance(args.broker, str): + broker_source = import_object(args.broker, app_dir=args.app_dir) + if inspect.isfunction(broker_source): + broker_source = broker_source() + else: + broker_source = args.broker + return normalize_broker_source(broker_source) + + def start_listen(args: WorkerArgs) -> None: """ This function starts actual listening process. @@ -133,19 +139,8 @@ def interrupt_handler(signum: int, _frame: Any) -> None: # We must set this field before importing tasks, # so broker will remember all tasks it's related to. - if isinstance(args.broker, AsyncBroker): - broker = args.broker - else: - broker = import_object(args.broker, app_dir=args.app_dir) - if inspect.isfunction(broker): - broker = broker() - if not isinstance(broker, AsyncBroker): - raise ValueError( - "Unknown broker type. Please use AsyncBroker instance " - "or pass broker factory function that returns an AsyncBroker instance.", - ) - - broker.is_worker_process = True + selection = get_broker_selection(args) + selection.mark_listener_brokers() import_tasks(args.modules, args.tasks_pattern, args.fs_discover) receiver_type = get_receiver_type(args) @@ -157,12 +152,13 @@ def interrupt_handler(signum: int, _frame: Any) -> None: else: executor = ThreadPoolExecutor(max_workers=args.max_threadpool_threads) - try: - logger.debug("Initialize receiver.") - with executor as pool: - receiver = receiver_type( - broker=broker, - executor=pool, + logger.debug("Initialize worker runtime.") + with executor as pool: + runtime = WorkerRuntime( + selection=selection, + receiver_type=receiver_type, + executor=pool, + options=WorkerRuntimeOptions( validate_params=not args.no_parse, max_async_tasks=args.max_async_tasks, max_async_tasks_jitter=args.max_async_tasks_jitter, @@ -171,11 +167,11 @@ def interrupt_handler(signum: int, _frame: Any) -> None: ack_type=args.ack_type, max_tasks_to_execute=args.max_tasks_per_child, wait_tasks_timeout=args.wait_tasks_timeout, - **receiver_kwargs, # type: ignore - ) - loop.run_until_complete(receiver.listen(shutdown_event)) - finally: - loop.run_until_complete(shutdown_broker(broker, args.shutdown_timeout)) + shutdown_timeout=args.shutdown_timeout, + receiver_kwargs=receiver_kwargs, + ), + ) + loop.run_until_complete(runtime.run(shutdown_event)) def run_worker(args: WorkerArgs) -> int | None: diff --git a/taskiq/context.py b/taskiq/context.py index 1e380cf0..365a46e6 100644 --- a/taskiq/context.py +++ b/taskiq/context.py @@ -2,12 +2,15 @@ from taskiq.abc.broker import AsyncBroker from taskiq.acks import AckController +from taskiq.compat import model_copy from taskiq.exceptions import NoResultError, TaskRejectedError from taskiq.message import TaskiqMessage if TYPE_CHECKING: # pragma: no cover from taskiq.state import TaskiqState +_REQUEUE_LABEL = "X-Taskiq-requeue" + class Context: """Context class.""" @@ -46,18 +49,39 @@ async def ack(self) -> None: async def requeue(self) -> None: """ - Requeue task. + Publish a replacement task and settle the current delivery. - This function creates a task with - the same message and sends it using - current broker. + The current message is acknowledged only after the replacement has + been published successfully. Non-ackable messages are still republished + but cannot be settled by Taskiq. :raises NoResultError: to not store result for current task. """ - requeue_count = int(self.message.labels.get("X-Taskiq-requeue", 0)) + requeue_count = int(self.message.labels.get(_REQUEUE_LABEL, 0)) requeue_count += 1 - self.message.labels["X-Taskiq-requeue"] = str(requeue_count) - await self.broker.router.requeue(self.message, broker=self.broker) + requeue_value = str(requeue_count) + requeued_message = model_copy( + self.message, + update={ + "labels": { + **self.message.labels, + _REQUEUE_LABEL: requeue_value, + }, + }, + ) + + if self._ack_controller is not None: + self._ack_controller.start_requeue() + try: + await self.broker.router.requeue(requeued_message, broker=self.broker) + except BaseException: + if self._ack_controller is not None: + self._ack_controller.fail_requeue() + raise + + self.message.labels[_REQUEUE_LABEL] = requeue_value + if self._ack_controller is not None: + self._ack_controller.complete_requeue() raise NoResultError def reject(self) -> None: diff --git a/taskiq/decor.py b/taskiq/decor.py index 8f90541c..b4ff0da2 100644 --- a/taskiq/decor.py +++ b/taskiq/decor.py @@ -13,6 +13,7 @@ overload, ) +from taskiq._task_function import is_task_binding_function from taskiq.kicker import AsyncKicker from taskiq.scheduler.created_schedule import CreatedSchedule from taskiq.task import AsyncTaskiqTask @@ -43,6 +44,10 @@ class AsyncTaskiqDecoratedTask(Generic[_FuncParams, _ReturnType]): current broker. """ + __name__: str + __qualname__: str + __wrapped__: Callable[_FuncParams, _ReturnType] + def __init__( self, broker: "AsyncBroker", @@ -57,6 +62,9 @@ def __init__( self.labels = labels self.return_type = return_type + if is_task_binding_function(original_func): + return + # This is a hack to make ProcessPoolExecutor work # with decorated functions. # diff --git a/taskiq/kicker.py b/taskiq/kicker.py index 8de79f90..cdef2328 100644 --- a/taskiq/kicker.py +++ b/taskiq/kicker.py @@ -11,20 +11,16 @@ Generic, ParamSpec, TypeVar, - cast, overload, ) -from taskiq.abc.middleware import TaskiqMiddleware -from taskiq.abc.result_backend import AsyncResultBackend -from taskiq.exceptions import SendTaskError from taskiq.flow import FlowProtocol from taskiq.message import TaskiqMessage, _build_taskiq_message from taskiq.router import TaskiqRoute, TaskiqRouter from taskiq.scheduler.created_schedule import CreatedSchedule from taskiq.scheduler.scheduled_task import CronSpec, ScheduledTask +from taskiq.sending import send_task_message from taskiq.task import AsyncTaskiqTask -from taskiq.utils import maybe_awaitable if TYPE_CHECKING: # pragma: no cover from taskiq.abc.broker import AsyncBroker @@ -132,7 +128,7 @@ def with_broker( broker: AsyncBroker, ) -> AsyncKicker[_FuncParams, _ReturnType]: """ - Replace broker for the function. + Replace broker for the function and clear earlier route/flow overrides. This method can be used with shared tasks. @@ -148,14 +144,22 @@ def with_broker( def with_flow( self, - flow: FlowProtocol | None, + flow: FlowProtocol, ) -> AsyncKicker[_FuncParams, _ReturnType]: """ Replace flow for the current invocation. - :param flow: flow to send message to. + Use an explicit `TaskiqRoute(broker=..., flow=None)` when an invocation + must bypass a broker default flow. + + :param flow: non-null flow to send message to. :return: Kicker with a route flow override. """ + if flow is None: + raise TypeError( + "flow cannot be None; pass an explicit TaskiqRoute with " + "flow=None to bypass the broker default flow.", + ) if self.route is not None: self.route = replace(self.route, flow=flow) self.route_flow = flow @@ -253,36 +257,31 @@ async def kiq_message( use_current_route: bool = True, ) -> AsyncTaskiqTask[_ReturnType]: """Send a prepared message.""" - try: - target_broker = broker or self.broker - if use_current_route: - target_route = self.route if route is None else route - target_flow = self.route_flow if flow is None else flow - else: - target_route = route - target_flow = flow - if target_route is not None: - if broker is not None and broker is not target_route.broker: - raise ValueError("Pass either route or broker override.") - target_broker = target_route.broker - target_flow = None - router = getattr(target_broker, "router", None) - if isinstance(router, TaskiqRouter): - broker_override = None - if target_route is None and ( - broker is not None or self._broker_overridden - ): - broker_override = target_broker - return await router.kiq( - message, - route=target_route, - broker=broker_override, - flow=target_flow, - return_type=self.return_type, - ) - return await self._legacy_kiq(target_broker, message) - except Exception as exc: - raise SendTaskError from exc + target_broker = broker or self.broker + if use_current_route: + target_route = self.route if route is None else route + target_flow = self.route_flow if flow is None else flow + else: + target_route = route + target_flow = flow + if target_route is not None: + if broker is not None and broker is not target_route.broker: + raise ValueError("Pass either route or broker override.") + target_broker = target_route.broker + target_flow = None + router = getattr(target_broker, "router", None) + if isinstance(router, TaskiqRouter): + broker_override = None + if target_route is None and (broker is not None or self._broker_overridden): + broker_override = target_broker + return await router.kiq( + message, + route=target_route, + broker=broker_override, + flow=target_flow, + return_type=self.return_type, + ) + return await self._legacy_kiq(target_broker, message) async def _legacy_kiq( self, @@ -295,26 +294,10 @@ async def _legacy_kiq( This keeps middleware tests and external broker-like mocks compatible while real AsyncBroker instances use TaskiqRouter. """ - middlewares = getattr(broker, "middlewares", []) - if not isinstance(middlewares, list): - middlewares = [] - - for middleware in middlewares: - if middleware.__class__.pre_send != TaskiqMiddleware.pre_send: - message = await maybe_awaitable(middleware.pre_send(message)) - - await broker.kick(broker.formatter.dumps(message)) - - for middleware in reversed(middlewares): - if middleware.__class__.post_send != TaskiqMiddleware.post_send: - await maybe_awaitable(middleware.post_send(message)) - - return AsyncTaskiqTask( - task_id=message.task_id, - result_backend=cast( - AsyncResultBackend[_ReturnType], - broker.result_backend, - ), + return await send_task_message( + broker, + message, + send=broker.kick, return_type=self.return_type, ) diff --git a/taskiq/receiver/receiver.py b/taskiq/receiver/receiver.py index 871d13d6..c0904325 100644 --- a/taskiq/receiver/receiver.py +++ b/taskiq/receiver/receiver.py @@ -4,8 +4,9 @@ import inspect import random import sys -from collections.abc import Callable +from collections.abc import AsyncGenerator, Callable from concurrent.futures import Executor, ProcessPoolExecutor +from dataclasses import dataclass from logging import getLogger from time import time from typing import Any, get_type_hints @@ -20,13 +21,30 @@ from taskiq.exceptions import NoResultError from taskiq.message import TaskiqMessage from taskiq.receiver.params_parser import parse_params +from taskiq.receiver.runtime import ReceiverRuntimeState from taskiq.result import TaskiqResult from taskiq.state import TaskiqState from taskiq.utils import maybe_awaitable logger = getLogger(__name__) PY_VERSION = sys.version_info -QUEUE_DONE = b"-1" +QUEUE_DONE = object() + + +@dataclass(frozen=True, slots=True) +class _PrefetchedMessage: + """A delivery admitted to the runner with its prefetch-slot ownership.""" + + data: bytes | AckableMessage + owns_prefetch_slot: bool + + +@dataclass(slots=True) +class _PrefetchState: + """Own one broker iterator and its pending read task.""" + + iterator: AsyncGenerator[bytes | AckableMessage, None] + current_message: asyncio.Task[bytes | AckableMessage] | None def _execute_sync_task_in_executor( @@ -78,6 +96,9 @@ def __init__( self.known_tasks: set[str] = set() self.max_tasks_to_execute = max_tasks_to_execute self.wait_tasks_timeout = wait_tasks_timeout + self._runtime_state: ReceiverRuntimeState | None = None + self._listen_error: BaseException | None = None + self._active_tasks: set[asyncio.Task[Any]] = set() for task in self.broker.get_all_tasks().values(): self._prepare_task(task.task_name, task.original_func) self.sem: asyncio.Semaphore | None = None @@ -96,9 +117,32 @@ def __init__( "Setting unlimited number of async tasks " "can result in undefined behavior", ) - self.sem_prefetch = asyncio.Semaphore(max_prefetch) + if max_prefetch < 0: + raise ValueError("max_prefetch cannot be negative.") + self.sem_prefetch = asyncio.Semaphore(max_prefetch + 1) self.is_process_pool = isinstance(executor, ProcessPoolExecutor) + def attach_runtime_state( + self, + runtime_state: ReceiverRuntimeState, + ) -> ReceiverRuntimeState: + """Attach process-wide limits supplied by a multi-broker runtime.""" + self._runtime_state = runtime_state + self.sem = runtime_state.execution_semaphore + self.sem_prefetch = runtime_state.prefetch_semaphore + return runtime_state + + def is_runtime_state_attached( + self, + runtime_state: ReceiverRuntimeState, + ) -> bool: + """Return whether this Receiver uses all supplied process-wide state.""" + return ( + self._runtime_state is runtime_state + and self.sem is runtime_state.execution_semaphore + and self.sem_prefetch is runtime_state.prefetch_semaphore + ) + async def callback( # noqa: C901, PLR0912 self, message: bytes | AckableMessage, @@ -167,7 +211,11 @@ async def callback( # noqa: C901, PLR0912 ack_controller=ack_controller, ) - if ack_time == AcknowledgeType.WHEN_EXECUTED and ack_controller.is_ackable: + if ( + ack_time == AcknowledgeType.WHEN_EXECUTED + and ack_controller.is_ackable + and not ack_controller.requeue_failed + ): await ack_controller.ack() for middleware in reversed(self.broker.middlewares): @@ -191,7 +239,15 @@ async def callback( # noqa: C901, PLR0912 if raise_err: raise exc - if ack_time == AcknowledgeType.WHEN_SAVED and ack_controller.is_ackable: + should_ack_saved = ( + ack_time == AcknowledgeType.WHEN_SAVED and not ack_controller.requeue_failed + ) + should_ack_manual_requeue = ( + ack_time == AcknowledgeType.MANUAL and ack_controller.requeue_published + ) + if ( + should_ack_saved or should_ack_manual_requeue + ) and ack_controller.is_ackable: await ack_controller.ack() def _get_ack_time(self, message: TaskiqMessage) -> AcknowledgeType: @@ -386,19 +442,23 @@ async def listen(self, finish_event: asyncio.Event) -> None: # pragma: no cover """ if self.run_startup: await self.broker.startup() + self._listen_error = None logger.info("Listening started.") - queue: asyncio.Queue[bytes | AckableMessage] = asyncio.Queue() + queue: asyncio.Queue[_PrefetchedMessage | object] = asyncio.Queue() async with anyio.create_task_group() as gr: gr.start_soon(self.prefetcher, queue, finish_event) gr.start_soon(self.runner, queue) + if self._listen_error is not None: + raise self._listen_error + if self.on_exit is not None: self.on_exit(self) async def prefetcher( self, - queue: "asyncio.Queue[bytes | AckableMessage]", + queue: "asyncio.Queue[_PrefetchedMessage | object]", finish_event: asyncio.Event, ) -> None: """ @@ -407,117 +467,301 @@ async def prefetcher( :param queue: queue for prefetched data. :param finish_event: event to indicate that we need to stop prefetching. """ - fetched_tasks: int = 0 - iterator = self.broker.listen() - current_message: asyncio.Task[bytes | AckableMessage] = asyncio.create_task( - iterator.__anext__(), # type: ignore - ) + try: + iterator = await self._open_broker_iterator() + state = _PrefetchState( + iterator=iterator, + current_message=None, + ) + except BaseException as exc: + self._record_listen_error(exc) + await queue.put(QUEUE_DONE) + return + fetched_tasks = 0 - while True: - if finish_event.is_set(): - break - try: - await self.sem_prefetch.acquire() - if ( - self.max_tasks_to_execute - and fetched_tasks >= self.max_tasks_to_execute - ): - logger.info("Max number of tasks executed.") - break - # Here we wait for the message to be fetched, - # but we make it with timeout so it can be interrupted - done, _ = await asyncio.wait({current_message}, timeout=0.3) - # If the message is not fetched, we release the semaphore - # and continue the loop. So it will check if finished event was set. - if not done: - self.sem_prefetch.release() + try: + while not self._should_stop_prefetch(finish_event, fetched_tasks): + message = await self._get_prefetched_message(state, finish_event) + if message is None: continue - # We're done, so now we need to check - # whether task has returned an error. - message = current_message.result() - current_message = asyncio.create_task(iterator.__anext__()) # type: ignore fetched_tasks += 1 - await queue.put(message) - # Custom hooks for OTel and any future instrumentations - for middleware in reversed(self.broker.middlewares): - if hasattr(middleware, "on_prefetch_queue_add"): - await maybe_awaitable( - middleware.on_prefetch_queue_add(), # type: ignore - ) - except (asyncio.CancelledError, StopAsyncIteration): - break - # We don't want to fetch new messages if we are shutting down. - logger.info("Stopping prefetching messages...") - current_message.cancel() - await queue.put(QUEUE_DONE) + if self._runtime_state is not None: + self._runtime_state.record_fetched_task() + await queue.put( + _PrefetchedMessage( + data=message, + owns_prefetch_slot=True, + ), + ) + await self._notify_prefetch_hook("on_prefetch_queue_add") + except StopAsyncIteration: + pass + except asyncio.CancelledError: + raise + except BaseException as exc: + self._record_listen_error(exc) + finally: + logger.info("Stopping prefetching messages...") + try: + await self._enqueue_late_prefetched_message(queue, state) + finally: + await queue.put(QUEUE_DONE) + + async def _enqueue_late_prefetched_message( + self, + queue: "asyncio.Queue[_PrefetchedMessage | object]", + state: _PrefetchState, + ) -> None: + """Drain a delivery that completed concurrently with listener stop.""" + late_message = await self._close_prefetch_state(state) + if late_message is None: + return + if self._runtime_state is not None: + self._runtime_state.record_fetched_task() + await queue.put( + _PrefetchedMessage( + data=late_message, + owns_prefetch_slot=False, + ), + ) + try: + await self._notify_prefetch_hook("on_prefetch_queue_add") + except BaseException as exc: + self._record_listen_error(exc) + + async def _open_broker_iterator( + self, + ) -> AsyncGenerator[bytes | AckableMessage, None]: + """Open and validate the transport listener owned by this Receiver.""" + iterator: object = self.broker.listen() + if inspect.isawaitable(iterator): + iterator = await iterator + if not isinstance(iterator, AsyncGenerator): + raise TypeError("Broker.listen() must return an async generator.") + return iterator + + def _should_stop_prefetch( + self, + finish_event: asyncio.Event, + fetched_tasks: int, + ) -> bool: + """Return whether this Receiver should stop requesting deliveries.""" + if finish_event.is_set(): + return True + if ( + self._runtime_state is None + and self.max_tasks_to_execute + and fetched_tasks >= self.max_tasks_to_execute + ): + logger.info("Max number of tasks executed.") + return True + return False + + async def _get_prefetched_message( + self, + state: _PrefetchState, + finish_event: asyncio.Event, + ) -> bytes | AckableMessage | None: + """Wait briefly for one delivery while holding a prefetch slot.""" + await self.sem_prefetch.acquire() + if state.current_message is None: + state.current_message = asyncio.create_task( + state.iterator.__anext__(), + ) + current_message = state.current_message + finish_waiter = asyncio.create_task(finish_event.wait()) + try: + try: + done, _ = await asyncio.wait( + {current_message, finish_waiter}, + timeout=0.3, + return_when=asyncio.FIRST_COMPLETED, + ) + finally: + finish_waiter.cancel() + await asyncio.gather(finish_waiter, return_exceptions=True) + except BaseException: + self.sem_prefetch.release() + raise + + if current_message in done: + state.current_message = None + try: + return current_message.result() + except BaseException: + self.sem_prefetch.release() + raise + + if not done: + self.sem_prefetch.release() + return None + self.sem_prefetch.release() + return None + + async def _close_prefetch_state( + self, + state: _PrefetchState, + ) -> bytes | AckableMessage | None: + """Close the broker iterator and return a concurrently finished read.""" + late_message: bytes | AckableMessage | None = None + current_message = state.current_message + state.current_message = None + if current_message is not None: + current_message.cancel() + try: + late_message = await current_message + except (asyncio.CancelledError, StopAsyncIteration): + pass + except BaseException as exc: + self._record_listen_error(exc) + + try: + await state.iterator.aclose() + except (asyncio.CancelledError, StopAsyncIteration): + pass + except BaseException as exc: + self._record_listen_error(exc) + return late_message + + async def _notify_prefetch_hook(self, hook_name: str) -> None: + """Run one optional prefetch instrumentation hook in reverse order.""" + for middleware in reversed(self.broker.middlewares): + hook = getattr(middleware, hook_name, None) + if hook is not None: + await maybe_awaitable(hook()) async def runner( self, - queue: "asyncio.Queue[bytes | AckableMessage]", + queue: "asyncio.Queue[_PrefetchedMessage | object]", ) -> None: """ Run tasks. :param queue: queue with prefetched data. """ - tasks: set[asyncio.Task[Any]] = set() - - def task_cb(task: "asyncio.Task[Any]") -> None: - """ - Callback for tasks. - - This function used to remove task - from the list of active tasks and release - the semaphore, so other tasks can use it. + cancelled = False + try: + while True: + queued_message = await queue.get() + if queued_message is QUEUE_DONE: + break + if not isinstance(queued_message, _PrefetchedMessage): + raise TypeError("Receiver queue contains an invalid message.") + await self._start_callback(queued_message) + + except asyncio.CancelledError: + cancelled = True + raise + finally: + await self._drain_active_tasks(cancel_immediately=cancelled) + logger.info("The runner is stopped.") - :param task: finished task - """ - tasks.discard(task) + async def _start_callback(self, message: _PrefetchedMessage) -> None: + """Transfer one prefetched delivery into an owned callback task.""" + owns_prefetch_slot = message.owns_prefetch_slot + owns_execution_slot = False + try: + await self._notify_prefetch_hook("on_prefetch_queue_remove") if self.sem is not None: + await self.sem.acquire() + owns_execution_slot = True + + if owns_prefetch_slot: + self.sem_prefetch.release() + owns_prefetch_slot = False + task = asyncio.create_task( + self.callback(message=message.data, raise_err=False), + ) + except BaseException: + if owns_prefetch_slot: + self.sem_prefetch.release() + if owns_execution_slot and self.sem is not None: self.sem.release() + raise - while True: - try: - # Waits for semaphore to be released. - if self.sem is not None: - await self.sem.acquire() + self._active_tasks.add(task) - self.sem_prefetch.release() - message = await queue.get() - if message is QUEUE_DONE: - # asyncio.wait will throw an error if there is nothing to wait for - if tasks: - logger.info( - f"Waiting for {len(tasks)} running tasks to complete...", - ) - await asyncio.wait(tasks, timeout=self.wait_tasks_timeout) - logger.info("No more tasks to wait for. Shutting down.") - break + # Keep a strong reference until completion. Otherwise Python may + # garbage-collect a detached task while it is still running. + task.add_done_callback(self._on_callback_done) - # Custom hooks for OTel and any future instrumentations - for middleware in reversed(self.broker.middlewares): - if hasattr(middleware, "on_prefetch_queue_remove"): - await maybe_awaitable( - middleware.on_prefetch_queue_remove(), # type: ignore - ) + def _on_callback_done(self, task: "asyncio.Task[Any]") -> None: + """Release shared capacity and retrieve callback failures.""" + self._active_tasks.discard(task) + if self.sem is not None: + self.sem.release() + if task.cancelled(): + return + task_exception = task.exception() + if task_exception is not None: + logger.error( + "Receiver callback failed outside task execution handling.", + exc_info=( + type(task_exception), + task_exception, + task_exception.__traceback__, + ), + ) - task = asyncio.create_task( - self.callback(message=message, raise_err=False), + async def _drain_active_tasks(self, *, cancel_immediately: bool) -> None: + """Wait for owned callbacks and cancel them after the graceful boundary.""" + tasks = set(self._active_tasks) + if not tasks: + return + + logger.info("Waiting for %d running tasks to complete...", len(tasks)) + cancellation: asyncio.CancelledError | None = None + if cancel_immediately: + pending = tasks + else: + try: + _, pending = await asyncio.wait( + tasks, + timeout=self.wait_tasks_timeout, ) - tasks.add(task) - - # We want the task to remove itself from the set when it's done. - # - # Because if we won't save it anywhere, - # python's GC can silently cancel task - # and this behaviour considered to be a Hisenbug. - # https://textual.textualize.io/blog/2023/02/11/the-heisenbug-lurking-in-your-async-code/ - task.add_done_callback(task_cb) - - except asyncio.CancelledError: - break - logger.info("The runner is stopped.") + except asyncio.CancelledError as exc: + cancellation = exc + pending = {task for task in tasks if not task.done()} + + if pending: + logger.warning( + "Cancelling %d tasks that exceeded the graceful wait boundary.", + len(pending), + ) + cleanup_cancellation = await self._cancel_callback_tasks(pending) + if cancellation is None: + cancellation = cleanup_cancellation + logger.info("No more tasks to wait for. Shutting down.") + if cancellation is not None: + raise cancellation + + @staticmethod + async def _cancel_callback_tasks( + tasks: set[asyncio.Task[Any]], + ) -> asyncio.CancelledError | None: + """Cancel callback tasks and await cleanup despite one outer cancellation.""" + for task in tasks: + task.cancel() + try: + await asyncio.gather(*tasks, return_exceptions=True) + except asyncio.CancelledError as exc: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + return exc + return None + + def _record_listen_error(self, error: BaseException) -> None: + """Preserve the first transport error while allowing iterator cleanup.""" + if self._listen_error is None: + self._listen_error = error + return + logger.error( + "Additional error while closing broker listener.", + exc_info=(type(error), error, error.__traceback__), + ) def _prepare_task(self, name: str, handler: Callable[..., Any]) -> None: """ diff --git a/taskiq/receiver/runtime.py b/taskiq/receiver/runtime.py new file mode 100644 index 00000000..9d4645e3 --- /dev/null +++ b/taskiq/receiver/runtime.py @@ -0,0 +1,567 @@ +"""Worker-process orchestration for one or several broker listeners.""" + +from __future__ import annotations + +import asyncio +import random +from collections.abc import Mapping, Sequence +from concurrent.futures import Executor +from dataclasses import dataclass, field +from logging import getLogger +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, cast + +from taskiq.abc.broker import AsyncBroker +from taskiq.acks import AcknowledgeType + +if TYPE_CHECKING: # pragma: no cover + from taskiq.receiver.receiver import Receiver + +logger = getLogger("taskiq.worker.runtime") + + +async def _call_broker_shutdown(broker: AsyncBroker) -> object | None: + """Call shutdown while tolerating legacy diagnostic return values.""" + # Third-party brokers historically returned diagnostics even though the + # public AsyncBroker contract is annotated as returning None. + return await broker.shutdown() # type: ignore[func-returns-value] + + +def _consume_late_shutdown_result( + shutdown_task: asyncio.Task[object | None], +) -> None: + """Retrieve the outcome of a shutdown task detached after its deadline.""" + try: + shutdown_task.result() + except asyncio.CancelledError: + pass + except BaseException: + logger.exception( + "Broker shutdown task %s failed after its cleanup deadline.", + shutdown_task.get_name(), + ) + + +async def _shutdown_broker( + broker: AsyncBroker, + timeout: float, +) -> asyncio.CancelledError | None: + """Shut down one broker without preventing later resource cleanup.""" + logger.info("Shutting down broker %s.", broker.broker_name) + shutdown_task = asyncio.create_task( + _call_broker_shutdown(broker), + name=f"taskiq-shutdown-{broker.broker_name}", + ) + try: + done, _ = await asyncio.wait( + (shutdown_task,), + timeout=timeout, + ) + except asyncio.CancelledError as exc: + shutdown_task.cancel() + shutdown_task.add_done_callback(_consume_late_shutdown_result) + logger.warning( + "Broker %s shutdown was interrupted by cancellation.", + broker.broker_name, + ) + return exc + + if shutdown_task not in done: + shutdown_task.cancel() + shutdown_task.add_done_callback(_consume_late_shutdown_result) + logger.warning( + "Broker %s did not shut down in %s seconds.", + broker.broker_name, + timeout, + ) + return None + + try: + result = shutdown_task.result() + except asyncio.CancelledError as exc: + logger.warning( + "Broker %s failed during shutdown: %s", + broker.broker_name, + type(exc).__name__, + exc_info=True, + ) + return exc + except Exception as exc: + logger.warning( + "Broker %s failed during shutdown: %s", + broker.broker_name, + type(exc).__name__, + exc_info=True, + ) + return None + + if result is not None: + logger.info( + "Broker %s returned a diagnostic value during shutdown.", + broker.broker_name, + ) + return None + + +class WorkerRuntimeConfigurationError(ValueError): + """Raised when a multi-broker worker configuration is inconsistent.""" + + +@dataclass(frozen=True, slots=True) +class BrokerSelection: + """Normalized listener selection resolved from the CLI import target.""" + + listeners: tuple[AsyncBroker, ...] + explicit_group: bool + + def mark_listener_brokers(self) -> None: + """Set worker mode before task modules and adapter resources are loaded.""" + if not self.explicit_group: + self.listeners[0].is_worker_process = True + return + + listener_ids = {id(broker) for broker in self.listeners} + for broker in self.listeners[0].router.brokers.values(): + broker.is_worker_process = id(broker) in listener_ids + + def managed_brokers(self) -> tuple[AsyncBroker, ...]: + """Return listener brokers plus outbound-only lifecycle participants.""" + if not self.explicit_group: + return self.listeners + + listener_ids = {id(broker) for broker in self.listeners} + outbound_brokers = tuple( + broker + for broker in self.listeners[0].router.brokers.values() + if id(broker) not in listener_ids + ) + for broker in outbound_brokers: + broker.is_worker_process = False + return (*self.listeners, *outbound_brokers) + + +def normalize_broker_source(source: object) -> BrokerSelection: + """Normalize one broker or an explicit listener sequence.""" + listeners, explicit_group = _coerce_listener_tuple(source) + _validate_listener_tuple(listeners) + return BrokerSelection(listeners=listeners, explicit_group=explicit_group) + + +def _coerce_listener_tuple( + source: object, +) -> tuple[tuple[AsyncBroker, ...], bool]: + """Coerce a supported import target into a listener tuple.""" + if isinstance(source, AsyncBroker): + return (source,), False + if isinstance(source, Sequence) and not isinstance( + source, + (str, bytes, bytearray), + ): + return cast(tuple[AsyncBroker, ...], tuple(source)), True + raise WorkerRuntimeConfigurationError( + "Worker target must be an AsyncBroker or a sequence of AsyncBroker " + "instances.", + ) + + +def _validate_listener_tuple(listeners: tuple[AsyncBroker, ...]) -> None: + """Validate listener identity and shared Router ownership.""" + if not listeners: + raise WorkerRuntimeConfigurationError("Worker broker sequence cannot be empty.") + + for index, broker in enumerate(listeners): + if not isinstance(broker, AsyncBroker): + raise WorkerRuntimeConfigurationError( + f"Worker broker at index {index} is not an AsyncBroker instance.", + ) + + if len({id(broker) for broker in listeners}) != len(listeners): + raise WorkerRuntimeConfigurationError( + "Worker broker sequence contains the same broker more than once.", + ) + + router = listeners[0].router + if any(broker.router is not router for broker in listeners[1:]): + raise WorkerRuntimeConfigurationError( + "Multi-broker listeners must use the same router instance.", + ) + + broker_names: set[str] = set() + for broker in listeners: + if broker.broker_name in broker_names: + raise WorkerRuntimeConfigurationError( + f"Worker broker name {broker.broker_name!r} is not unique.", + ) + broker_names.add(broker.broker_name) + if router.brokers.get(broker.broker_name) is not broker: + raise WorkerRuntimeConfigurationError( + f"Worker broker {broker.broker_name!r} is not owned by its router.", + ) + + +@dataclass(frozen=True, slots=True) +class WorkerRuntimeOptions: + """Receiver and lifecycle settings shared by a worker process.""" + + validate_params: bool + max_async_tasks: int | None + max_async_tasks_jitter: int + max_prefetch: int + propagate_exceptions: bool + ack_type: AcknowledgeType | None + max_tasks_to_execute: int | None + wait_tasks_timeout: float | None + shutdown_timeout: float + receiver_kwargs: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.max_prefetch < 0: + raise WorkerRuntimeConfigurationError("max_prefetch cannot be negative.") + if self.max_async_tasks_jitter < 0: + raise WorkerRuntimeConfigurationError( + "max_async_tasks_jitter cannot be negative.", + ) + if self.shutdown_timeout < 0: + raise WorkerRuntimeConfigurationError( + "shutdown_timeout cannot be negative.", + ) + if self.wait_tasks_timeout is not None and self.wait_tasks_timeout < 0: + raise WorkerRuntimeConfigurationError( + "wait_tasks_timeout cannot be negative.", + ) + object.__setattr__( + self, + "receiver_kwargs", + MappingProxyType(dict(self.receiver_kwargs)), + ) + + +class ReceiverRuntimeState: + """Process-wide budgets shared by child Receiver instances.""" + + def __init__( + self, + *, + finish_event: asyncio.Event, + max_async_tasks: int | None, + max_async_tasks_jitter: int, + max_prefetch: int, + max_tasks_to_execute: int | None, + ) -> None: + actual_limit = max_async_tasks + if actual_limit is not None and actual_limit > 0: + if max_async_tasks_jitter > 0: + actual_limit += random.randint( # noqa: S311 + 0, + max_async_tasks_jitter, + ) + self.execution_semaphore: asyncio.Semaphore | None = asyncio.Semaphore( + actual_limit, + ) + else: + self.execution_semaphore = None + + self.prefetch_semaphore = asyncio.Semaphore(max_prefetch + 1) + self.finish_event = finish_event + self.max_tasks_to_execute = max_tasks_to_execute + self.fetched_tasks = 0 + + def record_fetched_task(self) -> None: + """Count one fetched delivery and stop the process at its shared limit.""" + self.fetched_tasks += 1 + if ( + self.max_tasks_to_execute + and self.fetched_tasks >= self.max_tasks_to_execute + ): + self.finish_event.set() + + +class WorkerRuntime: + """Own broker lifecycle and concurrent Receiver listeners in one process.""" + + def __init__( + self, + *, + selection: BrokerSelection, + receiver_type: type[Receiver], + executor: Executor, + options: WorkerRuntimeOptions, + ) -> None: + selection.mark_listener_brokers() + self.selection = selection + self.managed_brokers = selection.managed_brokers() + self.receiver_type = receiver_type + self.executor = executor + self.options = options + self._validate_runtime_configuration() + + async def run(self, finish_event: asyncio.Event) -> None: + """Start resources, run every listener and clean up deterministically.""" + primary_error: BaseException | None = None + started_brokers: list[AsyncBroker] = [] + + try: + runtime_state = ReceiverRuntimeState( + finish_event=finish_event, + max_async_tasks=self.options.max_async_tasks, + max_async_tasks_jitter=self.options.max_async_tasks_jitter, + max_prefetch=self.options.max_prefetch, + max_tasks_to_execute=self.options.max_tasks_to_execute, + ) + receivers = self._build_receivers(runtime_state) + + for broker in self.managed_brokers: + started_brokers.append(broker) + logger.info("Starting broker %s.", broker.broker_name) + await broker.startup() + + await self._run_listeners(receivers, finish_event) + except BaseException as exc: + primary_error = exc + finally: + finish_event.set() + cleanup_cancellation = await self._shutdown_brokers(started_brokers) + if primary_error is None: + primary_error = cleanup_cancellation + + if primary_error is not None: + raise primary_error + + def _build_receivers( + self, + runtime_state: ReceiverRuntimeState, + ) -> tuple[Receiver, ...]: + receivers = [] + for broker in self.selection.listeners: + receiver = self.receiver_type( + broker=broker, + executor=self.executor, + run_startup=False, + validate_params=self.options.validate_params, + max_async_tasks=self.options.max_async_tasks, + max_async_tasks_jitter=self.options.max_async_tasks_jitter, + max_prefetch=self.options.max_prefetch, + propagate_exceptions=self.options.propagate_exceptions, + ack_type=self.options.ack_type, + max_tasks_to_execute=self.options.max_tasks_to_execute, + wait_tasks_timeout=self.options.wait_tasks_timeout, + **self.options.receiver_kwargs, + ) + attached_state = receiver.attach_runtime_state(runtime_state) + if ( + attached_state is not runtime_state + or not receiver.is_runtime_state_attached(runtime_state) + ): + raise WorkerRuntimeConfigurationError( + "Custom Receiver.attach_runtime_state() must preserve the " + "process-wide worker state.", + ) + receivers.append(receiver) + return tuple(receivers) + + async def _run_listeners( + self, + receivers: tuple[Receiver, ...], + finish_event: asyncio.Event, + ) -> None: + listener_tasks = tuple( + asyncio.create_task( + receiver.listen(finish_event), + name=f"taskiq-listener-{receiver.broker.broker_name}", + ) + for receiver in receivers + ) + finish_waiter = asyncio.create_task( + finish_event.wait(), + name="taskiq-worker-finish", + ) + primary_error: BaseException | None = None + + try: + first_done, _ = await asyncio.wait( + (*listener_tasks, finish_waiter), + return_when=asyncio.FIRST_COMPLETED, + ) + completed_listeners = { + listener_task + for listener_task in listener_tasks + if listener_task in first_done + } + primary_error = self._find_listener_error( + listener_tasks, + completed_listeners, + ) + except BaseException as exc: + primary_error = exc + finally: + finish_event.set() + if not finish_waiter.done(): + finish_waiter.cancel() + await asyncio.gather(finish_waiter, return_exceptions=True) + results, hard_cancelled = await self._settle_listener_tasks( + listener_tasks, + receivers, + ) + + if primary_error is None: + for listener_task, result in zip(listener_tasks, results, strict=True): + if listener_task in hard_cancelled and isinstance( + result, + asyncio.CancelledError, + ): + continue + if isinstance(result, BaseException): + primary_error = result + break + + if primary_error is not None: + raise primary_error + + @classmethod + def _find_listener_error( + cls, + listener_tasks: tuple[asyncio.Task[None], ...], + completed_tasks: set[asyncio.Task[None]], + ) -> BaseException | None: + """Select the first failed listener in deterministic broker order.""" + for listener_task in listener_tasks: + if listener_task not in completed_tasks: + continue + result = cls._get_listener_result(listener_task) + if isinstance(result, BaseException): + return result + return None + + async def _settle_listener_tasks( + self, + listener_tasks: tuple[asyncio.Task[None], ...], + receivers: tuple[Receiver, ...], + ) -> tuple[list[None | BaseException], set[asyncio.Task[None]]]: + graceful_timeout = None + if self.options.wait_tasks_timeout is not None: + graceful_timeout = ( + self.options.wait_tasks_timeout + self.options.shutdown_timeout + ) + _, pending = await asyncio.wait( + listener_tasks, + timeout=graceful_timeout, + ) + hard_cancelled = set(pending) + for listener_task, receiver in zip(listener_tasks, receivers, strict=True): + if listener_task not in hard_cancelled: + continue + logger.warning( + "Listener for broker %s did not stop in %s seconds; cancelling it.", + receiver.broker.broker_name, + graceful_timeout, + ) + listener_task.cancel() + + if hard_cancelled: + _, stubborn_tasks = await asyncio.wait( + hard_cancelled, + timeout=self.options.shutdown_timeout, + ) + for listener_task, receiver in zip( + listener_tasks, + receivers, + strict=True, + ): + if listener_task not in stubborn_tasks: + continue + logger.error( + "Listener for broker %s ignored cancellation; continuing " + "broker cleanup.", + receiver.broker.broker_name, + ) + listener_task.add_done_callback(self._consume_listener_result) + + results = [ + ( + self._get_listener_result(listener_task) + if listener_task.done() + else asyncio.CancelledError() + ) + for listener_task in listener_tasks + ] + return results, hard_cancelled + + @staticmethod + def _get_listener_result( + listener_task: asyncio.Task[None], + ) -> None | BaseException: + """Return one listener outcome without re-raising it.""" + try: + return listener_task.result() + except BaseException as exc: + return exc + + @staticmethod + def _consume_listener_result(listener_task: asyncio.Task[None]) -> None: + """Retrieve a late outcome from a cancellation-resistant listener.""" + try: + listener_task.result() + except asyncio.CancelledError: + pass + except BaseException: + logger.exception( + "Cancellation-resistant listener failed after runtime cleanup.", + ) + + async def _shutdown_brokers( + self, + started_brokers: list[AsyncBroker], + ) -> asyncio.CancelledError | None: + cleanup_cancellation: asyncio.CancelledError | None = None + for broker in reversed(started_brokers): + broker_cancellation = await _shutdown_broker( + broker, + self.options.shutdown_timeout, + ) + if cleanup_cancellation is None: + cleanup_cancellation = broker_cancellation + return cleanup_cancellation + + def _validate_runtime_configuration(self) -> None: + self._validate_unshared_resources() + self._validate_task_lookup() + + def _validate_unshared_resources(self) -> None: + middleware_owners: dict[int, str] = {} + backend_owners: dict[int, str] = {} + + for broker in self.managed_brokers: + for middleware in broker.middlewares: + previous_owner = middleware_owners.setdefault( + id(middleware), + broker.broker_name, + ) + if previous_owner != broker.broker_name: + raise WorkerRuntimeConfigurationError( + "One middleware instance cannot be lifecycle-owned by " + f"brokers {previous_owner!r} and {broker.broker_name!r}.", + ) + + previous_owner = backend_owners.setdefault( + id(broker.result_backend), + broker.broker_name, + ) + if previous_owner != broker.broker_name: + raise WorkerRuntimeConfigurationError( + "One result backend instance cannot be lifecycle-owned by " + f"brokers {previous_owner!r} and {broker.broker_name!r}.", + ) + + def _validate_task_lookup(self) -> None: + task_owners: dict[str, tuple[object, str]] = {} + for broker in self.selection.listeners: + for task_name, task in broker.get_all_tasks().items(): + existing = task_owners.setdefault( + task_name, + (task, broker.broker_name), + ) + if existing[0] is not task: + raise WorkerRuntimeConfigurationError( + f"Task {task_name!r} resolves differently for brokers " + f"{existing[1]!r} and {broker.broker_name!r}.", + ) diff --git a/taskiq/router.py b/taskiq/router.py index 8b526066..49cc31f1 100644 --- a/taskiq/router.py +++ b/taskiq/router.py @@ -1,7 +1,8 @@ from __future__ import annotations import warnings -from collections.abc import Callable +from collections.abc import Callable, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, overload from taskiq.flow import FlowProtocol @@ -44,9 +45,9 @@ def __init__(self) -> None: self._dispatcher = RouterDispatcher(self._routes) @property - def brokers(self) -> dict[str, AsyncBroker]: - """Return mutable broker registry for compatibility.""" - return self._brokers.brokers + def brokers(self) -> Mapping[str, AsyncBroker]: + """Return an immutable snapshot of registered brokers.""" + return MappingProxyType(self._brokers.get_all()) @property def default_broker(self) -> AsyncBroker | None: @@ -55,24 +56,24 @@ def default_broker(self) -> AsyncBroker | None: @default_broker.setter def default_broker(self, broker: AsyncBroker | None) -> None: - if broker is not None: - self._brokers.resolve(broker) - self._brokers.default_broker = broker + self._brokers.set_default(broker) @property - def task_registry(self) -> dict[str, AsyncTaskiqDecoratedTask[Any, Any]]: - """Return mutable task registry for compatibility.""" - return self._tasks.tasks + def task_registry( + self, + ) -> Mapping[str, AsyncTaskiqDecoratedTask[Any, Any]]: + """Return an immutable snapshot of registered tasks.""" + return MappingProxyType(self._tasks.get_all()) @property - def routes(self) -> dict[str, TaskiqRoute]: - """Return mutable route registry for compatibility.""" - return self._routes.routes + def routes(self) -> Mapping[str, TaskiqRoute]: + """Return an immutable snapshot of resolved task routes.""" + return MappingProxyType(self._routes.get_all()) @property - def subscriptions(self) -> list[TaskiqSubscription]: - """Return mutable subscription registry for compatibility.""" - return self._subscriptions.subscriptions + def subscriptions(self) -> tuple[TaskiqSubscription, ...]: + """Return an immutable snapshot of inbound subscriptions.""" + return self._subscriptions.get() @property def default_broker_name(self) -> str | None: @@ -123,17 +124,17 @@ def register_task( """Register a bound task or bind a task definition to a broker.""" if isinstance(task, TaskDefinition): target_broker = self._brokers.resolve(broker) + self._tasks.ensure_name_available(task.task_name) registered_task = target_broker.bind_task_definition( task, register=False, ) - self._register_bound_task( + target_broker.store_registered_task(registered_task) + return self._register_bound_task( registered_task, broker=target_broker, flow=flow, ) - target_broker.store_registered_task(registered_task) - return registered_task return self._register_bound_task(task, broker=broker, flow=flow) @@ -217,6 +218,20 @@ def resolve_route( """Resolve outbound route for a task invocation.""" return self._routes.resolve_route(task, broker=broker, flow=flow) + def has_route( + self, + task: str | AsyncTaskiqDecoratedTask[Any, Any], + ) -> bool: + """Return whether a task has an explicit outbound route policy.""" + return self._routes.has_route(task) + + def remove_route( + self, + task: str | AsyncTaskiqDecoratedTask[Any, Any], + ) -> TaskiqRoute | None: + """Remove and return a task's resolved outbound route, if present.""" + return self._routes.remove_route(task) + def subscribe( self, broker: AsyncBroker, @@ -231,6 +246,14 @@ def subscribe( task_names, ) + def unsubscribe( + self, + broker: AsyncBroker, + flow: FlowProtocol, + ) -> TaskiqSubscription | None: + """Remove and return one inbound broker/flow subscription.""" + return self._subscriptions.unsubscribe(broker, flow) + def get_subscriptions( self, broker: AsyncBroker | None = None, @@ -286,11 +309,11 @@ def _register_bound_task( if route_broker is None: route_broker = getattr(task, "broker", None) - route = None - if route_broker is not None or flow is not None: - route = self._routes.build_route(broker=route_broker, flow=flow) - self._tasks.register(task) - if route is not None: - self._routes.set_resolved_route(task.task_name, route) + if route_broker is not None or flow is not None: + self._routes.set_route( + task.task_name, + broker=route_broker, + flow=flow, + ) return task diff --git a/taskiq/routing/dispatcher.py b/taskiq/routing/dispatcher.py index d5b1ad43..60791c61 100644 --- a/taskiq/routing/dispatcher.py +++ b/taskiq/routing/dispatcher.py @@ -1,15 +1,13 @@ from __future__ import annotations -from typing import TYPE_CHECKING, TypeVar, cast +from typing import TYPE_CHECKING, TypeVar -from taskiq.abc.middleware import TaskiqMiddleware -from taskiq.abc.result_backend import AsyncResultBackend from taskiq.flow import FlowProtocol -from taskiq.message import TaskiqMessage +from taskiq.message import BrokerMessage, TaskiqMessage from taskiq.routing.models import TaskiqRoute from taskiq.routing.routes import RouteRegistry +from taskiq.sending import send_task_message from taskiq.task import AsyncTaskiqTask -from taskiq.utils import maybe_awaitable if TYPE_CHECKING: # pragma: no cover from taskiq.abc.broker import AsyncBroker @@ -23,7 +21,7 @@ class RouterDispatcher: """Send task messages through resolved router routes.""" def __init__(self, routes: RouteRegistry) -> None: - self.routes = routes + self._routes = routes async def kiq( self, @@ -35,7 +33,7 @@ async def kiq( return_type: type[_ReturnType] | None = None, ) -> AsyncTaskiqTask[_ReturnType]: """Send message through the resolved broker and flow.""" - target_route = self.routes.resolve_explicit_or_default_route( + target_route = self._routes.resolve_explicit_or_default_route( message.task_name, route=route, broker=broker, @@ -43,22 +41,13 @@ async def kiq( ) target_broker = target_route.broker - for middleware in target_broker.middlewares: - if middleware.__class__.pre_send != TaskiqMiddleware.pre_send: - message = await maybe_awaitable(middleware.pre_send(message)) - broker_message = target_broker.formatter.dumps(message) - await target_broker.kick_to_flow(broker_message, target_route.flow) + async def send(broker_message: BrokerMessage) -> None: + await target_broker.kick_to_flow(broker_message, target_route.flow) - for middleware in reversed(target_broker.middlewares): - if middleware.__class__.post_send != TaskiqMiddleware.post_send: - await maybe_awaitable(middleware.post_send(message)) - - return AsyncTaskiqTask( - task_id=message.task_id, - result_backend=cast( - AsyncResultBackend[_ReturnType], - target_broker.result_backend, - ), + return await send_task_message( + target_broker, + message, + send=send, return_type=return_type, ) @@ -71,7 +60,7 @@ async def requeue( flow: FlowProtocol | None = None, ) -> None: """Send an existing message again through the resolved route.""" - target_route = self.routes.resolve_explicit_or_default_route( + target_route = self._routes.resolve_explicit_or_default_route( message.task_name, route=route, broker=broker, diff --git a/taskiq/routing/registries.py b/taskiq/routing/registries.py index 9226cdaa..de3084bf 100644 --- a/taskiq/routing/registries.py +++ b/taskiq/routing/registries.py @@ -15,43 +15,58 @@ class BrokerRegistry: """Registry of brokers attached to one router.""" def __init__(self) -> None: - self.brokers: dict[str, AsyncBroker] = {} - self.default_broker: AsyncBroker | None = None + self._brokers: dict[str, AsyncBroker] = {} + self._default_broker: AsyncBroker | None = None + + @property + def default_broker(self) -> AsyncBroker | None: + """Return the configured default broker.""" + return self._default_broker @property def default_broker_name(self) -> str | None: """Return default broker name for compatibility and diagnostics.""" - if self.default_broker is None: + if self._default_broker is None: return None - return self.default_broker.broker_name + return self._default_broker.broker_name + + def get_all(self) -> dict[str, AsyncBroker]: + """Return a snapshot of registered brokers.""" + return dict(self._brokers) def register(self, broker: AsyncBroker, name: str | None = None) -> str: """Register broker as a transport in this registry.""" broker_name = name or broker.__class__.__name__ - registered = self.brokers.get(broker_name) + registered = self._brokers.get(broker_name) if registered is not None and registered is not broker: raise ValueError( f"Broker name {broker_name!r} is already registered. " "Please provide an explicit unique broker_name.", ) - for registered_name, registered_broker in self.brokers.items(): + for registered_name, registered_broker in self._brokers.items(): if registered_broker is broker and registered_name != broker_name: raise ValueError( f"Broker is already registered as {registered_name!r}.", ) - self.brokers[broker_name] = broker - if self.default_broker is None: - self.default_broker = broker + self._brokers[broker_name] = broker + if self._default_broker is None: + self._default_broker = broker return broker_name def get(self, name: str) -> AsyncBroker: """Return a broker by registered name.""" try: - return self.brokers[name] + return self._brokers[name] except KeyError as exc: raise ValueError(f"Unknown broker {name!r}.") from exc + def set_default(self, broker: AsyncBroker | None) -> None: + """Set a validated default broker or clear the default.""" + if broker is not None: + broker = self.resolve(broker) + self._default_broker = broker + def resolve(self, broker: AsyncBroker | None) -> AsyncBroker: """Resolve an explicit broker or return the default broker.""" if isinstance(broker, str): @@ -61,14 +76,14 @@ def resolve(self, broker: AsyncBroker | None) -> AsyncBroker: ) if broker is not None: - for registered_broker in self.brokers.values(): + for registered_broker in self._brokers.values(): if registered_broker is broker: return registered_broker raise ValueError("Broker is not registered in this router.") - if self.default_broker is None: + if self._default_broker is None: raise ValueError("Router doesn't have registered brokers.") - return self.default_broker + return self._default_broker def default_flow(self, broker: AsyncBroker) -> FlowProtocol | None: """Return broker-owned default flow.""" @@ -79,28 +94,43 @@ class TaskRegistry: """Registry of tasks attached to one router.""" def __init__(self) -> None: - self.tasks: dict[str, AsyncTaskiqDecoratedTask[Any, Any]] = {} + self._tasks: dict[str, AsyncTaskiqDecoratedTask[Any, Any]] = {} def register( self, task: AsyncTaskiqDecoratedTask[Any, Any], ) -> None: """Register a decorated task and reject name conflicts.""" - existing_task = self.tasks.get(task.task_name) + self.ensure_available(task) + self._tasks[task.task_name] = task + + def ensure_available( + self, + task: AsyncTaskiqDecoratedTask[Any, Any], + ) -> None: + """Validate that a task can be registered without mutating state.""" + existing_task = self._tasks.get(task.task_name) if existing_task is not None and existing_task is not task: - raise ValueError( - f"Task name {task.task_name!r} is already registered " - "in this router.", - ) - self.tasks[task.task_name] = task + self._raise_name_conflict(task.task_name) + + def ensure_name_available(self, task_name: str) -> None: + """Validate that an unbound task name is available.""" + if task_name in self._tasks: + self._raise_name_conflict(task_name) + + @staticmethod + def _raise_name_conflict(task_name: str) -> None: + raise ValueError( + f"Task name {task_name!r} is already registered in this router.", + ) def find( self, task_name: str, ) -> AsyncTaskiqDecoratedTask[Any, Any] | None: """Find a task by name.""" - return self.tasks.get(task_name) + return self._tasks.get(task_name) def get_all(self) -> dict[str, AsyncTaskiqDecoratedTask[Any, Any]]: """Return all tasks registered in this registry.""" - return dict(self.tasks) + return dict(self._tasks) diff --git a/taskiq/routing/routes.py b/taskiq/routing/routes.py index 96a866f8..f346f566 100644 --- a/taskiq/routing/routes.py +++ b/taskiq/routing/routes.py @@ -1,6 +1,6 @@ from __future__ import annotations -from dataclasses import replace +from dataclasses import dataclass from typing import TYPE_CHECKING, Any from taskiq.flow import FlowProtocol @@ -15,12 +15,44 @@ __all__ = ("RouteRegistry",) +@dataclass(frozen=True, slots=True) +class _TaskRoutePolicy: + """Stored outbound policy before broker default-flow resolution.""" + + broker: AsyncBroker + flow: FlowProtocol | None + + class RouteRegistry: """Outbound route policy for task invocations.""" def __init__(self, brokers: BrokerRegistry) -> None: - self.brokers = brokers - self.routes: dict[str, TaskiqRoute] = {} + self._brokers = brokers + self._routes: dict[str, _TaskRoutePolicy] = {} + + def get_all(self) -> dict[str, TaskiqRoute]: + """Return resolved route snapshots keyed by task name.""" + return { + task_name: self._resolve_policy(policy) + for task_name, policy in self._routes.items() + } + + def has_route( + self, + task: str | AsyncTaskiqDecoratedTask[Any, Any], + ) -> bool: + """Return whether a task has an explicit route policy.""" + return resolve_task_name(task) in self._routes + + def remove_route( + self, + task: str | AsyncTaskiqDecoratedTask[Any, Any], + ) -> TaskiqRoute | None: + """Remove and return a task's resolved route policy, if present.""" + policy = self._routes.pop(resolve_task_name(task), None) + if policy is None: + return None + return self._resolve_policy(policy) def build_route( self, @@ -28,10 +60,10 @@ def build_route( flow: FlowProtocol | None = None, ) -> TaskiqRoute: """Build a route without mutating route state.""" - target_broker = self.brokers.resolve(broker) + target_broker = self._brokers.resolve(broker) route_flow = flow if route_flow is None: - route_flow = self.brokers.default_flow(target_broker) + route_flow = self._brokers.default_flow(target_broker) return TaskiqRoute(broker=target_broker, flow=route_flow) def set_route( @@ -40,16 +72,17 @@ def set_route( broker: AsyncBroker | None = None, flow: FlowProtocol | None = None, ) -> TaskiqRoute: - """Set default outbound route for a task.""" - task_name = resolve_task_name(task) - route = self.build_route(broker=broker, flow=flow) - self.routes[task_name] = route - return route + """ + Set default outbound route for a task. - def set_resolved_route(self, task_name: str, route: TaskiqRoute) -> TaskiqRoute: - """Store an already resolved route for a task.""" - self.routes[task_name] = route - return route + A missing flow is stored as a dynamic broker-default policy. The + returned route is resolved against the broker's current default flow. + """ + task_name = resolve_task_name(task) + target_broker = self._brokers.resolve(broker) + policy = _TaskRoutePolicy(broker=target_broker, flow=flow) + self._routes[task_name] = policy + return self._resolve_policy(policy) def resolve_route( self, @@ -60,7 +93,7 @@ def resolve_route( """Resolve outbound route for a task invocation.""" task_name = resolve_task_name(task) if broker is not None: - target_broker = self.brokers.resolve(broker) + target_broker = self._brokers.resolve(broker) route_flow = self._resolve_flow_for_broker_override( task_name, target_broker, @@ -68,11 +101,11 @@ def resolve_route( ) return TaskiqRoute(broker=target_broker, flow=route_flow) - route = self.routes.get(task_name) - if route is not None: - if flow is None: - return route - return replace(route, flow=flow) + policy = self._routes.get(task_name) + if policy is not None: + if flow is not None: + return TaskiqRoute(broker=policy.broker, flow=flow) + return self._resolve_policy(policy) return self.build_route(flow=flow) @@ -88,7 +121,7 @@ def resolve_explicit_or_default_route( if route is not None: if broker is not None or flow is not None: raise ValueError("Pass either route or broker/flow overrides.") - self.brokers.resolve(route.broker) + self._brokers.resolve(route.broker) return route return self.resolve_route(task, broker=broker, flow=flow) @@ -100,7 +133,17 @@ def _resolve_flow_for_broker_override( ) -> FlowProtocol | None: if flow is not None: return flow - registered_route = self.routes.get(task_name) - if registered_route is not None and registered_route.broker is broker: - return registered_route.flow - return self.brokers.default_flow(broker) + registered_policy = self._routes.get(task_name) + if ( + registered_policy is not None + and registered_policy.broker is broker + and registered_policy.flow is not None + ): + return registered_policy.flow + return self._brokers.default_flow(broker) + + def _resolve_policy(self, policy: _TaskRoutePolicy) -> TaskiqRoute: + flow = policy.flow + if flow is None: + flow = self._brokers.default_flow(policy.broker) + return TaskiqRoute(broker=policy.broker, flow=flow) diff --git a/taskiq/routing/subscriptions.py b/taskiq/routing/subscriptions.py index 7f38ba2d..dadd9562 100644 --- a/taskiq/routing/subscriptions.py +++ b/taskiq/routing/subscriptions.py @@ -23,8 +23,8 @@ class SubscriptionPlan: """ def __init__(self, brokers: BrokerRegistry) -> None: - self.brokers = brokers - self.subscriptions: list[TaskiqSubscription] = [] + self._brokers = brokers + self._subscriptions: list[TaskiqSubscription] = [] def subscribe( self, @@ -33,9 +33,9 @@ def subscribe( task_names: Iterable[str], ) -> TaskiqSubscription: """Register an inbound flow subscription for a broker.""" - target_broker = self.brokers.resolve(broker) + target_broker = self._brokers.resolve(broker) resolved_task_names = frozenset(task_names) - default_flow = self.brokers.default_flow(target_broker) + default_flow = self._brokers.default_flow(target_broker) if default_flow is not None: self._ensure_compatible_flow( broker=target_broker, @@ -43,7 +43,7 @@ def subscribe( new_flow=flow, ) - for subscription_index, subscription in enumerate(self.subscriptions): + for subscription_index, subscription in enumerate(self._subscriptions): if subscription.broker is target_broker and self._is_same_flow( subscription.flow, flow, @@ -57,7 +57,7 @@ def subscribe( subscription, task_names=subscription.task_names | resolved_task_names, ) - self.subscriptions[subscription_index] = updated + self._subscriptions[subscription_index] = updated return updated subscription = TaskiqSubscription( @@ -65,28 +65,48 @@ def subscribe( flow=flow, task_names=resolved_task_names, ) - self.subscriptions.append(subscription) + self._subscriptions.append(subscription) return subscription + def unsubscribe( + self, + broker: AsyncBroker, + flow: FlowProtocol, + ) -> TaskiqSubscription | None: + """Remove and return one compatible broker/flow subscription.""" + target_broker = self._brokers.resolve(broker) + for index, subscription in enumerate(self._subscriptions): + if subscription.broker is target_broker and self._is_same_flow( + subscription.flow, + flow, + ): + self._ensure_compatible_flow( + broker=target_broker, + registered_flow=subscription.flow, + new_flow=flow, + ) + return self._subscriptions.pop(index) + return None + def get( self, broker: AsyncBroker | None = None, ) -> tuple[TaskiqSubscription, ...]: """Return registered inbound subscriptions.""" if broker is None: - return tuple(self.subscriptions) - target_broker = self.brokers.resolve(broker) + return tuple(self._subscriptions) + target_broker = self._brokers.resolve(broker) return tuple( subscription - for subscription in self.subscriptions + for subscription in self._subscriptions if subscription.broker is target_broker ) def get_broker_flows(self, broker: AsyncBroker) -> tuple[FlowProtocol, ...]: """Return flows a broker should subscribe to.""" - target_broker = self.brokers.resolve(broker) + target_broker = self._brokers.resolve(broker) flows = [subscription.flow for subscription in self.get(target_broker)] - default_flow = self.brokers.default_flow(target_broker) + default_flow = self._brokers.default_flow(target_broker) if default_flow is not None and not self._contains_flow_identity( flows, default_flow, diff --git a/taskiq/scheduler/scheduler.py b/taskiq/scheduler/scheduler.py index b0a1e00d..c65e74a5 100644 --- a/taskiq/scheduler/scheduler.py +++ b/taskiq/scheduler/scheduler.py @@ -74,7 +74,7 @@ def _apply_scheduled_route( send through the scheduler broker instead of the router default broker. """ router = getattr(self.broker, "router", None) - if isinstance(router, TaskiqRouter) and task_name in router.routes: + if isinstance(router, TaskiqRouter) and router.has_route(task_name): kicker.with_route(router.resolve_route(task_name)) return diff --git a/taskiq/sending.py b/taskiq/sending.py new file mode 100644 index 00000000..7349d3aa --- /dev/null +++ b/taskiq/sending.py @@ -0,0 +1,57 @@ +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING, TypeVar, cast + +from taskiq.abc.middleware import TaskiqMiddleware +from taskiq.abc.result_backend import AsyncResultBackend +from taskiq.exceptions import SendTaskError +from taskiq.message import BrokerMessage, TaskiqMessage +from taskiq.task import AsyncTaskiqTask +from taskiq.utils import maybe_awaitable + +if TYPE_CHECKING: # pragma: no cover + from taskiq.abc.broker import AsyncBroker + +_ReturnType = TypeVar("_ReturnType") +_TransportSend = Callable[[BrokerMessage], Awaitable[None]] + + +async def send_task_message( + broker: "AsyncBroker", + message: TaskiqMessage, + *, + send: _TransportSend, + return_type: type[_ReturnType] | None = None, +) -> AsyncTaskiqTask[_ReturnType]: + """ + Run the client-side send lifecycle through one broker. + + Middleware failures keep their original types. Only message formatting and + transport dispatch failures are mapped to `SendTaskError`, matching the + established Taskiq kicker contract. + """ + middlewares = getattr(broker, "middlewares", []) + if not isinstance(middlewares, list): + middlewares = [] + + for middleware in middlewares: + if middleware.__class__.pre_send != TaskiqMiddleware.pre_send: + message = await maybe_awaitable(middleware.pre_send(message)) + + try: + broker_message = broker.formatter.dumps(message) + await send(broker_message) + except Exception as exc: + raise SendTaskError from exc + + for middleware in reversed(middlewares): + if middleware.__class__.post_send != TaskiqMiddleware.post_send: + await maybe_awaitable(middleware.post_send(message)) + + return AsyncTaskiqTask( + task_id=message.task_id, + result_backend=cast( + AsyncResultBackend[_ReturnType], + broker.result_backend, + ), + return_type=return_type, + ) diff --git a/taskiq/task_builder.py b/taskiq/task_builder.py index eeec6ee2..367458f9 100644 --- a/taskiq/task_builder.py +++ b/taskiq/task_builder.py @@ -3,9 +3,11 @@ import inspect from collections.abc import Callable, Coroutine, Mapping from dataclasses import dataclass, field +from functools import WRAPPER_ASSIGNMENTS from types import CoroutineType, MappingProxyType from typing import TYPE_CHECKING, Any, Generic, ParamSpec, TypeVar, overload +from taskiq._task_function import ensure_task_binding_function from taskiq.message import TaskiqMessage, _build_taskiq_message if TYPE_CHECKING: # pragma: no cover @@ -18,8 +20,22 @@ _ReturnType = TypeVar("_ReturnType") +class _CallableMetadata(Generic[_FuncParams, _ReturnType]): + """Type declarations for function metadata populated on task definitions.""" + + __name__: str + __qualname__: str + __module__: str + __doc__: str | None + __annotations__: dict[str, Any] + __wrapped__: Callable[_FuncParams, _ReturnType] + + @dataclass(frozen=True, slots=True) -class TaskDefinition(Generic[_FuncParams, _ReturnType]): +class TaskDefinition( + _CallableMetadata[_FuncParams, _ReturnType], + Generic[_FuncParams, _ReturnType], +): """Unbound task declaration that can be registered later.""" task_name: str @@ -29,8 +45,20 @@ class TaskDefinition(Generic[_FuncParams, _ReturnType]): base_cls: type[AsyncTaskiqDecoratedTask[Any, Any]] | None = None def __post_init__(self) -> None: - """Freeze labels to make task definitions stable declarations.""" + """Freeze declaration state and preserve the wrapped callable contract.""" object.__setattr__(self, "labels", MappingProxyType(dict(self.labels))) + for attribute in WRAPPER_ASSIGNMENTS: + if hasattr(self.original_func, attribute): + object.__setattr__( + self, + attribute, + getattr(self.original_func, attribute), + ) + object.__setattr__(self, "__wrapped__", self.original_func) + ensure_task_binding_function( + self.original_func, + binding_key=self.task_name, + ) def __call__( self, @@ -88,6 +116,16 @@ def message( ) +def _get_task_binding_function( + task: TaskDefinition[_FuncParams, _ReturnType], +) -> Callable[_FuncParams, _ReturnType]: + """Return the framework-owned callable used for broker execution.""" + return ensure_task_binding_function( + task.original_func, + binding_key=task.task_name, + ) + + @overload def task_builder( task_name: Callable[_FuncParams, _ReturnType], diff --git a/tests/brokers/test_inmemory.py b/tests/brokers/test_inmemory.py index 9dc85019..ac8f9a0b 100644 --- a/tests/brokers/test_inmemory.py +++ b/tests/brokers/test_inmemory.py @@ -1,13 +1,66 @@ import asyncio import uuid +from typing import Any import pytest from taskiq import InMemoryBroker +from taskiq.abc.middleware import TaskiqMiddleware +from taskiq.brokers.inmemory_broker import InmemoryResultBackend from taskiq.events import TaskiqEvents from taskiq.state import TaskiqState +class LifecycleError(Exception): + """Marker exception for lifecycle fault tests.""" + + +class RecordingLifecycleMiddleware(TaskiqMiddleware): + """Record middleware resource lifecycle calls.""" + + def __init__( + self, + events: list[str], + shutdown_error: Exception | None = None, + ) -> None: + super().__init__() + self.events = events + self.shutdown_error = shutdown_error + + async def startup(self) -> None: + """Record startup.""" + self.events.append("middleware.startup") + + async def shutdown(self) -> None: + """Record shutdown.""" + self.events.append("middleware.shutdown") + if self.shutdown_error is not None: + raise self.shutdown_error + + +class RecordingResultBackend(InmemoryResultBackend[Any]): + """Record result backend lifecycle calls and optional startup failure.""" + + def __init__( + self, + events: list[str], + startup_error: Exception | None = None, + ) -> None: + super().__init__() + self.events = events + self.startup_error = startup_error + + async def startup(self) -> None: + """Record startup and raise the configured error.""" + self.events.append("backend.startup") + if self.startup_error is not None: + raise self.startup_error + + async def shutdown(self) -> None: + """Record shutdown.""" + self.events.append("backend.shutdown") + + async def test_inmemory_success() -> None: broker = InMemoryBroker() test_val = uuid.uuid4().hex @@ -48,6 +101,88 @@ async def _c_startup(state: TaskiqState) -> None: assert broker.state.from_client == test_value +async def test_lifecycle_runs_events_and_resources_once_in_order() -> None: + events: list[str] = [] + broker = InMemoryBroker() + broker.with_middlewares(RecordingLifecycleMiddleware(events)) + broker.with_result_backend(RecordingResultBackend(events)) + + for event, event_name in ( + (TaskiqEvents.CLIENT_STARTUP, "client.startup"), + (TaskiqEvents.WORKER_STARTUP, "worker.startup"), + (TaskiqEvents.CLIENT_SHUTDOWN, "client.shutdown"), + (TaskiqEvents.WORKER_SHUTDOWN, "worker.shutdown"), + ): + + @broker.on_event(event) + def record_event( + state: TaskiqState, + lifecycle_event: str = event_name, + ) -> None: + events.append(lifecycle_event) + + await broker.startup() + await broker.shutdown() + + assert events == [ + "client.startup", + "worker.startup", + "middleware.startup", + "backend.startup", + "client.shutdown", + "worker.shutdown", + "middleware.shutdown", + "backend.shutdown", + ] + + +async def test_shutdown_cleans_resources_after_startup_failure() -> None: + events: list[str] = [] + startup_error = LifecycleError("backend startup failed") + broker = InMemoryBroker() + broker.with_middlewares(RecordingLifecycleMiddleware(events)) + broker.with_result_backend(RecordingResultBackend(events, startup_error)) + + with pytest.raises(LifecycleError) as exc_info: + await broker.startup() + + assert exc_info.value is startup_error + + await broker.shutdown() + + assert events == [ + "middleware.startup", + "backend.startup", + "middleware.shutdown", + "backend.shutdown", + ] + + +async def test_shutdown_closes_remaining_resources_after_middleware_failure() -> None: + events: list[str] = [] + shutdown_error = LifecycleError("middleware shutdown failed") + broker = InMemoryBroker() + broker.with_middlewares( + RecordingLifecycleMiddleware(events, shutdown_error=shutdown_error), + ) + broker.with_result_backend(RecordingResultBackend(events)) + + await broker.startup() + + with pytest.raises(LifecycleError) as exc_info: + await broker.shutdown() + + assert exc_info.value is shutdown_error + assert events == [ + "middleware.startup", + "backend.startup", + "middleware.shutdown", + "backend.shutdown", + ] + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + broker.executor.submit(int) + + async def test_shutdown() -> None: broker = InMemoryBroker() test_value = uuid.uuid4().hex @@ -114,3 +249,46 @@ async def test_task() -> None: assert slept assert await task.is_ready() assert not broker._running_tasks + + +async def test_shutdown_waits_for_running_tasks_before_resource_cleanup() -> None: + task_started = asyncio.Event() + finish_task = asyncio.Event() + events: list[str] = [] + broker = InMemoryBroker() + broker.with_middlewares(RecordingLifecycleMiddleware(events)) + broker.with_result_backend(RecordingResultBackend(events)) + + @broker.on_event(TaskiqEvents.CLIENT_SHUTDOWN) + def record_client_shutdown(state: TaskiqState) -> None: + events.append("client.shutdown") + + @broker.task + async def running_task() -> None: + events.append("task.started") + task_started.set() + await finish_task.wait() + events.append("task.finished") + + await broker.startup() + await running_task.kiq() + await task_started.wait() + + shutdown_task = asyncio.create_task(broker.shutdown()) + await asyncio.sleep(0) + + assert not shutdown_task.done() + + finish_task.set() + await shutdown_task + + assert events == [ + "middleware.startup", + "backend.startup", + "task.started", + "task.finished", + "client.shutdown", + "middleware.shutdown", + "backend.shutdown", + ] + assert not broker._running_tasks diff --git a/tests/cli/worker/test_broker_selection.py b/tests/cli/worker/test_broker_selection.py new file mode 100644 index 00000000..5b03bed1 --- /dev/null +++ b/tests/cli/worker/test_broker_selection.py @@ -0,0 +1,108 @@ +import logging + +import pytest + +from taskiq import InMemoryBroker, TaskiqRouter +from taskiq.cli.worker.args import WorkerArgs +from taskiq.cli.worker.run import get_broker_selection, shutdown_broker +from taskiq.warnings import TaskiqDeprecationWarning + + +def test_get_broker_selection_preserves_direct_single_broker() -> None: + broker = InMemoryBroker() + + selection = get_broker_selection(WorkerArgs(broker=broker, modules=[])) + + assert selection.listeners == (broker,) + assert not selection.explicit_group + + +@pytest.mark.parametrize("use_factory", [False, True]) +def test_get_broker_selection_preserves_imported_single_broker( + monkeypatch: pytest.MonkeyPatch, + use_factory: bool, +) -> None: + broker = InMemoryBroker() + imported_source: object = broker + if use_factory: + + def broker_factory() -> InMemoryBroker: + return broker + + imported_source = broker_factory + + def fake_import_object(object_spec: str, app_dir: str | None = None) -> object: + assert object_spec == "application.worker:broker" + assert app_dir is None + return imported_source + + monkeypatch.setattr("taskiq.cli.worker.run.import_object", fake_import_object) + + selection = get_broker_selection( + WorkerArgs(broker="application.worker:broker", modules=[]), + ) + + assert selection.listeners == (broker,) + assert not selection.explicit_group + + +async def test_legacy_shutdown_broker_helper_is_deprecated_and_redacted( + caplog: pytest.LogCaptureFixture, +) -> None: + diagnostic = "sensitive-shutdown-diagnostic" + + class DiagnosticBroker(InMemoryBroker): + async def shutdown(self) -> None: + await super().shutdown() + return diagnostic # type: ignore[return-value] + + broker = DiagnosticBroker() + await broker.startup() + caplog.set_level(logging.INFO, logger="taskiq.worker.runtime") + + with pytest.warns(TaskiqDeprecationWarning, match="WorkerRuntime"): + await shutdown_broker(broker, timeout=1) + + assert "returned a diagnostic value" in caplog.text + assert diagnostic not in caplog.text + + +@pytest.mark.parametrize("use_factory", [False, True]) +def test_get_broker_selection_imports_explicit_listener_sequence( + monkeypatch: pytest.MonkeyPatch, + use_factory: bool, +) -> None: + router = TaskiqRouter() + first = InMemoryBroker(router=router, broker_name="first") + second = InMemoryBroker(router=router, broker_name="second") + listeners = (first, second) + imported_source: object = listeners + if use_factory: + + def listener_factory() -> tuple[InMemoryBroker, InMemoryBroker]: + return listeners + + imported_source = listener_factory + + import_calls: list[tuple[str, str | None]] = [] + + def fake_import_object(object_spec: str, app_dir: str | None = None) -> object: + import_calls.append((object_spec, app_dir)) + return imported_source + + monkeypatch.setattr( + "taskiq.cli.worker.run.import_object", + fake_import_object, + ) + + selection = get_broker_selection( + WorkerArgs( + broker="application.worker:listeners", + modules=[], + app_dir="/application", + ), + ) + + assert selection.listeners == listeners + assert selection.explicit_group + assert import_calls == [("application.worker:listeners", "/application")] diff --git a/tests/middlewares/admin_middleware/conftest.py b/tests/middlewares/admin_middleware/conftest.py index dbec77a4..9375a9d2 100644 --- a/tests/middlewares/admin_middleware/conftest.py +++ b/tests/middlewares/admin_middleware/conftest.py @@ -39,13 +39,12 @@ async def handle_executed(request: web.Request) -> web.Response: async def broker_with_admin_middleware( admin_api_server: TestServer, ) -> AsyncGenerator[InMemoryBroker, None]: - broker = InMemoryBroker(await_inplace=True).with_middlewares( - TaskiqAdminMiddleware( - str(admin_api_server.make_url("/")), # URL тестового сервера - "supersecret", - taskiq_broker_name="InMemory", - ), + middleware = TaskiqAdminMiddleware( + str(admin_api_server.make_url("/")), # URL тестового сервера + "supersecret", + taskiq_broker_name="InMemory", ) + broker = InMemoryBroker(await_inplace=True).with_middlewares(middleware) broker.register_task(task_with_dataclass, task_name="task_with_dataclass") broker.register_task(task_with_typed_dict, task_name="task_with_typed_dict") @@ -55,6 +54,8 @@ async def broker_with_admin_middleware( await broker.startup() yield broker await broker.shutdown() + assert middleware._client is not None + assert middleware._client.closed async def task_with_dataclass(dto: DataclassDTO) -> None: diff --git a/tests/receiver/test_runtime.py b/tests/receiver/test_runtime.py new file mode 100644 index 00000000..b14ed599 --- /dev/null +++ b/tests/receiver/test_runtime.py @@ -0,0 +1,1072 @@ +import asyncio +import logging +from collections.abc import AsyncGenerator, Iterator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace +from typing import Any + +import pytest + +from taskiq import AsyncBroker, BrokerMessage, TaskiqRouter, task_builder +from taskiq.abc.middleware import TaskiqMiddleware +from taskiq.receiver import Receiver +from taskiq.receiver.runtime import ( + ReceiverRuntimeState, + WorkerRuntime, + WorkerRuntimeConfigurationError, + WorkerRuntimeOptions, + normalize_broker_source, +) +from taskiq.result_backends.dummy import DummyResultBackend + + +class RuntimeTestError(Exception): + """Marker exception for worker runtime failures.""" + + +class FatalRuntimeTestError(BaseException): + """Marker fatal signal that runtime cleanup must not swallow.""" + + +class RuntimeBroker(AsyncBroker): + """Controllable broker used for multi-broker lifecycle tests.""" + + def __init__( + self, + *, + router: TaskiqRouter, + broker_name: str, + events: list[str], + startup_error: BaseException | None = None, + listen_error: BaseException | None = None, + shutdown_error: BaseException | None = None, + ) -> None: + self.events = events + self.startup_error = startup_error + self.listen_error = listen_error + self.shutdown_error = shutdown_error + self.listen_started = asyncio.Event() + self.release_listener = asyncio.Event() + self.iterator_closed = asyncio.Event() + self.started = False + self.queue: asyncio.Queue[bytes] = asyncio.Queue() + super().__init__(router=router, broker_name=broker_name) + + async def startup(self) -> None: + """Record startup and optionally fail after core resources start.""" + self.events.append( + f"{self.broker_name}.startup." + f"{'worker' if self.is_worker_process else 'client'}", + ) + await super().startup() + self.started = True + if self.startup_error is not None: + raise self.startup_error + + async def shutdown(self) -> None: + """Record shutdown and close core resources.""" + self.events.append(f"{self.broker_name}.shutdown") + await super().shutdown() + self.started = False + if self.shutdown_error is not None: + raise self.shutdown_error + + async def kick(self, message: BrokerMessage) -> None: + """Queue an encoded message for this listener.""" + if not self.started: + raise RuntimeTestError(f"{self.broker_name} kicked before startup") + self.events.append(f"{self.broker_name}.kick") + await self.queue.put(message.message) + + async def listen(self) -> AsyncGenerator[bytes, None]: + """Listen until cancelled or a configured failure is released.""" + self.events.append(f"{self.broker_name}.listen") + self.listen_started.set() + try: + if self.listen_error is not None: + await self.release_listener.wait() + raise self.listen_error + while True: + yield await self.queue.get() + finally: + self.events.append(f"{self.broker_name}.iterator_closed") + self.iterator_closed.set() + + +class PrefetchRecordingMiddleware(TaskiqMiddleware): + """Record broker deliveries admitted into the shared prefetch budget.""" + + def __init__(self, broker_name: str, prefetched: list[str]) -> None: + super().__init__() + self.broker_name = broker_name + self.prefetched = prefetched + + def on_prefetch_queue_add(self) -> None: + """Record one delivery accepted by a Receiver prefetch queue.""" + self.prefetched.append(self.broker_name) + + +@pytest.fixture +def runtime_executor() -> Iterator[ThreadPoolExecutor]: + """Provide the process-wide sync executor owned by a worker runtime.""" + with ThreadPoolExecutor(max_workers=2) as executor: + yield executor + + +def runtime_options(**overrides: Any) -> WorkerRuntimeOptions: + """Build runtime options with focused per-test overrides.""" + options = WorkerRuntimeOptions( + validate_params=True, + max_async_tasks=10, + max_async_tasks_jitter=0, + max_prefetch=0, + propagate_exceptions=True, + ack_type=None, + max_tasks_to_execute=None, + wait_tasks_timeout=1, + shutdown_timeout=1, + receiver_kwargs={}, + ) + return replace(options, **overrides) + + +def build_runtime( + brokers: object, + executor: ThreadPoolExecutor, + *, + options: WorkerRuntimeOptions | None = None, + receiver_type: type[Receiver] = Receiver, +) -> WorkerRuntime: + """Normalize broker configuration and build the runtime under test.""" + return WorkerRuntime( + selection=normalize_broker_source(brokers), + receiver_type=receiver_type, + executor=executor, + options=options or runtime_options(), + ) + + +async def test_runtime_listens_concurrently_and_starts_outbound_broker( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + first = RuntimeBroker(router=router, broker_name="first", events=events) + second = RuntimeBroker(router=router, broker_name="second", events=events) + outbound = RuntimeBroker(router=router, broker_name="outbound", events=events) + finish_event = asyncio.Event() + runtime = build_runtime((first, second), runtime_executor) + + runtime_task = asyncio.create_task(runtime.run(finish_event)) + await asyncio.gather(first.listen_started.wait(), second.listen_started.wait()) + + assert not runtime_task.done() + assert not outbound.listen_started.is_set() + assert first.is_worker_process + assert second.is_worker_process + assert not outbound.is_worker_process + assert events[:3] == [ + "first.startup.worker", + "second.startup.worker", + "outbound.startup.client", + ] + + finish_event.set() + await runtime_task + + assert first.iterator_closed.is_set() + assert second.iterator_closed.is_set() + assert [event for event in events if event.endswith(".shutdown")] == [ + "outbound.shutdown", + "second.shutdown", + "first.shutdown", + ] + + +async def test_single_broker_runtime_does_not_start_router_peers( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + listener = RuntimeBroker(router=router, broker_name="listener", events=events) + peer = RuntimeBroker(router=router, broker_name="peer", events=events) + finish_event = asyncio.Event() + runtime = build_runtime(listener, runtime_executor) + + runtime_task = asyncio.create_task(runtime.run(finish_event)) + await listener.listen_started.wait() + + assert listener.started + assert not peer.started + + finish_event.set() + await runtime_task + + assert [event for event in events if ".startup." in event] == [ + "listener.startup.worker", + ] + assert [event for event in events if event.endswith(".shutdown")] == [ + "listener.shutdown", + ] + + +async def test_listener_task_can_send_through_outbound_only_broker( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + first = RuntimeBroker(router=router, broker_name="first", events=events) + second = RuntimeBroker(router=router, broker_name="second", events=events) + outbound = RuntimeBroker(router=router, broker_name="outbound", events=events) + parent_finished = asyncio.Event() + + @first.task(task_name="runtime.outbound_child") + async def outbound_child() -> None: + return None + + @first.task(task_name="runtime.outbound_parent") + async def outbound_parent() -> None: + await outbound_child.kiq() + parent_finished.set() + + router.route_task(outbound_child, broker=outbound) + finish_event = asyncio.Event() + runtime = build_runtime((first, second), runtime_executor) + runtime_task = asyncio.create_task(runtime.run(finish_event)) + await asyncio.gather(first.listen_started.wait(), second.listen_started.wait()) + + await outbound_parent.kicker().with_broker(first).kiq() + await asyncio.wait_for(parent_finished.wait(), timeout=1) + finish_event.set() + await runtime_task + + assert "outbound.kick" in events + assert not outbound.listen_started.is_set() + assert events.index("outbound.startup.client") < events.index("outbound.kick") + assert events.index("outbound.kick") < events.index("outbound.shutdown") + + +async def test_runtime_rolls_back_partial_startup_in_reverse_order( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + startup_error = RuntimeTestError("second startup failed") + first = RuntimeBroker(router=router, broker_name="first", events=events) + second = RuntimeBroker( + router=router, + broker_name="second", + events=events, + startup_error=startup_error, + ) + runtime = build_runtime((first, second), runtime_executor) + + with pytest.raises(RuntimeTestError) as exc_info: + await runtime.run(asyncio.Event()) + + assert exc_info.value is startup_error + assert events == [ + "first.startup.worker", + "second.startup.worker", + "second.shutdown", + "first.shutdown", + ] + assert not first.listen_started.is_set() + assert not second.listen_started.is_set() + + +async def test_runtime_preserves_listener_failure_after_sibling_cleanup( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + listen_error = RuntimeTestError("first listener failed") + first = RuntimeBroker( + router=router, + broker_name="first", + events=events, + listen_error=listen_error, + ) + second = RuntimeBroker(router=router, broker_name="second", events=events) + runtime = build_runtime((first, second), runtime_executor) + + runtime_task = asyncio.create_task(runtime.run(asyncio.Event())) + await asyncio.gather(first.listen_started.wait(), second.listen_started.wait()) + first.release_listener.set() + + with pytest.raises(RuntimeTestError) as exc_info: + await runtime_task + + assert exc_info.value is listen_error + assert first.iterator_closed.is_set() + assert second.iterator_closed.is_set() + assert [event for event in events if event.endswith(".shutdown")] == [ + "second.shutdown", + "first.shutdown", + ] + + +async def test_runtime_keeps_first_listener_error_over_sibling_stop_error( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + first = RuntimeBroker(router=router, broker_name="first", events=events) + second = RuntimeBroker(router=router, broker_name="second", events=events) + first_started = asyncio.Event() + primary_error = RuntimeTestError("second listener failed") + sibling_error = RuntimeTestError("first listener failed while stopping") + + class OrderedFailureReceiver(Receiver): + async def listen(self, finish_event: asyncio.Event) -> None: + if self.broker is first: + first_started.set() + await finish_event.wait() + raise sibling_error + + await first_started.wait() + raise primary_error + + runtime = build_runtime( + (first, second), + runtime_executor, + receiver_type=OrderedFailureReceiver, + ) + + with pytest.raises(RuntimeTestError) as exc_info: + await runtime.run(asyncio.Event()) + + assert exc_info.value is primary_error + + +async def test_runtime_cancellation_closes_all_listeners_and_brokers( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + first = RuntimeBroker(router=router, broker_name="first", events=events) + second = RuntimeBroker(router=router, broker_name="second", events=events) + runtime = build_runtime((first, second), runtime_executor) + + runtime_task = asyncio.create_task(runtime.run(asyncio.Event())) + await asyncio.gather(first.listen_started.wait(), second.listen_started.wait()) + runtime_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await runtime_task + + assert first.iterator_closed.is_set() + assert second.iterator_closed.is_set() + assert [event for event in events if event.endswith(".shutdown")] == [ + "second.shutdown", + "first.shutdown", + ] + + +async def test_runtime_shares_async_execution_limit_across_brokers( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + first = RuntimeBroker(router=router, broker_name="first", events=events) + second = RuntimeBroker(router=router, broker_name="second", events=events) + finish_event = asyncio.Event() + release_tasks = asyncio.Event() + first_started = asyncio.Event() + all_finished = asyncio.Event() + active_tasks = 0 + max_active_tasks = 0 + started_tasks = 0 + finished_tasks = 0 + + @first.task(task_name="runtime.shared_limit") + async def limited_task() -> None: + nonlocal active_tasks, finished_tasks, max_active_tasks, started_tasks + active_tasks += 1 + started_tasks += 1 + max_active_tasks = max(max_active_tasks, active_tasks) + first_started.set() + try: + await release_tasks.wait() + finally: + active_tasks -= 1 + finished_tasks += 1 + if finished_tasks == 2: + all_finished.set() + + runtime = build_runtime( + (first, second), + runtime_executor, + options=runtime_options(max_async_tasks=1), + ) + runtime_task = asyncio.create_task(runtime.run(finish_event)) + await asyncio.gather(first.listen_started.wait(), second.listen_started.wait()) + + await limited_task.kicker().with_broker(first).kiq() + await limited_task.kicker().with_broker(second).kiq() + await first_started.wait() + await asyncio.sleep(0.05) + + assert started_tasks == 1 + assert max_active_tasks == 1 + + release_tasks.set() + await asyncio.wait_for(all_finished.wait(), timeout=1) + finish_event.set() + await runtime_task + + assert max_active_tasks == 1 + + +async def test_runtime_shares_prefetch_limit_across_brokers( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + first = RuntimeBroker(router=router, broker_name="first", events=events) + second = RuntimeBroker(router=router, broker_name="second", events=events) + prefetched: list[str] = [] + first.with_middlewares(PrefetchRecordingMiddleware("first", prefetched)) + second.with_middlewares(PrefetchRecordingMiddleware("second", prefetched)) + first_task_started = asyncio.Event() + all_tasks_finished = asyncio.Event() + release_tasks = asyncio.Event() + finished_tasks = 0 + + @first.task(task_name="runtime.shared_prefetch") + async def limited_task() -> None: + nonlocal finished_tasks + first_task_started.set() + await release_tasks.wait() + finished_tasks += 1 + if finished_tasks == 4: + all_tasks_finished.set() + + finish_event = asyncio.Event() + runtime = build_runtime( + (first, second), + runtime_executor, + options=runtime_options(max_async_tasks=1, max_prefetch=0), + ) + runtime_task = asyncio.create_task(runtime.run(finish_event)) + await asyncio.gather(first.listen_started.wait(), second.listen_started.wait()) + + for target_broker in (first, second, first, second): + await limited_task.kicker().with_broker(target_broker).kiq() + await first_task_started.wait() + + async def wait_for_two_prefetched_tasks() -> None: + while len(prefetched) < 2: + await asyncio.sleep(0) + + await asyncio.wait_for(wait_for_two_prefetched_tasks(), timeout=1) + await asyncio.sleep(0.05) + + assert len(prefetched) == 2 + + release_tasks.set() + await asyncio.wait_for(all_tasks_finished.wait(), timeout=1) + finish_event.set() + await runtime_task + + +async def test_runtime_shares_max_tasks_per_child_across_brokers( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + first = RuntimeBroker(router=router, broker_name="first", events=events) + second = RuntimeBroker(router=router, broker_name="second", events=events) + first_finished = asyncio.Event() + all_finished = asyncio.Event() + execution_count = 0 + + @first.task(task_name="runtime.shared_task_count") + async def counted_task() -> None: + nonlocal execution_count + execution_count += 1 + if execution_count == 1: + first_finished.set() + elif execution_count == 2: + all_finished.set() + + finish_event = asyncio.Event() + runtime = build_runtime( + (first, second), + runtime_executor, + options=runtime_options(max_tasks_to_execute=2), + ) + runtime_task = asyncio.create_task(runtime.run(finish_event)) + await asyncio.gather(first.listen_started.wait(), second.listen_started.wait()) + + await counted_task.kicker().with_broker(first).kiq() + await first_finished.wait() + await counted_task.kicker().with_broker(second).kiq() + + await asyncio.wait_for(all_finished.wait(), timeout=1) + await asyncio.wait_for(runtime_task, timeout=1) + + assert execution_count == 2 + assert finish_event.is_set() + + +async def test_task_count_stop_does_not_discard_eager_transport_delivery( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + broker = RuntimeBroker( + router=TaskiqRouter(), + broker_name="listener", + events=events, + ) + prefetch_hook_started = asyncio.Event() + release_prefetch_hook = asyncio.Event() + execution_count = 0 + + class PausingPrefetchMiddleware(TaskiqMiddleware): + async def on_prefetch_queue_add(self) -> None: + prefetch_hook_started.set() + await release_prefetch_hook.wait() + + broker.with_middlewares(PausingPrefetchMiddleware()) + + @broker.task(task_name="runtime.no_eager_delivery_loss") + async def counted_task() -> None: + nonlocal execution_count + execution_count += 1 + + runtime = build_runtime( + (broker,), + runtime_executor, + options=runtime_options(max_tasks_to_execute=1), + ) + runtime_task = asyncio.create_task(runtime.run(asyncio.Event())) + await broker.listen_started.wait() + + await counted_task.kiq() + await prefetch_hook_started.wait() + await counted_task.kiq() + release_prefetch_hook.set() + await asyncio.wait_for(runtime_task, timeout=1) + + assert execution_count == 1 + assert broker.queue.qsize() == 1 + + +async def test_simultaneous_task_count_deliveries_have_bounded_overshoot( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + first = RuntimeBroker(router=router, broker_name="first", events=events) + second = RuntimeBroker(router=router, broker_name="second", events=events) + execution_count = 0 + + @first.task(task_name="runtime.bounded_task_count_overshoot") + async def counted_task() -> None: + nonlocal execution_count + execution_count += 1 + + runtime = build_runtime( + (first, second), + runtime_executor, + options=runtime_options( + max_prefetch=1, + max_tasks_to_execute=1, + ), + ) + runtime_task = asyncio.create_task(runtime.run(asyncio.Event())) + await asyncio.gather(first.listen_started.wait(), second.listen_started.wait()) + + await asyncio.gather( + counted_task.kicker().with_broker(first).kiq(), + counted_task.kicker().with_broker(second).kiq(), + ) + await asyncio.wait_for(runtime_task, timeout=1) + + assert execution_count == 2 + assert first.queue.empty() + assert second.queue.empty() + + +async def test_runtime_cancels_callback_after_graceful_timeout( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + broker = RuntimeBroker(router=router, broker_name="listener", events=events) + task_started = asyncio.Event() + task_cancelled = asyncio.Event() + + @broker.task(task_name="runtime.callback_timeout") + async def blocked_task() -> None: + task_started.set() + try: + await asyncio.Event().wait() + finally: + task_cancelled.set() + + finish_event = asyncio.Event() + runtime = build_runtime( + (broker,), + runtime_executor, + options=runtime_options( + wait_tasks_timeout=0.01, + shutdown_timeout=0.1, + ), + ) + runtime_task = asyncio.create_task(runtime.run(finish_event)) + await broker.listen_started.wait() + await blocked_task.kiq() + await task_started.wait() + + finish_event.set() + await asyncio.wait_for(runtime_task, timeout=1) + + assert task_cancelled.is_set() + assert broker.iterator_closed.is_set() + + +async def test_runtime_waits_for_callbacks_without_configured_task_timeout( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + broker = RuntimeBroker( + router=TaskiqRouter(), + broker_name="listener", + events=events, + ) + task_started = asyncio.Event() + release_task = asyncio.Event() + task_cancelled = False + + @broker.task(task_name="runtime.unbounded_graceful_wait") + async def blocked_task() -> None: + nonlocal task_cancelled + task_started.set() + try: + await release_task.wait() + except asyncio.CancelledError: + task_cancelled = True + raise + + finish_event = asyncio.Event() + runtime = build_runtime( + (broker,), + runtime_executor, + options=runtime_options( + wait_tasks_timeout=None, + shutdown_timeout=0.01, + ), + ) + runtime_task = asyncio.create_task(runtime.run(finish_event)) + await broker.listen_started.wait() + await blocked_task.kiq() + await task_started.wait() + + finish_event.set() + await asyncio.sleep(0.05) + finished_before_release = runtime_task.done() + release_task.set() + await asyncio.wait_for(runtime_task, timeout=1) + + assert not finished_before_release + assert not task_cancelled + + +async def test_runtime_cancellation_during_drain_awaits_callback_cleanup( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + broker = RuntimeBroker( + router=TaskiqRouter(), + broker_name="listener", + events=events, + ) + task_started = asyncio.Event() + release_task = asyncio.Event() + task_cleaned = asyncio.Event() + + @broker.task(task_name="runtime.cancel_during_drain") + async def blocked_task() -> None: + task_started.set() + try: + await release_task.wait() + finally: + task_cleaned.set() + + finish_event = asyncio.Event() + runtime = build_runtime( + (broker,), + runtime_executor, + options=runtime_options( + wait_tasks_timeout=0.4, + shutdown_timeout=0.05, + ), + ) + runtime_task = asyncio.create_task(runtime.run(finish_event)) + await broker.listen_started.wait() + await blocked_task.kiq() + await task_started.wait() + + finish_event.set() + await asyncio.wait_for(runtime_task, timeout=1) + cleaned_before_runtime_returned = task_cleaned.is_set() + + release_task.set() + await asyncio.wait_for(task_cleaned.wait(), timeout=1) + + assert cleaned_before_runtime_returned + + +async def test_runtime_continues_cleanup_after_broker_shutdown_failure( + runtime_executor: ThreadPoolExecutor, + caplog: pytest.LogCaptureFixture, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + first = RuntimeBroker(router=router, broker_name="first", events=events) + shutdown_error = RuntimeTestError("second shutdown failed") + second = RuntimeBroker( + router=router, + broker_name="second", + events=events, + shutdown_error=shutdown_error, + ) + finish_event = asyncio.Event() + runtime = build_runtime((first, second), runtime_executor) + caplog.set_level(logging.WARNING, logger="taskiq.worker.runtime") + + runtime_task = asyncio.create_task(runtime.run(finish_event)) + await asyncio.gather(first.listen_started.wait(), second.listen_started.wait()) + finish_event.set() + await runtime_task + + assert [event for event in events if event.endswith(".shutdown")] == [ + "second.shutdown", + "first.shutdown", + ] + assert "second failed during shutdown" in caplog.text + + +async def test_runtime_bounds_cancellation_resistant_broker_shutdown( + runtime_executor: ThreadPoolExecutor, + caplog: pytest.LogCaptureFixture, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + release_shutdown = asyncio.Event() + shutdown_cancelled = asyncio.Event() + shutdown_finished = asyncio.Event() + + class CancellationResistantShutdownBroker(RuntimeBroker): + async def shutdown(self) -> None: + self.events.append(f"{self.broker_name}.shutdown") + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + shutdown_cancelled.set() + await release_shutdown.wait() + finally: + shutdown_finished.set() + + first = RuntimeBroker(router=router, broker_name="first", events=events) + second = CancellationResistantShutdownBroker( + router=router, + broker_name="second", + events=events, + ) + runtime = build_runtime( + (first, second), + runtime_executor, + options=runtime_options(shutdown_timeout=0.01), + ) + caplog.set_level(logging.WARNING, logger="taskiq.worker.runtime") + finish_event = asyncio.Event() + runtime_task = asyncio.create_task(runtime.run(finish_event)) + await asyncio.gather(first.listen_started.wait(), second.listen_started.wait()) + + finish_event.set() + completed, _ = await asyncio.wait((runtime_task,), timeout=1) + + release_shutdown.set() + if runtime_task not in completed: + await asyncio.wait_for(runtime_task, timeout=0.2) + else: + await runtime_task + await asyncio.wait_for(shutdown_finished.wait(), timeout=0.2) + + assert shutdown_cancelled.is_set() + assert runtime_task in completed + assert [event for event in events if event.endswith(".shutdown")] == [ + "second.shutdown", + "first.shutdown", + ] + assert "second did not shut down" in caplog.text + + +async def test_runtime_cancellation_during_broker_shutdown_cleans_remaining_brokers( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + shutdown_started = asyncio.Event() + + class BlockingShutdownBroker(RuntimeBroker): + async def shutdown(self) -> None: + self.events.append(f"{self.broker_name}.shutdown") + shutdown_started.set() + await asyncio.Event().wait() + + first = RuntimeBroker(router=router, broker_name="first", events=events) + second = BlockingShutdownBroker( + router=router, + broker_name="second", + events=events, + ) + runtime = build_runtime((first, second), runtime_executor) + finish_event = asyncio.Event() + runtime_task = asyncio.create_task(runtime.run(finish_event)) + await asyncio.gather(first.listen_started.wait(), second.listen_started.wait()) + + finish_event.set() + await shutdown_started.wait() + runtime_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await runtime_task + + assert [event for event in events if event.endswith(".shutdown")] == [ + "second.shutdown", + "first.shutdown", + ] + + +async def test_runtime_does_not_swallow_fatal_broker_shutdown_error( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + shutdown_error = FatalRuntimeTestError("fatal shutdown") + broker = RuntimeBroker( + router=TaskiqRouter(), + broker_name="listener", + events=events, + shutdown_error=shutdown_error, + ) + runtime = build_runtime((broker,), runtime_executor) + finish_event = asyncio.Event() + finish_event.set() + + with pytest.raises(FatalRuntimeTestError) as exc_info: + await runtime.run(finish_event) + + assert exc_info.value is shutdown_error + + +async def test_runtime_bounds_cancellation_resistant_listener_cleanup( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + first = RuntimeBroker(router=router, broker_name="first", events=events) + second = RuntimeBroker(router=router, broker_name="second", events=events) + listeners_started = { + "first": asyncio.Event(), + "second": asyncio.Event(), + } + cancellation_seen = asyncio.Event() + release_listener = asyncio.Event() + listener_tasks: list[asyncio.Task[None]] = [] + + class CancellationResistantReceiver(Receiver): + async def listen(self, finish_event: asyncio.Event) -> None: + listeners_started[self.broker.broker_name].set() + listener_task = asyncio.current_task() + assert listener_task is not None + listener_tasks.append(listener_task) + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancellation_seen.set() + await release_listener.wait() + + runtime = build_runtime( + (first, second), + runtime_executor, + options=runtime_options( + wait_tasks_timeout=0.01, + shutdown_timeout=0.01, + ), + receiver_type=CancellationResistantReceiver, + ) + + finish_event = asyncio.Event() + runtime_task = asyncio.create_task(runtime.run(finish_event)) + await asyncio.gather( + listeners_started["first"].wait(), + listeners_started["second"].wait(), + ) + finish_event.set() + await asyncio.wait_for(runtime_task, timeout=1) + + assert cancellation_seen.is_set() + assert [event for event in events if event.endswith(".shutdown")] == [ + "second.shutdown", + "first.shutdown", + ] + + release_listener.set() + await asyncio.wait_for(asyncio.gather(*listener_tasks), timeout=0.2) + + +def test_runtime_rejects_mixed_routers_before_startup( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + first = RuntimeBroker( + router=TaskiqRouter(), + broker_name="first", + events=events, + ) + second = RuntimeBroker( + router=TaskiqRouter(), + broker_name="second", + events=events, + ) + + with pytest.raises(WorkerRuntimeConfigurationError, match="same router"): + build_runtime((first, second), runtime_executor) + + assert events == [] + + +@pytest.mark.parametrize( + ("source", "error_match"), + [ + ((), "cannot be empty"), + ((object(),), "index 0"), + ("module:broker", "must be an AsyncBroker"), + ], +) +def test_runtime_rejects_invalid_listener_sources( + source: object, + error_match: str, +) -> None: + with pytest.raises(WorkerRuntimeConfigurationError, match=error_match): + normalize_broker_source(source) + + +def test_runtime_rejects_duplicate_listener_before_startup() -> None: + events: list[str] = [] + broker = RuntimeBroker( + router=TaskiqRouter(), + broker_name="listener", + events=events, + ) + + with pytest.raises(WorkerRuntimeConfigurationError, match="more than once"): + normalize_broker_source((broker, broker)) + + assert events == [] + + +def test_single_broker_selection_does_not_change_router_peer_mode() -> None: + events: list[str] = [] + router = TaskiqRouter() + listener = RuntimeBroker(router=router, broker_name="listener", events=events) + peer = RuntimeBroker(router=router, broker_name="peer", events=events) + peer.is_worker_process = True + + normalize_broker_source(listener).mark_listener_brokers() + + assert listener.is_worker_process + assert peer.is_worker_process + + +def test_runtime_rejects_shared_middleware_before_startup( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + first = RuntimeBroker(router=router, broker_name="first", events=events) + second = RuntimeBroker(router=router, broker_name="second", events=events) + middleware = TaskiqMiddleware() + first.with_middlewares(middleware) + second.with_middlewares(middleware) + + with pytest.raises(WorkerRuntimeConfigurationError, match="middleware"): + build_runtime((first, second), runtime_executor) + + assert events == [] + + +def test_runtime_rejects_shared_result_backend_before_startup( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + first = RuntimeBroker(router=router, broker_name="first", events=events) + second = RuntimeBroker(router=router, broker_name="second", events=events) + backend: DummyResultBackend[Any] = DummyResultBackend() + first.with_result_backend(backend) + second.with_result_backend(backend) + + with pytest.raises(WorkerRuntimeConfigurationError, match="result backend"): + build_runtime((first, second), runtime_executor) + + assert events == [] + + +def test_runtime_rejects_conflicting_effective_task_before_startup( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + router = TaskiqRouter() + first = RuntimeBroker(router=router, broker_name="first", events=events) + second = RuntimeBroker(router=router, broker_name="second", events=events) + + @first.task(task_name="runtime.conflict") + async def first_task() -> str: + return "first" + + @task_builder("runtime.conflict") + async def second_task() -> str: + return "second" + + second.store_registered_task( + second.bind_task_definition(second_task, register=False), + ) + + with pytest.raises(WorkerRuntimeConfigurationError, match="resolves differently"): + build_runtime((first, second), runtime_executor) + + assert events == [] + + +async def test_runtime_rejects_custom_receiver_without_shared_state( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + broker = RuntimeBroker( + router=TaskiqRouter(), + broker_name="listener", + events=events, + ) + + class IncompatibleReceiver(Receiver): + def attach_runtime_state( + self, + runtime_state: ReceiverRuntimeState, + ) -> ReceiverRuntimeState: + return runtime_state + + runtime = build_runtime( + (broker,), + runtime_executor, + receiver_type=IncompatibleReceiver, + ) + + with pytest.raises(WorkerRuntimeConfigurationError, match="Custom Receiver"): + await runtime.run(asyncio.Event()) + + assert events == [] diff --git a/tests/routing/test_dispatch.py b/tests/routing/test_dispatch.py index 0317f4e3..fa2e74b5 100644 --- a/tests/routing/test_dispatch.py +++ b/tests/routing/test_dispatch.py @@ -1,4 +1,6 @@ -from taskiq import Flow, FlowProtocol, InMemoryBroker, TaskiqRouter +import pytest + +from taskiq import Flow, FlowProtocol, InMemoryBroker, TaskiqRoute, TaskiqRouter from tests.routing.models import ( BrokerQueue, OldStyleRecordingBroker, @@ -39,6 +41,49 @@ async def demo_task() -> None: assert broker.get_subscribed_flows() == (flow,) +async def test_registered_task_resolves_current_default_flow_at_send_time() -> None: + first_flow = Flow("first") + second_flow = Flow("second") + broker = RecordingBroker(default_flow=first_flow) + + @broker.task(task_name="demo.task") + async def demo_task() -> None: + return None + + initial_route = broker.router.resolve_route(demo_task) + broker.default_flow = second_flow + + await demo_task.kiq() + await demo_task.kicker().with_route(initial_route).kiq() + + assert initial_route.flow == first_flow + assert broker.router.resolve_route(demo_task).flow == second_flow + assert broker.router.routes[demo_task.task_name].flow == second_flow + assert [flow for _, flow in broker.sent] == [second_flow, first_flow] + + +async def test_explicit_task_route_does_not_follow_default_flow_changes() -> None: + router = TaskiqRouter() + broker = RecordingBroker( + router=router, + broker_name="broker", + default_flow=Flow("default.first"), + ) + explicit_flow = Flow("explicit") + + @broker.task(task_name="demo.task") + async def demo_task() -> None: + return None + + router.route_task(demo_task, broker=broker, flow=explicit_flow) + broker.default_flow = Flow("default.second") + + await demo_task.kiq() + + assert router.resolve_route(demo_task).flow == explicit_flow + assert broker.sent[0][1] == explicit_flow + + def test_router_set_broker_can_configure_default_flow() -> None: router = TaskiqRouter() broker = RecordingBroker(router=router, broker_name="broker") @@ -138,6 +183,69 @@ async def demo_task() -> None: assert second.sent[0][1] == second_flow +async def test_kicker_broker_override_clears_earlier_flow_override() -> None: + router = TaskiqRouter() + source = RecordingBroker(router=router, broker_name="source") + target_flow = Flow("target.default") + target = RecordingBroker( + router=router, + broker_name="target", + default_flow=target_flow, + ) + + @source.task(task_name="demo.task") + async def demo_task() -> None: + return None + + await demo_task.kicker().with_flow(Flow("stale")).with_broker(target).kiq() + + assert source.sent == [] + assert target.sent[0][1] == target_flow + + +async def test_kicker_flow_override_after_broker_override_wins() -> None: + router = TaskiqRouter() + source = RecordingBroker(router=router, broker_name="source") + target = RecordingBroker( + router=router, + broker_name="target", + default_flow=Flow("target.default"), + ) + explicit_flow = Flow("target.explicit") + + @source.task(task_name="demo.task") + async def demo_task() -> None: + return None + + await demo_task.kicker().with_broker(target).with_flow(explicit_flow).kiq() + + assert source.sent == [] + assert target.sent[0][1] == explicit_flow + + +def test_kicker_rejects_ambiguous_none_flow_override() -> None: + broker = RecordingBroker(default_flow=Flow("default")) + + @broker.task(task_name="demo.task") + async def demo_task() -> None: + return None + + with pytest.raises(TypeError, match="flow cannot be None"): + demo_task.kicker().with_flow(None) # type: ignore[arg-type] + + +async def test_explicit_route_can_bypass_broker_default_flow() -> None: + broker = RecordingBroker(default_flow=Flow("default")) + + @broker.task(task_name="demo.task") + async def demo_task() -> None: + return None + + await demo_task.kicker().with_route(TaskiqRoute(broker, flow=None)).kiq() + + assert broker.sent[0][1] is None + + async def test_router_uses_explicit_broker_lookup_for_config_names() -> None: router = TaskiqRouter() source = RecordingBroker(router=router, broker_name="source") @@ -260,6 +368,24 @@ async def demo_task(value: int) -> None: assert second.sent == [] +async def test_prepared_invocation_snapshots_current_default_flow() -> None: + first_flow = Flow("default.first") + second_flow = Flow("default.second") + broker = RecordingBroker(default_flow=first_flow) + + @broker.task(task_name="demo.task") + async def demo_task() -> None: + return None + + prepared = demo_task.kicker().prepare() + broker.default_flow = second_flow + + await prepared.kiq() + await demo_task.kiq() + + assert [flow for _, flow in broker.sent] == [first_flow, second_flow] + + async def test_router_task_decorator_can_choose_broker_and_flow() -> None: router = TaskiqRouter() target = RecordingBroker(router=router, broker_name="target") diff --git a/tests/routing/test_router_core.py b/tests/routing/test_router_core.py index d2b92301..79f7f818 100644 --- a/tests/routing/test_router_core.py +++ b/tests/routing/test_router_core.py @@ -51,6 +51,46 @@ def test_router_rejects_string_broker_references() -> None: router.route_task("demo.task", broker="broker") # type: ignore[arg-type] +def test_router_routes_are_immutable_resolved_snapshots() -> None: + router = TaskiqRouter() + broker = RecordingBroker(router=router, broker_name="broker") + route = router.route_task("demo.task", broker=broker) + + routes = router.routes + + assert routes == {"demo.task": route} + with pytest.raises(TypeError): + routes["other.task"] = route # type: ignore[index] + + assert router.has_route("demo.task") + assert router.remove_route("demo.task") == route + assert not router.has_route("demo.task") + assert router.remove_route("demo.task") is None + + +def test_broker_and_task_registries_are_immutable_snapshots() -> None: + router = TaskiqRouter() + first = RecordingBroker(router=router, broker_name="first") + broker_snapshot = router.brokers + task_snapshot = router.task_registry + + @first.task(task_name="demo.task") + async def demo_task() -> None: + return None + + second = RecordingBroker(router=router, broker_name="second") + + assert broker_snapshot == {"first": first} + assert task_snapshot == {} + assert router.brokers == {"first": first, "second": second} + assert router.task_registry == {"demo.task": demo_task} + + with pytest.raises(TypeError): + broker_snapshot["other"] = first # type: ignore[index] + with pytest.raises(TypeError): + task_snapshot["other.task"] = demo_task # type: ignore[index] + + def test_router_rejects_duplicate_task_names() -> None: router = TaskiqRouter() first = RecordingBroker(router=router, broker_name="first") diff --git a/tests/routing/test_send_lifecycle.py b/tests/routing/test_send_lifecycle.py new file mode 100644 index 00000000..33b65c6c --- /dev/null +++ b/tests/routing/test_send_lifecycle.py @@ -0,0 +1,258 @@ +import asyncio +from typing import Literal + +import pytest + +from taskiq import TaskiqRoute +from taskiq.abc.formatter import TaskiqFormatter +from taskiq.abc.middleware import TaskiqMiddleware +from taskiq.exceptions import SendTaskError +from taskiq.message import BrokerMessage, TaskiqMessage +from tests.utils import RecordingBroker + + +class LifecycleError(Exception): + """Marker error for send lifecycle tests.""" + + +class LifecycleMiddleware(TaskiqMiddleware): + """Record send hooks and optionally fail one phase.""" + + def __init__( + self, + events: list[str], + *, + error: LifecycleError | None = None, + fail_at: Literal["pre_send", "post_send"] | None = None, + ) -> None: + super().__init__() + self.events = events + self.error = error + self.fail_at = fail_at + + def pre_send(self, message: TaskiqMessage) -> TaskiqMessage: + """Record and optionally fail before transport dispatch.""" + self.events.append("pre_send") + if self.fail_at == "pre_send" and self.error is not None: + raise self.error + return message + + def post_send(self, message: TaskiqMessage) -> None: + """Record and optionally fail after transport dispatch.""" + self.events.append("post_send") + if self.fail_at == "post_send" and self.error is not None: + raise self.error + + +class FailingFormatter(TaskiqFormatter): + """Formatter that raises a configured error on serialization.""" + + def __init__(self, error: LifecycleError) -> None: + self.error = error + + def dumps(self, message: TaskiqMessage) -> BrokerMessage: + """Fail message serialization.""" + raise self.error + + def loads(self, message: bytes) -> TaskiqMessage: + """Incoming deserialization is not used by these tests.""" + raise AssertionError("loads() must not be called") + + +class AsyncOrderedMiddleware(TaskiqMiddleware): + """Record async send hooks with a middleware name.""" + + def __init__(self, name: str, events: list[str]) -> None: + super().__init__() + self.name = name + self.events = events + + async def pre_send(self, message: TaskiqMessage) -> TaskiqMessage: + """Record async pre-send order.""" + await asyncio.sleep(0) + self.events.append(f"{self.name}.pre_send") + return message + + async def post_send(self, message: TaskiqMessage) -> None: + """Record async post-send order.""" + await asyncio.sleep(0) + self.events.append(f"{self.name}.post_send") + + +class FailingSendBroker(RecordingBroker): + """Broker that records transport entry and then fails.""" + + def __init__(self, events: list[str], error: BaseException) -> None: + self.events = events + self.send_error = error + super().__init__() + + async def kick(self, message: BrokerMessage) -> None: + """Fail old-style transport dispatch.""" + self.events.append("kick") + raise self.send_error + + async def kick_to_flow( + self, + message: BrokerMessage, + flow: object | None = None, + ) -> None: + """Fail flow-aware transport dispatch.""" + self.events.append("kick") + raise self.send_error + + +async def test_send_middleware_order_supports_async_hooks() -> None: + events: list[str] = [] + broker = RecordingBroker() + broker.with_middlewares( + AsyncOrderedMiddleware("first", events), + AsyncOrderedMiddleware("second", events), + ) + + @broker.task(task_name="lifecycle.order") + async def demo_task() -> None: + return None + + await demo_task.kiq() + + assert events == [ + "first.pre_send", + "second.pre_send", + "second.post_send", + "first.post_send", + ] + + +async def test_pre_send_error_is_not_wrapped() -> None: + events: list[str] = [] + error = LifecycleError("pre_send") + broker = RecordingBroker() + broker.with_middlewares( + LifecycleMiddleware(events, error=error, fail_at="pre_send"), + ) + + @broker.task(task_name="lifecycle.pre_send") + async def demo_task() -> None: + return None + + with pytest.raises(LifecycleError) as exc_info: + await demo_task.kiq() + + assert exc_info.value is error + assert events == ["pre_send"] + assert broker.sent == [] + + +async def test_post_send_error_is_not_wrapped() -> None: + events: list[str] = [] + error = LifecycleError("post_send") + broker = RecordingBroker() + broker.with_middlewares( + LifecycleMiddleware(events, error=error, fail_at="post_send"), + ) + + @broker.task(task_name="lifecycle.post_send") + async def demo_task() -> None: + return None + + with pytest.raises(LifecycleError) as exc_info: + await demo_task.kiq() + + assert exc_info.value is error + assert events == ["pre_send", "post_send"] + assert len(broker.sent) == 1 + + +async def test_formatter_error_is_wrapped_and_skips_post_send() -> None: + events: list[str] = [] + error = LifecycleError("formatter") + broker = RecordingBroker() + broker.with_middlewares(LifecycleMiddleware(events)) + broker.with_formatter(FailingFormatter(error)) + + @broker.task(task_name="lifecycle.formatter") + async def demo_task() -> None: + return None + + with pytest.raises(SendTaskError) as exc_info: + await demo_task.kiq() + + assert exc_info.value.__cause__ is error + assert events == ["pre_send"] + assert broker.sent == [] + + +async def test_transport_error_is_wrapped_and_skips_post_send() -> None: + events: list[str] = [] + error = LifecycleError("transport") + broker = FailingSendBroker(events, error) + broker.with_middlewares(LifecycleMiddleware(events)) + + @broker.task(task_name="lifecycle.transport") + async def demo_task() -> None: + return None + + with pytest.raises(SendTaskError) as exc_info: + await demo_task.kiq() + + assert exc_info.value.__cause__ is error + assert events == ["pre_send", "kick"] + + +async def test_transport_cancellation_is_not_wrapped() -> None: + events: list[str] = [] + cancellation = asyncio.CancelledError() + broker = FailingSendBroker(events, cancellation) + broker.with_middlewares(LifecycleMiddleware(events)) + + @broker.task(task_name="lifecycle.cancel") + async def demo_task() -> None: + return None + + with pytest.raises(asyncio.CancelledError) as exc_info: + await demo_task.kiq() + + assert exc_info.value is cancellation + assert events == ["pre_send", "kick"] + + +async def test_route_configuration_error_is_not_wrapped() -> None: + source = RecordingBroker() + other = RecordingBroker(router=source.router, broker_name="other") + + @source.task(task_name="lifecycle.route") + async def demo_task() -> None: + return None + + kicker = demo_task.kicker() + message = kicker.prepare().message + + with pytest.raises(ValueError, match="either route or broker override"): + await kicker.kiq_message( + message, + broker=other, + route=TaskiqRoute(source), + ) + + +async def test_legacy_fallback_keeps_middleware_error_unwrapped() -> None: + events: list[str] = [] + error = LifecycleError("legacy.pre_send") + broker = RecordingBroker() + broker.with_middlewares( + LifecycleMiddleware(events, error=error, fail_at="pre_send"), + ) + + @broker.task(task_name="lifecycle.legacy") + async def demo_task() -> None: + return None + + del broker.router + + with pytest.raises(LifecycleError) as exc_info: + await demo_task.kiq() + + assert exc_info.value is error + assert events == ["pre_send"] + assert broker.sent == [] diff --git a/tests/routing/test_shared_tasks.py b/tests/routing/test_shared_tasks.py index e7aa596c..23220499 100644 --- a/tests/routing/test_shared_tasks.py +++ b/tests/routing/test_shared_tasks.py @@ -1,10 +1,28 @@ +from typing import Any + import pytest from taskiq import AsyncTaskiqDecoratedTask, Flow, task_builder +from taskiq.brokers.shared_broker import AsyncSharedBroker, SharedDecoratedTask from tests.routing.models import RecordingMiddleware, TracingTask from tests.utils import RecordingBroker +class StorageError(Exception): + """Marker error for broker-owned registration storage tests.""" + + +class FailingStorageBroker(RecordingBroker): + """Broker whose registration storage fails before publishing a task.""" + + def store_registered_task( + self, + task: AsyncTaskiqDecoratedTask[Any, Any], + ) -> None: + """Fail broker-owned task storage.""" + raise StorageError("storage failed") + + async def test_task_builder_can_be_registered_later() -> None: broker = RecordingBroker() @@ -114,3 +132,62 @@ def add(left: int, right: int) -> int: with pytest.raises(ValueError, match="TaskDefinition already defines"): broker.register_task(add, queue="other") + + +async def test_task_definition_uses_shared_broker_registration_contract() -> None: + task_name = "shared.global.definition" + broker = AsyncSharedBroker() + target = RecordingBroker() + broker.default_broker(target) + broker.global_task_registry.pop(task_name, None) + + @task_builder(task_name) + async def shared_task() -> None: + return None + + try: + registered = broker.register_task(shared_task) + + assert isinstance(registered, SharedDecoratedTask) + assert broker.global_task_registry[task_name] is registered + assert task_name not in broker.local_task_registry + assert broker.router.find_task(task_name) is registered + + await registered.kiq() + + assert target.sent[0][0].task_name == task_name + finally: + broker.global_task_registry.pop(task_name, None) + + +def test_storage_failure_does_not_publish_partial_router_state() -> None: + broker = FailingStorageBroker() + + @task_builder("shared.storage.failure") + def shared_task() -> None: + return None + + with pytest.raises(StorageError, match="storage failed"): + broker.register_task(shared_task) + + assert broker.router.find_task(shared_task.task_name) is None + assert not broker.router.has_route(shared_task.task_name) + assert shared_task.task_name not in broker.local_task_registry + + +def test_repeated_task_definition_binding_keeps_first_registration() -> None: + broker = RecordingBroker() + + @task_builder("shared.duplicate") + def shared_task() -> None: + return None + + registered = broker.register_task(shared_task) + bound_function_name = shared_task.original_func.__name__ + + with pytest.raises(ValueError, match="already registered"): + broker.register_task(shared_task) + + assert shared_task.original_func.__name__ == bound_function_name + assert broker.router.find_task(shared_task.task_name) is registered + assert broker.local_task_registry[shared_task.task_name] is registered diff --git a/tests/routing/test_subscriptions.py b/tests/routing/test_subscriptions.py index cc82603b..45d22a1c 100644 --- a/tests/routing/test_subscriptions.py +++ b/tests/routing/test_subscriptions.py @@ -33,6 +33,37 @@ def test_route_task_does_not_register_subscription_by_default() -> None: assert broker.get_subscribed_flows() == () +def test_subscription_snapshot_and_explicit_unsubscribe() -> None: + router = TaskiqRouter() + broker = RecordingBroker(router=router, broker_name="broker") + first_flow = Flow("events.first") + second_flow = Flow("events.second") + + first = router.subscribe(broker, first_flow, "first.task") + snapshot = router.subscriptions + second = router.subscribe(broker, second_flow, "second.task") + + assert snapshot == (first,) + assert router.subscriptions == (first, second) + assert router.unsubscribe(broker, first_flow) == first + assert router.subscriptions == (second,) + assert router.unsubscribe(broker, first_flow) is None + + +def test_unsubscribe_rejects_conflicting_flow_declaration() -> None: + router = TaskiqRouter() + broker = RecordingBroker(router=router, broker_name="broker") + registered_flow = Flow("events").with_options(durable=True) + conflicting_flow = Flow("events").with_options(durable=False) + + subscription = router.subscribe(broker, registered_flow) + + with pytest.raises(ValueError, match="different broker options"): + router.unsubscribe(broker, conflicting_flow) + + assert router.subscriptions == (subscription,) + + def test_route_task_subscribe_true_is_deprecated_shim() -> None: router = TaskiqRouter() broker = RecordingBroker(router=router, broker_name="broker") diff --git a/tests/scheduler/test_scheduler.py b/tests/scheduler/test_scheduler.py index bc37effc..a5ed5eb5 100644 --- a/tests/scheduler/test_scheduler.py +++ b/tests/scheduler/test_scheduler.py @@ -138,7 +138,10 @@ async def demo_task() -> None: kwargs={}, cron="* * * * *", ) - del router.routes["demo.task"] + removed_route = router.remove_route("demo.task") + + assert removed_route is not None + assert removed_route.broker is target_broker await scheduler.on_ready(source, scheduled_task) diff --git a/tests/test_exceptions_flow.py b/tests/test_exceptions_flow.py index 8d65194a..7676937d 100644 --- a/tests/test_exceptions_flow.py +++ b/tests/test_exceptions_flow.py @@ -1,3 +1,4 @@ +import asyncio import re import pytest @@ -11,6 +12,7 @@ UnknownTaskError, ) from taskiq.message import BrokerMessage +from taskiq.receiver import Receiver def _broker_message(task_name: str) -> BrokerMessage: @@ -51,6 +53,14 @@ async def test_shared_broker_raises_listen_error() -> None: await broker.listen() +async def test_receiver_preserves_awaitable_listen_error() -> None: + broker = AsyncSharedBroker() + receiver = Receiver(broker, max_async_tasks=1, run_startup=False) + + with pytest.raises(SharedBrokerListenError, match="Shared broker cannot listen"): + await receiver.listen(asyncio.Event()) + + def test_registering_task_in_another_broker_raises_mismatch_error() -> None: first_broker = InMemoryBroker() second_broker = InMemoryBroker() diff --git a/tests/test_requeue.py b/tests/test_requeue.py index a23f98c3..cf1fc14a 100644 --- a/tests/test_requeue.py +++ b/tests/test_requeue.py @@ -1,11 +1,59 @@ +import asyncio + import pytest -from taskiq import Context, Flow, InMemoryBroker, TaskiqDepends, TaskiqRouter +from taskiq import ( + Context, + Flow, + FlowProtocol, + InMemoryBroker, + TaskiqDepends, + TaskiqRouter, +) +from taskiq.acks import AckableMessage, AckController from taskiq.exceptions import NoResultError -from taskiq.message import TaskiqMessage +from taskiq.message import BrokerMessage, TaskiqMessage +from taskiq.receiver import Receiver from tests.utils import RecordingBroker +class RequeueError(Exception): + """Marker error for requeue lifecycle tests.""" + + +class RequeueRecordingBroker(RecordingBroker): + """Record publication order for receiver acknowledgement tests.""" + + def __init__(self, events: list[str]) -> None: + self.events = events + super().__init__() + + async def kick_to_flow( + self, + message: BrokerMessage, + flow: FlowProtocol | None = None, + ) -> None: + """Record successful requeue publication.""" + self.events.append("publish") + await super().kick_to_flow(message, flow) + + +class FailingRequeueBroker(RecordingBroker): + """Fail requeue publication with a configured exception.""" + + def __init__(self, error: BaseException) -> None: + self.error = error + super().__init__() + + async def kick_to_flow( + self, + message: BrokerMessage, + flow: FlowProtocol | None = None, + ) -> None: + """Raise the configured publication error.""" + raise self.error + + def build_message(task_name: str = "demo.task") -> TaskiqMessage: """Build a message for context requeue tests.""" return TaskiqMessage( @@ -134,3 +182,136 @@ async def test_router_requeue_accepts_explicit_flow_override() -> None: await router.requeue(message, broker=broker, flow=explicit_flow) assert broker.sent[0][1] == explicit_flow + + +@pytest.mark.parametrize( + ("ack_type", "expected_events"), + [ + pytest.param("when_received", ["ack", "publish"], id="when-received"), + pytest.param("when_executed", ["publish", "ack"], id="when-executed"), + pytest.param("when_saved", ["publish", "ack"], id="when-saved"), + pytest.param("manual", ["publish", "ack"], id="manual"), + ], +) +async def test_receiver_requeue_settles_original_after_publication( + ack_type: str, + expected_events: list[str], +) -> None: + events: list[str] = [] + broker = RequeueRecordingBroker(events) + + @broker.task(task_name="demo.requeue") + async def demo_task(context: Context = TaskiqDepends()) -> None: + await context.requeue() + + message = demo_task.kicker().with_labels(ack_type=ack_type).prepare().message + broker_message = broker.formatter.dumps(message) + receiver = Receiver(broker, max_async_tasks=1) + + await receiver.callback( + AckableMessage( + data=broker_message.message, + ack=lambda: events.append("ack"), + ), + ) + + assert events == expected_events + assert broker.sent[0][0].labels["X-Taskiq-requeue"] == "1" + + +async def test_requeue_publish_failure_keeps_original_unacked_and_unchanged() -> None: + events: list[str] = [] + broker = FailingRequeueBroker(RequeueError("publish failed")) + message = build_message() + controller = AckController(lambda: events.append("ack")) + context = Context(message, broker, controller) + + with pytest.raises(RequeueError, match="publish failed"): + await context.requeue() + + assert events == [] + assert not context.is_acked + assert controller.requeue_failed + assert "X-Taskiq-requeue" not in message.labels + + +async def test_requeue_cancellation_keeps_original_unacked_and_unchanged() -> None: + events: list[str] = [] + cancellation = asyncio.CancelledError() + broker = FailingRequeueBroker(cancellation) + message = build_message() + controller = AckController(lambda: events.append("ack")) + context = Context(message, broker, controller) + + with pytest.raises(asyncio.CancelledError) as exc_info: + await context.requeue() + + assert exc_info.value is cancellation + assert events == [] + assert not context.is_acked + assert controller.requeue_failed + assert "X-Taskiq-requeue" not in message.labels + + +@pytest.mark.parametrize("ack_type", ["when_executed", "when_saved", "manual"]) +@pytest.mark.parametrize( + "error", + [ + pytest.param(RequeueError("publish failed"), id="error"), + pytest.param(asyncio.CancelledError(), id="cancelled"), + ], +) +async def test_receiver_does_not_ack_failed_requeue( + ack_type: str, + error: BaseException, +) -> None: + events: list[str] = [] + broker = FailingRequeueBroker(error) + + @broker.task(task_name="demo.failed-requeue") + async def demo_task(context: Context = TaskiqDepends()) -> None: + await context.requeue() + + message = demo_task.kicker().with_labels(ack_type=ack_type).prepare().message + broker_message = broker.formatter.dumps(message) + receiver = Receiver(broker, max_async_tasks=1) + + await receiver.callback( + AckableMessage( + data=broker_message.message, + ack=lambda: events.append("ack"), + ), + ) + + context = broker.custom_dependency_context[Context] + assert events == [] + assert not context.is_acked + assert "X-Taskiq-requeue" not in context.message.labels + + +async def test_requeue_ack_failure_propagates_after_successful_publication() -> None: + events: list[str] = [] + broker = RequeueRecordingBroker(events) + + def fail_ack() -> None: + events.append("ack") + raise RequeueError("ack failed") + + @broker.task(task_name="demo.ack-failure") + async def demo_task(context: Context = TaskiqDepends()) -> None: + await context.requeue() + + message = demo_task.kicker().with_labels(ack_type="manual").prepare().message + broker_message = broker.formatter.dumps(message) + receiver = Receiver(broker, max_async_tasks=1) + + with pytest.raises(RequeueError, match="ack failed"): + await receiver.callback( + AckableMessage(data=broker_message.message, ack=fail_ack), + ) + + context = broker.custom_dependency_context[Context] + assert events == ["publish", "ack"] + assert len(broker.sent) == 1 + assert context.message.labels["X-Taskiq-requeue"] == "1" + assert not context.is_acked diff --git a/tests/test_task_builder.py b/tests/test_task_builder.py index 0acf3066..9dbefb95 100644 --- a/tests/test_task_builder.py +++ b/tests/test_task_builder.py @@ -1,16 +1,32 @@ +import inspect +from concurrent.futures import ProcessPoolExecutor from dataclasses import dataclass +from multiprocessing import get_context from typing import Any import pytest from pydantic import BaseModel -from taskiq import AsyncTaskiqDecoratedTask, InMemoryBroker, task_builder +from taskiq import ( + AsyncTaskiqDecoratedTask, + InMemoryBroker, + TaskDefinition, + task_builder, +) from taskiq.labels import LabelType class CustomBrokerTask(AsyncTaskiqDecoratedTask[Any, Any]): """Custom broker-level task class used by old decorator tests.""" + def broker_extension(self) -> str: + """Return a marker exposed by the broker-specific task class.""" + return self.task_name + + +class ExplicitTask(AsyncTaskiqDecoratedTask[Any, Any]): + """Task class explicitly selected by a shared task declaration.""" + @dataclass(frozen=True, slots=True) class Payload: @@ -25,6 +41,35 @@ class User(BaseModel): name: str +@task_builder("shared.process_pool") +def shared_process_pool_add(left: int, right: int) -> int: + """Add values in a spawned worker process.""" + return left + right + + +def build_process_pool_task( + task_name: str, + increment: int, +) -> TaskDefinition[..., int]: + """Build one shared closure from a reusable declaration factory.""" + + @task_builder(task_name) + def generated(value: int) -> int: + return value + increment + + return generated + + +factory_process_pool_add_one = build_process_pool_task( + "shared.factory_add_one", + 1, +) +factory_process_pool_add_two = build_process_pool_task( + "shared.factory_add_two", + 2, +) + + async def test_task_definition_call_executes_sync_and_async_functions() -> None: @task_builder("shared.add") def add(left: int, right: int) -> int: @@ -39,6 +84,105 @@ async def multiply(left: int, right: int) -> int: assert await multiply(2, 3) == 6 assert await multiply.call(2, 3) == 6 + broker = InMemoryBroker() + registered_add = broker.register_task(add) + registered_multiply = broker.register_task(multiply) + + assert not inspect.iscoroutinefunction(registered_add.original_func) + assert inspect.iscoroutinefunction(registered_multiply.original_func) + + +def test_task_definition_preserves_metadata_across_broker_bindings() -> None: + @task_builder("shared.metadata") + def calculate(left: int, right: int = 2) -> int: + """Calculate a shared result.""" + return left * right + + source = calculate.original_func + source_metadata = ( + source.__name__, + source.__qualname__, + source.__module__, + source.__doc__, + inspect.signature(source), + ) + + first_registered = InMemoryBroker().register_task(calculate) + second_registered = InMemoryBroker().register_task(calculate) + + assert calculate.__name__ == source.__name__ + assert calculate.__qualname__ == source.__qualname__ + assert calculate.__module__ == source.__module__ + assert calculate.__doc__ == source.__doc__ + assert calculate.__wrapped__ is source + assert inspect.signature(calculate) == inspect.signature(source) + assert ( + source.__name__, + source.__qualname__, + source.__module__, + source.__doc__, + inspect.signature(source), + ) == source_metadata + + for registered in (first_registered, second_registered): + assert registered.__name__ == source.__name__ + assert registered.__qualname__ == source.__qualname__ + assert registered.__module__ == source.__module__ + assert registered.__doc__ == source.__doc__ + assert registered.__wrapped__ is source + assert inspect.signature(registered) == inspect.signature(source) + + assert first_registered.original_func is second_registered.original_func + + +def test_task_definition_binding_runs_in_spawned_process_pool() -> None: + registered = InMemoryBroker().register_task(shared_process_pool_add) + + with ProcessPoolExecutor( + max_workers=1, + mp_context=get_context("spawn"), + ) as executor: + result = executor.submit(registered.original_func, 2, 3).result(timeout=10) + + assert result == 5 + + +def test_factory_task_bindings_run_in_spawned_process_pool() -> None: + first_registered = InMemoryBroker().register_task(factory_process_pool_add_one) + second_registered = InMemoryBroker().register_task(factory_process_pool_add_two) + + assert first_registered.original_func is not second_registered.original_func + assert ( + first_registered.original_func.__name__ + != second_registered.original_func.__name__ + ) + + with ProcessPoolExecutor( + max_workers=1, + mp_context=get_context("spawn"), + ) as executor: + first_result = executor.submit( + first_registered.original_func, + 10, + ).result(timeout=10) + second_result = executor.submit( + second_registered.original_func, + 10, + ).result(timeout=10) + + assert first_result == 11 + assert second_result == 12 + + +def test_factory_task_bindings_reject_duplicate_task_names() -> None: + build_process_pool_task("shared.factory_duplicate", 1) + + with pytest.raises( + ValueError, + match="Factory-generated TaskDefinitions must use unique task names", + ): + build_process_pool_task("shared.factory_duplicate", 2) + def test_task_definition_message_matches_registered_kicker_message() -> None: @task_builder( @@ -91,7 +235,7 @@ def process(payload: Payload) -> None: process.message("task-id", Payload) # type: ignore[arg-type] -def test_shared_task_without_base_cls_uses_native_task_class() -> None: +def test_shared_task_without_base_cls_uses_broker_task_class() -> None: broker = InMemoryBroker() broker.decorator_class = CustomBrokerTask @@ -106,7 +250,22 @@ async def shared_task() -> None: registered = broker.register_task(shared_task) assert isinstance(old_task, CustomBrokerTask) - assert type(registered) is AsyncTaskiqDecoratedTask + assert isinstance(registered, CustomBrokerTask) + assert registered.broker_extension() == "shared.default" + + +def test_explicit_base_cls_overrides_broker_task_class() -> None: + broker = InMemoryBroker() + broker.decorator_class = CustomBrokerTask + + @task_builder("shared.explicit", base_cls=ExplicitTask) + async def shared_task() -> None: + return None + + registered = broker.register_task(shared_task) + + assert isinstance(registered, ExplicitTask) + assert not isinstance(registered, CustomBrokerTask) @pytest.mark.parametrize( From 4a94eda2e1f9f97d0258042be0c42a71a6740b29 Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Sat, 11 Jul 2026 15:37:55 +0200 Subject: [PATCH 07/10] feat: harden multi-broker routing compatibility - manage shared router brokers through the scheduler lifecycle - reject unsupported explicit flows instead of silently discarding them - preserve legacy no-flow and in-memory dispatch behavior - add routing, lifecycle, compatibility tests, and documentation --- docs/guide/routing-and-flows.md | 17 ++++ taskiq/__init__.py | 2 + taskiq/abc/broker.py | 12 ++- taskiq/brokers/inmemory_broker.py | 8 ++ taskiq/exceptions.py | 11 +++ taskiq/scheduler/scheduler.py | 81 ++++++++++++++-- taskiq/sending.py | 7 +- tests/brokers/test_inmemory.py | 16 +++- tests/routing/test_dispatch.py | 38 +++++++- tests/scheduler/test_scheduler.py | 151 ++++++++++++++++++++++++++++++ 10 files changed, 324 insertions(+), 19 deletions(-) diff --git a/docs/guide/routing-and-flows.md b/docs/guide/routing-and-flows.md index 72119e93..7107b00b 100644 --- a/docs/guide/routing-and-flows.md +++ b/docs/guide/routing-and-flows.md @@ -10,6 +10,8 @@ This page describes the routing and flow contract on the `experiment/separate_broker` branch. The old task declaration API remains compatible. Broker-specific packages may keep using their existing send/listen implementation until they opt into flow-aware dispatch and listen-plan support. +Legacy sends without a selected flow keep working; selecting an explicit flow +requires the broker adapter to implement that capability. ::: Taskiq keeps the old task declaration API: @@ -308,6 +310,13 @@ before `on_ready(...)` runs affect that scheduled dispatch. If no router route exists for the scheduled task, Taskiq keeps the old behavior and sends through the scheduler broker. +Scheduler startup owns producer lifecycle for every broker already attached to +that router. Brokers start in registration order and shut down in reverse +order. A partial startup failure still shuts down every broker whose startup +was attempted. Configure the router before scheduler startup; brokers attached +later are not part of that lifecycle snapshot. A scheduler whose broker has no +other router members keeps the legacy single-broker lifecycle. + `with_broker(...)`, `with_route(...)` and `with_flow(...)` on a kicker are not persisted schedule route metadata. They can affect `CreatedSchedule.kiq()`, which is an immediate queued invocation helper, but they do not change the @@ -476,8 +485,13 @@ Migration guidance for broker packages: - Existing brokers that implement only `kick(message)` and `listen()` remain valid. The default `kick_to_flow(message, flow=None)` calls `kick(message)`. + If an explicit flow reaches that fallback, it raises + `UnsupportedFlowError` instead of silently sending to the broker's old + default destination. - Flow-aware brokers can override `kick_to_flow(...)` and map `flow.name` to a queue, topic, subject or stream. +- The built-in `InMemoryBroker` explicitly accepts every Flow as local routing + metadata because it has one in-process execution destination. - A successful `kick_to_flow(...)` return is the publication boundary used by `Context.requeue()` before Receiver acknowledges the original delivery. Flow-aware adapters must not return success before their normal send contract @@ -500,6 +514,9 @@ Scheduler and requeue compatibility: broker-specific flow options in this iteration. - Scheduled dispatch resolves the route late in the scheduler process through the scheduler broker's router. +- Scheduler startup starts all brokers already attached to that router and + shuts them down in reverse order, including cleanup after partial startup. + Configure broker membership before scheduler startup. - `Context.requeue()` sends through the broker currently executing the task. It does not re-resolve a task route across brokers. - Taskiq does not preserve inbound source-flow provenance yet because that data diff --git a/taskiq/__init__.py b/taskiq/__init__.py index 6ebb34ab..6d8a172d 100644 --- a/taskiq/__init__.py +++ b/taskiq/__init__.py @@ -23,6 +23,7 @@ SendTaskError, TaskiqError, TaskiqResultTimeoutError, + UnsupportedFlowError, ) from taskiq.flow import Flow, FlowIdentity, FlowProtocol, get_flow_identity from taskiq.funcs import gather @@ -81,6 +82,7 @@ "TaskiqScheduler", "TaskiqState", "TaskiqSubscription", + "UnsupportedFlowError", "ZeroMQBroker", "__version__", "async_shared_broker", diff --git a/taskiq/abc/broker.py b/taskiq/abc/broker.py index 55e4d4c9..cf37a3c8 100644 --- a/taskiq/abc/broker.py +++ b/taskiq/abc/broker.py @@ -23,7 +23,7 @@ from taskiq.acks import AckableMessage from taskiq.decor import AsyncTaskiqDecoratedTask from taskiq.events import TaskiqEvents -from taskiq.exceptions import TaskBrokerMismatchError +from taskiq.exceptions import TaskBrokerMismatchError, UnsupportedFlowError from taskiq.flow import FlowProtocol from taskiq.formatters.proxy_formatter import ProxyFormatter from taskiq.message import BrokerMessage @@ -320,11 +320,19 @@ async def kick_to_flow( 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. + stream or any other transport address. The compatibility fallback is + valid only when no explicit flow was selected. :param message: message to send. :param flow: optional transport-neutral flow. + :raises UnsupportedFlowError: if an explicit flow requires adapter + support that this broker has not implemented. """ + if flow is not None: + raise UnsupportedFlowError( + broker_type=type(self).__name__, + flow_name=flow.name, + ) await self.kick(message) @abstractmethod diff --git a/taskiq/brokers/inmemory_broker.py b/taskiq/brokers/inmemory_broker.py index d37bdd75..2ff0ca43 100644 --- a/taskiq/brokers/inmemory_broker.py +++ b/taskiq/brokers/inmemory_broker.py @@ -179,6 +179,14 @@ async def kick(self, message: BrokerMessage) -> None: self._running_tasks.add(task) task.add_done_callback(self._running_tasks.discard) + async def kick_to_flow( + self, + message: BrokerMessage, + flow: FlowProtocol | None = None, + ) -> None: + """Execute locally because every flow shares one in-memory runtime.""" + await self.kick(message) + def listen(self) -> AsyncGenerator[bytes, None]: """ Inmemory broker cannot listen. diff --git a/taskiq/exceptions.py b/taskiq/exceptions.py index 1e6fbb83..8817d3a1 100644 --- a/taskiq/exceptions.py +++ b/taskiq/exceptions.py @@ -38,6 +38,17 @@ class SendTaskError(BrokerError): __template__ = "Cannot send task to the queue" +class UnsupportedFlowError(SendTaskError): + """Error raised when a broker cannot dispatch to an explicit flow.""" + + __template__ = ( + "Broker {broker_type} does not support explicit flow {flow_name!r}. " + "The broker adapter must override kick_to_flow()." + ) + broker_type: str + flow_name: str + + class SharedBrokerSendTaskError(SendTaskError): """Error when someone tries to send task with shared broker.""" diff --git a/taskiq/scheduler/scheduler.py b/taskiq/scheduler/scheduler.py index c65e74a5..298b845b 100644 --- a/taskiq/scheduler/scheduler.py +++ b/taskiq/scheduler/scheduler.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from logging import getLogger from typing import TYPE_CHECKING, Any @@ -19,22 +21,47 @@ class TaskiqScheduler: def __init__( self, - broker: "AsyncBroker", - sources: list["ScheduleSource"], + broker: AsyncBroker, + sources: list[ScheduleSource], ) -> None: # pragma: no cover self.broker = broker self.sources = sources + self._started_brokers: list[AsyncBroker] = [] - async def startup(self) -> None: # pragma: no cover + async def startup(self) -> None: """ - This method is called on startup. + Start every producer broker the scheduler may route through. - Here you can do stuff, like creating - connections or anything you'd like. + A scheduler with a standalone broker preserves the legacy one-broker + lifecycle. When brokers share a router, the scheduler owns all of their + producer-side resources because late route resolution may select any of + them. """ - await self.broker.startup() + if self._started_brokers: + raise RuntimeError("TaskiqScheduler is already started.") + + try: + for broker in self._managed_brokers(): + broker.is_scheduler_process = True + self._started_brokers.append(broker) + await broker.startup() + except BaseException: + cleanup_error = await self._shutdown_brokers( + tuple(self._started_brokers), + ) + self._started_brokers.clear() + if cleanup_error is not None: + logger.error( + "Error while cleaning up partially started scheduler brokers.", + exc_info=( + type(cleanup_error), + cleanup_error, + cleanup_error.__traceback__, + ), + ) + raise - async def on_ready(self, source: "ScheduleSource", task: ScheduledTask) -> None: + async def on_ready(self, source: ScheduleSource, task: ScheduledTask) -> None: """ This method is called when task is ready to be enqueued. @@ -58,8 +85,42 @@ async def on_ready(self, source: "ScheduleSource", task: ScheduledTask) -> None: await maybe_awaitable(source.post_send(task)) async def shutdown(self) -> None: - """Shutdown the scheduler process.""" - await self.broker.shutdown() + """Shut down scheduler brokers in reverse startup order.""" + brokers = tuple(self._started_brokers) + if not brokers: + brokers = (self.broker,) + self._started_brokers.clear() + shutdown_error = await self._shutdown_brokers(brokers) + if shutdown_error is not None: + raise shutdown_error + + def _managed_brokers(self) -> tuple[AsyncBroker, ...]: + """Return scheduler-owned producer brokers in startup order.""" + router = getattr(self.broker, "router", None) + if isinstance(router, TaskiqRouter): + brokers = tuple(router.brokers.values()) + if brokers: + return brokers + return (self.broker,) + + @staticmethod + async def _shutdown_brokers( + brokers: tuple[AsyncBroker, ...], + ) -> BaseException | None: + """Close all brokers while preserving the first shutdown failure.""" + first_error: BaseException | None = None + for broker in reversed(brokers): + try: + await broker.shutdown() + except BaseException as exc: + if first_error is None: + first_error = exc + else: + logger.error( + "Additional error while shutting down scheduler brokers.", + exc_info=(type(exc), exc, exc.__traceback__), + ) + return first_error def _apply_scheduled_route( self, diff --git a/taskiq/sending.py b/taskiq/sending.py index 7349d3aa..9c634a32 100644 --- a/taskiq/sending.py +++ b/taskiq/sending.py @@ -3,7 +3,7 @@ from taskiq.abc.middleware import TaskiqMiddleware from taskiq.abc.result_backend import AsyncResultBackend -from taskiq.exceptions import SendTaskError +from taskiq.exceptions import SendTaskError, UnsupportedFlowError from taskiq.message import BrokerMessage, TaskiqMessage from taskiq.task import AsyncTaskiqTask from taskiq.utils import maybe_awaitable @@ -27,7 +27,8 @@ async def send_task_message( Middleware failures keep their original types. Only message formatting and transport dispatch failures are mapped to `SendTaskError`, matching the - established Taskiq kicker contract. + established Taskiq kicker contract. `UnsupportedFlowError` keeps its + specific public meaning so missing adapter capability is actionable. """ middlewares = getattr(broker, "middlewares", []) if not isinstance(middlewares, list): @@ -40,6 +41,8 @@ async def send_task_message( try: broker_message = broker.formatter.dumps(message) await send(broker_message) + except UnsupportedFlowError: + raise except Exception as exc: raise SendTaskError from exc diff --git a/tests/brokers/test_inmemory.py b/tests/brokers/test_inmemory.py index ac8f9a0b..7983aa8b 100644 --- a/tests/brokers/test_inmemory.py +++ b/tests/brokers/test_inmemory.py @@ -4,7 +4,7 @@ import pytest -from taskiq import InMemoryBroker +from taskiq import Flow, InMemoryBroker from taskiq.abc.middleware import TaskiqMiddleware from taskiq.brokers.inmemory_broker import InmemoryResultBackend from taskiq.events import TaskiqEvents @@ -75,6 +75,20 @@ async def task() -> str: assert not broker._running_tasks +async def test_inmemory_accepts_explicit_flow_as_local_metadata() -> None: + flow = Flow("local.priority") + broker = InMemoryBroker(default_flow=flow, await_inplace=True) + + @broker.task + async def task() -> str: + return flow.name + + kicked = await task.kiq() + result = await kicked.wait_result() + + assert result.return_value == "local.priority" + + async def test_cannot_listen() -> None: broker = InMemoryBroker() diff --git a/tests/routing/test_dispatch.py b/tests/routing/test_dispatch.py index fa2e74b5..b0ee4e32 100644 --- a/tests/routing/test_dispatch.py +++ b/tests/routing/test_dispatch.py @@ -1,6 +1,13 @@ import pytest -from taskiq import Flow, FlowProtocol, InMemoryBroker, TaskiqRoute, TaskiqRouter +from taskiq import ( + Flow, + FlowProtocol, + InMemoryBroker, + TaskiqRoute, + TaskiqRouter, + UnsupportedFlowError, +) from tests.routing.models import ( BrokerQueue, OldStyleRecordingBroker, @@ -121,17 +128,16 @@ async def demo_task() -> None: assert target.sent[0][1] == flow -async def test_old_style_broker_uses_kick_to_flow_fallback() -> None: +async def test_old_style_broker_uses_fallback_without_flow() -> None: router = TaskiqRouter() source = RecordingBroker(router=router, broker_name="source") target = OldStyleRecordingBroker(router=router, broker_name="target") - flow = Flow("legacy") @source.task(task_name="demo.task") async def demo_task() -> None: return None - router.route_task(demo_task, broker=target, flow=flow) + router.route_task(demo_task, broker=target) await demo_task.kiq() @@ -139,6 +145,30 @@ async def demo_task() -> None: assert target.sent[0].task_name == "demo.task" +async def test_old_style_broker_rejects_explicit_flow() -> None: + router = TaskiqRouter() + source = RecordingBroker(router=router, broker_name="source") + target = OldStyleRecordingBroker(router=router, broker_name="target") + flow = Flow("legacy") + + @source.task(task_name="demo.task") + async def demo_task() -> None: + return None + + router.route_task(demo_task, broker=target, flow=flow) + + with pytest.raises( + UnsupportedFlowError, + match="OldStyleRecordingBroker does not support explicit flow 'legacy'", + ) as exc_info: + await demo_task.kiq() + + assert exc_info.value.broker_type == "OldStyleRecordingBroker" + assert exc_info.value.flow_name == "legacy" + assert source.sent == [] + assert target.sent == [] + + async def test_kicker_route_override_wins_over_registered_route() -> None: router = TaskiqRouter() first = RecordingBroker(router=router, broker_name="first") diff --git a/tests/scheduler/test_scheduler.py b/tests/scheduler/test_scheduler.py index a5ed5eb5..f6c44210 100644 --- a/tests/scheduler/test_scheduler.py +++ b/tests/scheduler/test_scheduler.py @@ -1,5 +1,7 @@ from datetime import datetime, timezone +import pytest + from taskiq import Flow, TaskiqRouter from taskiq.abc.schedule_source import ScheduleSource from taskiq.brokers.inmemory_broker import InMemoryBroker @@ -10,6 +12,42 @@ from tests.utils import RecordingBroker +class LifecycleError(RuntimeError): + """Marker error for scheduler broker lifecycle tests.""" + + +class LifecycleBroker(RecordingBroker): + """Recording broker with observable startup and shutdown behavior.""" + + def __init__( + self, + events: list[str], + *, + router: TaskiqRouter, + broker_name: str, + startup_error: BaseException | None = None, + shutdown_error: BaseException | None = None, + ) -> None: + self.events = events + self.startup_error = startup_error + self.shutdown_error = shutdown_error + super().__init__(router=router, broker_name=broker_name) + + async def startup(self) -> None: + """Record startup and raise the configured failure.""" + self.events.append(f"startup:{self.broker_name}") + if self.startup_error is not None: + raise self.startup_error + await super().startup() + + async def shutdown(self) -> None: + """Record shutdown and raise the configured failure.""" + self.events.append(f"shutdown:{self.broker_name}") + await super().shutdown() + if self.shutdown_error is not None: + raise self.shutdown_error + + class RecordingScheduleSource(ScheduleSource): """Schedule source that records dynamically added schedules.""" @@ -38,6 +76,119 @@ def pre_send( raise ScheduledTaskCancelledError +async def test_scheduler_owns_all_router_broker_lifecycles() -> None: + events: list[str] = [] + router = TaskiqRouter() + first = LifecycleBroker(events, router=router, broker_name="first") + second = LifecycleBroker(events, router=router, broker_name="second") + scheduler = TaskiqScheduler(broker=first, sources=[]) + + await scheduler.startup() + + assert first.is_scheduler_process + assert second.is_scheduler_process + assert events == ["startup:first", "startup:second"] + + await scheduler.shutdown() + + assert events == [ + "startup:first", + "startup:second", + "shutdown:second", + "shutdown:first", + ] + + +async def test_scheduler_cleans_up_partial_broker_startup() -> None: + events: list[str] = [] + router = TaskiqRouter() + cleanup_error = LifecycleError("first cleanup failed") + first = LifecycleBroker( + events, + router=router, + broker_name="first", + shutdown_error=cleanup_error, + ) + startup_error = LifecycleError("second startup failed") + LifecycleBroker( + events, + router=router, + broker_name="second", + startup_error=startup_error, + ) + LifecycleBroker(events, router=router, broker_name="not-started") + scheduler = TaskiqScheduler(broker=first, sources=[]) + + with pytest.raises(LifecycleError) as exc_info: + await scheduler.startup() + + assert exc_info.value is startup_error + + assert events == [ + "startup:first", + "startup:second", + "shutdown:second", + "shutdown:first", + ] + + +async def test_scheduler_shutdown_closes_all_brokers_after_failure() -> None: + events: list[str] = [] + router = TaskiqRouter() + first = LifecycleBroker( + events, + router=router, + broker_name="first", + shutdown_error=LifecycleError("first shutdown failed"), + ) + second_error = LifecycleError("second shutdown failed") + LifecycleBroker( + events, + router=router, + broker_name="second", + shutdown_error=second_error, + ) + LifecycleBroker(events, router=router, broker_name="third") + scheduler = TaskiqScheduler(broker=first, sources=[]) + await scheduler.startup() + + with pytest.raises(LifecycleError) as exc_info: + await scheduler.shutdown() + + assert exc_info.value is second_error + + assert events[-3:] == [ + "shutdown:third", + "shutdown:second", + "shutdown:first", + ] + + +async def test_scheduler_rejects_duplicate_startup() -> None: + events: list[str] = [] + router = TaskiqRouter() + broker = LifecycleBroker(events, router=router, broker_name="broker") + scheduler = TaskiqScheduler(broker=broker, sources=[]) + await scheduler.startup() + + with pytest.raises(RuntimeError, match="already started"): + await scheduler.startup() + + await scheduler.shutdown() + assert events == ["startup:broker", "shutdown:broker"] + + +async def test_scheduler_shutdown_before_startup_preserves_legacy_behavior() -> None: + events: list[str] = [] + router = TaskiqRouter() + broker = LifecycleBroker(events, router=router, broker_name="broker") + scheduler = TaskiqScheduler(broker=broker, sources=[]) + + await scheduler.shutdown() + + assert events == ["shutdown:broker"] + + async def test_scheduled_task_cancelled() -> None: broker = InMemoryBroker() source = CancellingScheduleSource() From a7eb7a59fbbabf3ab98b0fdc96a633d3338e93d0 Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Sat, 11 Jul 2026 17:32:33 +0200 Subject: [PATCH 08/10] fix(routing): harden retry and lifecycle ownership - keep immediate retries on the receiving broker - protect callback cleanup from repeated cancellation - make scheduler rollback cleanup idempotent - document unsupported-flow requeue behavior --- docs/guide/routing-and-flows.md | 9 +++ taskiq/middlewares/simple_retry_middleware.py | 14 ++-- taskiq/middlewares/smart_retry_middleware.py | 1 + taskiq/receiver/receiver.py | 55 +++++++++----- taskiq/scheduler/scheduler.py | 35 ++++++++- tests/middlewares/test_simple_retry.py | 48 ++++++++++++ tests/receiver/test_runtime.py | 74 +++++++++++++++++++ tests/scheduler/test_scheduler.py | 23 ++++++ tests/test_requeue.py | 24 +++++- 9 files changed, 255 insertions(+), 28 deletions(-) diff --git a/docs/guide/routing-and-flows.md b/docs/guide/routing-and-flows.md index 7107b00b..9e9bfe3d 100644 --- a/docs/guide/routing-and-flows.md +++ b/docs/guide/routing-and-flows.md @@ -341,6 +341,9 @@ contract: - A failed or cancelled publication remains unacknowledged for `when_executed`, `when_saved` and `manual`, allowing broker redelivery. The original label is not incremented. +- `UnsupportedFlowError` follows the same failure contract. A legacy adapter + with an explicit same-broker Flow can therefore redeliver the original until + the route is corrected or the transport's dead-letter policy intervenes. - `when_received` is intentionally different: the original was acknowledged before task execution, so a later publication failure cannot restore it. - If acknowledgement itself fails after successful publication, that error is @@ -519,6 +522,12 @@ Scheduler and requeue compatibility: Configure broker membership before scheduler startup. - `Context.requeue()` sends through the broker currently executing the task. It does not re-resolve a task route across brokers. +- Immediate retries published by `SimpleRetryMiddleware` or + `SmartRetryMiddleware` stay on the broker currently executing the task, so + the retry keeps that broker's send lifecycle and result backend. When + `SmartRetryMiddleware` uses a `ScheduleSource`, it creates a normal + `ScheduledTask`; the scheduler then applies the late route-resolution rules + above. - Taskiq does not preserve inbound source-flow provenance yet because that data is not stored in `TaskiqMessage` or `Context`. diff --git a/taskiq/middlewares/simple_retry_middleware.py b/taskiq/middlewares/simple_retry_middleware.py index 9bd80f7b..d26b1bbc 100644 --- a/taskiq/middlewares/simple_retry_middleware.py +++ b/taskiq/middlewares/simple_retry_middleware.py @@ -65,11 +65,15 @@ async def on_error( if not retry_on_error: return - kicker: AsyncKicker[Any, Any] = AsyncKicker( - task_name=message.task_name, - broker=self.broker, - labels=message.labels, - ).with_task_id(message.task_id) + kicker: AsyncKicker[Any, Any] = ( + AsyncKicker( + task_name=message.task_name, + broker=self.broker, + labels=message.labels, + ) + .with_broker(self.broker) + .with_task_id(message.task_id) + ) # Getting number of previous retries. retries = int(message.labels.get("_retries", 0)) + 1 diff --git a/taskiq/middlewares/smart_retry_middleware.py b/taskiq/middlewares/smart_retry_middleware.py index 66874504..8a21d869 100644 --- a/taskiq/middlewares/smart_retry_middleware.py +++ b/taskiq/middlewares/smart_retry_middleware.py @@ -176,6 +176,7 @@ async def on_error( broker=self.broker, labels=message.labels, ) + .with_broker(self.broker) .with_task_id(message.task_id) .with_labels(_retries=retries) ) diff --git a/taskiq/receiver/receiver.py b/taskiq/receiver/receiver.py index c0904325..83abc0f7 100644 --- a/taskiq/receiver/receiver.py +++ b/taskiq/receiver/receiver.py @@ -640,7 +640,7 @@ async def runner( :param queue: queue with prefetched data. """ - cancelled = False + cancellation: asyncio.CancelledError | None = None try: while True: queued_message = await queue.get() @@ -650,13 +650,25 @@ async def runner( raise TypeError("Receiver queue contains an invalid message.") await self._start_callback(queued_message) - except asyncio.CancelledError: - cancelled = True - raise + except asyncio.CancelledError as exc: + cancellation = exc finally: - await self._drain_active_tasks(cancel_immediately=cancelled) + with anyio.CancelScope(shield=True): + cleanup_cancellation = await self._drain_active_tasks( + cancel_immediately=cancellation is not None, + ) + if cancellation is None: + cancellation = cleanup_cancellation + + if cancellation is not None: + raise cancellation logger.info("The runner is stopped.") + async def _run_owned_callback(self, message: bytes | AckableMessage) -> None: + """Run a callback outside repeated listener-scope cancellation.""" + with anyio.CancelScope(shield=True): + await self.callback(message=message, raise_err=False) + async def _start_callback(self, message: _PrefetchedMessage) -> None: """Transfer one prefetched delivery into an owned callback task.""" owns_prefetch_slot = message.owns_prefetch_slot @@ -671,7 +683,7 @@ async def _start_callback(self, message: _PrefetchedMessage) -> None: self.sem_prefetch.release() owns_prefetch_slot = False task = asyncio.create_task( - self.callback(message=message.data, raise_err=False), + self._run_owned_callback(message.data), ) except BaseException: if owns_prefetch_slot: @@ -704,11 +716,15 @@ def _on_callback_done(self, task: "asyncio.Task[Any]") -> None: ), ) - async def _drain_active_tasks(self, *, cancel_immediately: bool) -> None: + async def _drain_active_tasks( + self, + *, + cancel_immediately: bool, + ) -> asyncio.CancelledError | None: """Wait for owned callbacks and cancel them after the graceful boundary.""" tasks = set(self._active_tasks) if not tasks: - return + return None logger.info("Waiting for %d running tasks to complete...", len(tasks)) cancellation: asyncio.CancelledError | None = None @@ -733,25 +749,24 @@ async def _drain_active_tasks(self, *, cancel_immediately: bool) -> None: if cancellation is None: cancellation = cleanup_cancellation logger.info("No more tasks to wait for. Shutting down.") - if cancellation is not None: - raise cancellation + return cancellation @staticmethod async def _cancel_callback_tasks( tasks: set[asyncio.Task[Any]], ) -> asyncio.CancelledError | None: - """Cancel callback tasks and await cleanup despite one outer cancellation.""" + """Cancel callbacks once and await cleanup despite outer cancellation.""" for task in tasks: task.cancel() - try: - await asyncio.gather(*tasks, return_exceptions=True) - except asyncio.CancelledError as exc: - for task in tasks: - if not task.done(): - task.cancel() - await asyncio.gather(*tasks, return_exceptions=True) - return exc - return None + waiter = asyncio.gather(*tasks, return_exceptions=True) + cancellation: asyncio.CancelledError | None = None + while not waiter.done(): + try: + await asyncio.shield(waiter) + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc + return cancellation def _record_listen_error(self, error: BaseException) -> None: """Preserve the first transport error while allowing iterator cleanup.""" diff --git a/taskiq/scheduler/scheduler.py b/taskiq/scheduler/scheduler.py index 298b845b..19d9ab52 100644 --- a/taskiq/scheduler/scheduler.py +++ b/taskiq/scheduler/scheduler.py @@ -1,5 +1,6 @@ from __future__ import annotations +from enum import Enum, auto from logging import getLogger from typing import TYPE_CHECKING, Any @@ -16,6 +17,16 @@ logger = getLogger(__name__) +class _SchedulerLifecycle(Enum): + """Internal scheduler resource lifecycle.""" + + CREATED = auto() + STARTING = auto() + STARTED = auto() + STOPPING = auto() + STOPPED = auto() + + class TaskiqScheduler: """Scheduler class.""" @@ -27,6 +38,7 @@ def __init__( self.broker = broker self.sources = sources self._started_brokers: list[AsyncBroker] = [] + self._lifecycle = _SchedulerLifecycle.CREATED async def startup(self) -> None: """ @@ -37,9 +49,15 @@ async def startup(self) -> None: producer-side resources because late route resolution may select any of them. """ - if self._started_brokers: + if self._lifecycle in { + _SchedulerLifecycle.STARTING, + _SchedulerLifecycle.STARTED, + }: raise RuntimeError("TaskiqScheduler is already started.") + if self._lifecycle is _SchedulerLifecycle.STOPPING: + raise RuntimeError("TaskiqScheduler is shutting down.") + self._lifecycle = _SchedulerLifecycle.STARTING try: for broker in self._managed_brokers(): broker.is_scheduler_process = True @@ -50,6 +68,7 @@ async def startup(self) -> None: tuple(self._started_brokers), ) self._started_brokers.clear() + self._lifecycle = _SchedulerLifecycle.STOPPED if cleanup_error is not None: logger.error( "Error while cleaning up partially started scheduler brokers.", @@ -60,6 +79,7 @@ async def startup(self) -> None: ), ) raise + self._lifecycle = _SchedulerLifecycle.STARTED async def on_ready(self, source: ScheduleSource, task: ScheduledTask) -> None: """ @@ -86,11 +106,22 @@ async def on_ready(self, source: ScheduleSource, task: ScheduledTask) -> None: async def shutdown(self) -> None: """Shut down scheduler brokers in reverse startup order.""" + if self._lifecycle is _SchedulerLifecycle.STOPPED: + return + if self._lifecycle is _SchedulerLifecycle.STARTING: + raise RuntimeError("TaskiqScheduler is still starting.") + if self._lifecycle is _SchedulerLifecycle.STOPPING: + raise RuntimeError("TaskiqScheduler is already shutting down.") + brokers = tuple(self._started_brokers) if not brokers: brokers = (self.broker,) self._started_brokers.clear() - shutdown_error = await self._shutdown_brokers(brokers) + self._lifecycle = _SchedulerLifecycle.STOPPING + try: + shutdown_error = await self._shutdown_brokers(brokers) + finally: + self._lifecycle = _SchedulerLifecycle.STOPPED if shutdown_error is not None: raise shutdown_error diff --git a/tests/middlewares/test_simple_retry.py b/tests/middlewares/test_simple_retry.py index 783d98b7..5fd57739 100644 --- a/tests/middlewares/test_simple_retry.py +++ b/tests/middlewares/test_simple_retry.py @@ -3,10 +3,12 @@ import pytest +from taskiq import SmartRetryMiddleware, TaskiqRouter from taskiq.formatters.json_formatter import JSONFormatter from taskiq.message import TaskiqMessage from taskiq.middlewares.simple_retry_middleware import SimpleRetryMiddleware from taskiq.result import TaskiqResult +from tests.utils import RecordingBroker @pytest.fixture @@ -73,3 +75,49 @@ async def test_max_retries(broker: AsyncMock) -> None: Exception(), ) broker.kick.assert_not_called() + + +@pytest.mark.parametrize( + "middleware", + [ + pytest.param(SimpleRetryMiddleware(), id="simple"), + pytest.param(SmartRetryMiddleware(default_delay=0), id="smart"), + ], +) +async def test_retry_keeps_current_broker_when_route_points_elsewhere( + middleware: SimpleRetryMiddleware | SmartRetryMiddleware, +) -> None: + router = TaskiqRouter() + current_broker = RecordingBroker( + router=router, + broker_name="current", + ) + routed_broker = RecordingBroker( + router=router, + broker_name="routed", + ) + current_broker.add_middlewares(middleware) + router.route_task("meme", broker=routed_broker) + result = TaskiqResult(is_err=True, return_value=None, execution_time=0.0) + + await middleware.on_error( + TaskiqMessage( + task_id="test_id", + task_name="meme", + labels={ + "max_retries": 2, + "retry_on_error": True, + }, + args=[], + kwargs={}, + ), + result, + RuntimeError("retry"), + ) + + assert routed_broker.sent == [] + assert len(current_broker.sent) == 1 + retried_message, retried_flow = current_broker.sent[0] + assert retried_message.task_name == "meme" + assert retried_message.labels["_retries"] == "1" + assert retried_flow is None diff --git a/tests/receiver/test_runtime.py b/tests/receiver/test_runtime.py index b14ed599..8c3d5c7c 100644 --- a/tests/receiver/test_runtime.py +++ b/tests/receiver/test_runtime.py @@ -719,6 +719,80 @@ async def blocked_task() -> None: assert cleaned_before_runtime_returned +async def test_runtime_hard_cancel_awaits_multiphase_callback_cleanup( + runtime_executor: ThreadPoolExecutor, + caplog: pytest.LogCaptureFixture, +) -> None: + events: list[str] = [] + cleanup_started = asyncio.Event() + cleanup_finished = asyncio.Event() + cleanup_interrupted = asyncio.Event() + release_cleanup = asyncio.Event() + callback_started = asyncio.Event() + + class CleanupOrderBroker(RuntimeBroker): + async def shutdown(self) -> None: + self.events.append("broker.shutdown.started") + await super().shutdown() + + broker = CleanupOrderBroker( + router=TaskiqRouter(), + broker_name="listener", + events=events, + ) + + @broker.task(task_name="runtime.multiphase_callback_cleanup") + async def blocked_task() -> None: + callback_started.set() + try: + await asyncio.Event().wait() + finally: + cleanup_started.set() + try: + await release_cleanup.wait() + for _ in range(20): + await asyncio.sleep(0) + except asyncio.CancelledError: + cleanup_interrupted.set() + raise + events.append("callback.cleanup.finished") + cleanup_finished.set() + + finish_event = asyncio.Event() + runtime = build_runtime( + (broker,), + runtime_executor, + options=runtime_options( + wait_tasks_timeout=0.01, + shutdown_timeout=0.05, + ), + ) + caplog.set_level(logging.WARNING, logger="taskiq.worker.runtime") + runtime_task = asyncio.create_task(runtime.run(finish_event)) + await broker.listen_started.wait() + await blocked_task.kiq() + await callback_started.wait() + + finish_event.set() + await cleanup_started.wait() + + async def wait_for_hard_cancel() -> None: + while not any( + "did not stop" in record.getMessage() for record in caplog.records + ): + await asyncio.sleep(0) + + await asyncio.wait_for(wait_for_hard_cancel(), timeout=1) + release_cleanup.set() + await asyncio.wait_for(runtime_task, timeout=1) + + assert cleanup_finished.is_set() + assert not cleanup_interrupted.is_set() + assert events.index("callback.cleanup.finished") < events.index( + "broker.shutdown.started", + ) + + async def test_runtime_continues_cleanup_after_broker_shutdown_failure( runtime_executor: ThreadPoolExecutor, caplog: pytest.LogCaptureFixture, diff --git a/tests/scheduler/test_scheduler.py b/tests/scheduler/test_scheduler.py index f6c44210..0df22be1 100644 --- a/tests/scheduler/test_scheduler.py +++ b/tests/scheduler/test_scheduler.py @@ -132,6 +132,29 @@ async def test_scheduler_cleans_up_partial_broker_startup() -> None: ] +async def test_scheduler_shutdown_after_failed_startup_does_not_repeat_cleanup() -> ( + None +): + events: list[str] = [] + router = TaskiqRouter() + first = LifecycleBroker(events, router=router, broker_name="first") + LifecycleBroker( + events, + router=router, + broker_name="second", + startup_error=LifecycleError("second startup failed"), + ) + scheduler = TaskiqScheduler(broker=first, sources=[]) + + with pytest.raises(LifecycleError, match="second startup failed"): + await scheduler.startup() + + events_after_rollback = list(events) + await scheduler.shutdown() + + assert events == events_after_rollback + + async def test_scheduler_shutdown_closes_all_brokers_after_failure() -> None: events: list[str] = [] router = TaskiqRouter() diff --git a/tests/test_requeue.py b/tests/test_requeue.py index cf1fc14a..52afe0ce 100644 --- a/tests/test_requeue.py +++ b/tests/test_requeue.py @@ -11,9 +11,10 @@ TaskiqRouter, ) from taskiq.acks import AckableMessage, AckController -from taskiq.exceptions import NoResultError +from taskiq.exceptions import NoResultError, UnsupportedFlowError from taskiq.message import BrokerMessage, TaskiqMessage from taskiq.receiver import Receiver +from tests.routing.models import OldStyleRecordingBroker from tests.utils import RecordingBroker @@ -184,6 +185,27 @@ async def test_router_requeue_accepts_explicit_flow_override() -> None: assert broker.sent[0][1] == explicit_flow +async def test_requeue_unsupported_flow_keeps_original_unacked() -> None: + router = TaskiqRouter() + broker = OldStyleRecordingBroker(router=router, broker_name="legacy") + flow = Flow("unsupported") + message = build_message() + ack_events: list[str] = [] + controller = AckController(lambda: ack_events.append("ack")) + context = Context(message, broker, controller) + router.route_task(message.task_name, broker=broker, flow=flow) + + with pytest.raises(UnsupportedFlowError) as exc_info: + await context.requeue() + + assert exc_info.value.flow_name == flow.name + assert broker.sent == [] + assert ack_events == [] + assert not context.is_acked + assert controller.requeue_failed + assert "X-Taskiq-requeue" not in message.labels + + @pytest.mark.parametrize( ("ack_type", "expected_events"), [ From d89d18692f290a8af8507c77692fa5e51ccbee34 Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Sat, 11 Jul 2026 21:17:39 +0200 Subject: [PATCH 09/10] fix(routing): harden lifecycle coverage and callable task binding - cover receiver and worker runtime cancellation and shutdown paths - validate scheduler, router, registry, CLI, and broker lifecycle contracts - support callable-object shared tasks during broker binding - organize runtime fault tests into focused modules --- taskiq/abc/broker.py | 17 +- tests/abc/test_broker.py | 12 + tests/brokers/test_inmemory.py | 111 +++++- tests/cli/worker/test_broker_selection.py | 123 ++++++- tests/receiver/conftest.py | 11 + tests/receiver/runtime_support.py | 133 +++++++ tests/receiver/test_lifecycle_faults.py | 422 ++++++++++++++++++++++ tests/receiver/test_runtime.py | 147 +------- tests/receiver/test_runtime_faults.py | 262 ++++++++++++++ tests/routing/test_dispatch.py | 47 +++ tests/routing/test_router_core.py | 52 +++ tests/routing/test_subscriptions.py | 16 + tests/scheduler/test_scheduler.py | 53 +++ tests/test_task_builder.py | 127 +++++++ 14 files changed, 1395 insertions(+), 138 deletions(-) create mode 100644 tests/receiver/conftest.py create mode 100644 tests/receiver/runtime_support.py create mode 100644 tests/receiver/test_lifecycle_faults.py create mode 100644 tests/receiver/test_runtime_faults.py diff --git a/taskiq/abc/broker.py b/taskiq/abc/broker.py index cf37a3c8..45da7890 100644 --- a/taskiq/abc/broker.py +++ b/taskiq/abc/broker.py @@ -4,7 +4,7 @@ from abc import ABC, abstractmethod from collections import defaultdict from collections.abc import AsyncGenerator, Awaitable, Callable -from functools import wraps +from functools import WRAPPER_ASSIGNMENTS, update_wrapper from logging import getLogger from typing import ( TYPE_CHECKING, @@ -13,6 +13,7 @@ ParamSpec, TypeAlias, TypeVar, + cast, get_type_hints, overload, ) @@ -532,12 +533,20 @@ def _decorate_task( ): raise TypeError("base_cls must be a subclass of AsyncTaskiqDecoratedTask.") - sign = get_type_hints(metadata_source) + hint_source = metadata_source + if not hasattr(hint_source, "__annotations__"): + hint_source = cast(Any, metadata_source).__call__ + sign = get_type_hints(hint_source) return_type = None if "return" in sign: return_type = sign["return"] - decorated_task = wraps(metadata_source)( + assigned_metadata = tuple( + attribute + for attribute in WRAPPER_ASSIGNMENTS + if hasattr(metadata_source, attribute) + ) + decorated_task = update_wrapper( task_cls( broker=self, original_func=func, @@ -545,6 +554,8 @@ def _decorate_task( task_name=real_task_name, return_type=return_type, # type: ignore ), + metadata_source, + assigned=assigned_metadata, ) if register: diff --git a/tests/abc/test_broker.py b/tests/abc/test_broker.py index 636f9576..83ab9d47 100644 --- a/tests/abc/test_broker.py +++ b/tests/abc/test_broker.py @@ -68,6 +68,18 @@ async def test_func() -> None: } +def test_decorator_generates_unique_lambda_task_names() -> None: + """Anonymous tasks must not collide in the broker task registry.""" + broker = _TestBroker() + + first_task = broker.task(lambda: None) + second_task = broker.task(lambda: None) + + assert ":lambda_" in first_task.task_name + assert ":lambda_" in second_task.task_name + assert first_task.task_name != second_task.task_name + + def test_kicker_labels_modification() -> None: """Test that using kicker.with_labels doesn't modify task's labels globally.""" broker = _TestBroker() diff --git a/tests/brokers/test_inmemory.py b/tests/brokers/test_inmemory.py index 7983aa8b..11a2e096 100644 --- a/tests/brokers/test_inmemory.py +++ b/tests/brokers/test_inmemory.py @@ -1,10 +1,11 @@ import asyncio +import logging import uuid from typing import Any import pytest -from taskiq import Flow, InMemoryBroker +from taskiq import Flow, InMemoryBroker, TaskiqMessage from taskiq.abc.middleware import TaskiqMiddleware from taskiq.brokers.inmemory_broker import InmemoryResultBackend from taskiq.events import TaskiqEvents @@ -39,16 +40,18 @@ async def shutdown(self) -> None: class RecordingResultBackend(InmemoryResultBackend[Any]): - """Record result backend lifecycle calls and optional startup failure.""" + """Record result backend lifecycle calls and optional failures.""" def __init__( self, events: list[str], startup_error: Exception | None = None, + shutdown_error: Exception | None = None, ) -> None: super().__init__() self.events = events self.startup_error = startup_error + self.shutdown_error = shutdown_error async def startup(self) -> None: """Record startup and raise the configured error.""" @@ -59,6 +62,20 @@ async def startup(self) -> None: async def shutdown(self) -> None: """Record shutdown.""" self.events.append("backend.shutdown") + if self.shutdown_error is not None: + raise self.shutdown_error + + +class FailingExecutionMiddleware(TaskiqMiddleware): + """Raise before execution to expose background callback failures.""" + + def __init__(self, error: Exception) -> None: + super().__init__() + self.error = error + + def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage: + """Raise the configured callback failure.""" + raise self.error async def test_inmemory_success() -> None: @@ -197,6 +214,48 @@ async def test_shutdown_closes_remaining_resources_after_middleware_failure() -> broker.executor.submit(int) +async def test_shutdown_keeps_first_failure_and_closes_every_resource( + caplog: pytest.LogCaptureFixture, +) -> None: + events: list[str] = [] + event_error = LifecycleError("event shutdown failed") + middleware_error = LifecycleError("middleware shutdown failed") + backend_error = LifecycleError("backend shutdown failed") + broker = InMemoryBroker() + broker.with_middlewares( + RecordingLifecycleMiddleware(events, shutdown_error=middleware_error), + ) + broker.with_result_backend( + RecordingResultBackend(events, shutdown_error=backend_error), + ) + + @broker.on_event(TaskiqEvents.CLIENT_SHUTDOWN) + def fail_client_shutdown(state: TaskiqState) -> None: + events.append("client.shutdown") + raise event_error + + @broker.on_event(TaskiqEvents.WORKER_SHUTDOWN) + def record_worker_shutdown(state: TaskiqState) -> None: + events.append("worker.shutdown") + + caplog.set_level(logging.ERROR, logger="taskiq") + await broker.startup() + + with pytest.raises(LifecycleError) as exc_info: + await broker.shutdown() + + assert exc_info.value is event_error + assert events == [ + "middleware.startup", + "backend.startup", + "client.shutdown", + "worker.shutdown", + "middleware.shutdown", + "backend.shutdown", + ] + assert caplog.text.count("Additional error while shutting down") == 2 + + async def test_shutdown() -> None: broker = InMemoryBroker() test_value = uuid.uuid4().hex @@ -265,6 +324,54 @@ async def test_task() -> None: assert not broker._running_tasks +async def test_wait_all_propagates_task_failure() -> None: + task_error = LifecycleError("in-memory task failed") + broker = InMemoryBroker() + broker.with_middlewares(FailingExecutionMiddleware(task_error)) + + @broker.task + async def failing_task() -> None: + return None + + await failing_task.kiq() + + with pytest.raises(LifecycleError) as exc_info: + await broker.wait_all() + + assert exc_info.value is task_error + assert not broker._running_tasks + + +async def test_shutdown_cleans_resources_after_task_failure() -> None: + events: list[str] = [] + task_error = LifecycleError("in-memory task failed") + broker = InMemoryBroker() + broker.with_middlewares(RecordingLifecycleMiddleware(events)) + broker.with_middlewares(FailingExecutionMiddleware(task_error)) + broker.with_result_backend(RecordingResultBackend(events)) + + @broker.task + async def failing_task() -> None: + return None + + await broker.startup() + await failing_task.kiq() + + with pytest.raises(LifecycleError) as exc_info: + await broker.shutdown() + + assert exc_info.value is task_error + assert events == [ + "middleware.startup", + "backend.startup", + "middleware.shutdown", + "backend.shutdown", + ] + assert not broker._running_tasks + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + broker.executor.submit(int) + + async def test_shutdown_waits_for_running_tasks_before_resource_cleanup() -> None: task_started = asyncio.Event() finish_task = asyncio.Event() diff --git a/tests/cli/worker/test_broker_selection.py b/tests/cli/worker/test_broker_selection.py index 5b03bed1..fcce7936 100644 --- a/tests/cli/worker/test_broker_selection.py +++ b/tests/cli/worker/test_broker_selection.py @@ -1,10 +1,15 @@ +import asyncio import logging +from concurrent.futures import Executor import pytest from taskiq import InMemoryBroker, TaskiqRouter +from taskiq.acks import AcknowledgeType from taskiq.cli.worker.args import WorkerArgs -from taskiq.cli.worker.run import get_broker_selection, shutdown_broker +from taskiq.cli.worker.run import get_broker_selection, shutdown_broker, start_listen +from taskiq.receiver import Receiver +from taskiq.receiver.runtime import BrokerSelection, WorkerRuntimeOptions from taskiq.warnings import TaskiqDeprecationWarning @@ -17,6 +22,103 @@ def test_get_broker_selection_preserves_direct_single_broker() -> None: assert not selection.explicit_group +def test_start_listen_builds_runtime_from_worker_arguments( + monkeypatch: pytest.MonkeyPatch, +) -> None: + broker = InMemoryBroker() + args = WorkerArgs( + broker=broker, + modules=["application.tasks"], + tasks_pattern=("**/jobs.py",), + fs_discover=True, + configure_logging=False, + max_threadpool_threads=2, + no_parse=True, + shutdown_timeout=0.7, + max_async_tasks=7, + max_async_tasks_jitter=2, + receiver_arg=[("dependency", "value")], + max_prefetch=3, + no_propagate_errors=True, + ack_type=AcknowledgeType.WHEN_RECEIVED, + max_tasks_per_child=11, + wait_tasks_timeout=0.5, + ) + loop = asyncio.new_event_loop() + captured_selection: BrokerSelection | None = None + captured_receiver_type: type[Receiver] | None = None + captured_executor: Executor | None = None + captured_options: WorkerRuntimeOptions | None = None + runtime_ran = False + import_calls: list[tuple[list[str], tuple[str, ...], bool]] = [] + + class RecordingRuntime: + def __init__( + self, + *, + selection: BrokerSelection, + receiver_type: type[Receiver], + executor: Executor, + options: WorkerRuntimeOptions, + ) -> None: + nonlocal captured_selection + nonlocal captured_receiver_type + nonlocal captured_executor + nonlocal captured_options + captured_selection = selection + captured_receiver_type = receiver_type + captured_executor = executor + captured_options = options + + async def run(self, finish_event: asyncio.Event) -> None: + nonlocal runtime_ran + runtime_ran = True + finish_event.set() + + def record_import_tasks( + modules: list[str], + task_pattern: tuple[str, ...], + fs_discover: bool, + ) -> None: + import_calls.append((modules, task_pattern, fs_discover)) + + monkeypatch.setattr("taskiq.cli.worker.run.uvloop", None) + monkeypatch.setattr("taskiq.cli.worker.run.signal.signal", lambda *args: None) + monkeypatch.setattr("taskiq.cli.worker.run.asyncio.new_event_loop", lambda: loop) + monkeypatch.setattr("taskiq.cli.worker.run.import_tasks", record_import_tasks) + monkeypatch.setattr( + "taskiq.cli.worker.run.get_receiver_type", + lambda args: Receiver, + ) + monkeypatch.setattr("taskiq.cli.worker.run.WorkerRuntime", RecordingRuntime) + + try: + start_listen(args) + finally: + asyncio.set_event_loop(None) + loop.close() + + assert captured_selection is not None + assert captured_selection.listeners == (broker,) + assert broker.is_worker_process + assert captured_receiver_type is Receiver + assert captured_executor is not None + assert captured_options == WorkerRuntimeOptions( + validate_params=False, + max_async_tasks=7, + max_async_tasks_jitter=2, + max_prefetch=3, + propagate_exceptions=False, + ack_type=AcknowledgeType.WHEN_RECEIVED, + max_tasks_to_execute=11, + wait_tasks_timeout=0.5, + shutdown_timeout=0.7, + receiver_kwargs={"dependency": "value"}, + ) + assert import_calls == [(["application.tasks"], ("**/jobs.py",), True)] + assert runtime_ran + + @pytest.mark.parametrize("use_factory", [False, True]) def test_get_broker_selection_preserves_imported_single_broker( monkeypatch: pytest.MonkeyPatch, @@ -67,6 +169,25 @@ async def shutdown(self) -> None: assert diagnostic not in caplog.text +async def test_legacy_shutdown_broker_helper_preserves_cancellation() -> None: + shutdown_error = asyncio.CancelledError("adapter shutdown cancelled") + + class CancelledShutdownBroker(InMemoryBroker): + async def shutdown(self) -> None: + raise shutdown_error + + broker = CancelledShutdownBroker() + + with ( + pytest.warns( + TaskiqDeprecationWarning, + match="WorkerRuntime", + ), + pytest.raises(asyncio.CancelledError), + ): + await shutdown_broker(broker, timeout=1) + + @pytest.mark.parametrize("use_factory", [False, True]) def test_get_broker_selection_imports_explicit_listener_sequence( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/receiver/conftest.py b/tests/receiver/conftest.py new file mode 100644 index 00000000..168b7dce --- /dev/null +++ b/tests/receiver/conftest.py @@ -0,0 +1,11 @@ +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor + +import pytest + + +@pytest.fixture +def runtime_executor() -> Iterator[ThreadPoolExecutor]: + """Provide the process-wide sync executor owned by a worker runtime.""" + with ThreadPoolExecutor(max_workers=2) as executor: + yield executor diff --git a/tests/receiver/runtime_support.py b/tests/receiver/runtime_support.py new file mode 100644 index 00000000..31a0bd0a --- /dev/null +++ b/tests/receiver/runtime_support.py @@ -0,0 +1,133 @@ +import asyncio +from collections.abc import AsyncGenerator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace +from typing import Any + +from taskiq import AsyncBroker, BrokerMessage, TaskiqRouter +from taskiq.abc.middleware import TaskiqMiddleware +from taskiq.receiver import Receiver +from taskiq.receiver.runtime import ( + WorkerRuntime, + WorkerRuntimeOptions, + normalize_broker_source, +) + + +class RuntimeTestError(Exception): + """Marker exception for worker runtime failures.""" + + +class FatalRuntimeTestError(BaseException): + """Marker fatal signal that runtime cleanup must not swallow.""" + + +class RuntimeBroker(AsyncBroker): + """Controllable broker used for multi-broker lifecycle tests.""" + + def __init__( + self, + *, + router: TaskiqRouter, + broker_name: str, + events: list[str], + startup_error: BaseException | None = None, + listen_error: BaseException | None = None, + shutdown_error: BaseException | None = None, + ) -> None: + self.events = events + self.startup_error = startup_error + self.listen_error = listen_error + self.shutdown_error = shutdown_error + self.listen_started = asyncio.Event() + self.release_listener = asyncio.Event() + self.iterator_closed = asyncio.Event() + self.started = False + self.queue: asyncio.Queue[bytes] = asyncio.Queue() + super().__init__(router=router, broker_name=broker_name) + + async def startup(self) -> None: + """Record startup and optionally fail after core resources start.""" + self.events.append( + f"{self.broker_name}.startup." + f"{'worker' if self.is_worker_process else 'client'}", + ) + await super().startup() + self.started = True + if self.startup_error is not None: + raise self.startup_error + + async def shutdown(self) -> None: + """Record shutdown and close core resources.""" + self.events.append(f"{self.broker_name}.shutdown") + await super().shutdown() + self.started = False + if self.shutdown_error is not None: + raise self.shutdown_error + + async def kick(self, message: BrokerMessage) -> None: + """Queue an encoded message for this listener.""" + if not self.started: + raise RuntimeTestError(f"{self.broker_name} kicked before startup") + self.events.append(f"{self.broker_name}.kick") + await self.queue.put(message.message) + + async def listen(self) -> AsyncGenerator[bytes, None]: + """Listen until cancelled or a configured failure is released.""" + self.events.append(f"{self.broker_name}.listen") + self.listen_started.set() + try: + if self.listen_error is not None: + await self.release_listener.wait() + raise self.listen_error + while True: + yield await self.queue.get() + finally: + self.events.append(f"{self.broker_name}.iterator_closed") + self.iterator_closed.set() + + +class PrefetchRecordingMiddleware(TaskiqMiddleware): + """Record broker deliveries admitted into the shared prefetch budget.""" + + def __init__(self, broker_name: str, prefetched: list[str]) -> None: + super().__init__() + self.broker_name = broker_name + self.prefetched = prefetched + + def on_prefetch_queue_add(self) -> None: + """Record one delivery accepted by a Receiver prefetch queue.""" + self.prefetched.append(self.broker_name) + + +def runtime_options(**overrides: Any) -> WorkerRuntimeOptions: + """Build runtime options with focused per-test overrides.""" + options = WorkerRuntimeOptions( + validate_params=True, + max_async_tasks=10, + max_async_tasks_jitter=0, + max_prefetch=0, + propagate_exceptions=True, + ack_type=None, + max_tasks_to_execute=None, + wait_tasks_timeout=1, + shutdown_timeout=1, + receiver_kwargs={}, + ) + return replace(options, **overrides) + + +def build_runtime( + brokers: object, + executor: ThreadPoolExecutor, + *, + options: WorkerRuntimeOptions | None = None, + receiver_type: type[Receiver] = Receiver, +) -> WorkerRuntime: + """Normalize broker configuration and build the runtime under test.""" + return WorkerRuntime( + selection=normalize_broker_source(brokers), + receiver_type=receiver_type, + executor=executor, + options=options or runtime_options(), + ) diff --git a/tests/receiver/test_lifecycle_faults.py b/tests/receiver/test_lifecycle_faults.py new file mode 100644 index 00000000..87739676 --- /dev/null +++ b/tests/receiver/test_lifecycle_faults.py @@ -0,0 +1,422 @@ +import asyncio +import logging +from collections.abc import AsyncGenerator, Callable +from concurrent.futures import ThreadPoolExecutor +from typing import cast + +import pytest + +from taskiq import BrokerMessage, TaskiqMessage +from taskiq.abc.broker import AckableMessage, AsyncBroker +from taskiq.abc.middleware import TaskiqMiddleware +from taskiq.receiver import Receiver +from taskiq.receiver.receiver import _PrefetchedMessage +from taskiq.receiver.runtime import ReceiverRuntimeState + + +class ReceiverLifecycleError(RuntimeError): + """Marker error for listener and middleware lifecycle failures.""" + + +class ListenerBroker(AsyncBroker): + """Broker adapter with a controllable listener implementation.""" + + def __init__(self, listener_factory: Callable[[], object]) -> None: + self.listener_factory = listener_factory + super().__init__() + + async def kick(self, message: BrokerMessage) -> None: + """Ignore outbound messages in listener-only tests.""" + + def listen(self) -> AsyncGenerator[bytes, None]: + """Return the configured listener object for Receiver validation.""" + return cast(AsyncGenerator[bytes, None], self.listener_factory()) + + +def encoded_message(broker: AsyncBroker, task_name: str) -> bytes: + """Build one valid transport payload for a registered task.""" + return broker.formatter.dumps( + TaskiqMessage( + task_id="task-id", + task_name=task_name, + labels={}, + args=[], + kwargs={}, + ), + ).message + + +async def empty_listener() -> AsyncGenerator[bytes, None]: + """Return a valid listener that reaches the end of its stream.""" + if False: # pragma: no branch + yield b"" + + +async def test_receiver_accepts_finite_broker_listener() -> None: + broker = ListenerBroker(empty_listener) + with ThreadPoolExecutor(max_workers=1) as executor: + receiver = Receiver(broker, executor=executor, run_startup=False) + + await receiver.listen(asyncio.Event()) + + +async def test_single_receiver_stops_after_legacy_task_limit() -> None: + executed = 0 + + async def finite_listener() -> AsyncGenerator[bytes, None]: + yield encoded_message(broker, "receiver.legacy-limit") + yield encoded_message(broker, "receiver.legacy-limit") + + broker = ListenerBroker(finite_listener) + + @broker.task(task_name="receiver.legacy-limit") + async def limited_task() -> None: + nonlocal executed + executed += 1 + + with ThreadPoolExecutor(max_workers=1) as executor: + receiver = Receiver( + broker, + executor=executor, + run_startup=False, + max_tasks_to_execute=1, + ) + + await receiver.listen(asyncio.Event()) + + assert executed == 1 + + +async def test_receiver_rejects_non_generator_listener() -> None: + broker = ListenerBroker(lambda: ()) + with ThreadPoolExecutor(max_workers=1) as executor: + receiver = Receiver(broker, executor=executor, run_startup=False) + + with pytest.raises( + TypeError, + match=r"Broker\.listen\(\) must return an async generator", + ): + await receiver.listen(asyncio.Event()) + + +def test_receiver_rejects_negative_prefetch_limit() -> None: + broker = ListenerBroker(empty_listener) + with ( + ThreadPoolExecutor(max_workers=1) as executor, + pytest.raises( + ValueError, + match="max_prefetch cannot be negative", + ), + ): + Receiver( + broker, + executor=executor, + run_startup=False, + max_prefetch=-1, + ) + + +async def test_receiver_drains_delivery_completed_during_listener_stop() -> None: + read_started = asyncio.Event() + finish_event = asyncio.Event() + executed = asyncio.Event() + hook_error = ReceiverLifecycleError("late prefetch hook failed") + + async def late_listener() -> AsyncGenerator[bytes, None]: + read_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + yield encoded_message(broker, "receiver.late-delivery") + + class FailingPrefetchMiddleware(TaskiqMiddleware): + def on_prefetch_queue_add(self) -> None: + raise hook_error + + broker = ListenerBroker(late_listener) + broker.with_middlewares(FailingPrefetchMiddleware()) + + @broker.task(task_name="receiver.late-delivery") + async def late_delivery_task() -> None: + executed.set() + + with ThreadPoolExecutor(max_workers=1) as executor: + receiver = Receiver(broker, executor=executor, run_startup=False) + runtime_state = ReceiverRuntimeState( + finish_event=finish_event, + max_async_tasks=1, + max_async_tasks_jitter=0, + max_prefetch=0, + max_tasks_to_execute=None, + ) + receiver.attach_runtime_state(runtime_state) + listener_task = asyncio.create_task(receiver.listen(finish_event)) + await read_started.wait() + + finish_event.set() + with pytest.raises(ReceiverLifecycleError) as exc_info: + await listener_task + + assert exc_info.value is hook_error + assert executed.is_set() + assert runtime_state.fetched_tasks == 1 + + +async def test_receiver_preserves_pending_read_failure_during_close() -> None: + read_started = asyncio.Event() + finish_event = asyncio.Event() + read_error = ReceiverLifecycleError("pending transport read failed") + + async def failing_listener() -> AsyncGenerator[bytes, None]: + read_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + raise read_error from None + if False: # pragma: no branch + yield b"" + + broker = ListenerBroker(failing_listener) + with ThreadPoolExecutor(max_workers=1) as executor: + receiver = Receiver(broker, executor=executor, run_startup=False) + listener_task = asyncio.create_task(receiver.listen(finish_event)) + await read_started.wait() + + finish_event.set() + with pytest.raises(ReceiverLifecycleError) as exc_info: + await listener_task + + assert exc_info.value is read_error + + +async def test_receiver_preserves_iterator_close_failure_after_delivery() -> None: + finish_event = asyncio.Event() + executed = asyncio.Event() + close_error = ReceiverLifecycleError("listener close failed") + + class StopAfterPrefetchMiddleware(TaskiqMiddleware): + def on_prefetch_queue_add(self) -> None: + finish_event.set() + + async def close_failing_listener() -> AsyncGenerator[bytes, None]: + try: + yield encoded_message(broker, "receiver.close-failure") + finally: + raise close_error + + broker = ListenerBroker(close_failing_listener) + broker.with_middlewares(StopAfterPrefetchMiddleware()) + + @broker.task(task_name="receiver.close-failure") + async def delivered_task() -> None: + executed.set() + + with ThreadPoolExecutor(max_workers=1) as executor: + receiver = Receiver(broker, executor=executor, run_startup=False) + + with pytest.raises(ReceiverLifecycleError) as exc_info: + await receiver.listen(finish_event) + + assert exc_info.value is close_error + assert executed.is_set() + + +async def test_prefetch_hook_failure_releases_owned_capacity() -> None: + hook_error = ReceiverLifecycleError("prefetch remove failed") + + class FailingRemoveMiddleware(TaskiqMiddleware): + def on_prefetch_queue_remove(self) -> None: + raise hook_error + + broker = ListenerBroker(empty_listener) + broker.with_middlewares(FailingRemoveMiddleware()) + with ThreadPoolExecutor(max_workers=1) as executor: + receiver = Receiver( + broker, + executor=executor, + run_startup=False, + max_prefetch=0, + ) + await receiver.sem_prefetch.acquire() + + with pytest.raises(ReceiverLifecycleError) as exc_info: + await receiver._start_callback( + _PrefetchedMessage(data=b"message", owns_prefetch_slot=True), + ) + + await asyncio.wait_for(receiver.sem_prefetch.acquire(), timeout=0.1) + + assert exc_info.value is hook_error + + +async def test_callback_cancelled_before_start_releases_execution_capacity() -> None: + broker = ListenerBroker(empty_listener) + + @broker.task(task_name="receiver.cancel-before-start") + async def callback_task() -> None: + return None + + with ThreadPoolExecutor(max_workers=1) as executor: + receiver = Receiver( + broker, + executor=executor, + run_startup=False, + max_async_tasks=1, + max_prefetch=0, + ) + await receiver.sem_prefetch.acquire() + await receiver._start_callback( + _PrefetchedMessage( + data=encoded_message(broker, callback_task.task_name), + owns_prefetch_slot=True, + ), + ) + owned_task = next(iter(receiver._active_tasks)) + + owned_task.cancel() + await asyncio.gather(owned_task, return_exceptions=True) + await asyncio.sleep(0) + + assert not receiver._active_tasks + assert receiver.sem is not None + await asyncio.wait_for(receiver.sem.acquire(), timeout=0.1) + + +async def test_unexpected_owned_callback_failure_is_retrieved_and_logged( + caplog: pytest.LogCaptureFixture, +) -> None: + callback_error = ReceiverLifecycleError("owned callback failed") + + class FailingCallbackReceiver(Receiver): + async def _run_owned_callback( + self, + message: bytes | AckableMessage, + ) -> None: + raise callback_error + + broker = ListenerBroker(empty_listener) + caplog.set_level(logging.ERROR, logger="taskiq.receiver.receiver") + with ThreadPoolExecutor(max_workers=1) as executor: + receiver = FailingCallbackReceiver( + broker, + executor=executor, + run_startup=False, + max_async_tasks=1, + max_prefetch=0, + ) + await receiver.sem_prefetch.acquire() + await receiver._start_callback( + _PrefetchedMessage(data=b"message", owns_prefetch_slot=True), + ) + owned_task = next(iter(receiver._active_tasks)) + + await asyncio.gather(owned_task, return_exceptions=True) + await asyncio.sleep(0) + + assert not receiver._active_tasks + assert "Receiver callback failed outside task execution handling" in caplog.text + + +async def test_callback_cleanup_survives_outer_cancellation() -> None: + cleanup_started = asyncio.Event() + release_cleanup = asyncio.Event() + + async def callback() -> None: + try: + await asyncio.Event().wait() + finally: + cleanup_started.set() + await release_cleanup.wait() + + callback_task = asyncio.create_task(callback()) + await asyncio.sleep(0) + cleanup_task = asyncio.create_task( + Receiver._cancel_callback_tasks({callback_task}), + ) + await cleanup_started.wait() + + cleanup_task.cancel() + await asyncio.sleep(0) + release_cleanup.set() + cancellation = await cleanup_task + + assert isinstance(cancellation, asyncio.CancelledError) + assert callback_task.done() + + +async def test_active_callbacks_can_be_cancelled_without_graceful_wait() -> None: + callback_started = asyncio.Event() + callback_cleaned = asyncio.Event() + + async def callback() -> None: + callback_started.set() + try: + await asyncio.Event().wait() + finally: + callback_cleaned.set() + + broker = ListenerBroker(empty_listener) + with ThreadPoolExecutor(max_workers=1) as executor: + receiver = Receiver(broker, executor=executor, run_startup=False) + callback_task = asyncio.create_task(callback()) + receiver._active_tasks.add(callback_task) + await callback_started.wait() + + cancellation = await receiver._drain_active_tasks(cancel_immediately=True) + + assert cancellation is None + assert callback_cleaned.is_set() + assert callback_task.cancelled() + + +async def test_graceful_drain_preserves_outer_cancellation() -> None: + callback_started = asyncio.Event() + callback_cleaned = asyncio.Event() + + async def callback() -> None: + callback_started.set() + try: + await asyncio.Event().wait() + finally: + callback_cleaned.set() + + broker = ListenerBroker(empty_listener) + with ThreadPoolExecutor(max_workers=1) as executor: + receiver = Receiver( + broker, + executor=executor, + run_startup=False, + wait_tasks_timeout=None, + ) + callback_task = asyncio.create_task(callback()) + receiver._active_tasks.add(callback_task) + await callback_started.wait() + drain_task = asyncio.create_task( + receiver._drain_active_tasks(cancel_immediately=False), + ) + await asyncio.sleep(0) + assert not drain_task.done() + + drain_task.cancel() + cancellation = await drain_task + + assert isinstance(cancellation, asyncio.CancelledError) + assert callback_cleaned.is_set() + assert callback_task.cancelled() + + +async def test_additional_listener_close_failure_keeps_primary_error( + caplog: pytest.LogCaptureFixture, +) -> None: + broker = ListenerBroker(empty_listener) + first_error = ReceiverLifecycleError("primary listener failure") + second_error = ReceiverLifecycleError("additional close failure") + caplog.set_level(logging.ERROR, logger="taskiq.receiver.receiver") + with ThreadPoolExecutor(max_workers=1) as executor: + receiver = Receiver(broker, executor=executor, run_startup=False) + + receiver._record_listen_error(first_error) + receiver._record_listen_error(second_error) + + assert receiver._listen_error is first_error + assert "Additional error while closing broker listener" in caplog.text diff --git a/tests/receiver/test_runtime.py b/tests/receiver/test_runtime.py index 8c3d5c7c..810516c3 100644 --- a/tests/receiver/test_runtime.py +++ b/tests/receiver/test_runtime.py @@ -1,149 +1,27 @@ import asyncio import logging -from collections.abc import AsyncGenerator, Iterator from concurrent.futures import ThreadPoolExecutor -from dataclasses import replace from typing import Any import pytest -from taskiq import AsyncBroker, BrokerMessage, TaskiqRouter, task_builder +from taskiq import TaskiqRouter, task_builder from taskiq.abc.middleware import TaskiqMiddleware from taskiq.receiver import Receiver from taskiq.receiver.runtime import ( ReceiverRuntimeState, - WorkerRuntime, WorkerRuntimeConfigurationError, - WorkerRuntimeOptions, normalize_broker_source, ) from taskiq.result_backends.dummy import DummyResultBackend - - -class RuntimeTestError(Exception): - """Marker exception for worker runtime failures.""" - - -class FatalRuntimeTestError(BaseException): - """Marker fatal signal that runtime cleanup must not swallow.""" - - -class RuntimeBroker(AsyncBroker): - """Controllable broker used for multi-broker lifecycle tests.""" - - def __init__( - self, - *, - router: TaskiqRouter, - broker_name: str, - events: list[str], - startup_error: BaseException | None = None, - listen_error: BaseException | None = None, - shutdown_error: BaseException | None = None, - ) -> None: - self.events = events - self.startup_error = startup_error - self.listen_error = listen_error - self.shutdown_error = shutdown_error - self.listen_started = asyncio.Event() - self.release_listener = asyncio.Event() - self.iterator_closed = asyncio.Event() - self.started = False - self.queue: asyncio.Queue[bytes] = asyncio.Queue() - super().__init__(router=router, broker_name=broker_name) - - async def startup(self) -> None: - """Record startup and optionally fail after core resources start.""" - self.events.append( - f"{self.broker_name}.startup." - f"{'worker' if self.is_worker_process else 'client'}", - ) - await super().startup() - self.started = True - if self.startup_error is not None: - raise self.startup_error - - async def shutdown(self) -> None: - """Record shutdown and close core resources.""" - self.events.append(f"{self.broker_name}.shutdown") - await super().shutdown() - self.started = False - if self.shutdown_error is not None: - raise self.shutdown_error - - async def kick(self, message: BrokerMessage) -> None: - """Queue an encoded message for this listener.""" - if not self.started: - raise RuntimeTestError(f"{self.broker_name} kicked before startup") - self.events.append(f"{self.broker_name}.kick") - await self.queue.put(message.message) - - async def listen(self) -> AsyncGenerator[bytes, None]: - """Listen until cancelled or a configured failure is released.""" - self.events.append(f"{self.broker_name}.listen") - self.listen_started.set() - try: - if self.listen_error is not None: - await self.release_listener.wait() - raise self.listen_error - while True: - yield await self.queue.get() - finally: - self.events.append(f"{self.broker_name}.iterator_closed") - self.iterator_closed.set() - - -class PrefetchRecordingMiddleware(TaskiqMiddleware): - """Record broker deliveries admitted into the shared prefetch budget.""" - - def __init__(self, broker_name: str, prefetched: list[str]) -> None: - super().__init__() - self.broker_name = broker_name - self.prefetched = prefetched - - def on_prefetch_queue_add(self) -> None: - """Record one delivery accepted by a Receiver prefetch queue.""" - self.prefetched.append(self.broker_name) - - -@pytest.fixture -def runtime_executor() -> Iterator[ThreadPoolExecutor]: - """Provide the process-wide sync executor owned by a worker runtime.""" - with ThreadPoolExecutor(max_workers=2) as executor: - yield executor - - -def runtime_options(**overrides: Any) -> WorkerRuntimeOptions: - """Build runtime options with focused per-test overrides.""" - options = WorkerRuntimeOptions( - validate_params=True, - max_async_tasks=10, - max_async_tasks_jitter=0, - max_prefetch=0, - propagate_exceptions=True, - ack_type=None, - max_tasks_to_execute=None, - wait_tasks_timeout=1, - shutdown_timeout=1, - receiver_kwargs={}, - ) - return replace(options, **overrides) - - -def build_runtime( - brokers: object, - executor: ThreadPoolExecutor, - *, - options: WorkerRuntimeOptions | None = None, - receiver_type: type[Receiver] = Receiver, -) -> WorkerRuntime: - """Normalize broker configuration and build the runtime under test.""" - return WorkerRuntime( - selection=normalize_broker_source(brokers), - receiver_type=receiver_type, - executor=executor, - options=options or runtime_options(), - ) +from tests.receiver.runtime_support import ( + FatalRuntimeTestError, + PrefetchRecordingMiddleware, + RuntimeBroker, + RuntimeTestError, + build_runtime, + runtime_options, +) async def test_runtime_listens_concurrently_and_starts_outbound_broker( @@ -963,6 +841,7 @@ async def listen(self, finish_event: asyncio.Event) -> None: except asyncio.CancelledError: cancellation_seen.set() await release_listener.wait() + raise runtime = build_runtime( (first, second), @@ -990,7 +869,11 @@ async def listen(self, finish_event: asyncio.Event) -> None: ] release_listener.set() - await asyncio.wait_for(asyncio.gather(*listener_tasks), timeout=0.2) + results = await asyncio.wait_for( + asyncio.gather(*listener_tasks, return_exceptions=True), + timeout=0.2, + ) + assert all(isinstance(result, asyncio.CancelledError) for result in results) def test_runtime_rejects_mixed_routers_before_startup( diff --git a/tests/receiver/test_runtime_faults.py b/tests/receiver/test_runtime_faults.py new file mode 100644 index 00000000..8f659108 --- /dev/null +++ b/tests/receiver/test_runtime_faults.py @@ -0,0 +1,262 @@ +import asyncio +import logging +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from taskiq import TaskiqRouter +from taskiq.receiver import Receiver +from taskiq.receiver.runtime import ( + ReceiverRuntimeState, + WorkerRuntime, + WorkerRuntimeConfigurationError, + normalize_broker_source, +) +from tests.receiver.runtime_support import ( + RuntimeBroker, + RuntimeTestError, + build_runtime, + runtime_options, +) + + +async def test_runtime_propagates_adapter_shutdown_cancellation( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + shutdown_error = asyncio.CancelledError("adapter shutdown cancelled") + broker = RuntimeBroker( + router=TaskiqRouter(), + broker_name="listener", + events=events, + shutdown_error=shutdown_error, + ) + runtime = build_runtime((broker,), runtime_executor) + finish_event = asyncio.Event() + finish_event.set() + + with pytest.raises(asyncio.CancelledError): + await runtime.run(finish_event) + assert events[-1] == "listener.shutdown" + + +async def test_runtime_logs_shutdown_failure_after_cleanup_deadline( + runtime_executor: ThreadPoolExecutor, + caplog: pytest.LogCaptureFixture, +) -> None: + events: list[str] = [] + release_shutdown = asyncio.Event() + shutdown_finished = asyncio.Event() + late_error = RuntimeTestError("late shutdown failure") + + class LateFailureShutdownBroker(RuntimeBroker): + async def shutdown(self) -> None: + self.events.append(f"{self.broker_name}.shutdown") + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + await release_shutdown.wait() + raise late_error from None + finally: + shutdown_finished.set() + + broker = LateFailureShutdownBroker( + router=TaskiqRouter(), + broker_name="listener", + events=events, + ) + runtime = build_runtime( + (broker,), + runtime_executor, + options=runtime_options(shutdown_timeout=0.01), + ) + finish_event = asyncio.Event() + caplog.set_level(logging.ERROR, logger="taskiq.worker.runtime") + runtime_task = asyncio.create_task(runtime.run(finish_event)) + await broker.listen_started.wait() + + finish_event.set() + await asyncio.wait_for(runtime_task, timeout=1) + release_shutdown.set() + await asyncio.wait_for(shutdown_finished.wait(), timeout=1) + await asyncio.sleep(0) + + assert "failed after its cleanup deadline" in caplog.text + + +async def test_runtime_preserves_listener_failure_discovered_during_settlement( + runtime_executor: ThreadPoolExecutor, +) -> None: + events: list[str] = [] + broker = RuntimeBroker( + router=TaskiqRouter(), + broker_name="listener", + events=events, + ) + finish_observed = asyncio.Event() + settlement_started = asyncio.Event() + release_failure = asyncio.Event() + late_error = RuntimeTestError("listener failed while settling") + + class LateFailureReceiver(Receiver): + async def listen(self, finish_event: asyncio.Event) -> None: + await finish_event.wait() + finish_observed.set() + await release_failure.wait() + raise late_error + + class ObservedSettlementRuntime(WorkerRuntime): + async def _settle_listener_tasks( + self, + listener_tasks: tuple[asyncio.Task[None], ...], + receivers: tuple[Receiver, ...], + ) -> tuple[list[None | BaseException], set[asyncio.Task[None]]]: + settlement_started.set() + return await super()._settle_listener_tasks(listener_tasks, receivers) + + finish_event = asyncio.Event() + runtime = ObservedSettlementRuntime( + selection=normalize_broker_source((broker,)), + receiver_type=LateFailureReceiver, + executor=runtime_executor, + options=runtime_options(), + ) + runtime_task = asyncio.create_task(runtime.run(finish_event)) + while not broker.started: + await asyncio.sleep(0) + + finish_event.set() + await finish_observed.wait() + await settlement_started.wait() + release_failure.set() + + with pytest.raises(RuntimeTestError) as exc_info: + await runtime_task + + assert exc_info.value is late_error + + +async def test_runtime_logs_late_cancellation_resistant_listener_failure( + runtime_executor: ThreadPoolExecutor, + caplog: pytest.LogCaptureFixture, +) -> None: + events: list[str] = [] + broker = RuntimeBroker( + router=TaskiqRouter(), + broker_name="listener", + events=events, + ) + listener_started = asyncio.Event() + cancellation_seen = asyncio.Event() + release_listener = asyncio.Event() + late_error = RuntimeTestError("late listener failure") + listener_tasks: list[asyncio.Task[None]] = [] + + class LateFailureReceiver(Receiver): + async def listen(self, finish_event: asyncio.Event) -> None: + listener_task = asyncio.current_task() + assert listener_task is not None + listener_tasks.append(listener_task) + listener_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancellation_seen.set() + await release_listener.wait() + raise late_error from None + + runtime = build_runtime( + (broker,), + runtime_executor, + options=runtime_options( + wait_tasks_timeout=0.01, + shutdown_timeout=0.01, + ), + receiver_type=LateFailureReceiver, + ) + finish_event = asyncio.Event() + caplog.set_level(logging.ERROR, logger="taskiq.worker.runtime") + runtime_task = asyncio.create_task(runtime.run(finish_event)) + await listener_started.wait() + + finish_event.set() + await asyncio.wait_for(runtime_task, timeout=1) + assert cancellation_seen.is_set() + + release_listener.set() + results = await asyncio.wait_for( + asyncio.gather(*listener_tasks, return_exceptions=True), + timeout=1, + ) + await asyncio.sleep(0) + + assert results == [late_error] + assert "failed after runtime cleanup" in caplog.text + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + pytest.param( + {"max_prefetch": -1}, + "max_prefetch cannot be negative", + id="negative-prefetch", + ), + pytest.param( + {"max_async_tasks_jitter": -1}, + "max_async_tasks_jitter cannot be negative", + id="negative-jitter", + ), + pytest.param( + {"shutdown_timeout": -1}, + "shutdown_timeout cannot be negative", + id="negative-shutdown-timeout", + ), + pytest.param( + {"wait_tasks_timeout": -1}, + "wait_tasks_timeout cannot be negative", + id="negative-task-timeout", + ), + ], +) +def test_runtime_options_reject_negative_limits( + overrides: dict[str, int], + message: str, +) -> None: + with pytest.raises(WorkerRuntimeConfigurationError, match=message): + runtime_options(**overrides) + + +async def test_runtime_state_applies_shared_execution_jitter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + "taskiq.receiver.runtime.random.randint", + lambda start, end: 2, + ) + runtime_state = ReceiverRuntimeState( + finish_event=asyncio.Event(), + max_async_tasks=1, + max_async_tasks_jitter=2, + max_prefetch=0, + max_tasks_to_execute=None, + ) + semaphore = runtime_state.execution_semaphore + assert semaphore is not None + + for _ in range(3): + await semaphore.acquire() + + assert semaphore.locked() + + +def test_runtime_state_preserves_unlimited_async_execution() -> None: + runtime_state = ReceiverRuntimeState( + finish_event=asyncio.Event(), + max_async_tasks=None, + max_async_tasks_jitter=0, + max_prefetch=0, + max_tasks_to_execute=None, + ) + + assert runtime_state.execution_semaphore is None diff --git a/tests/routing/test_dispatch.py b/tests/routing/test_dispatch.py index b0ee4e32..9235820b 100644 --- a/tests/routing/test_dispatch.py +++ b/tests/routing/test_dispatch.py @@ -4,6 +4,7 @@ Flow, FlowProtocol, InMemoryBroker, + TaskiqMessage, TaskiqRoute, TaskiqRouter, UnsupportedFlowError, @@ -431,6 +432,52 @@ async def demo_task() -> None: assert target.sent[0][1] == flow +async def test_router_bare_task_decorator_uses_default_broker() -> None: + router = TaskiqRouter() + broker = RecordingBroker(router=router, broker_name="default") + + @router.task + async def demo_task() -> None: + return None + + await demo_task.kiq() + + assert demo_task.broker is broker + assert broker.sent[0][0].task_name == f"{__name__}:demo_task" + + +async def test_explicit_route_rejects_additional_broker_override() -> None: + router = TaskiqRouter() + broker = RecordingBroker(router=router, broker_name="broker") + route = TaskiqRoute(broker=broker, flow=Flow("events")) + message = TaskiqMessage( + task_id="task-id", + task_name="demo.task", + labels={}, + args=[], + kwargs={}, + ) + + with pytest.raises(ValueError, match="either route or broker/flow"): + await router.kiq(message, route=route, broker=broker) + + +async def test_flow_override_replaces_selected_route_flow() -> None: + router = TaskiqRouter() + broker = RecordingBroker(router=router, broker_name="broker") + + @broker.task(task_name="demo.task") + async def demo_task() -> None: + return None + + route = TaskiqRoute(broker=broker, flow=Flow("initial")) + override = Flow("override") + + await demo_task.kicker().with_route(route).with_flow(override).kiq() + + assert broker.sent[0][1] == override + + async def test_routed_dispatch_uses_target_middleware_and_result_backend() -> None: router = TaskiqRouter() source = RecordingBroker(router=router, broker_name="source") diff --git a/tests/routing/test_router_core.py b/tests/routing/test_router_core.py index 79f7f818..4b39b6b1 100644 --- a/tests/routing/test_router_core.py +++ b/tests/routing/test_router_core.py @@ -9,6 +9,7 @@ def test_broker_creates_default_router() -> None: broker = RecordingBroker() assert broker.router.brokers[broker.broker_name] is broker + assert broker.router.default_broker is broker assert broker.router.default_broker_name == broker.broker_name @@ -25,6 +26,33 @@ def test_router_rejects_duplicate_broker_names() -> None: RecordingBroker(router=router, broker_name="broker") +def test_router_rejects_registering_same_broker_under_another_name() -> None: + router = TaskiqRouter() + broker = RecordingBroker(router=router, broker_name="broker") + + with pytest.raises(ValueError, match="already registered as 'broker'"): + router.set_broker(broker, name="alias") + + +def test_router_rejects_unknown_broker_lookup() -> None: + router = TaskiqRouter() + + with pytest.raises(ValueError, match="Unknown broker 'missing'"): + router.get_broker("missing") + + +def test_router_can_clear_default_broker_explicitly() -> None: + router = TaskiqRouter() + RecordingBroker(router=router, broker_name="broker") + + router.default_broker = None + + assert router.default_broker is None + assert router.default_broker_name is None + with pytest.raises(ValueError, match="doesn't have registered brokers"): + router.resolve_route("demo.task") + + def test_router_rejects_broker_attached_to_another_router() -> None: first_router = TaskiqRouter() second_router = TaskiqRouter() @@ -51,6 +79,14 @@ def test_router_rejects_string_broker_references() -> None: router.route_task("demo.task", broker="broker") # type: ignore[arg-type] +def test_router_rejects_invalid_task_reference() -> None: + router = TaskiqRouter() + broker = RecordingBroker(router=router, broker_name="broker") + + with pytest.raises(TypeError, match="task name or decorated task"): + router.route_task(object(), broker=broker) # type: ignore[arg-type] + + def test_router_routes_are_immutable_resolved_snapshots() -> None: router = TaskiqRouter() broker = RecordingBroker(router=router, broker_name="broker") @@ -121,3 +157,19 @@ def shared_task() -> None: assert router.register_task_calls == 1 assert broker.router.find_task("shared.once") is registered assert broker.find_task("shared.once") is registered + + +def test_router_infers_broker_from_prebound_task() -> None: + router = TaskiqRouter() + broker = RecordingBroker(router=router, broker_name="broker") + + @task_builder("shared.prebound") + def shared_task() -> None: + return None + + bound_task = broker.bind_task_definition(shared_task, register=False) + + registered = router.register_task(bound_task) + + assert registered is bound_task + assert router.resolve_route(registered).broker is broker diff --git a/tests/routing/test_subscriptions.py b/tests/routing/test_subscriptions.py index 45d22a1c..068117b8 100644 --- a/tests/routing/test_subscriptions.py +++ b/tests/routing/test_subscriptions.py @@ -45,6 +45,7 @@ def test_subscription_snapshot_and_explicit_unsubscribe() -> None: assert snapshot == (first,) assert router.subscriptions == (first, second) + assert first.broker_name == "broker" assert router.unsubscribe(broker, first_flow) == first assert router.subscriptions == (second,) assert router.unsubscribe(broker, first_flow) is None @@ -179,6 +180,21 @@ def test_subscribe_rejects_default_flow_option_conflict() -> None: ) +def test_subscribe_preserves_distinct_default_flow() -> None: + router = TaskiqRouter() + default_flow = Flow("default") + subscribed_flow = Flow("events") + broker = RecordingBroker( + router=router, + broker_name="broker", + default_flow=default_flow, + ) + + router.subscribe(broker, subscribed_flow, "demo.task") + + assert broker.get_subscribed_flows() == (default_flow, subscribed_flow) + + def test_subscribe_allows_same_identity_on_different_brokers() -> None: router = TaskiqRouter() first_broker = RecordingBroker(router=router, broker_name="first") diff --git a/tests/scheduler/test_scheduler.py b/tests/scheduler/test_scheduler.py index 0df22be1..67323f5d 100644 --- a/tests/scheduler/test_scheduler.py +++ b/tests/scheduler/test_scheduler.py @@ -1,3 +1,4 @@ +import asyncio from datetime import datetime, timezone import pytest @@ -201,6 +202,58 @@ async def test_scheduler_rejects_duplicate_startup() -> None: assert events == ["startup:broker", "shutdown:broker"] +async def test_scheduler_rejects_shutdown_while_startup_is_in_progress() -> None: + events: list[str] = [] + startup_started = asyncio.Event() + release_startup = asyncio.Event() + + class BlockingStartupBroker(LifecycleBroker): + async def startup(self) -> None: + startup_started.set() + await release_startup.wait() + await super().startup() + + router = TaskiqRouter() + broker = BlockingStartupBroker(events, router=router, broker_name="broker") + scheduler = TaskiqScheduler(broker=broker, sources=[]) + startup_task = asyncio.create_task(scheduler.startup()) + await startup_started.wait() + + with pytest.raises(RuntimeError, match="still starting"): + await scheduler.shutdown() + + release_startup.set() + await startup_task + await scheduler.shutdown() + + +async def test_scheduler_rejects_operations_while_shutdown_is_in_progress() -> None: + events: list[str] = [] + shutdown_started = asyncio.Event() + release_shutdown = asyncio.Event() + + class BlockingShutdownBroker(LifecycleBroker): + async def shutdown(self) -> None: + shutdown_started.set() + await release_shutdown.wait() + await super().shutdown() + + router = TaskiqRouter() + broker = BlockingShutdownBroker(events, router=router, broker_name="broker") + scheduler = TaskiqScheduler(broker=broker, sources=[]) + await scheduler.startup() + shutdown_task = asyncio.create_task(scheduler.shutdown()) + await shutdown_started.wait() + + with pytest.raises(RuntimeError, match="shutting down"): + await scheduler.startup() + with pytest.raises(RuntimeError, match="already shutting down"): + await scheduler.shutdown() + + release_shutdown.set() + await shutdown_task + + async def test_scheduler_shutdown_before_startup_preserves_legacy_behavior() -> None: events: list[str] = [] router = TaskiqRouter() diff --git a/tests/test_task_builder.py b/tests/test_task_builder.py index 9dbefb95..ae071809 100644 --- a/tests/test_task_builder.py +++ b/tests/test_task_builder.py @@ -41,6 +41,17 @@ class User(BaseModel): name: str +class CallableIncrement: + """Callable object used to prove non-function shared task binding.""" + + def __init__(self, increment: int) -> None: + self.increment = increment + + def __call__(self, value: int) -> int: + """Increment one value.""" + return value + self.increment + + @task_builder("shared.process_pool") def shared_process_pool_add(left: int, right: int) -> int: """Add values in a spawned worker process.""" @@ -60,6 +71,25 @@ def generated(value: int) -> int: return generated +def build_stateful_task( + task_name: str, + captured: object, + positional_default: Any = 1, + keyword_default: Any = 1, +) -> TaskDefinition[..., int]: + """Build declarations that share code but vary execution state.""" + + @task_builder(task_name) + def generated( + value: int = positional_default, + *, + scale: int = keyword_default, + ) -> int: + return value * scale + int(bool(captured)) + + return generated + + factory_process_pool_add_one = build_process_pool_task( "shared.factory_add_one", 1, @@ -92,6 +122,103 @@ async def multiply(left: int, right: int) -> int: assert inspect.iscoroutinefunction(registered_multiply.original_func) +def test_bare_task_builder_uses_qualified_function_name() -> None: + @task_builder + def implicit_task(value: int) -> int: + return value + 1 + + assert implicit_task.task_name == f"{__name__}:implicit_task" + assert implicit_task(1) == 2 + + +async def test_callable_object_task_definition_can_be_bound_and_sent() -> None: + definition = task_builder("shared.callable-object")(CallableIncrement(2)) + broker = InMemoryBroker(await_inplace=True) + + registered = broker.register_task(definition) + sent_task = await registered.kiq(3) + result = await sent_task.wait_result() + await broker.shutdown() + + assert definition(3) == 5 + assert registered.return_type is int + assert registered.__wrapped__ is definition.original_func + assert result.return_value == 5 + + +def test_callable_object_binding_rejects_ambiguous_redeclaration() -> None: + task_builder("shared.callable-collision")(CallableIncrement(2)) + + with pytest.raises( + ValueError, + match="Factory-generated TaskDefinitions must use unique task names", + ): + task_builder("shared.callable-collision")(CallableIncrement(2)) + + +def test_equivalent_immutable_factory_state_reuses_binding() -> None: + first = build_stateful_task( + "shared.equivalent-state", + frozenset({1, 2}), + positional_default=(1, 2), + keyword_default=2, + ) + second = build_stateful_task( + "shared.equivalent-state", + frozenset({1, 2}), + positional_default=(1, 2), + keyword_default=2, + ) + + first_registered = InMemoryBroker().register_task(first) + second_registered = InMemoryBroker().register_task(second) + + assert first_registered.original_func is second_registered.original_func + + +@pytest.mark.parametrize( + ("task_name", "first_state", "second_state"), + [ + pytest.param( + "shared.mutable-state", + ([1], 1, 1), + ([1], 1, 1), + id="equal-mutable-closure", + ), + pytest.param( + "shared.positional-state", + (frozenset({1}), 1, 1), + (frozenset({1}), 1.0, 1), + id="different-positional-default-type", + ), + pytest.param( + "shared.keyword-state", + (frozenset({1}), 1, 1), + (frozenset({1}), 1, 2), + id="different-keyword-default", + ), + pytest.param( + "shared.frozenset-state", + (frozenset({1}), 1, 1), + (frozenset({2}), 1, 1), + id="different-frozen-closure", + ), + ], +) +def test_different_factory_execution_state_rejects_binding_reuse( + task_name: str, + first_state: tuple[object, Any, Any], + second_state: tuple[object, Any, Any], +) -> None: + build_stateful_task(task_name, *first_state) + + with pytest.raises( + ValueError, + match="Factory-generated TaskDefinitions must use unique task names", + ): + build_stateful_task(task_name, *second_state) + + def test_task_definition_preserves_metadata_across_broker_bindings() -> None: @task_builder("shared.metadata") def calculate(left: int, right: int = 2) -> int: From 9efb05ff148d566f06bed4b74729323b98d2ad3a Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Sat, 11 Jul 2026 22:42:29 +0200 Subject: [PATCH 10/10] fix(tasks): protect bound task state from callable metadata - prevent callable __dict__ from overriding task identity and broker state - preserve non-conflicting custom metadata for backward compatibility - cover callable-object and legacy decorator metadata collisions --- taskiq/abc/broker.py | 41 +++++++++++++++++++++++------ tests/abc/test_broker.py | 32 ++++++++++++++++++++++ tests/test_task_builder.py | 54 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 119 insertions(+), 8 deletions(-) diff --git a/taskiq/abc/broker.py b/taskiq/abc/broker.py index 45da7890..5dae7990 100644 --- a/taskiq/abc/broker.py +++ b/taskiq/abc/broker.py @@ -3,7 +3,7 @@ import warnings from abc import ABC, abstractmethod from collections import defaultdict -from collections.abc import AsyncGenerator, Awaitable, Callable +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping from functools import WRAPPER_ASSIGNMENTS, update_wrapper from logging import getLogger from typing import ( @@ -49,12 +49,43 @@ _T = TypeVar("_T") _FuncParams = ParamSpec("_FuncParams") _ReturnType = TypeVar("_ReturnType") +_DecoratedTask = TypeVar( + "_DecoratedTask", + bound=AsyncTaskiqDecoratedTask[Any, Any], +) EventHandler: TypeAlias = Callable[[TaskiqState], Awaitable[None] | None] logger = getLogger("taskiq") +def _copy_callable_metadata( + task: _DecoratedTask, + source: Callable[..., Any], +) -> _DecoratedTask: + """Preserve callable metadata while keeping task-owned state authoritative.""" + assigned_metadata = tuple( + attribute for attribute in WRAPPER_ASSIGNMENTS if hasattr(source, attribute) + ) + update_wrapper( + task, + source, + assigned=assigned_metadata, + updated=(), + ) + + protected_attributes = set(vars(task)) + for task_type in type(task).__mro__: + protected_attributes.update(vars(task_type)) + + source_attributes = getattr(source, "__dict__", None) + if isinstance(source_attributes, Mapping): + for attribute, value in source_attributes.items(): + if attribute not in protected_attributes: + setattr(task, attribute, value) + return task + + def default_id_generator() -> str: """ Default task_id generator. @@ -541,12 +572,7 @@ def _decorate_task( if "return" in sign: return_type = sign["return"] - assigned_metadata = tuple( - attribute - for attribute in WRAPPER_ASSIGNMENTS - if hasattr(metadata_source, attribute) - ) - decorated_task = update_wrapper( + decorated_task = _copy_callable_metadata( task_cls( broker=self, original_func=func, @@ -555,7 +581,6 @@ def _decorate_task( return_type=return_type, # type: ignore ), metadata_source, - assigned=assigned_metadata, ) if register: diff --git a/tests/abc/test_broker.py b/tests/abc/test_broker.py index 83ab9d47..d4f81ac1 100644 --- a/tests/abc/test_broker.py +++ b/tests/abc/test_broker.py @@ -68,6 +68,38 @@ async def test_func() -> None: } +def test_decorator_metadata_cannot_replace_task_contract() -> None: + """Task fields must win over colliding wrapped-function attributes.""" + broker = _TestBroker() + + def source(value: int) -> int: + return value + 1 + + vars(source).update( + broker="source-broker", + custom_metadata="preserved", + kiq="source-kiq", + labels={"source": "function"}, + original_func="source-original-func", + return_type=str, + task_name="source.task-name", + ) + + registered = broker.task( + task_name="expected.task-name", + queue="expected", + )(source) + + assert registered.broker is broker + assert registered.task_name == "expected.task-name" + assert registered.labels == {"queue": "expected"} + assert registered.original_func is source + assert registered.return_type is int + assert callable(registered.kiq) + assert vars(registered)["custom_metadata"] == "preserved" + assert registered(2) == 3 + + def test_decorator_generates_unique_lambda_task_names() -> None: """Anonymous tasks must not collide in the broker task registry.""" broker = _TestBroker() diff --git a/tests/test_task_builder.py b/tests/test_task_builder.py index ae071809..9c6448c3 100644 --- a/tests/test_task_builder.py +++ b/tests/test_task_builder.py @@ -52,6 +52,23 @@ def __call__(self, value: int) -> int: return value + self.increment +class CollidingCallable: + """Callable whose application state overlaps with the task API.""" + + def __init__(self, increment: int) -> None: + self.broker = "source-broker" + self.increment = increment + self.kiq = "source-kiq" + self.labels = {"source": "callable-object"} + self.original_func = "source-original-func" + self.return_type = str + self.task_name = "source.task-name" + + def __call__(self, value: int) -> int: + """Increment one value without exposing state as task metadata.""" + return value + self.increment + + @task_builder("shared.process_pool") def shared_process_pool_add(left: int, right: int) -> int: """Add values in a spawned worker process.""" @@ -146,6 +163,43 @@ async def test_callable_object_task_definition_can_be_bound_and_sent() -> None: assert result.return_value == 5 +async def test_callable_object_state_cannot_replace_task_contract() -> None: + source = CollidingCallable(2) + definition = task_builder( + "shared.callable-owned-state", + queue="expected", + )(source) + broker = InMemoryBroker(await_inplace=True) + + registered = broker.register_task(definition) + sent_task = await registered.kiq(3) + result = await sent_task.wait_result() + await broker.shutdown() + + assert registered.broker is broker + assert registered.task_name == "shared.callable-owned-state" + assert registered.labels == {"queue": "expected"} + assert registered.return_type is int + assert registered.__wrapped__ is source + assert registered.original_func(3) == 5 + assert callable(registered.kiq) + assert vars(registered)["increment"] == 2 + assert broker.local_task_registry[registered.task_name] is registered + assert result.return_value == 5 + + +def test_task_binding_preserves_non_conflicting_custom_metadata() -> None: + def source(value: int) -> int: + return value + 1 + + vars(source)["custom_metadata"] = "preserved" + definition = task_builder("shared.custom-metadata")(source) + + registered = InMemoryBroker().register_task(definition) + + assert vars(registered)["custom_metadata"] == "preserved" + + def test_callable_object_binding_rejects_ambiguous_redeclaration() -> None: task_builder("shared.callable-collision")(CallableIncrement(2))