fix(inmemory): harden accepted task lifecycle

Preserve accepted work across hook failures and cancellation.
Drain broker-owned tasks before closing lifecycle resources.
Move executor shutdown off the event loop and retain first failures.
Add deterministic ordering, concurrency, and shutdown coverage.

Refs: #586

Author: GefMar

docs/available-components/brokers.md

@@ -8,14 +8,20 @@ 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 waits for local background executions before
-closing those resources and the synchronous task executor.
+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
 

docs/guide/testing-taskiq.md

@@ -165,6 +165,9 @@ 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
@@ -191,9 +194,16 @@ 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()` performs the same drain before closing middleware, result
-backend, and executor resources. If shutdown is cancelled while draining, it
-finishes the drain and resource cleanup before propagating the cancellation.
+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.

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,
@@ -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:
         """
@@ -209,35 +212,37 @@ async def shutdown(self) -> None:
         This method is called,
         when broker is closing.
         """
-        shutdown_error: BaseException | None = None
+        shutdown_errors: list[BaseException] = []
+        await self._shutdown_resources(shutdown_errors)
 
+        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]:
-                try:
-                    await maybe_awaitable(handler(self.state))
-                except BaseException as exc:
-                    shutdown_error = self._remember_shutdown_error(
-                        shutdown_error,
-                        exc,
-                    )
+                yield partial(handler, self.state)
 
         for middleware in self.middlewares:
             if middleware.__class__.shutdown != TaskiqMiddleware.shutdown:
-                try:
-                    await maybe_awaitable(middleware.shutdown())
-                except BaseException as exc:
-                    shutdown_error = self._remember_shutdown_error(
-                        shutdown_error,
-                        exc,
-                    )
-
-        try:
-            await self.result_backend.shutdown()
-        except BaseException as exc:
-            shutdown_error = self._remember_shutdown_error(shutdown_error, exc)
+                yield middleware.shutdown
 
-        if shutdown_error is not None:
-            raise shutdown_error
+        yield self.result_backend.shutdown
 
     def _get_shutdown_events(self) -> tuple[TaskiqEvents, ...]:
         """Return event phases owned by this broker shutdown."""
@@ -259,6 +264,23 @@ def _remember_shutdown_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,