fix(routing): harden retry and lifecycle ownership

GefMar · GefMar · commit a7eb7a59fbba · 2026-07-11T17:32:33.000+02:00
- keep immediate retries on the receiving broker
- protect callback cleanup from repeated cancellation
- make scheduler rollback cleanup idempotent
- document unsupported-flow requeue behavior
diff --git a/docs/guide/routing-and-flows.md b/docs/guide/routing-and-flows.md
@@ -341,6 +341,9 @@ contract:
 - A failed or cancelled publication remains unacknowledged for
   `when_executed`, `when_saved` and `manual`, allowing broker redelivery. The
   original label is not incremented.
+- `UnsupportedFlowError` follows the same failure contract. A legacy adapter
+  with an explicit same-broker Flow can therefore redeliver the original until
+  the route is corrected or the transport's dead-letter policy intervenes.
 - `when_received` is intentionally different: the original was acknowledged
   before task execution, so a later publication failure cannot restore it.
 - If acknowledgement itself fails after successful publication, that error is
@@ -519,6 +522,12 @@ Scheduler and requeue compatibility:
   Configure broker membership before scheduler startup.
 - `Context.requeue()` sends through the broker currently executing the task. It
   does not re-resolve a task route across brokers.
+- Immediate retries published by `SimpleRetryMiddleware` or
+  `SmartRetryMiddleware` stay on the broker currently executing the task, so
+  the retry keeps that broker's send lifecycle and result backend. When
+  `SmartRetryMiddleware` uses a `ScheduleSource`, it creates a normal
+  `ScheduledTask`; the scheduler then applies the late route-resolution rules
+  above.
 - Taskiq does not preserve inbound source-flow provenance yet because that data
   is not stored in `TaskiqMessage` or `Context`.
 
diff --git a/taskiq/middlewares/simple_retry_middleware.py b/taskiq/middlewares/simple_retry_middleware.py
@@ -65,11 +65,15 @@ async def on_error(
         if not retry_on_error:
             return
 
-        kicker: AsyncKicker[Any, Any] = AsyncKicker(
-            task_name=message.task_name,
-            broker=self.broker,
-            labels=message.labels,
-        ).with_task_id(message.task_id)
+        kicker: AsyncKicker[Any, Any] = (
+            AsyncKicker(
+                task_name=message.task_name,
+                broker=self.broker,
+                labels=message.labels,
+            )
+            .with_broker(self.broker)
+            .with_task_id(message.task_id)
+        )
 
         # Getting number of previous retries.
         retries = int(message.labels.get("_retries", 0)) + 1
diff --git a/taskiq/middlewares/smart_retry_middleware.py b/taskiq/middlewares/smart_retry_middleware.py
@@ -176,6 +176,7 @@ async def on_error(
                     broker=self.broker,
                     labels=message.labels,
                 )
+                .with_broker(self.broker)
                 .with_task_id(message.task_id)
                 .with_labels(_retries=retries)
             )
