fix(task): tolerate transient result backend errors in wait_result

taskiq/task.py

@@ -78,6 +78,7 @@ async def wait_result(
         check_interval: float = 0.2,
         timeout: float = -1.0,
         with_logs: bool = False,
+        max_poll_failures: int = 3,
     ) -> "TaskiqResult[_ReturnType]":
         """
         Waits until result is ready.
@@ -89,15 +90,43 @@ async def wait_result(
         task didn't become ready in provided
         period of time.
 
+        Failing readiness checks are tolerated while they look
+        transient. The counter is reset by every successful check, so
+        the wait is only aborted after max_poll_failures checks have
+        failed in a row.
+
         :param check_interval: How often checks are performed.
         :param timeout: timeout for the result.
         :param with_logs: whether you want to fetch logs from worker.
+        :param max_poll_failures: how many readiness checks may fail
+            in a row before giving up. Pass 0 to give up on the first
+            failed check.
         :raises TaskiqResultTimeoutError: if task didn't
             become ready in provided period of time.
+        :raises ResultIsReadyError: if readiness checks kept failing.
         :return: task's return value.
         """
         start_time = time()
-        while not await self.is_ready():
+        failures = 0
+        while True:
+            try:
+                ready = await self.is_ready()
+            except ResultIsReadyError:
+                failures += 1
+                if failures > max_poll_failures:
+                    raise
+                logger.warning(
+                    "Cannot check readiness of task %s, "
+                    "retrying (%d of %d attempts used).",
+                    self.task_id,
+                    failures,
+                    max_poll_failures,
+                    exc_info=True,
+                )
+            else:
+                if ready:
+                    break
+                failures = 0
             if 0 < timeout < time() - start_time:
                 raise TaskiqResultTimeoutError(timeout=timeout)
             await asyncio.sleep(check_interval)

tests/test_task.py

@@ -8,6 +8,7 @@
 from taskiq.abc import AsyncResultBackend
 from taskiq.abc.serializer import TaskiqSerializer
 from taskiq.compat import model_dump, model_validate
+from taskiq.exceptions import ResultIsReadyError, TaskiqResultTimeoutError
 from taskiq.result import TaskiqResult
 from taskiq.task import AsyncTaskiqTask
 
@@ -69,3 +70,101 @@ class MyResult(BaseModel):
     sent_task = AsyncTaskiqTask(test_id, res_back, MyResult)
     parsed = await sent_task.wait_result()
     assert isinstance(parsed.return_value, MyResult)
+
+
+class ScriptedBackend(AsyncResultBackend[str]):
+    """Backend whose readiness checks follow a fixed script.
+
+    Each item of the script is either a value to return or an exception
+    to raise. The last item repeats once the script runs out.
+    """
+
+    def __init__(self, script: list[bool | Exception]) -> None:
+        self.script = script
+        self.polls = 0
+
+    async def set_result(
+        self,
+        task_id: str,
+        result: TaskiqResult[str],
+    ) -> None:
+        """Results are produced by the script, so nothing is stored."""
+
+    async def is_result_ready(self, task_id: str) -> bool:
+        """Perform the next scripted readiness check."""
+        step = self.script[min(self.polls, len(self.script) - 1)]
+        self.polls += 1
+        if isinstance(step, Exception):
+            raise step
+        return step
+
+    async def get_result(
+        self,
+        task_id: str,
+        with_logs: bool = False,
+    ) -> TaskiqResult[str]:
+        """Return a fixed successful result."""
+        return TaskiqResult(is_err=False, return_value="done", execution_time=0.0)
+
+
+async def test_wait_result_survives_transient_error() -> None:
+    backend = ScriptedBackend([ConnectionError("connection reset"), True])
+    task: AsyncTaskiqTask[str] = AsyncTaskiqTask("task-id", backend)
+
+    result = await task.wait_result(check_interval=0.0)
+
+    assert result.return_value == "done"
+    assert backend.polls == 2
+
+
+async def test_wait_result_raises_if_backend_never_recovers() -> None:
+    backend = ScriptedBackend([ConnectionError("backend is down")])
+    task: AsyncTaskiqTask[str] = AsyncTaskiqTask("task-id", backend)
+
+    with pytest.raises(ResultIsReadyError):
+        await task.wait_result(check_interval=0.0, max_poll_failures=2)
+
+    # Two tolerated failures, then the third one is fatal.
+    assert backend.polls == 3
+
+
+async def test_wait_result_counts_failures_consecutively() -> None:
+    backend = ScriptedBackend(
+        [
+            ConnectionError("connection reset"),
+            False,
+            ConnectionError("connection reset"),
+            False,
+            ConnectionError("connection reset"),
+            True,
+        ],
+    )
+    task: AsyncTaskiqTask[str] = AsyncTaskiqTask("task-id", backend)
+
+    # Three failures in total, but never two in a row.
+    result = await task.wait_result(check_interval=0.0, max_poll_failures=1)
+
+    assert result.return_value == "done"
+    assert backend.polls == 6
+
+
+async def test_wait_result_no_retries_by_configuration() -> None:
+    backend = ScriptedBackend([ConnectionError("connection reset"), True])
+    task: AsyncTaskiqTask[str] = AsyncTaskiqTask("task-id", backend)
+
+    with pytest.raises(ResultIsReadyError):
+        await task.wait_result(check_interval=0.0, max_poll_failures=0)
+
+    assert backend.polls == 1
+
+
+async def test_wait_result_timeout_while_retrying() -> None:
+    backend = ScriptedBackend([ConnectionError("backend is down")])
+    task: AsyncTaskiqTask[str] = AsyncTaskiqTask("task-id", backend)
+
+    with pytest.raises(TaskiqResultTimeoutError):
+        await task.wait_result(
+            check_interval=0.05,
+            timeout=0.1,
+            max_poll_failures=1000,
+        )