diff --git a/taskiq/task.py b/taskiq/task.py index 550dabee..63938d03 100644 --- a/taskiq/task.py +++ b/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) diff --git a/tests/test_task.py b/tests/test_task.py index 5f47102b..b69e4afa 100644 --- a/tests/test_task.py +++ b/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, + )