diff --git a/taskiq/receiver/receiver.py b/taskiq/receiver/receiver.py
@@ -640,7 +640,7 @@ async def runner(
 
         :param queue: queue with prefetched data.
         """
-        cancelled = False
+        cancellation: asyncio.CancelledError | None = None
         try:
             while True:
                 queued_message = await queue.get()
@@ -650,13 +650,25 @@ async def runner(
                     raise TypeError("Receiver queue contains an invalid message.")
                 await self._start_callback(queued_message)
 
-        except asyncio.CancelledError:
-            cancelled = True
-            raise
+        except asyncio.CancelledError as exc:
+            cancellation = exc
         finally:
-            await self._drain_active_tasks(cancel_immediately=cancelled)
+            with anyio.CancelScope(shield=True):
+                cleanup_cancellation = await self._drain_active_tasks(
+                    cancel_immediately=cancellation is not None,
+                )
+            if cancellation is None:
+                cancellation = cleanup_cancellation
+
+        if cancellation is not None:
+            raise cancellation
         logger.info("The runner is stopped.")
 
+    async def _run_owned_callback(self, message: bytes | AckableMessage) -> None:
+        """Run a callback outside repeated listener-scope cancellation."""
+        with anyio.CancelScope(shield=True):
+            await self.callback(message=message, raise_err=False)
+
     async def _start_callback(self, message: _PrefetchedMessage) -> None:
         """Transfer one prefetched delivery into an owned callback task."""
         owns_prefetch_slot = message.owns_prefetch_slot
@@ -671,7 +683,7 @@ async def _start_callback(self, message: _PrefetchedMessage) -> None:
                 self.sem_prefetch.release()
                 owns_prefetch_slot = False
             task = asyncio.create_task(
-                self.callback(message=message.data, raise_err=False),
+                self._run_owned_callback(message.data),
             )
         except BaseException:
             if owns_prefetch_slot:
@@ -704,11 +716,15 @@ def _on_callback_done(self, task: "asyncio.Task[Any]") -> None:
                 ),
             )
 
-    async def _drain_active_tasks(self, *, cancel_immediately: bool) -> None:
+    async def _drain_active_tasks(
+        self,
+        *,
+        cancel_immediately: bool,
+    ) -> asyncio.CancelledError | None:
         """Wait for owned callbacks and cancel them after the graceful boundary."""
         tasks = set(self._active_tasks)
         if not tasks:
-            return
+            return None
 
         logger.info("Waiting for %d running tasks to complete...", len(tasks))
         cancellation: asyncio.CancelledError | None = None
@@ -733,25 +749,24 @@ async def _drain_active_tasks(self, *, cancel_immediately: bool) -> None:
             if cancellation is None:
                 cancellation = cleanup_cancellation
         logger.info("No more tasks to wait for. Shutting down.")
-        if cancellation is not None:
-            raise cancellation
+        return cancellation
 
     @staticmethod
     async def _cancel_callback_tasks(
         tasks: set[asyncio.Task[Any]],
     ) -> asyncio.CancelledError | None:
-        """Cancel callback tasks and await cleanup despite one outer cancellation."""
+        """Cancel callbacks once and await cleanup despite outer cancellation."""
         for task in tasks:
             task.cancel()
-        try:
-            await asyncio.gather(*tasks, return_exceptions=True)
-        except asyncio.CancelledError as exc:
-            for task in tasks:
-                if not task.done():
-                    task.cancel()
-            await asyncio.gather(*tasks, return_exceptions=True)
-            return exc
-        return None
+        waiter = asyncio.gather(*tasks, return_exceptions=True)
+        cancellation: asyncio.CancelledError | None = None
+        while not waiter.done():
+            try:
+                await asyncio.shield(waiter)
+            except asyncio.CancelledError as exc:
+                if cancellation is None:
+                    cancellation = exc
+        return cancellation
 
     def _record_listen_error(self, error: BaseException) -> None:
         """Preserve the first transport error while allowing iterator cleanup."""
diff --git a/taskiq/scheduler/scheduler.py b/taskiq/scheduler/scheduler.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from enum import Enum, auto
 from logging import getLogger
 from typing import TYPE_CHECKING, Any
 
@@ -16,6 +17,16 @@
 logger = getLogger(__name__)
 
 
+class _SchedulerLifecycle(Enum):
+    """Internal scheduler resource lifecycle."""
+
+    CREATED = auto()
+    STARTING = auto()
+    STARTED = auto()
+    STOPPING = auto()
+    STOPPED = auto()
+
+
 class TaskiqScheduler:
     """Scheduler class."""
 
@@ -27,6 +38,7 @@ def __init__(
         self.broker = broker
         self.sources = sources
         self._started_brokers: list[AsyncBroker] = []
+        self._lifecycle = _SchedulerLifecycle.CREATED
 
     async def startup(self) -> None:
         """
@@ -37,9 +49,15 @@ async def startup(self) -> None:
         producer-side resources because late route resolution may select any of
         them.
         """
-        if self._started_brokers:
+        if self._lifecycle in {
+            _SchedulerLifecycle.STARTING,
+            _SchedulerLifecycle.STARTED,
+        }:
             raise RuntimeError("TaskiqScheduler is already started.")
+        if self._lifecycle is _SchedulerLifecycle.STOPPING:
+            raise RuntimeError("TaskiqScheduler is shutting down.")
 
+        self._lifecycle = _SchedulerLifecycle.STARTING
         try:
             for broker in self._managed_brokers():
                 broker.is_scheduler_process = True
@@ -50,6 +68,7 @@ async def startup(self) -> None:
                 tuple(self._started_brokers),
             )
             self._started_brokers.clear()
+            self._lifecycle = _SchedulerLifecycle.STOPPED
             if cleanup_error is not None:
                 logger.error(
                     "Error while cleaning up partially started scheduler brokers.",
@@ -60,6 +79,7 @@ async def startup(self) -> None:
                     ),
                 )
             raise
+        self._lifecycle = _SchedulerLifecycle.STARTED
 
     async def on_ready(self, source: ScheduleSource, task: ScheduledTask) -> None:
         """
@@ -86,11 +106,22 @@ async def on_ready(self, source: ScheduleSource, task: ScheduledTask) -> None:
 
     async def shutdown(self) -> None:
         """Shut down scheduler brokers in reverse startup order."""
+        if self._lifecycle is _SchedulerLifecycle.STOPPED:
+            return
+        if self._lifecycle is _SchedulerLifecycle.STARTING:
+            raise RuntimeError("TaskiqScheduler is still starting.")
+        if self._lifecycle is _SchedulerLifecycle.STOPPING:
+            raise RuntimeError("TaskiqScheduler is already shutting down.")
+
         brokers = tuple(self._started_brokers)
         if not brokers:
             brokers = (self.broker,)
         self._started_brokers.clear()
-        shutdown_error = await self._shutdown_brokers(brokers)
+        self._lifecycle = _SchedulerLifecycle.STOPPING
+        try:
+            shutdown_error = await self._shutdown_brokers(brokers)
+        finally:
+            self._lifecycle = _SchedulerLifecycle.STOPPED
         if shutdown_error is not None:
             raise shutdown_error
 
diff --git a/tests/middlewares/test_simple_retry.py b/tests/middlewares/test_simple_retry.py
@@ -3,10 +3,12 @@
 
 import pytest
 
+from taskiq import SmartRetryMiddleware, TaskiqRouter
 from taskiq.formatters.json_formatter import JSONFormatter
 from taskiq.message import TaskiqMessage
 from taskiq.middlewares.simple_retry_middleware import SimpleRetryMiddleware
 from taskiq.result import TaskiqResult
+from tests.utils import RecordingBroker
 
 
 @pytest.fixture
@@ -73,3 +75,49 @@ async def test_max_retries(broker: AsyncMock) -> None:
         Exception(),
     )
     broker.kick.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "middleware",
