diff --git a/taskiq/abc/broker.py b/taskiq/abc/broker.py index ea2e86c0..ccd880c0 100644 --- a/taskiq/abc/broker.py +++ b/taskiq/abc/broker.py @@ -237,6 +237,18 @@ async def kick( :param message: name of a task. """ + async def _finish_kick(self) -> None: # noqa: B027 + """ + Hook that runs after a message is sent and ``post_send`` fired. + + Most brokers send messages elsewhere and have nothing to do here, + so this is a no-op by default. Brokers that execute tasks in-place + (like :class:`~taskiq.InMemoryBroker` with ``await_inplace``) can + override this method to await the in-place execution *after* the + ``post_send`` middleware hook has run, preserving the expected + ``pre_send -> post_send -> pre_execute -> ...`` ordering. + """ + @abstractmethod def listen(self) -> AsyncGenerator[bytes | AckableMessage, None]: """ diff --git a/taskiq/brokers/inmemory_broker.py b/taskiq/brokers/inmemory_broker.py index 0a7cc98e..b0e9f493 100644 --- a/taskiq/brokers/inmemory_broker.py +++ b/taskiq/brokers/inmemory_broker.py @@ -146,6 +146,7 @@ def __init__( ) self.await_inplace = await_inplace self._running_tasks: set[asyncio.Task[Any]] = set() + self._inplace_tasks: set[asyncio.Task[Any]] = set() async def kick(self, message: BrokerMessage) -> None: """ @@ -162,14 +163,37 @@ async def kick(self, message: BrokerMessage) -> None: raise UnknownTaskError(task_name=message.task_name) receiver_cb = self.receiver.callback(message=message.message) + task = asyncio.create_task(receiver_cb) + if self.await_inplace: - await receiver_cb + # Schedule the receiver as a task instead of awaiting it here so + # that control returns to the sender and the ``post_send`` + # middleware hook fires *before* the task starts executing. The + # task is then awaited in `_finish_kick`, right after `post_send`, + # keeping the in-place execution while fixing the hook ordering. + self._inplace_tasks.add(task) + task.add_done_callback(self._inplace_tasks.discard) return - task = asyncio.create_task(receiver_cb) self._running_tasks.add(task) task.add_done_callback(self._running_tasks.discard) + async def _finish_kick(self) -> None: + """ + Await in-place task execution scheduled by `kick`. + + This runs after the `post_send` middleware hook, so with + ``await_inplace=True`` the hook ordering becomes + ``pre_send -> post_send -> pre_execute -> task -> post_execute`` + while tasks are still fully executed before `kiq` returns. + """ + if not self._inplace_tasks: + return + to_await = list(self._inplace_tasks) + self._inplace_tasks.clear() + for task in to_await: + await task + def listen(self) -> AsyncGenerator[bytes, None]: """ Inmemory broker cannot listen. diff --git a/taskiq/kicker.py b/taskiq/kicker.py index dc113a7e..31176dc9 100644 --- a/taskiq/kicker.py +++ b/taskiq/kicker.py @@ -169,6 +169,8 @@ async def kiq( if middleware.__class__.post_send != TaskiqMiddleware.post_send: await maybe_awaitable(middleware.post_send(message)) + await self.broker._finish_kick() # noqa: SLF001 + return AsyncTaskiqTask( task_id=message.task_id, result_backend=self.broker.result_backend, diff --git a/tests/brokers/test_inmemory.py b/tests/brokers/test_inmemory.py index 9dc85019..250a0902 100644 --- a/tests/brokers/test_inmemory.py +++ b/tests/brokers/test_inmemory.py @@ -4,7 +4,10 @@ import pytest from taskiq import InMemoryBroker +from taskiq.abc.middleware import TaskiqMiddleware from taskiq.events import TaskiqEvents +from taskiq.message import TaskiqMessage +from taskiq.result import TaskiqResult from taskiq.state import TaskiqState @@ -114,3 +117,46 @@ async def test_task() -> None: assert slept assert await task.is_ready() assert not broker._running_tasks + + +async def test_inplace_middleware_order() -> None: + """post_send must fire before task execution with await_inplace=True.""" + events: list[str] = [] + + class OrderMiddleware(TaskiqMiddleware): + async def pre_send(self, message: TaskiqMessage) -> TaskiqMessage: + events.append("pre_send") + return message + + async def post_send(self, message: TaskiqMessage) -> None: + events.append("post_send") + + async def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage: + events.append("pre_execute") + return message + + async def post_execute( + self, + message: TaskiqMessage, + result: "TaskiqResult[object]", + ) -> None: + events.append("post_execute") + + broker = InMemoryBroker(await_inplace=True).with_middlewares(OrderMiddleware()) + + @broker.task + async def test_task() -> None: + events.append("task") + + task = await test_task.kiq() + + # The task must be fully executed by the time kiq returns. + assert await task.is_ready() + assert not broker._inplace_tasks + assert events == [ + "pre_send", + "post_send", + "pre_execute", + "task", + "post_execute", + ]