Use monotonic clocks for result timeouts

Author: IlyaKhalafi

taskiq/funcs.py

@@ -1,5 +1,5 @@
 import asyncio
-from time import time
+from time import monotonic
 from typing import Any
 
 from taskiq.exceptions import TaskiqResultTimeoutError
@@ -30,7 +30,7 @@ async def gather(
     :return: list of TaskiqResults.
     """
     loop = asyncio.get_event_loop()
-    start_time = time()
+    start_time = monotonic()
     ordered_ids = [task.task_id for task in tasks]
     task_ids = {task.task_id for task in tasks}
     results = {}
@@ -46,7 +46,7 @@ async def check_task(task: AsyncTaskiqTask[Any]) -> None:
             results[task.task_id] = await task.get_result(with_logs=with_logs)
 
     while task_ids:
-        if 0 < timeout < time() - start_time:
+        if 0 < timeout < monotonic() - start_time:
             raise TaskiqResultTimeoutError(timeout=timeout)
         check_tasks = []
         for task in tasks:

taskiq/task.py

@@ -1,6 +1,6 @@
 import asyncio
 from logging import getLogger
-from time import time
+from time import monotonic
 from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
 from taskiq.compat import parse_obj_as
@@ -96,9 +96,9 @@ async def wait_result(
             become ready in provided period of time.
         :return: task's return value.
         """
-        start_time = time()
+        start_time = monotonic()
         while not await self.is_ready():
-            if 0 < timeout < time() - start_time:
+            if 0 < timeout < monotonic() - start_time:
                 raise TaskiqResultTimeoutError(timeout=timeout)
             await asyncio.sleep(check_interval)
         return await self.get_result(with_logs=with_logs)

tests/test_task_timeout.py

@@ -0,0 +1,128 @@
+from types import ModuleType
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+
+import taskiq.funcs as funcs_module
+import taskiq.task as task_module
+from taskiq.abc.result_backend import AsyncResultBackend
+from taskiq.exceptions import TaskiqResultTimeoutError
+from taskiq.result import TaskiqResult
+from taskiq.task import AsyncTaskiqTask
+
+
+async def wait_for_result(
+    task: AsyncTaskiqTask[int],
+    module: ModuleType,
+    timeout: float,
+) -> TaskiqResult[int]:
+    """Exercise either public result-waiting API."""
+    if module is funcs_module:
+        return (await funcs_module.gather(task, timeout=timeout, periodicity=0))[0]
+    return await task.wait_result(timeout=timeout, check_interval=0)
+
+
+@pytest.mark.parametrize("module", [task_module, funcs_module])
+async def test_result_timeout_survives_backward_wall_clock(
+    monkeypatch: pytest.MonkeyPatch,
+    module: ModuleType,
+) -> None:
+    """A backwards wall-clock jump must not postpone the timeout."""
+    # Keep the old clock patched so this also reproduces on the unfixed source.
+    monkeypatch.setattr(
+        module,
+        "time",
+        Mock(side_effect=[100, 99, 98, 97]),
+        raising=False,
+    )
+    monkeypatch.setattr(
+        module,
+        "monotonic",
+        Mock(side_effect=[0, 0.5, 1.1]),
+        raising=False,
+    )
+    backend = AsyncMock(spec=AsyncResultBackend)
+    backend.is_result_ready.side_effect = [False, False, True]
+    backend.get_result.return_value = TaskiqResult(
+        is_err=False,
+        return_value=42,
+        execution_time=0,
+    )
+    task = AsyncTaskiqTask[int]("task", backend)
+
+    with pytest.raises(TaskiqResultTimeoutError):
+        await wait_for_result(task, module, timeout=1)
+
+    assert backend.is_result_ready.await_count == (1 if module is funcs_module else 2)
+    backend.get_result.assert_not_awaited()
+
+
+@pytest.mark.parametrize("module", [task_module, funcs_module])
+async def test_forward_wall_clock_does_not_expire_result_wait(
+    monkeypatch: pytest.MonkeyPatch,
+    module: ModuleType,
+) -> None:
+    """A forward wall-clock jump must not end an otherwise timely wait."""
+    monkeypatch.setattr(module, "time", Mock(side_effect=[100, 10000]), raising=False)
+    monkeypatch.setattr(
+        module,
+        "monotonic",
+        Mock(side_effect=[0, 0.25, 0.5]),
+        raising=False,
+    )
+    backend = AsyncMock(spec=AsyncResultBackend)
+    backend.is_result_ready.side_effect = [False, True]
+    result = TaskiqResult(is_err=False, return_value=42, execution_time=0)
+    backend.get_result.return_value = result
+    task = AsyncTaskiqTask[int]("task", backend)
+
+    assert await wait_for_result(task, module, timeout=1) is result
+
+
+@pytest.mark.parametrize("module", [task_module, funcs_module])
+@pytest.mark.parametrize("timeout", [-1, 0])
+async def test_result_wait_keeps_unlimited_timeout(
+    monkeypatch: pytest.MonkeyPatch,
+    module: ModuleType,
+    timeout: float,
+) -> None:
+    """Nonpositive timeouts keep polling until the result is available."""
+    monkeypatch.setattr(
+        module,
+        "time",
+        Mock(side_effect=[100, 99, 98, 97]),
+        raising=False,
+    )
+    monkeypatch.setattr(
+        module,
+        "monotonic",
+        Mock(side_effect=[0, 100, 200]),
+        raising=False,
+    )
+    backend = AsyncMock(spec=AsyncResultBackend)
+    backend.is_result_ready.side_effect = [False, False, True]
+    result = TaskiqResult(is_err=False, return_value=42, execution_time=0)
+    backend.get_result.return_value = result
+    task = AsyncTaskiqTask[int]("task", backend)
+
+    assert await wait_for_result(task, module, timeout=timeout) is result
+    assert backend.is_result_ready.await_count == 3
+
+
+@pytest.mark.parametrize("module", [task_module, funcs_module])
+async def test_result_wait_keeps_strict_timeout_boundary(
+    monkeypatch: pytest.MonkeyPatch,
+    module: ModuleType,
+) -> None:
+    """Elapsed time equal to the timeout does not expire the wait."""
+    monkeypatch.setattr(module, "time", Mock(side_effect=[0, 1]), raising=False)
+    monkeypatch.setattr(module, "monotonic", Mock(side_effect=[0, 1]), raising=False)
+    backend = AsyncMock(spec=AsyncResultBackend)
+    backend.is_result_ready.side_effect = (
+        [False, True] if module is task_module else [True]
+    )
+    result = TaskiqResult(is_err=False, return_value=42, execution_time=0)
+    backend.get_result.return_value = result
+    task = AsyncTaskiqTask[int]("task", backend)
+
+    assert await wait_for_result(task, module, timeout=1) is result