+    [
+        pytest.param(SimpleRetryMiddleware(), id="simple"),
+        pytest.param(SmartRetryMiddleware(default_delay=0), id="smart"),
+    ],
+)
+async def test_retry_keeps_current_broker_when_route_points_elsewhere(
+    middleware: SimpleRetryMiddleware | SmartRetryMiddleware,
+) -> None:
+    router = TaskiqRouter()
+    current_broker = RecordingBroker(
+        router=router,
+        broker_name="current",
+    )
+    routed_broker = RecordingBroker(
+        router=router,
+        broker_name="routed",
+    )
+    current_broker.add_middlewares(middleware)
+    router.route_task("meme", broker=routed_broker)
+    result = TaskiqResult(is_err=True, return_value=None, execution_time=0.0)
+
+    await middleware.on_error(
+        TaskiqMessage(
+            task_id="test_id",
+            task_name="meme",
+            labels={
+                "max_retries": 2,
+                "retry_on_error": True,
+            },
+            args=[],
+            kwargs={},
+        ),
+        result,
+        RuntimeError("retry"),
+    )
+
+    assert routed_broker.sent == []
+    assert len(current_broker.sent) == 1
+    retried_message, retried_flow = current_broker.sent[0]
+    assert retried_message.task_name == "meme"
+    assert retried_message.labels["_retries"] == "1"
+    assert retried_flow is None
diff --git a/tests/receiver/test_runtime.py b/tests/receiver/test_runtime.py
@@ -719,6 +719,80 @@ async def blocked_task() -> None:
     assert cleaned_before_runtime_returned
 
 
