don't check timeout on each yield by default

daniel-sanche · daniel-sanche · commit 9cadd63ad125 · 2023-08-15T10:47:42.000-07:00
diff --git a/google/api_core/retry_streaming.py b/google/api_core/retry_streaming.py
@@ -16,7 +16,6 @@
 
 from typing import Callable, Optional, Iterable, Iterator, Generator, TypeVar, Any, cast
 
-import datetime
 import logging
 import time
 
@@ -104,6 +103,7 @@ def __init__(
         sleep_generator: Iterable[float],
         timeout: Optional[float] = None,
         on_error: Optional[Callable[[Exception], None]] = None,
+        check_timeout_on_yield=False,
     ):
         """
         Args:
@@ -119,6 +119,11 @@ def __init__(
             on_error: A function to call while processing a
                 retryable exception.  Any error raised by this function will *not*
                 be caught.
+            check_timeout_on_yield: If True, the timeout value will be checked
+                after each yield. If the timeout has been exceeded, the generator
+                will raise a RetryError. Note that this adds an overhead to each
+                yield, so it is preferred to add the timeout logic to the wrapped 
+                stream when possible.
         """
         self.target_fn = target
         self.active_target: Iterator[T] = self.target_fn().__iter__()
@@ -130,6 +135,7 @@ def __init__(
             self.deadline = time.monotonic() + self.timeout
         else:
             self.deadline = None
+        self._check_timeout_on_yield = check_timeout_on_yield
 
     def __iter__(self) -> Generator[T, Any, None]:
         """
@@ -194,7 +200,8 @@ def __next__(self) -> T:
           - the next value of the active_target iterator
         """
         # check for expired timeouts before attempting to iterate
-        self._check_timeout(time.monotonic())
+        if self._check_timeout_on_yield:
+            self._check_timeout(time.monotonic())
         try:
             return next(self.active_target)
         except Exception as exc:
@@ -233,7 +240,8 @@ def send(self, *args, **kwargs) -> T:
           - AttributeError if the active_target does not have a send() method
         """
         # check for expired timeouts before attempting to iterate
-        self._check_timeout(time.monotonic())
+        if self._check_timeout_on_yield:
+            self._check_timeout(time.monotonic())
         if getattr(self.active_target, "send", None):
             casted_target = cast(Generator, self.active_target)
             try:
diff --git a/google/api_core/retry_streaming_async.py b/google/api_core/retry_streaming_async.py
@@ -118,6 +118,7 @@ def __init__(
         sleep_generator: Iterable[float],
         timeout: Optional[float] = None,
         on_error: Optional[Callable[[Exception], None]] = None,
+        check_timeout_on_yield=False,
     ):
         """
         Args:
@@ -133,6 +134,11 @@ def __init__(
             on_error: A function to call while processing a
                 retryable exception.  Any error raised by this function will *not*
                 be caught.
+            check_timeout_on_yield: If True, the timeout value will be checked
+                after each yield. If the timeout has been exceeded, the generator
+                will raise a RetryError. Note that this adds an overhead to each
+                yield, so it is preferred to add the timeout logic to the wrapped 
+                stream when possible.
         """
         self.target_fn = target
         # active target must be populated in an async context
@@ -146,6 +152,7 @@ def __init__(
             self.deadline = time.monotonic() + self.timeout
         else:
             self.deadline = None
+        self._check_timeout_on_yield = check_timeout_on_yield
 
     def _check_timeout(
         self, current_time: float, source_exception: Optional[Exception] = None
@@ -176,7 +183,7 @@ async def _ensure_active_target(self) -> AsyncIterator[T]:
         Returns:
           - The active_target iterable
         """
-        if not self.active_target:
+        if self.active_target is None:
             new_iterable = self.target_fn()
             if isinstance(new_iterable, Awaitable):
                 new_iterable = await new_iterable
@@ -226,7 +233,8 @@ async def _iteration_helper(self, iteration_routine: Awaitable) -> T:
           - The next value from the active_target iterator.
         """
         # check for expired timeouts before attempting to iterate
-        self._check_timeout(time.monotonic())
+        if self._check_timeout_on_yield:
+            self._check_timeout(time.monotonic())
         try:
             # grab the next value from the active_target
             # Note: here would be a good place to add a timeout, like asyncio.wait_for.
diff --git a/tests/asyncio/test_retry_async.py b/tests/asyncio/test_retry_async.py
@@ -746,15 +746,38 @@ async def __anext__(self):
         with pytest.raises(StopAsyncIteration):
             await retryable.__anext__()
 
+    @mock.patch("random.uniform", autospec=True, side_effect=lambda m, n: n)
+    @mock.patch("asyncio.sleep", autospec=True)
     @pytest.mark.asyncio
-    async def test_iterate_stream_after_deadline(self):
+    async def test_yield_stream_after_deadline(self, sleep, uniform):
         """
-        Streaming retries should raise RetryError when calling next or send after deadline has passed
+        By default, if the deadline is hit between yields, the generator will continue.
+
+        There is a flag that should cause the wrapper to test for the deadline after
+        each yield.
         """
-        retry_ = retry_async.AsyncRetry(is_stream=True, deadline=0.01)
-        decorated = retry_(self._generator_mock)
-        generator = decorated(10)
-        await generator.__anext__()
-        await asyncio.sleep(0.02)
-        with pytest.raises(exceptions.RetryError):
-            await generator.__anext__()
+        import time
+        from google.api_core.retry_streaming_async import AsyncRetryableGenerator
+        timeout = 2
+        time_now = time.monotonic()
+        now_patcher = mock.patch(
+            "time.monotonic", return_value=time_now,
+        )
+
+        with now_patcher as patched_now:
+            no_check = AsyncRetryableGenerator(self._generator_mock, None, [], timeout=timeout, check_timeout_on_yield=False)
+            assert no_check._check_timeout_on_yield is False
+            check = AsyncRetryableGenerator(self._generator_mock, None, [], timeout=timeout, check_timeout_on_yield=True)
+            assert check._check_timeout_on_yield is True
+
+            # first yield should be fine
+            await check.__anext__()
+            await no_check.__anext__()
+
+            # simulate a delay before next yield
+            patched_now.return_value += timeout + 1
+
+            # second yield should raise when check_timeout_on_yield is True
+            with pytest.raises(exceptions.RetryError):
+                await check.__anext__()
+            await no_check.__anext__()
diff --git a/tests/unit/test_retry.py b/tests/unit/test_retry.py
@@ -772,39 +772,38 @@ def test___call___with_is_stream(self, sleep):
         unpacked = [next(gen) for i in range(10)]
         assert unpacked == [0, 1, 2, 3, 4, 5, 0, 1, 2, 3]
 
-    def test_iterate_stream_after_deadline(self):
-        """
-        Streaming retries should raise RetryError when calling next after deadline has passed
+    @mock.patch("random.uniform", autospec=True, side_effect=lambda m, n: n)
+    @mock.patch("asyncio.sleep", autospec=True)
+    @pytest.mark.asyncio
+    async def test_yield_stream_after_deadline(self, sleep, uniform):
         """
-        from time import sleep
+        By default, if the deadline is hit between yields, the generator will continue.
 
-        retry_ = retry.Retry(
-            predicate=retry.if_exception_type(ValueError),
-            is_stream=True,
-            deadline=0.01,
+        There is a flag that should cause the wrapper to test for the deadline after
+        each yield.
+        """
+        import time
+        from google.api_core.retry_streaming import RetryableGenerator
+        timeout = 2
+        time_now = time.monotonic()
+        now_patcher = mock.patch(
+            "time.monotonic", return_value=time_now,
         )
-        decorated = retry_(self._generator_mock)
-        generator = decorated(10)
-        next(generator)
-        sleep(0.02)
-        with pytest.raises(exceptions.RetryError):
-            next(generator)
 
-    def test_iterate_stream_send_after_deadline(self):
-        """
-        Streaming retries should raise RetryError when calling send after deadline has passed
-        """
-        from time import sleep
+        with now_patcher as patched_now:
+            no_check = RetryableGenerator(self._generator_mock, None, [], timeout=timeout, check_timeout_on_yield=False)
+            assert no_check._check_timeout_on_yield is False
+            check = RetryableGenerator(self._generator_mock, None, [], timeout=timeout, check_timeout_on_yield=True)
+            assert check._check_timeout_on_yield is True
 
-        retry_ = retry.Retry(
-            predicate=retry.if_exception_type(ValueError),
-            is_stream=True,
-            deadline=0.01,
-        )
-        decorated = retry_(self._generator_mock)
-        generator = decorated(10)
-        next(generator)
-        generator.send("test")
-        sleep(0.02)
-        with pytest.raises(exceptions.RetryError):
-            generator.send("test")
+            # first yield should be fine
+            next(check)
+            next(no_check)
+
+            # simulate a delay before next yield
+            patched_now.return_value += timeout + 1
+
+            # second yield should raise when check_timeout_on_yield is True
+            with pytest.raises(exceptions.RetryError):
+                next(check)
+            next(no_check)
