diff --git a/docs/available-components/brokers.md b/docs/available-components/brokers.md index c4dad6df..a4b60a59 100644 --- a/docs/available-components/brokers.md +++ b/docs/available-components/brokers.md @@ -8,10 +8,21 @@ In this section we'll list officially supported brokers. ## InMemoryBroker -This is a special broker for local development. It uses the same functions to execute tasks, -but all tasks are executed locally in the current thread. +This is a special broker for local development. It uses the same functions to +execute tasks in the current process. Async task functions run on the event +loop, while sync task functions use the broker's thread pool. By default it uses `InMemoryResultBackend` but this can be overridden. +`startup()` and `shutdown()` run both the client and worker event phases because +the broker performs both roles in one process. They also own middleware and +result backend lifecycle. Shutdown stops accepting new local executions, waits +for already accepted work (including active send middleware), then closes those +resources and the synchronous task executor. A send started after shutdown is +rejected before its `pre_send` hooks run. +Once cleanup completes, later `shutdown()` calls do not close those resources +again. If cancellation interrupts a lifecycle hook, it propagates immediately +and a later call resumes cleanup from that hook. + ## ZeroMQBroker This broker uses [ZMQ](https://zeromq.org/) to communicate between worker and client processes. diff --git a/docs/guide/testing-taskiq.md b/docs/guide/testing-taskiq.md index 14ae47fe..4f81db59 100644 --- a/docs/guide/testing-taskiq.md +++ b/docs/guide/testing-taskiq.md @@ -165,6 +165,13 @@ broker = InMemoryBroker(await_inplace=True) With this setup all `await function.kiq()` calls will behave similarly to `await function()`, but with dependency injection and all taskiq-related functionality. +For async task functions, successful inline execution remains in the caller's +asyncio task and context. Sync task functions still run in the broker's thread +pool with a copy of the current `ContextVars`. +The normal middleware boundary is preserved in both execution modes: all +`post_send` hooks finish before `pre_execute` and the task body begin. +With `await_inplace=True`, the `kiq()` call also returns only after inline +execution completes. 2. Alternatively, you can manually await all tasks after invoking the target function by using the `wait_all` method. @@ -182,6 +189,25 @@ async def test_add_one(): # have been completed and do all the assertions. ``` +`wait_all()` also surfaces the first pending background callback failure even +if that callback completed before `wait_all()` was called. The failure is +consumed after it is raised, so a later `wait_all()` only observes newer work. +Cancelling a `wait_all()` call cancels only that waiter; accepted executions +remain tracked and can be drained by a later call. +Calling `shutdown()` rejects new sends, performs the same drain, then closes +middleware, result backend, and executor resources. If shutdown is cancelled +while draining accepted work, it finishes that drain and cleanup before +propagating the cancellation. +Cancellation from a shutdown event, middleware, or result backend hook instead +propagates immediately. A later `shutdown()` resumes at the interrupted hook +without repeating completed cleanup, while calls after completed cleanup are +no-ops. +An invocation already running `pre_send` is part of that drain; a later +invocation is rejected before `pre_send` can produce side effects. +Both drain methods must be called by the external test or application lifecycle +owner. Calling `wait_all()` or `shutdown()` from a task or `post_send` hook +managed by the same broker raises `RuntimeError` instead of waiting on itself. + ## Dependency injection If you use dependencies in your tasks, you may think that this can become a problem for tests. But it's not. diff --git a/taskiq/abc/broker.py b/taskiq/abc/broker.py index ea2e86c0..85c27e3a 100644 --- a/taskiq/abc/broker.py +++ b/taskiq/abc/broker.py @@ -3,8 +3,9 @@ import warnings from abc import ABC, abstractmethod from collections import defaultdict -from collections.abc import AsyncGenerator, Awaitable, Callable -from functools import wraps +from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator +from contextlib import contextmanager +from functools import partial, wraps from logging import getLogger from typing import ( TYPE_CHECKING, @@ -23,7 +24,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 SendTaskError, TaskBrokerMismatchError from taskiq.formatters.proxy_formatter import ProxyFormatter from taskiq.message import BrokerMessage from taskiq.result_backends.dummy import DummyResultBackend @@ -47,6 +48,7 @@ _ReturnType = TypeVar("_ReturnType") EventHandler: TypeAlias = Callable[[TaskiqState], Awaitable[None] | None] +ShutdownHook: TypeAlias = Callable[[], Awaitable[None] | None] logger = getLogger("taskiq") @@ -118,6 +120,7 @@ def __init__( self.is_worker_process = False # True only if broker runs in scheduler process. self.is_scheduler_process = False + self._shutdown_resource_index = 0 def find_task(self, task_name: str) -> AsyncTaskiqDecoratedTask[Any, Any] | None: """ @@ -186,12 +189,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: @@ -199,6 +199,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. @@ -206,19 +212,86 @@ 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_errors: list[BaseException] = [] + await self._shutdown_resources(shutdown_errors) - # Call all shutdown events. - for handler in self.event_handlers[event]: - await maybe_awaitable(handler(self.state)) + if shutdown_errors: + raise shutdown_errors[0] + + async def _shutdown_resources( + self, + shutdown_errors: list[BaseException], + ) -> None: + """Close every registered resource and record failures in order.""" + for index, shutdown_hook in enumerate(self._iter_shutdown_hooks()): + if index < self._shutdown_resource_index: + continue + try: + await maybe_awaitable(shutdown_hook()) + except Exception as exc: + self._record_shutdown_error(shutdown_errors, exc) + self._shutdown_resource_index = index + 1 + + def _iter_shutdown_hooks(self) -> Iterator[ShutdownHook]: + """Yield lifecycle hooks in their shutdown order.""" + for event in self._get_shutdown_events(): + for handler in self.event_handlers[event]: + yield partial(handler, self.state) for middleware in self.middlewares: if middleware.__class__.shutdown != TaskiqMiddleware.shutdown: - await maybe_awaitable(middleware.shutdown()) + yield middleware.shutdown - await self.result_backend.shutdown() + yield self.result_backend.shutdown + + 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 cleanup continues.""" + if first_error is None: + return current_error + logger.error( + "Additional error while shutting down broker resources.", + exc_info=current_error, + ) + return first_error + + @classmethod + def _record_shutdown_error( + cls, + shutdown_errors: list[BaseException], + current_error: BaseException, + ) -> None: + """Record the first failure when it occurs and log later failures.""" + first_error = shutdown_errors[0] if shutdown_errors else None + remembered_error = cls._remember_shutdown_error(first_error, current_error) + if first_error is None: + shutdown_errors.append(remembered_error) + + @contextmanager + def _send_lifecycle(self) -> Iterator[None]: + """Own package-internal client send work through broker handoff.""" + yield + + async def _kick_with_post_send( + self, + message: BrokerMessage, + post_send: Callable[[], Awaitable[None]], + ) -> None: + """Run the package-internal send boundary used by AsyncKicker.""" + try: + await self.kick(message) + except Exception as exc: + raise SendTaskError from exc + await post_send() @abstractmethod async def kick( diff --git a/taskiq/brokers/inmemory_broker.py b/taskiq/brokers/inmemory_broker.py index 0a7cc98e..6656e882 100644 --- a/taskiq/brokers/inmemory_broker.py +++ b/taskiq/brokers/inmemory_broker.py @@ -1,17 +1,20 @@ import asyncio from collections import OrderedDict -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator from concurrent.futures import ThreadPoolExecutor -from typing import Any, TypeVar +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Any, NoReturn, TypeVar + +import anyio from taskiq.abc.broker import AsyncBroker from taskiq.abc.result_backend import AsyncResultBackend, TaskiqResult from taskiq.depends.progress_tracker import TaskProgress from taskiq.events import TaskiqEvents -from taskiq.exceptions import UnknownTaskError +from taskiq.exceptions import SendTaskError, UnknownTaskError from taskiq.message import BrokerMessage from taskiq.receiver import Receiver -from taskiq.utils import maybe_awaitable _ReturnType = TypeVar("_ReturnType") @@ -145,7 +148,40 @@ def __init__( propagate_exceptions=propagate_exceptions, ) self.await_inplace = await_inplace + self._inflight_tasks: dict[ + asyncio.Future[None], + asyncio.Task[Any] | None, + ] = {} self._running_tasks: set[asyncio.Task[Any]] = set() + self._running_task_error: BaseException | None = None + self._accepting_tasks = True + self._shutdown_lock = asyncio.Lock() + self._shutdown_complete = False + self._executor_shutdown_complete = False + self._pending_shutdown_error: BaseException | None = None + self._managed_context: ContextVar[asyncio.Future[None] | None] = ContextVar( + "taskiq_inmemory_managed_context", + default=None, + ) + self._deferred_execution: ContextVar[ + asyncio.Future[BrokerMessage | None] | None + ] = ContextVar("taskiq_inmemory_deferred_execution", default=None) + + @contextmanager + def _send_lifecycle(self) -> Iterator[None]: + """Register a client send before it starts using broker resources.""" + try: + self._ensure_accepting_tasks() + except Exception as exc: + raise SendTaskError from exc + + send_done = self._track_inflight_send() + context_token = self._managed_context.set(send_done) + try: + yield + finally: + self._managed_context.reset(context_token) + self._finish_inflight_send(send_done) async def kick(self, message: BrokerMessage) -> None: """ @@ -157,18 +193,201 @@ async def kick(self, message: BrokerMessage) -> None: :raises TaskiqError: if someone wants to kick unknown task. """ - target_task = self.find_task(message.task_name) - if target_task is None: + deferred_execution = self._deferred_execution.get() + if deferred_execution is not None and not deferred_execution.done(): + self._ensure_known_task(message) + deferred_execution.set_result(message) + return + + self._ensure_accepting_tasks() + self._ensure_known_task(message) + send_done = self._track_inflight_send() + context_token = self._managed_context.set(send_done) + try: + await self._dispatch_message(message) + finally: + self._managed_context.reset(context_token) + self._finish_inflight_send(send_done) + + async def _kick_with_post_send( + self, + message: BrokerMessage, + post_send: Callable[[], Awaitable[None]], + ) -> None: + """Run client post-send hooks before local task execution starts.""" + accepted_message, kick_error = await self._accept_message(message) + + if kick_error is not None: + self._raise_kick_error(accepted_message, kick_error, post_send) + + if accepted_message is None: + await post_send() + return + + message = accepted_message + + try: + await post_send() + except asyncio.CancelledError: + self._start_background_task(message) + raise + except Exception as post_send_error: + if not self.await_inplace: + self._start_background_task(message) + raise + await self._finish_after_post_send_error( + message, + post_send_error, + ) + except BaseException: + self._start_background_task(message) + raise + else: + try: + await self._dispatch_message(message) + except Exception as exc: + raise SendTaskError from exc + + def _raise_kick_error( + self, + accepted_message: BrokerMessage | None, + kick_error: BaseException, + post_send: Callable[[], Awaitable[None]], + ) -> NoReturn: + """Preserve accepted work and expose the public kick failure.""" + if accepted_message is not None: + self._start_accepted_send(accepted_message, post_send) + if isinstance(kick_error, Exception): + raise SendTaskError from kick_error + raise kick_error + + async def _finish_after_post_send_error( + self, + message: BrokerMessage, + post_send_error: Exception, + ) -> NoReturn: + """Finish accepted inline work without replacing its post-send error.""" + task = self._create_managed_task( + self.receiver.callback(message=message.message), + ) + try: + await asyncio.wait((task,)) + except asyncio.CancelledError: + self._track_running_task(task) + raise + try: + task.result() + except BaseException as execution_error: + raise post_send_error from execution_error + raise post_send_error + + def _ensure_known_task(self, message: BrokerMessage) -> None: + """Reject direct sends for tasks that are not registered locally.""" + if self.find_task(message.task_name) is None: raise UnknownTaskError(task_name=message.task_name) - receiver_cb = self.receiver.callback(message=message.message) + def _ensure_accepting_tasks(self) -> None: + """Reject work once broker shutdown has started.""" + if not self._accepting_tasks: + raise RuntimeError("InMemoryBroker is shutting down.") + + async def _dispatch_message(self, message: BrokerMessage) -> None: + """Execute a validated message inline or track it in the background.""" if self.await_inplace: - await receiver_cb + await self.receiver.callback(message=message.message) return - task = asyncio.create_task(receiver_cb) + self._start_background_task(message) + + async def _accept_message( + self, + message: BrokerMessage, + ) -> tuple[BrokerMessage | None, BaseException | None]: + """Run the public send extension point while deferring local execution.""" + deferred_execution = asyncio.get_running_loop().create_future() + dispatch_token = self._deferred_execution.set(deferred_execution) + kick_error: BaseException | None = None + try: + await self.kick(message) + except BaseException as exc: + kick_error = exc + finally: + self._deferred_execution.reset(dispatch_token) + if not deferred_execution.done(): + deferred_execution.set_result(None) + return deferred_execution.result(), kick_error + + def _start_background_task(self, message: BrokerMessage) -> None: + """Start and track one local receiver callback.""" + task = self._create_managed_task( + self.receiver.callback(message=message.message), + ) + self._track_running_task(task) + + def _start_accepted_send( + self, + message: BrokerMessage, + post_send: Callable[[], Awaitable[None]], + ) -> None: + """Finish an accepted send in order after its public override fails.""" + task = self._create_managed_task( + self._run_accepted_send(message, post_send), + ) + self._track_running_task(task) + + async def _run_accepted_send( + self, + message: BrokerMessage, + post_send: Callable[[], Awaitable[None]], + ) -> None: + """Run post-send and accepted execution while preserving first failure.""" + try: + await post_send() + except BaseException as post_send_error: + try: + await self.receiver.callback(message=message.message) + except BaseException as execution_error: + raise post_send_error from execution_error + raise + await self.receiver.callback(message=message.message) + + def _create_managed_task( + self, + awaitable: Awaitable[None], + ) -> asyncio.Task[None]: + """Create a broker-owned task with a bounded managed context.""" + managed_done = asyncio.get_running_loop().create_future() + return asyncio.create_task(self._run_managed(awaitable, managed_done)) + + async def _run_managed( + self, + awaitable: Awaitable[None], + managed_done: asyncio.Future[None], + ) -> None: + """Run one awaitable and release its inherited managed context.""" + context_token = self._managed_context.set(managed_done) + try: + await awaitable + finally: + self._managed_context.reset(context_token) + managed_done.set_result(None) + + def _track_running_task(self, task: asyncio.Task[Any]) -> None: + """Track one local receiver callback until it finishes.""" self._running_tasks.add(task) - task.add_done_callback(self._running_tasks.discard) + task.add_done_callback(self._on_running_task_done) + + def _track_inflight_send(self) -> asyncio.Future[None]: + """Track one accepted send until execution ownership is settled.""" + send_done = asyncio.get_running_loop().create_future() + self._inflight_tasks[send_done] = asyncio.current_task() + return send_done + + def _finish_inflight_send(self, send_done: asyncio.Future[None]) -> None: + """Release an accepted send after execution ownership is settled.""" + self._inflight_tasks.pop(send_done, None) + if not send_done.done(): + send_done.set_result(None) def listen(self) -> AsyncGenerator[bytes, None]: """ @@ -187,20 +406,149 @@ async def wait_all(self) -> None: Useful when used in testing and you need to await all sent tasks before asserting results. + Cancelling this waiter does not cancel accepted executions. """ - to_await = list(self._running_tasks) - for task in to_await: - await task + self._ensure_external_drain("wait_all") + + while self._inflight_tasks or self._running_tasks: + inflight_sends = tuple(self._inflight_tasks) + running_tasks = tuple(self._running_tasks) + await asyncio.gather( + *(asyncio.shield(task) for task in (*inflight_sends, *running_tasks)), + return_exceptions=True, + ) + for send_done in inflight_sends: + self._inflight_tasks.pop(send_done, None) + for task in running_tasks: + self._on_running_task_done(task) + + if self._running_task_error is not None: + task_error = self._running_task_error + self._running_task_error = None + raise task_error + + def _on_running_task_done(self, task: asyncio.Task[Any]) -> None: + """Release a completed task and retain its first failure until drain.""" + if task not in self._running_tasks: + return - 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)) + self._running_tasks.remove(task) + try: + task.result() + except BaseException as exc: + if self._running_task_error is None: + self._running_task_error = exc + + def _ensure_external_drain(self, operation: str) -> None: + """Reject drain calls made by work that the drain must await.""" + managed_done = self._managed_context.get() + if managed_done is not None and not managed_done.done(): + raise RuntimeError( + f"InMemoryBroker.{operation}() cannot be called from " + "a task managed by this broker.", + ) + + current_task = asyncio.current_task() + if current_task is None: + return + + is_managed_execution = ( + current_task in self._running_tasks + or current_task in self._inflight_tasks + or current_task in self._inflight_tasks.values() + ) + if is_managed_execution: + raise RuntimeError( + f"InMemoryBroker.{operation}() cannot be called from " + "a task managed by this broker.", + ) + + def _get_startup_events(self) -> tuple[TaskiqEvents, ...]: + """Run both sides because tasks execute in the client process.""" + return (TaskiqEvents.CLIENT_STARTUP, TaskiqEvents.WORKER_STARTUP) + + def _get_shutdown_events(self) -> tuple[TaskiqEvents, ...]: + """Shut down both 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 resources and stop the executor. + + Cancellation during accepted-work drain is propagated after cleanup. + """ + self._ensure_external_drain("shutdown") + self._accepting_tasks = False + + async with self._shutdown_lock: + if self._shutdown_complete: + return + await self._shutdown_once() + + async def _shutdown_once(self) -> None: + """Run one retryable shutdown attempt while holding the lifecycle lock.""" + shutdown_error = self._pending_shutdown_error + drain_cancelled = False + + while True: + try: + with anyio.CancelScope(shield=drain_cancelled): + await self.wait_all() + except asyncio.CancelledError as exc: + shutdown_error = self._remember_shutdown_error( + shutdown_error, + exc, + ) + drain_cancelled = True + continue + except BaseException as exc: + shutdown_error = self._remember_shutdown_error( + shutdown_error, + exc, + ) + break + + shutdown_errors = [] if shutdown_error is None else [shutdown_error] + drain_cancelled = await self._shutdown_executor( + shutdown_errors, + drain_cancelled, + ) + + try: + with anyio.CancelScope(shield=drain_cancelled): + await super()._shutdown_resources(shutdown_errors) + except BaseException: + if shutdown_errors: + self._pending_shutdown_error = shutdown_errors[0] + raise + + self._pending_shutdown_error = None + self._shutdown_complete = True + if shutdown_errors: + raise shutdown_errors[0] + + async def _shutdown_executor( + self, + shutdown_errors: list[BaseException], + drain_cancelled: bool, + ) -> bool: + """Close the sync executor without blocking the event loop.""" + if self._executor_shutdown_complete: + return drain_cancelled + + while True: + try: + with anyio.CancelScope(shield=drain_cancelled): + await anyio.to_thread.run_sync( + self.executor.shutdown, + abandon_on_cancel=False, + ) + except asyncio.CancelledError as exc: + self._record_shutdown_error(shutdown_errors, exc) + drain_cancelled = True + continue + except BaseException as exc: # pragma: no cover + self._record_shutdown_error(shutdown_errors, exc) + break + + self._executor_shutdown_complete = True + return drain_cancelled diff --git a/taskiq/kicker.py b/taskiq/kicker.py index b8b9c073..d73d9840 100644 --- a/taskiq/kicker.py +++ b/taskiq/kicker.py @@ -1,6 +1,7 @@ from collections.abc import Coroutine from dataclasses import asdict, is_dataclass from datetime import datetime, timedelta +from functools import partial from logging import getLogger from types import CoroutineType from typing import ( @@ -141,10 +142,10 @@ async def kiq( **kwargs: _FuncParams.kwargs, ) -> Any: """ - This method sends function call over the network. + Prepare and dispatch a task invocation through the current broker. - It gets current broker and calls it's kick method, - returning what it returns. + Send middleware runs around broker dispatch before the task handle is + returned. :param args: function's arguments. :param kwargs: function's key word arguments. @@ -156,18 +157,22 @@ 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)) - try: - await self.broker.kick(self.broker.formatter.dumps(message)) - except Exception as exc: - raise SendTaskError from exc - - for middleware in reversed(self.broker.middlewares): - if middleware.__class__.post_send != TaskiqMiddleware.post_send: - await maybe_awaitable(middleware.post_send(message)) + with self.broker._send_lifecycle(): # noqa: SLF001 + 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)) + + try: + broker_message = self.broker.formatter.dumps(message) + except Exception as exc: + raise SendTaskError from exc + + # AsyncKicker and AsyncBroker share this package-internal send boundary. + await self.broker._kick_with_post_send( # noqa: SLF001 + broker_message, + partial(self._run_post_send, message), + ) return AsyncTaskiqTask( task_id=message.task_id, @@ -175,6 +180,12 @@ async def kiq( return_type=self.return_type, # type: ignore # (pyright issue) ) + async def _run_post_send(self, message: TaskiqMessage) -> None: + """Run post-send middleware in reverse registration order.""" + for middleware in reversed(self.broker.middlewares): + if middleware.__class__.post_send != TaskiqMiddleware.post_send: + await maybe_awaitable(middleware.post_send(message)) + async def schedule_by_cron( self, source: "ScheduleSource", diff --git a/tests/abc/test_broker.py b/tests/abc/test_broker.py index 636f9576..6c2bcb83 100644 --- a/tests/abc/test_broker.py +++ b/tests/abc/test_broker.py @@ -1,12 +1,17 @@ +import asyncio from collections.abc import AsyncGenerator +from contextvars import ContextVar, Token from copy import copy import pytest from taskiq.abc.broker import AsyncBroker +from taskiq.abc.middleware import TaskiqMiddleware +from taskiq.cli.worker.run import shutdown_broker from taskiq.decor import AsyncTaskiqDecoratedTask from taskiq.events import TaskiqEvents -from taskiq.message import BrokerMessage +from taskiq.exceptions import SendTaskError +from taskiq.message import BrokerMessage, TaskiqMessage from taskiq.state import TaskiqState @@ -30,6 +35,79 @@ async def listen(self) -> AsyncGenerator[BrokerMessage, None]: # type: ignore """ +class _RecordingSendBroker(_TestBroker): + """Record transport dispatch and raise an optional send failure.""" + + def __init__( + self, + events: list[str], + kick_error: Exception | None = None, + ) -> None: + super().__init__() + self.events = events + self.kick_error = kick_error + + async def kick(self, message: BrokerMessage) -> None: + self.events.append("kick") + if self.kick_error is not None: + raise self.kick_error + + +class _RecordingSendMiddleware(TaskiqMiddleware): + """Record client send hooks and raise an optional post-send failure.""" + + def __init__( + self, + events: list[str], + post_send_error: Exception | None = None, + ) -> None: + super().__init__() + self.events = events + self.post_send_error = post_send_error + + def pre_send(self, message: TaskiqMessage) -> TaskiqMessage: + self.events.append("pre_send") + return message + + def post_send(self, message: TaskiqMessage) -> None: + self.events.append("post_send") + if self.post_send_error is not None: + raise self.post_send_error + + +class _BlockingShutdownMiddleware(TaskiqMiddleware): + """Hold base broker shutdown until cancelled or explicitly released.""" + + def __init__( + self, + started: asyncio.Event, + release: asyncio.Event, + ) -> None: + super().__init__() + self.started = started + self.release = release + + async def shutdown(self) -> None: + self.started.set() + await self.release.wait() + + +class _ContextLifecycleMiddleware(TaskiqMiddleware): + """Pair one ContextVar token across startup and shutdown.""" + + def __init__(self, context: ContextVar[str]) -> None: + super().__init__() + self.context = context + self.token: Token[str] | None = None + + async def startup(self) -> None: + self.token = self.context.set("started") + + async def shutdown(self) -> None: + assert self.token is not None + self.context.reset(self.token) + + def test_decorator_success() -> None: """Test that decoration without parameters works.""" tbrok = _TestBroker() @@ -82,6 +160,87 @@ async def test_task() -> None: ... assert test_task.labels == old_labels +async def test_kicker_preserves_external_broker_send_order() -> None: + events: list[str] = [] + broker = _RecordingSendBroker(events) + broker.with_middlewares(_RecordingSendMiddleware(events)) + + @broker.task + async def task() -> None: + return None + + await task.kiq() + + assert events == ["pre_send", "kick", "post_send"] + + +async def test_kicker_wraps_external_broker_transport_error() -> None: + events: list[str] = [] + kick_error = RuntimeError("transport failed") + broker = _RecordingSendBroker(events, kick_error=kick_error) + broker.with_middlewares(_RecordingSendMiddleware(events)) + + @broker.task + async def task() -> None: + return None + + with pytest.raises(SendTaskError) as exc_info: + await task.kiq() + + assert exc_info.value.__cause__ is kick_error + assert events == ["pre_send", "kick"] + + +async def test_kicker_preserves_post_send_error_type() -> None: + events: list[str] = [] + post_send_error = RuntimeError("post-send failed") + broker = _RecordingSendBroker(events) + broker.with_middlewares( + _RecordingSendMiddleware(events, post_send_error=post_send_error), + ) + + @broker.task + async def task() -> None: + return None + + with pytest.raises(RuntimeError) as exc_info: + await task.kiq() + + assert exc_info.value is post_send_error + assert events == ["pre_send", "kick", "post_send"] + + +async def test_base_shutdown_preserves_lifecycle_context() -> None: + lifecycle_context = ContextVar("lifecycle_context", default="outside") + broker = _TestBroker().with_middlewares( + _ContextLifecycleMiddleware(lifecycle_context), + ) + + async with broker: + assert lifecycle_context.get() == "started" + + assert lifecycle_context.get() == "outside" + + +async def test_worker_timeout_cancels_base_broker_shutdown() -> None: + shutdown_started = asyncio.Event() + release_shutdown = asyncio.Event() + broker = _TestBroker().with_middlewares( + _BlockingShutdownMiddleware(shutdown_started, release_shutdown), + ) + loop = asyncio.get_running_loop() + release_handle = loop.call_later(0.5, release_shutdown.set) + started_at = loop.time() + + try: + await shutdown_broker(broker, timeout=0.01) + finally: + release_handle.cancel() + + assert shutdown_started.is_set() + assert loop.time() - started_at < 0.2 + + @pytest.mark.anyio @pytest.mark.parametrize( ("is_worker_process", "startup", "shutdown"), diff --git a/tests/brokers/inmemory_contract_support.py b/tests/brokers/inmemory_contract_support.py new file mode 100644 index 00000000..99e9a704 --- /dev/null +++ b/tests/brokers/inmemory_contract_support.py @@ -0,0 +1,372 @@ +import asyncio +from typing import Any, cast + +from taskiq import InMemoryBroker, TaskiqMessage, TaskiqResult +from taskiq.abc.middleware import TaskiqMiddleware +from taskiq.brokers.inmemory_broker import InmemoryResultBackend +from taskiq.message import BrokerMessage + + +class LifecycleError(Exception): + """Marker exception for lifecycle fault tests.""" + + +class PostSendError(Exception): + """Marker exception for post-send compatibility tests.""" + + +class RecordingLifecycleMiddleware(TaskiqMiddleware): + """Record middleware lifecycle calls and optional shutdown failure.""" + + 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: + self.events.append("middleware.startup") + + async def shutdown(self) -> None: + 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 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: + self.events.append("backend.startup") + if self.startup_error is not None: + raise self.startup_error + + async def shutdown(self) -> None: + 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: BaseException) -> None: + super().__init__() + self.error = error + + def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage: + raise self.error + + +class BlockingPostSendMiddleware(TaskiqMiddleware): + """Pause post-send to prove execution cannot overlap it.""" + + def __init__( + self, + events: list[str], + post_send_started: asyncio.Event, + release_post_send: asyncio.Event, + ) -> None: + super().__init__() + self.events = events + self.post_send_started = post_send_started + self.release_post_send = release_post_send + + def pre_send(self, message: TaskiqMessage) -> TaskiqMessage: + self.events.append("pre_send") + return message + + async def post_send(self, message: TaskiqMessage) -> None: + self.events.append("post_send.started") + self.post_send_started.set() + await self.release_post_send.wait() + self.events.append("post_send.finished") + + def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage: + self.events.append("pre_execute") + return message + + def post_execute( + self, + message: TaskiqMessage, + result: TaskiqResult[Any], + ) -> None: + self.events.append("post_execute") + + +class BlockingPreSendMiddleware(TaskiqMiddleware): + """Pause pre-send and expose subsequent resource shutdown.""" + + def __init__( + self, + pre_send_started: asyncio.Event, + release_pre_send: asyncio.Event, + shutdown_started: asyncio.Event, + ) -> None: + super().__init__() + self.pre_send_started = pre_send_started + self.release_pre_send = release_pre_send + self.shutdown_started = shutdown_started + self.pre_send_calls = 0 + + async def pre_send(self, message: TaskiqMessage) -> TaskiqMessage: + self.pre_send_calls += 1 + self.pre_send_started.set() + await self.release_pre_send.wait() + return message + + async def shutdown(self) -> None: + self.shutdown_started.set() + + +class BlockingShutdownMiddleware(TaskiqMiddleware): + """Block middleware shutdown until the test releases it.""" + + def __init__( + self, + shutdown_started: asyncio.Event, + release_shutdown: asyncio.Event, + shutdown_finished: asyncio.Event, + ) -> None: + super().__init__() + self.shutdown_started = shutdown_started + self.release_shutdown = release_shutdown + self.shutdown_finished = shutdown_finished + self.shutdown_calls = 0 + + async def shutdown(self) -> None: + self.shutdown_calls += 1 + self.shutdown_started.set() + await self.release_shutdown.wait() + self.shutdown_finished.set() + + +class BlockingShutdownResultBackend(InmemoryResultBackend[Any]): + """Block backend shutdown until the test releases it.""" + + def __init__( + self, + shutdown_started: asyncio.Event, + release_shutdown: asyncio.Event, + shutdown_finished: asyncio.Event, + ) -> None: + super().__init__() + self.shutdown_started = shutdown_started + self.release_shutdown = release_shutdown + self.shutdown_finished = shutdown_finished + + async def shutdown(self) -> None: + self.shutdown_started.set() + await self.release_shutdown.wait() + self.shutdown_finished.set() + + +class CoordinatedPostSendMiddleware(TaskiqMiddleware): + """Hold two post-send hooks at different deterministic boundaries.""" + + def __init__( + self, + first_task_name: str, + second_task_name: str, + first_post_send_started: asyncio.Event, + second_post_send_started: asyncio.Event, + release_second_post_send: asyncio.Event, + ) -> None: + super().__init__() + self.first_task_name = first_task_name + self.second_task_name = second_task_name + self.first_post_send_started = first_post_send_started + self.second_post_send_started = second_post_send_started + self.release_second_post_send = release_second_post_send + + async def post_send(self, message: TaskiqMessage) -> None: + if message.task_name == self.first_task_name: + self.first_post_send_started.set() + await self.second_post_send_started.wait() + return + + if message.task_name == self.second_task_name: + self.second_post_send_started.set() + await self.release_second_post_send.wait() + + +class FailingPostSendMiddleware(TaskiqMiddleware): + """Fail after transport acceptance.""" + + def __init__(self, error: Exception) -> None: + super().__init__() + self.error = error + + def post_send(self, message: TaskiqMessage) -> None: + raise self.error + + +class ReentrantDrainMiddleware(TaskiqMiddleware): + """Attempt a broker drain from the active send lifecycle.""" + + def __init__(self, operation: str) -> None: + super().__init__() + self.operation = operation + + async def post_send(self, message: TaskiqMessage) -> None: + await getattr(self.broker, self.operation)() + + +class ChildReentrantDrainMiddleware(TaskiqMiddleware): + """Attempt a broker drain from a child of the active send lifecycle.""" + + def __init__(self, operation: str) -> None: + super().__init__() + self.operation = operation + + async def post_send(self, message: TaskiqMessage) -> None: + drain = asyncio.create_task(getattr(self.broker, self.operation)()) + await drain + + +class DeferredChildDrainMiddleware(TaskiqMiddleware): + """Drain from a child after its originating send lifecycle has ended.""" + + def __init__(self, release: asyncio.Event) -> None: + super().__init__() + self.release = release + self.drain_task: asyncio.Task[BaseException | None] | None = None + + async def post_send(self, message: TaskiqMessage) -> None: + self.drain_task = asyncio.create_task(self._drain_after_release()) + + async def _drain_after_release(self) -> BaseException | None: + await self.release.wait() + try: + await cast(InMemoryBroker, self.broker).wait_all() + except BaseException as exc: + return exc + return None + + +class DrainSignallingInMemoryBroker(InMemoryBroker): + """Expose drain attempts for deterministic overlap tests.""" + + def __init__( + self, + drain_started: asyncio.Event, + drain_restarted: asyncio.Event | None = None, + *, + await_inplace: bool = False, + ) -> None: + super().__init__(await_inplace=await_inplace) + self.drain_started = drain_started + self.drain_restarted = drain_restarted + self.drain_calls = 0 + + async def wait_all(self) -> None: + self.drain_calls += 1 + if self.drain_calls == 1: + self.drain_started.set() + elif self.drain_restarted is not None: + self.drain_restarted.set() + await super().wait_all() + + +class RecordingKickInMemoryBroker(InMemoryBroker): + """Record calls through the public kick extension point.""" + + def __init__(self, *, await_inplace: bool) -> None: + super().__init__(await_inplace=await_inplace) + self.kick_calls = 0 + + async def kick(self, message: BrokerMessage) -> None: + self.kick_calls += 1 + await super().kick(message) + + +class BlockingAfterAcceptKickInMemoryBroker(InMemoryBroker): + """Block a public kick override after its base send was accepted.""" + + def __init__(self, *, await_inplace: bool) -> None: + super().__init__(await_inplace=await_inplace) + self.accepted = asyncio.Event() + self.release = asyncio.Event() + + async def kick(self, message: BrokerMessage) -> None: + await super().kick(message) + self.accepted.set() + await self.release.wait() + + +class FailingAfterAcceptKickInMemoryBroker(InMemoryBroker): + """Fail a public kick override after its base send was accepted.""" + + def __init__( + self, + kick_error: Exception, + *, + await_inplace: bool, + ) -> None: + super().__init__(await_inplace=await_inplace) + self.kick_error = kick_error + + async def kick(self, message: BrokerMessage) -> None: + await super().kick(message) + raise self.kick_error + + +class TransformingKickInMemoryBroker(InMemoryBroker): + """Replace the locally accepted message through the public kick seam.""" + + def __init__( + self, + replacement_task_name: str, + *, + await_inplace: bool, + ) -> None: + super().__init__(await_inplace=await_inplace) + self.replacement_task_name = replacement_task_name + + async def kick(self, message: BrokerMessage) -> None: + accepted_message = self.formatter.loads(message.message) + accepted_message.task_name = self.replacement_task_name + await super().kick(self.formatter.dumps(accepted_message)) + + +class RejectingKickInMemoryBroker(InMemoryBroker): + """Reject one public send before post-send middleware can observe it.""" + + def __init__( + self, + kick_error: Exception, + *, + await_inplace: bool, + ) -> None: + super().__init__(await_inplace=await_inplace) + self.kick_error = kick_error + self.kick_calls = 0 + + async def kick(self, message: BrokerMessage) -> None: + self.kick_calls += 1 + raise self.kick_error + + +class ReplacingKickInMemoryBroker(InMemoryBroker): + """Replace local dispatch while preserving the public send pipeline.""" + + def __init__(self, *, await_inplace: bool) -> None: + super().__init__(await_inplace=await_inplace) + self.kick_calls = 0 + + async def kick(self, message: BrokerMessage) -> None: + self.kick_calls += 1 diff --git a/tests/brokers/test_inmemory_lifecycle.py b/tests/brokers/test_inmemory_lifecycle.py new file mode 100644 index 00000000..b28cde8f --- /dev/null +++ b/tests/brokers/test_inmemory_lifecycle.py @@ -0,0 +1,788 @@ +import asyncio +import logging +import threading +from typing import Any + +import anyio +import pytest + +from taskiq import InMemoryBroker, TaskiqMessage +from taskiq.events import TaskiqEvents +from taskiq.exceptions import SendTaskError +from taskiq.state import TaskiqState +from tests.brokers.inmemory_contract_support import ( + BlockingPreSendMiddleware, + BlockingShutdownMiddleware, + BlockingShutdownResultBackend, + DrainSignallingInMemoryBroker, + FailingExecutionMiddleware, + LifecycleError, + RecordingLifecycleMiddleware, + RecordingResultBackend, +) + + +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)) + + @broker.on_event(TaskiqEvents.CLIENT_STARTUP) + def record_client_startup(state: TaskiqState) -> None: + events.append("client.startup") + + @broker.on_event(TaskiqEvents.WORKER_STARTUP) + def record_worker_startup(state: TaskiqState) -> None: + events.append("worker.startup") + + @broker.on_event(TaskiqEvents.CLIENT_SHUTDOWN) + def record_client_shutdown(state: TaskiqState) -> None: + events.append("client.shutdown") + + @broker.on_event(TaskiqEvents.WORKER_SHUTDOWN) + def record_worker_shutdown(state: TaskiqState) -> None: + events.append("worker.shutdown") + + await broker.startup() + await broker.shutdown() + 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_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 + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + broker.executor.submit(int) + + await broker.shutdown() + assert events == [ + "middleware.startup", + "backend.startup", + "client.shutdown", + "worker.shutdown", + "middleware.shutdown", + "backend.shutdown", + ] + + +async def test_shutdown_cancellation_during_event_is_retryable() -> None: + events: list[str] = [] + shutdown_started = asyncio.Event() + keep_shutdown_blocked = asyncio.Event() + broker = InMemoryBroker() + broker.with_middlewares(RecordingLifecycleMiddleware(events)) + broker.with_result_backend(RecordingResultBackend(events)) + + @broker.on_event(TaskiqEvents.CLIENT_SHUTDOWN) + async def block_client_shutdown(state: TaskiqState) -> None: + events.append("client.shutdown") + shutdown_started.set() + await keep_shutdown_blocked.wait() + + await broker.startup() + shutdown_task = asyncio.create_task(broker.shutdown()) + await asyncio.wait_for(shutdown_started.wait(), timeout=1) + shutdown_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(shutdown_task, timeout=1) + + assert events == [ + "middleware.startup", + "backend.startup", + "client.shutdown", + ] + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + broker.executor.submit(int) + + keep_shutdown_blocked.set() + await broker.shutdown() + + assert events == [ + "middleware.startup", + "backend.startup", + "client.shutdown", + "client.shutdown", + "middleware.shutdown", + "backend.shutdown", + ] + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + broker.executor.submit(int) + + +@pytest.mark.parametrize("resource_kind", ["middleware", "backend"]) +async def test_shutdown_cancellation_during_resource_is_retryable( + resource_kind: str, +) -> None: + shutdown_started = asyncio.Event() + release_shutdown = asyncio.Event() + shutdown_finished = asyncio.Event() + broker = InMemoryBroker() + + if resource_kind == "middleware": + broker.with_middlewares( + BlockingShutdownMiddleware( + shutdown_started, + release_shutdown, + shutdown_finished, + ), + ) + else: + broker.with_result_backend( + BlockingShutdownResultBackend( + shutdown_started, + release_shutdown, + shutdown_finished, + ), + ) + + await broker.startup() + shutdown_task = asyncio.create_task(broker.shutdown()) + await asyncio.wait_for(shutdown_started.wait(), timeout=1) + shutdown_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(shutdown_task, timeout=1) + + assert not shutdown_finished.is_set() + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + broker.executor.submit(int) + + release_shutdown.set() + await broker.shutdown() + + assert shutdown_finished.is_set() + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + broker.executor.submit(int) + + +async def test_shutdown_preserves_failure_that_precedes_cancellation( + caplog: pytest.LogCaptureFixture, +) -> None: + shutdown_started = asyncio.Event() + release_shutdown = asyncio.Event() + shutdown_finished = asyncio.Event() + lifecycle_error = LifecycleError("shutdown handler failed") + shutdown_event_calls = 0 + broker = InMemoryBroker().with_middlewares( + BlockingShutdownMiddleware( + shutdown_started, + release_shutdown, + shutdown_finished, + ), + ) + + @broker.on_event(TaskiqEvents.CLIENT_SHUTDOWN) + def fail_shutdown(state: TaskiqState) -> None: + nonlocal shutdown_event_calls + shutdown_event_calls += 1 + raise lifecycle_error + + caplog.set_level(logging.ERROR, logger="taskiq") + shutdown_task = asyncio.create_task(broker.shutdown()) + await asyncio.wait_for(shutdown_started.wait(), timeout=1) + shutdown_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(shutdown_task, timeout=1) + + assert not shutdown_finished.is_set() + assert caplog.text.count("Additional error while shutting down") == 0 + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + broker.executor.submit(int) + + release_shutdown.set() + with pytest.raises(LifecycleError) as retry_exc_info: + await broker.shutdown() + + assert retry_exc_info.value is lifecycle_error + assert shutdown_finished.is_set() + assert shutdown_event_calls == 1 + assert caplog.text.count("Additional error while shutting down") == 0 + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + broker.executor.submit(int) + + +async def test_concurrent_shutdown_closes_resources_once() -> None: + shutdown_started = asyncio.Event() + release_shutdown = asyncio.Event() + shutdown_finished = asyncio.Event() + middleware = BlockingShutdownMiddleware( + shutdown_started, + release_shutdown, + shutdown_finished, + ) + broker = InMemoryBroker().with_middlewares(middleware) + + first_shutdown = asyncio.create_task(broker.shutdown()) + await asyncio.wait_for(shutdown_started.wait(), timeout=1) + second_shutdown = asyncio.create_task(broker.shutdown()) + await asyncio.sleep(0) + + assert not second_shutdown.done() + release_shutdown.set() + await asyncio.wait_for( + asyncio.gather(first_shutdown, second_shutdown), + timeout=1, + ) + + assert middleware.shutdown_calls == 1 + assert shutdown_finished.is_set() + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + broker.executor.submit(int) + + +async def test_shutdown_retry_skips_completed_resources() -> None: + events: list[str] = [] + shutdown_started = asyncio.Event() + release_shutdown = asyncio.Event() + shutdown_finished = asyncio.Event() + completed_middleware = RecordingLifecycleMiddleware(events) + blocking_middleware = BlockingShutdownMiddleware( + shutdown_started, + release_shutdown, + shutdown_finished, + ) + broker = InMemoryBroker().with_middlewares( + completed_middleware, + blocking_middleware, + ) + + first_attempt = asyncio.create_task(broker.shutdown()) + await asyncio.wait_for(shutdown_started.wait(), timeout=1) + first_attempt.cancel() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(first_attempt, timeout=1) + + assert events == ["middleware.shutdown"] + assert blocking_middleware.shutdown_calls == 1 + + release_shutdown.set() + await broker.shutdown() + + assert events == ["middleware.shutdown"] + assert blocking_middleware.shutdown_calls == 2 + assert shutdown_finished.is_set() + + +async def test_task_failure_does_not_shield_resource_cancellation() -> None: + shutdown_started = asyncio.Event() + release_shutdown = asyncio.Event() + shutdown_finished = asyncio.Event() + shutdown_exited = asyncio.Event() + task_error = LifecycleError("in-memory task failed") + broker = InMemoryBroker() + broker.with_middlewares(FailingExecutionMiddleware(task_error)) + broker.with_middlewares( + BlockingShutdownMiddleware( + shutdown_started, + release_shutdown, + shutdown_finished, + ), + ) + + @broker.task + async def failing_task() -> None: + return None + + await failing_task.kiq() + running_task = next(iter(broker._running_tasks)) + await asyncio.wait((running_task,)) + assert not broker._running_tasks + + cancel_scope = anyio.CancelScope() + + async def run_shutdown() -> None: + with cancel_scope: + await broker.shutdown() + shutdown_exited.set() + + shutdown_task = asyncio.create_task(run_shutdown()) + await asyncio.wait_for(shutdown_started.wait(), timeout=1) + cancel_scope.cancel() + + try: + await asyncio.wait_for(shutdown_exited.wait(), timeout=1) + finally: + release_shutdown.set() + await asyncio.wait_for(shutdown_task, timeout=1) + + assert not shutdown_finished.is_set() + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + broker.executor.submit(int) + + with pytest.raises(LifecycleError) as exc_info: + await broker.shutdown() + + assert exc_info.value is task_error + assert shutdown_finished.is_set() + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + broker.executor.submit(int) + + +async def test_cancelled_shutdown_drains_before_resource_cleanup() -> None: + task_started = asyncio.Event() + finish_task = asyncio.Event() + drain_started = asyncio.Event() + drain_restarted = asyncio.Event() + events: list[str] = [] + broker = DrainSignallingInMemoryBroker(drain_started, drain_restarted) + broker.with_middlewares(RecordingLifecycleMiddleware(events)) + broker.with_result_backend(RecordingResultBackend(events)) + + @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 asyncio.wait_for(task_started.wait(), timeout=1) + + shutdown_task = asyncio.create_task(broker.shutdown()) + await asyncio.wait_for(drain_started.wait(), timeout=1) + shutdown_task.cancel() + await asyncio.wait_for(drain_restarted.wait(), timeout=1) + + assert not shutdown_task.done() + finish_task.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(asyncio.shield(shutdown_task), timeout=1) + + assert events == [ + "middleware.startup", + "backend.startup", + "task.started", + "task.finished", + "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_anyio_cancellation_does_not_spin_while_draining() -> None: + task_started = asyncio.Event() + finish_task = asyncio.Event() + drain_started = asyncio.Event() + broker = DrainSignallingInMemoryBroker(drain_started) + + @broker.task + async def running_task() -> None: + task_started.set() + await finish_task.wait() + + await running_task.kiq() + await asyncio.wait_for(task_started.wait(), timeout=1) + + cancel_scope = anyio.CancelScope() + + async def cancel_then_release() -> None: + await drain_started.wait() + cancel_scope.cancel() + await asyncio.sleep(0) + finish_task.set() + + cancellation_driver = asyncio.create_task(cancel_then_release()) + with cancel_scope: + await broker.shutdown() + await asyncio.wait_for(cancellation_driver, timeout=1) + + assert broker.drain_calls == 2 + 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() + drain_started = asyncio.Event() + events: list[str] = [] + broker = DrainSignallingInMemoryBroker(drain_started) + 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 asyncio.wait_for(task_started.wait(), timeout=1) + + shutdown_task = asyncio.create_task(broker.shutdown()) + await asyncio.wait_for(drain_started.wait(), timeout=1) + assert not shutdown_task.done() + + finish_task.set() + await asyncio.wait_for(shutdown_task, timeout=1) + + assert events == [ + "middleware.startup", + "backend.startup", + "task.started", + "task.finished", + "client.shutdown", + "middleware.shutdown", + "backend.shutdown", + ] + assert not broker._running_tasks + + +async def test_shutdown_waits_for_cancelled_inline_sync_execution() -> None: + events: list[str] = [] + task_started = threading.Event() + finish_task = threading.Event() + broker = InMemoryBroker(await_inplace=True) + broker.with_middlewares(RecordingLifecycleMiddleware(events)) + + @broker.task + def sync_task() -> None: + events.append("task.started") + task_started.set() + finish_task.wait() + events.append("task.finished") + + await broker.startup() + sender = asyncio.create_task(sync_task.kiq()) + assert await asyncio.to_thread(task_started.wait, 1) + sender.cancel() + await sender + + release_timer = threading.Timer(0.05, finish_task.set) + release_timer.start() + try: + await broker.shutdown() + finally: + release_timer.cancel() + + assert events == [ + "middleware.startup", + "task.started", + "task.finished", + "middleware.shutdown", + ] + + +async def test_shutdown_keeps_loop_responsive_for_cancelled_sync_execution() -> None: + loop = asyncio.get_running_loop() + task_started = threading.Event() + finish_task = threading.Event() + loop_callback_finished = asyncio.Event() + broker = InMemoryBroker(await_inplace=True) + + async def finish_on_loop() -> None: + loop_callback_finished.set() + + @broker.task + def sync_task() -> None: + task_started.set() + finish_task.wait() + callback = asyncio.run_coroutine_threadsafe(finish_on_loop(), loop) + callback.result(timeout=1) + + sender = asyncio.create_task(sync_task.kiq()) + assert await asyncio.to_thread(task_started.wait, 1) + sender.cancel() + await sender + + finish_task.set() + await asyncio.wait_for(broker.shutdown(), timeout=1) + + assert loop_callback_finished.is_set() + + +async def test_shutdown_defers_cancellation_while_closing_sync_executor() -> None: + events: list[str] = [] + task_started = threading.Event() + finish_task = threading.Event() + broker = InMemoryBroker(await_inplace=True) + broker.with_middlewares(RecordingLifecycleMiddleware(events)) + + @broker.task + def sync_task() -> None: + events.append("task.started") + task_started.set() + finish_task.wait() + events.append("task.finished") + + await broker.startup() + sender = asyncio.create_task(sync_task.kiq()) + assert await asyncio.to_thread(task_started.wait, 1) + sender.cancel() + await sender + + shutdown_task = asyncio.create_task(broker.shutdown()) + await asyncio.sleep(0) + assert not shutdown_task.done() + shutdown_task.cancel() + release_timer = threading.Timer(0.05, finish_task.set) + release_timer.start() + try: + with pytest.raises(asyncio.CancelledError): + await shutdown_task + finally: + release_timer.cancel() + + assert events == [ + "middleware.startup", + "task.started", + "task.finished", + "middleware.shutdown", + ] + + +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_shutdown_waits_for_send_blocked_in_pre_send( + await_inplace: bool, +) -> None: + pre_send_started = asyncio.Event() + release_pre_send = asyncio.Event() + resource_shutdown_started = asyncio.Event() + drain_started = asyncio.Event() + task_executed = asyncio.Event() + middleware = BlockingPreSendMiddleware( + pre_send_started, + release_pre_send, + resource_shutdown_started, + ) + broker = DrainSignallingInMemoryBroker( + drain_started, + await_inplace=await_inplace, + ).with_middlewares(middleware) + + @broker.task + async def task() -> None: + task_executed.set() + + sender = asyncio.create_task(task.kiq()) + await asyncio.wait_for(pre_send_started.wait(), timeout=1) + shutdown_task = asyncio.create_task(broker.shutdown()) + await asyncio.wait_for(drain_started.wait(), timeout=1) + + assert not shutdown_task.done() + assert not resource_shutdown_started.is_set() + + release_pre_send.set() + await asyncio.wait_for(sender, timeout=1) + await asyncio.wait_for(shutdown_task, timeout=1) + + assert task_executed.is_set() + assert resource_shutdown_started.is_set() + assert middleware.pre_send_calls == 1 + assert not broker._inflight_tasks + assert not broker._running_tasks + + +async def test_shutdown_rejects_send_before_pre_send_side_effects() -> None: + pre_send_started = asyncio.Event() + release_pre_send = asyncio.Event() + release_pre_send.set() + resource_shutdown_started = asyncio.Event() + middleware = BlockingPreSendMiddleware( + pre_send_started, + release_pre_send, + resource_shutdown_started, + ) + broker = InMemoryBroker().with_middlewares(middleware) + + @broker.task + async def task() -> None: + return None + + await broker.shutdown() + + with pytest.raises(SendTaskError) as exc_info: + await task.kiq() + + assert isinstance(exc_info.value.__cause__, RuntimeError) + assert middleware.pre_send_calls == 0 + assert not pre_send_started.is_set() + + +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_shutdown_rejects_work_after_drain_starts( + await_inplace: bool, +) -> None: + shutdown_started = asyncio.Event() + finish_shutdown = asyncio.Event() + task_executed = False + broker = InMemoryBroker(await_inplace=await_inplace) + + @broker.on_event(TaskiqEvents.CLIENT_SHUTDOWN) + async def block_client_shutdown(state: TaskiqState) -> None: + shutdown_started.set() + await finish_shutdown.wait() + + @broker.task + async def task() -> None: + nonlocal task_executed + task_executed = True + + shutdown_task = asyncio.create_task(broker.shutdown()) + await asyncio.wait_for(shutdown_started.wait(), timeout=1) + + with pytest.raises(SendTaskError) as exc_info: + await task.kiq() + + assert isinstance(exc_info.value.__cause__, RuntimeError) + message = TaskiqMessage( + task_id="late-direct-task-id", + task_name=task.task_name, + labels={}, + labels_types={}, + args=[], + kwargs={}, + ) + with pytest.raises(RuntimeError, match="shutting down"): + await broker.kick(broker.formatter.dumps(message)) + + assert not task_executed + finish_shutdown.set() + await asyncio.wait_for(shutdown_task, timeout=1) + assert not broker._inflight_tasks + assert not broker._running_tasks + + +async def test_wait_all_retains_failure_from_already_completed_task() -> None: + task_error = LifecycleError("completed 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() + running_task = next(iter(broker._running_tasks)) + task_completed = asyncio.Event() + + def mark_completed(completed_task: asyncio.Task[Any]) -> None: + assert completed_task.done() + task_completed.set() + + running_task.add_done_callback(mark_completed) + await asyncio.wait_for(task_completed.wait(), timeout=1) + assert not broker._running_tasks + + with pytest.raises(LifecycleError) as exc_info: + await broker.wait_all() + + assert exc_info.value is task_error + await broker.wait_all() + await broker.shutdown() + + +async def test_shutdown_cleans_resources_after_task_failure() -> None: + events: list[str] = [] + task_error = LifecycleError("in-memory task failed") + middleware_error = LifecycleError("middleware shutdown failed") + broker = InMemoryBroker() + broker.with_middlewares(FailingExecutionMiddleware(task_error)) + broker.with_middlewares( + RecordingLifecycleMiddleware(events, shutdown_error=middleware_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) diff --git a/tests/brokers/test_inmemory_send.py b/tests/brokers/test_inmemory_send.py new file mode 100644 index 00000000..4b2bced7 --- /dev/null +++ b/tests/brokers/test_inmemory_send.py @@ -0,0 +1,744 @@ +import asyncio +from contextvars import ContextVar +from typing import Any + +import pytest + +from taskiq import InMemoryBroker, TaskiqMessage +from taskiq.exceptions import SendTaskError, UnknownTaskError +from taskiq.kicker import AsyncKicker +from tests.brokers.inmemory_contract_support import ( + BlockingAfterAcceptKickInMemoryBroker, + BlockingPostSendMiddleware, + ChildReentrantDrainMiddleware, + CoordinatedPostSendMiddleware, + DeferredChildDrainMiddleware, + DrainSignallingInMemoryBroker, + FailingAfterAcceptKickInMemoryBroker, + FailingExecutionMiddleware, + FailingPostSendMiddleware, + LifecycleError, + PostSendError, + RecordingKickInMemoryBroker, + ReentrantDrainMiddleware, + RejectingKickInMemoryBroker, + ReplacingKickInMemoryBroker, + TransformingKickInMemoryBroker, +) + + +async def test_direct_kick_preserves_inline_execution() -> None: + task_executed = False + broker = InMemoryBroker(await_inplace=True) + + @broker.task(task_name="direct.task") + async def task() -> None: + nonlocal task_executed + task_executed = True + + message = TaskiqMessage( + task_id="direct-task-id", + task_name=task.task_name, + labels={}, + labels_types={}, + args=[], + kwargs={}, + ) + await broker.kick(broker.formatter.dumps(message)) + + assert task_executed + await broker.wait_all() + await broker.shutdown() + + +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_kicker_uses_inmemory_public_kick_extension( + await_inplace: bool, +) -> None: + task_executed = False + broker = RecordingKickInMemoryBroker(await_inplace=await_inplace) + + @broker.task + async def task() -> None: + nonlocal task_executed + task_executed = True + + await task.kiq() + await broker.wait_all() + + assert broker.kick_calls == 1 + assert task_executed + await broker.shutdown() + + +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_cancelled_public_kick_preserves_accepted_execution( + await_inplace: bool, +) -> None: + events: list[str] = [] + post_send_started = asyncio.Event() + release_post_send = asyncio.Event() + task_executed = asyncio.Event() + broker = BlockingAfterAcceptKickInMemoryBroker( + await_inplace=await_inplace, + ).with_middlewares( + BlockingPostSendMiddleware( + events, + post_send_started, + release_post_send, + ), + ) + + @broker.task + async def task() -> None: + task_executed.set() + + send_task = asyncio.create_task(task.kiq()) + await asyncio.wait_for(broker.accepted.wait(), timeout=1) + send_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await send_task + + await asyncio.wait_for(post_send_started.wait(), timeout=1) + assert not task_executed.is_set() + release_post_send.set() + await asyncio.wait_for(broker.wait_all(), timeout=1) + + assert task_executed.is_set() + assert events == [ + "pre_send", + "post_send.started", + "post_send.finished", + "pre_execute", + "post_execute", + ] + await broker.shutdown() + + +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_failed_public_kick_preserves_accepted_execution( + await_inplace: bool, +) -> None: + events: list[str] = [] + post_send_started = asyncio.Event() + release_post_send = asyncio.Event() + kick_error = LifecycleError("public kick failed after acceptance") + task_executed = asyncio.Event() + broker = FailingAfterAcceptKickInMemoryBroker( + kick_error, + await_inplace=await_inplace, + ).with_middlewares( + BlockingPostSendMiddleware( + events, + post_send_started, + release_post_send, + ), + ) + + @broker.task + async def task() -> None: + task_executed.set() + + with pytest.raises(SendTaskError) as exc_info: + await task.kiq() + + await asyncio.wait_for(post_send_started.wait(), timeout=1) + assert not task_executed.is_set() + release_post_send.set() + await asyncio.wait_for(broker.wait_all(), timeout=1) + + assert exc_info.value.__cause__ is kick_error + assert task_executed.is_set() + assert events == [ + "pre_send", + "post_send.started", + "post_send.finished", + "pre_execute", + "post_execute", + ] + await broker.shutdown() + + +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_public_kick_dispatches_the_message_accepted_by_base( + await_inplace: bool, +) -> None: + executed_tasks: list[str] = [] + broker = TransformingKickInMemoryBroker( + "replacement", + await_inplace=await_inplace, + ) + + @broker.task(task_name="original") + async def original() -> None: + executed_tasks.append("original") + + @broker.task(task_name="replacement") + async def replacement() -> None: + executed_tasks.append("replacement") + + await original.kiq() + await broker.wait_all() + + assert executed_tasks == ["replacement"] + await broker.shutdown() + + +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_async_post_send_finishes_before_execution( + await_inplace: bool, +) -> None: + events: list[str] = [] + post_send_started = asyncio.Event() + release_post_send = asyncio.Event() + task_started = asyncio.Event() + broker = RecordingKickInMemoryBroker(await_inplace=await_inplace) + broker.with_middlewares( + BlockingPostSendMiddleware( + events, + post_send_started, + release_post_send, + ), + ) + + @broker.task + async def task() -> None: + events.append("task") + task_started.set() + + kiq_task = asyncio.create_task(task.kiq()) + await asyncio.wait_for(post_send_started.wait(), timeout=1) + + assert broker.kick_calls == 1 + assert not task_started.is_set() + assert events == ["pre_send", "post_send.started"] + + release_post_send.set() + await asyncio.wait_for(kiq_task, timeout=1) + await asyncio.wait_for(task_started.wait(), timeout=1) + await broker.shutdown() + + assert events == [ + "pre_send", + "post_send.started", + "post_send.finished", + "pre_execute", + "task", + "post_execute", + ] + + +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_rejected_public_kick_skips_post_send( + await_inplace: bool, +) -> None: + events: list[str] = [] + post_send_started = asyncio.Event() + release_post_send = asyncio.Event() + release_post_send.set() + kick_error = LifecycleError("public kick rejected the send") + task_executed = False + broker = RejectingKickInMemoryBroker( + kick_error, + await_inplace=await_inplace, + ).with_middlewares( + BlockingPostSendMiddleware( + events, + post_send_started, + release_post_send, + ), + ) + + @broker.task + async def task() -> None: + nonlocal task_executed + task_executed = True + + with pytest.raises(SendTaskError) as exc_info: + await task.kiq() + + assert exc_info.value.__cause__ is kick_error + assert broker.kick_calls == 1 + assert events == ["pre_send"] + assert not post_send_started.is_set() + assert not task_executed + await broker.shutdown() + + +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_replaced_public_kick_skips_local_execution( + await_inplace: bool, +) -> None: + events: list[str] = [] + post_send_started = asyncio.Event() + release_post_send = asyncio.Event() + release_post_send.set() + task_executed = False + broker = ReplacingKickInMemoryBroker( + await_inplace=await_inplace, + ).with_middlewares( + BlockingPostSendMiddleware( + events, + post_send_started, + release_post_send, + ), + ) + + @broker.task + async def task() -> None: + nonlocal task_executed + task_executed = True + + await task.kiq() + + assert broker.kick_calls == 1 + assert events == [ + "pre_send", + "post_send.started", + "post_send.finished", + ] + assert post_send_started.is_set() + assert not task_executed + assert not broker._running_tasks + await broker.shutdown() + + +async def test_concurrent_inline_sends_keep_per_invocation_ownership() -> None: + first_post_send_started = asyncio.Event() + second_post_send_started = asyncio.Event() + release_second_post_send = asyncio.Event() + first_task_finished = asyncio.Event() + second_task_started = asyncio.Event() + finish_second_task = asyncio.Event() + second_task_finished = asyncio.Event() + broker = InMemoryBroker(await_inplace=True) + + @broker.task(task_name="first.task") + async def first_task() -> None: + first_task_finished.set() + + @broker.task(task_name="second.task") + async def second_task() -> None: + second_task_started.set() + await finish_second_task.wait() + second_task_finished.set() + + broker.with_middlewares( + CoordinatedPostSendMiddleware( + first_task.task_name, + second_task.task_name, + first_post_send_started, + second_post_send_started, + release_second_post_send, + ), + ) + + first_sender = asyncio.create_task(first_task.kiq()) + await asyncio.wait_for(first_post_send_started.wait(), timeout=1) + second_sender = asyncio.create_task(second_task.kiq()) + await asyncio.wait_for(second_post_send_started.wait(), timeout=1) + + await asyncio.wait_for(first_sender, timeout=1) + assert first_task_finished.is_set() + assert not second_task_started.is_set() + assert not second_sender.done() + + release_second_post_send.set() + await asyncio.wait_for(second_task_started.wait(), timeout=1) + assert not second_sender.done() + + finish_second_task.set() + await asyncio.wait_for(second_sender, timeout=1) + + assert second_task_finished.is_set() + await broker.shutdown() + + +async def test_inline_send_preserves_caller_task_context() -> None: + context_value: ContextVar[str] = ContextVar( + "inmemory_inline_context", + default="caller", + ) + caller_task = asyncio.current_task() + execution_task: asyncio.Task[Any] | None = None + broker = InMemoryBroker(await_inplace=True) + + @broker.task + async def task() -> None: + nonlocal execution_task + execution_task = asyncio.current_task() + context_value.set("execution") + + await task.kiq() + + assert execution_task is caller_task + assert context_value.get() == "execution" + await broker.shutdown() + + +async def test_post_send_failure_does_not_retract_accepted_inline_task() -> None: + post_send_error = PostSendError("post-send failed") + task_executed = False + broker = InMemoryBroker(await_inplace=True) + broker.with_middlewares(FailingPostSendMiddleware(post_send_error)) + + @broker.task + async def task() -> None: + nonlocal task_executed + task_executed = True + + with pytest.raises(PostSendError) as exc_info: + await task.kiq() + + assert exc_info.value is post_send_error + assert task_executed + await broker.shutdown() + + +async def test_sender_cancellation_after_post_send_failure_keeps_execution() -> None: + post_send_error = PostSendError("post-send failed") + task_started = asyncio.Event() + finish_task = asyncio.Event() + task_cancelled = False + task_finished = False + broker = InMemoryBroker(await_inplace=True) + broker.with_middlewares(FailingPostSendMiddleware(post_send_error)) + + @broker.task + async def task() -> None: + nonlocal task_cancelled, task_finished + task_started.set() + try: + await finish_task.wait() + except asyncio.CancelledError: + task_cancelled = True + raise + task_finished = True + + sender = asyncio.create_task(task.kiq()) + await asyncio.wait_for(task_started.wait(), timeout=1) + sender.cancel() + + with pytest.raises(asyncio.CancelledError): + await sender + + assert not task_cancelled + finish_task.set() + await asyncio.wait_for(broker.wait_all(), timeout=1) + assert task_finished + await broker.shutdown() + + +async def test_inline_execution_failure_remains_a_send_error() -> None: + execution_error = LifecycleError("execution failed") + broker = InMemoryBroker(await_inplace=True) + broker.with_middlewares(FailingExecutionMiddleware(execution_error)) + + @broker.task + async def task() -> None: + return None + + with pytest.raises(SendTaskError) as exc_info: + await task.kiq() + + assert exc_info.value.__cause__ is execution_error + await broker.wait_all() + await broker.shutdown() + + +async def test_kicker_wraps_unknown_inmemory_task() -> None: + broker = InMemoryBroker() + kicker: AsyncKicker[Any, Any] = AsyncKicker("missing.task", broker, {}) + + with pytest.raises(SendTaskError) as exc_info: + await kicker.kiq() + + assert isinstance(exc_info.value.__cause__, UnknownTaskError) + await broker.shutdown() + + +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_post_send_failure_preserves_execution_error_ownership( + await_inplace: bool, +) -> None: + post_send_error = PostSendError("post-send failed") + execution_error = LifecycleError("execution failed") + broker = InMemoryBroker(await_inplace=await_inplace) + broker.with_middlewares( + FailingPostSendMiddleware(post_send_error), + FailingExecutionMiddleware(execution_error), + ) + + @broker.task + async def task() -> None: + return None + + with pytest.raises(PostSendError) as exc_info: + await task.kiq() + + assert exc_info.value is post_send_error + if await_inplace: + assert exc_info.value.__cause__ is execution_error + await broker.wait_all() + else: + with pytest.raises(LifecycleError) as execution_exc_info: + await broker.wait_all() + assert execution_exc_info.value is execution_error + + await broker.shutdown() + + +async def test_post_send_failure_remains_primary_when_execution_is_cancelled() -> None: + post_send_error = PostSendError("post-send failed") + execution_cancel = asyncio.CancelledError("execution cancelled") + broker = InMemoryBroker(await_inplace=True) + broker.with_middlewares( + FailingPostSendMiddleware(post_send_error), + FailingExecutionMiddleware(execution_cancel), + ) + + @broker.task + async def task() -> None: + return None + + with pytest.raises(PostSendError) as exc_info: + await task.kiq() + + assert exc_info.value is post_send_error + assert isinstance(exc_info.value.__cause__, asyncio.CancelledError) + await broker.wait_all() + await broker.shutdown() + + +async def test_post_send_cancellation_does_not_retract_accepted_inline_task() -> None: + events: list[str] = [] + post_send_started = asyncio.Event() + release_post_send = asyncio.Event() + task_executed = asyncio.Event() + broker = InMemoryBroker(await_inplace=True) + broker.with_middlewares( + BlockingPostSendMiddleware( + events, + post_send_started, + release_post_send, + ), + ) + + @broker.task + async def task() -> None: + task_executed.set() + + kiq_task = asyncio.create_task(task.kiq()) + await asyncio.wait_for(post_send_started.wait(), timeout=1) + kiq_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await kiq_task + + await asyncio.wait_for(task_executed.wait(), timeout=1) + await asyncio.wait_for(broker.wait_all(), timeout=1) + await broker.shutdown() + + +async def test_sender_cancellation_preserves_inline_execution_cancellation() -> None: + task_started = asyncio.Event() + task_cancelled = asyncio.Event() + keep_running = asyncio.Event() + broker = InMemoryBroker(await_inplace=True) + + @broker.task + async def task() -> None: + task_started.set() + try: + await keep_running.wait() + except asyncio.CancelledError: + task_cancelled.set() + raise + + kiq_task = asyncio.create_task(task.kiq()) + await asyncio.wait_for(task_started.wait(), timeout=1) + kiq_task.cancel() + + await asyncio.wait_for(asyncio.shield(kiq_task), timeout=1) + + assert not kiq_task.cancelled() + assert task_cancelled.is_set() + await broker.wait_all() + await broker.shutdown() + + +async def test_shutdown_waits_for_send_blocked_in_post_send() -> None: + events: list[str] = [] + post_send_started = asyncio.Event() + release_post_send = asyncio.Event() + task_executed = asyncio.Event() + drain_started = asyncio.Event() + broker = DrainSignallingInMemoryBroker(drain_started) + broker.with_middlewares( + BlockingPostSendMiddleware( + events, + post_send_started, + release_post_send, + ), + ) + + @broker.task + async def task() -> None: + task_executed.set() + + kiq_task = asyncio.create_task(task.kiq()) + await asyncio.wait_for(post_send_started.wait(), timeout=1) + shutdown_task = asyncio.create_task(broker.shutdown()) + await asyncio.wait_for(drain_started.wait(), timeout=1) + + assert not shutdown_task.done() + assert not task_executed.is_set() + + release_post_send.set() + await kiq_task + await asyncio.wait_for(shutdown_task, timeout=1) + + assert task_executed.is_set() + + +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_cancelled_wait_all_does_not_cancel_accepted_send( + await_inplace: bool, +) -> None: + events: list[str] = [] + post_send_started = asyncio.Event() + release_post_send = asyncio.Event() + task_executed = asyncio.Event() + drain_started = asyncio.Event() + broker = DrainSignallingInMemoryBroker( + drain_started, + await_inplace=await_inplace, + ) + broker.with_middlewares( + BlockingPostSendMiddleware( + events, + post_send_started, + release_post_send, + ), + ) + + @broker.task + async def task() -> None: + task_executed.set() + + sender = asyncio.create_task(task.kiq()) + await asyncio.wait_for(post_send_started.wait(), timeout=1) + drain = asyncio.create_task(broker.wait_all()) + await asyncio.wait_for(drain_started.wait(), timeout=1) + drain.cancel() + + with pytest.raises(asyncio.CancelledError): + await drain + + assert not sender.done() + assert not task_executed.is_set() + + release_post_send.set() + await asyncio.wait_for(sender, timeout=1) + await asyncio.wait_for(task_executed.wait(), timeout=1) + await asyncio.wait_for(broker.wait_all(), timeout=1) + await broker.shutdown() + + +@pytest.mark.parametrize("operation", ["wait_all", "shutdown"]) +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_execution_cannot_reenter_broker_drain( + operation: str, + await_inplace: bool, +) -> None: + broker = InMemoryBroker(await_inplace=await_inplace) + drain_error: RuntimeError | None = None + + @broker.task + async def task() -> None: + nonlocal drain_error + try: + await getattr(broker, operation)() + except RuntimeError as exc: + drain_error = exc + + await asyncio.wait_for(task.kiq(), timeout=1) + await asyncio.wait_for(broker.wait_all(), timeout=1) + + assert drain_error is not None + assert str(drain_error).startswith(f"InMemoryBroker.{operation}()") + assert broker.executor.submit(int).result() == 0 + await broker.shutdown() + + +@pytest.mark.parametrize("operation", ["wait_all", "shutdown"]) +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_post_send_cannot_reenter_broker_drain( + operation: str, + await_inplace: bool, +) -> None: + task_executed = False + broker = InMemoryBroker(await_inplace=await_inplace) + broker.with_middlewares(ReentrantDrainMiddleware(operation)) + + @broker.task + async def task() -> None: + nonlocal task_executed + task_executed = True + + with pytest.raises(RuntimeError) as exc_info: + await asyncio.wait_for(task.kiq(), timeout=1) + + await asyncio.wait_for(broker.wait_all(), timeout=1) + assert str(exc_info.value).startswith(f"InMemoryBroker.{operation}()") + assert task_executed + assert broker.executor.submit(int).result() == 0 + await broker.shutdown() + + +@pytest.mark.parametrize("operation", ["wait_all", "shutdown"]) +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_post_send_child_cannot_reenter_broker_drain( + operation: str, + await_inplace: bool, +) -> None: + task_executed = False + broker = InMemoryBroker(await_inplace=await_inplace) + broker.with_middlewares(ChildReentrantDrainMiddleware(operation)) + + @broker.task + async def task() -> None: + nonlocal task_executed + task_executed = True + + with pytest.raises(RuntimeError) as exc_info: + await asyncio.wait_for(task.kiq(), timeout=1) + + await asyncio.wait_for(broker.wait_all(), timeout=1) + assert str(exc_info.value).startswith(f"InMemoryBroker.{operation}()") + assert task_executed + assert broker.executor.submit(int).result() == 0 + await broker.shutdown() + + +@pytest.mark.parametrize("await_inplace", [False, True]) +async def test_post_send_child_can_drain_after_send_lifecycle( + await_inplace: bool, +) -> None: + release = asyncio.Event() + broker = InMemoryBroker(await_inplace=await_inplace) + middleware = DeferredChildDrainMiddleware(release) + broker.with_middlewares(middleware) + + @broker.task + async def task() -> None: + return None + + await asyncio.wait_for(task.kiq(), timeout=1) + await asyncio.wait_for(broker.wait_all(), timeout=1) + release.set() + + assert middleware.drain_task is not None + assert await asyncio.wait_for(middleware.drain_task, timeout=1) is None + await broker.shutdown() 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/middlewares/admin_middleware/test_arguments_formatting.py b/tests/middlewares/admin_middleware/test_arguments_formatting.py index a65b5669..0d021625 100644 --- a/tests/middlewares/admin_middleware/test_arguments_formatting.py +++ b/tests/middlewares/admin_middleware/test_arguments_formatting.py @@ -22,7 +22,6 @@ class PydanticDTOFactory(ModelFactory[PydanticDTO]): __model__ = PydanticDTO -# @pytest.mark.skip class TestArgumentsFormattingInAdminMiddleware: @pytest.mark.parametrize( "dto_factory, task_name", diff --git a/tests/middlewares/test_simple_retry.py b/tests/middlewares/test_simple_retry.py index 783d98b7..1e7c7754 100644 --- a/tests/middlewares/test_simple_retry.py +++ b/tests/middlewares/test_simple_retry.py @@ -1,23 +1,11 @@ -import uuid -from unittest.mock import AsyncMock - -import pytest - -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 AsyncQueueBroker -@pytest.fixture -def broker() -> AsyncMock: - mocked_broker = AsyncMock() - mocked_broker.id_generator = lambda: uuid.uuid4().hex - mocked_broker.formatter = JSONFormatter() - return mocked_broker - - -async def test_successful_retry(broker: AsyncMock) -> None: +async def test_successful_retry() -> None: + broker = AsyncQueueBroker() middleware = SimpleRetryMiddleware() middleware.set_broker(broker) await middleware.on_error( @@ -33,12 +21,13 @@ async def test_successful_retry(broker: AsyncMock) -> None: TaskiqResult(is_err=True, return_value=None, execution_time=0.0), Exception(), ) - resend: TaskiqMessage = broker.kick.await_args.args[0] + resend = broker.formatter.loads(broker.queue.get_nowait()) assert resend.task_name == "meme" assert resend.labels["_retries"] == "1" -async def test_no_retry(broker: AsyncMock) -> None: +async def test_no_retry() -> None: + broker = AsyncQueueBroker() middleware = SimpleRetryMiddleware() middleware.set_broker(broker) await middleware.on_error( @@ -52,10 +41,11 @@ async def test_no_retry(broker: AsyncMock) -> None: TaskiqResult(is_err=True, return_value=None, execution_time=0.0), Exception(), ) - broker.kick.assert_not_called() + assert broker.queue.empty() -async def test_max_retries(broker: AsyncMock) -> None: +async def test_max_retries() -> None: + broker = AsyncQueueBroker() middleware = SimpleRetryMiddleware(default_retry_count=3) middleware.set_broker(broker) await middleware.on_error( @@ -72,4 +62,4 @@ async def test_max_retries(broker: AsyncMock) -> None: TaskiqResult(is_err=True, return_value=None, execution_time=0.0), Exception(), ) - broker.kick.assert_not_called() + assert broker.queue.empty() diff --git a/tests/opentelemetry/test_auto_instrumentation.py b/tests/opentelemetry/test_auto_instrumentation.py index c7290869..f6350330 100644 --- a/tests/opentelemetry/test_auto_instrumentation.py +++ b/tests/opentelemetry/test_auto_instrumentation.py @@ -23,10 +23,16 @@ async def task_add(a: float, b: float) -> float: asyncio.run(test()) - spans = self.sorted_spans(self.memory_exporter.get_finished_spans()) + spans = self.memory_exporter.get_finished_spans() self.assertEqual(len(spans), 2) - consumer, producer = spans + spans_by_kind = {span.kind: span for span in spans} + self.assertEqual( + set(spans_by_kind), + {SpanKind.CONSUMER, SpanKind.PRODUCER}, + ) + consumer = spans_by_kind[SpanKind.CONSUMER] + producer = spans_by_kind[SpanKind.PRODUCER] self.assertEqual( consumer.name, diff --git a/tests/opentelemetry/test_tasks.py b/tests/opentelemetry/test_tasks.py index ab443f96..cd889de3 100644 --- a/tests/opentelemetry/test_tasks.py +++ b/tests/opentelemetry/test_tasks.py @@ -36,10 +36,16 @@ async def test() -> None: asyncio.run(test()) - spans = self.sorted_spans(self.memory_exporter.get_finished_spans()) + spans = self.memory_exporter.get_finished_spans() self.assertEqual(len(spans), 2) - consumer, producer = spans + spans_by_kind = {span.kind: span for span in spans} + self.assertEqual( + set(spans_by_kind), + {SpanKind.CONSUMER, SpanKind.PRODUCER}, + ) + consumer = spans_by_kind[SpanKind.CONSUMER] + producer = spans_by_kind[SpanKind.PRODUCER] self.assertEqual( consumer.name, @@ -85,10 +91,16 @@ async def test() -> None: asyncio.run(test()) - spans = self.sorted_spans(self.memory_exporter.get_finished_spans()) + spans = self.memory_exporter.get_finished_spans() self.assertEqual(len(spans), 2) - consumer, producer = spans + spans_by_kind = {span.kind: span for span in spans} + self.assertEqual( + set(spans_by_kind), + {SpanKind.CONSUMER, SpanKind.PRODUCER}, + ) + consumer = spans_by_kind[SpanKind.CONSUMER] + producer = spans_by_kind[SpanKind.PRODUCER] self.assertEqual( consumer.name,