+async def test_runtime_hard_cancel_awaits_multiphase_callback_cleanup(
+    runtime_executor: ThreadPoolExecutor,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    events: list[str] = []
+    cleanup_started = asyncio.Event()
+    cleanup_finished = asyncio.Event()
+    cleanup_interrupted = asyncio.Event()
+    release_cleanup = asyncio.Event()
+    callback_started = asyncio.Event()
+
+    class CleanupOrderBroker(RuntimeBroker):
+        async def shutdown(self) -> None:
+            self.events.append("broker.shutdown.started")
+            await super().shutdown()
+
+    broker = CleanupOrderBroker(
+        router=TaskiqRouter(),
+        broker_name="listener",
+        events=events,
+    )
+
+    @broker.task(task_name="runtime.multiphase_callback_cleanup")
+    async def blocked_task() -> None:
+        callback_started.set()
+        try:
+            await asyncio.Event().wait()
+        finally:
+            cleanup_started.set()
+            try:
+                await release_cleanup.wait()
+                for _ in range(20):
+                    await asyncio.sleep(0)
+            except asyncio.CancelledError:
+                cleanup_interrupted.set()
+                raise
+            events.append("callback.cleanup.finished")
+            cleanup_finished.set()
+
+    finish_event = asyncio.Event()
+    runtime = build_runtime(
+        (broker,),
+        runtime_executor,
+        options=runtime_options(
+            wait_tasks_timeout=0.01,
+            shutdown_timeout=0.05,
+        ),
+    )
+    caplog.set_level(logging.WARNING, logger="taskiq.worker.runtime")
+    runtime_task = asyncio.create_task(runtime.run(finish_event))
+    await broker.listen_started.wait()
+    await blocked_task.kiq()
+    await callback_started.wait()
+
+    finish_event.set()
+    await cleanup_started.wait()
+
+    async def wait_for_hard_cancel() -> None:
+        while not any(
+            "did not stop" in record.getMessage() for record in caplog.records
+        ):
+            await asyncio.sleep(0)
+
+    await asyncio.wait_for(wait_for_hard_cancel(), timeout=1)
+    release_cleanup.set()
+    await asyncio.wait_for(runtime_task, timeout=1)
+
+    assert cleanup_finished.is_set()
+    assert not cleanup_interrupted.is_set()
+    assert events.index("callback.cleanup.finished") < events.index(
+        "broker.shutdown.started",
+    )
+
+
 async def test_runtime_continues_cleanup_after_broker_shutdown_failure(
     runtime_executor: ThreadPoolExecutor,
     caplog: pytest.LogCaptureFixture,
diff --git a/tests/scheduler/test_scheduler.py b/tests/scheduler/test_scheduler.py
diff --git a/tests/test_requeue.py b/tests/test_requeue.py

Original file line number	Diff line number	Diff line change
`@@ -176,6 +176,7 @@ async def on_error(`
`176`	`176`	`broker=self.broker,`
`177`	`177`	`labels=message.labels,`
`178`	`178`	`)`
	`179`	`+ .with_broker(self.broker)`
`179`	`180`	`.with_task_id(message.task_id)`
`180`	`181`	`.with_labels(_retries=retries)`
`181`	`182`	`)`