remove wait_for in async streaming for perf reasons

Author: daniel-sanche

google/api_core/retry_streaming.py

@@ -97,24 +97,39 @@ def _handle_exception(self, exc) -> None:
                 next_attempt = datetime_helpers.utcnow() + datetime.timedelta(
                     seconds=next_sleep
                 )
-                if self.deadline < next_attempt:
-                    raise exceptions.RetryError(
-                        f"Deadline of {self.timeout:.1f} seconds exceeded", exc
-                    ) from exc
+                self._check_timeout(next_attempt, exc)
             # sleep before retrying
             _LOGGER.debug(
                 "Retrying due to {}, sleeping {:.1f}s ...".format(exc, next_sleep)
             )
             time.sleep(next_sleep)
             self.active_target = self.target_fn().__iter__()
 
+    def _check_timeout(self, current_time:float, source_exception: Optional[Exception] = None) -> None:
+        """
+        Helper function to check if the timeout has been exceeded, and raise a RetryError if so.
+
+        Args:
+          - current_time: the timestamp to check against the deadline
+          - source_exception: the exception that triggered the timeout check, if any
+        Raises:
+          - RetryError if the deadline has been exceeded
+        """
+        if self.deadline is not None and self.deadline < current_time:
+            raise exceptions.RetryError(
+                "Timeout of {:.1f}s exceeded".format(self.timeout),
+                source_exception,
+            ) from source_exception
+
     def __next__(self) -> T:
         """
         Implement the iterator protocol.
 
         Returns:
           - the next value of the active_target iterator
         """
+        # check for expired timeouts before attempting to iterate
+        self._check_timeout(datetime_helpers.utcnow())
         try:
             return next(self.active_target)
         except Exception as exc:
@@ -152,6 +167,8 @@ def send(self, *args, **kwargs) -> T:
         Raises:
           - AttributeError if the active_target does not have a send() method
         """
+        # check for expired timeouts before attempting to iterate
+        self._check_timeout(datetime_helpers.utcnow())
         if getattr(self.active_target, "send", None):
             casted_target = cast(Generator, self.active_target)
             try:

google/api_core/retry_streaming_async.py

@@ -170,10 +170,9 @@ async def _iteration_helper(self, iteration_routine: Awaitable) -> T:
             # start the timer for the current operation
             start_timestamp = datetime_helpers.utcnow()
             # grab the next value from the active_target
-            next_val_routine = asyncio.wait_for(
-                iteration_routine, self.remaining_timeout_budget
-            )
-            next_val = await next_val_routine
+            # Note: interrupting with asyncio.wait_for is expensive,
+            # so we only check  for timeouts at the start of each iteration
+            next_val = await iteration_routine
             # subtract the time spent waiting for the next value from the
             # remaining timeout budget
             self._subtract_time_from_budget(start_timestamp)

tests/asyncio/test_retry_async.py

@@ -540,44 +540,6 @@ def increase_time(sleep_delay):
         assert last_wait == 4
         assert total_wait == 7
 
-    @pytest.mark.asyncio
-    async def test___call___generator_timeout_cancellations(self):
-        """
-        Tests that a retry-decorated generator will throw a RetryError
-        after using its time budget
-        """
-        on_error = mock.Mock(return_value=None)
-        retry_ = retry_async.AsyncRetry(
-            predicate=retry_async.if_exception_type(ValueError),
-            deadline=0.2,
-            is_stream=True,
-        )
-        utcnow = datetime.datetime.utcnow()
-        utcnow_patcher = mock.patch(
-            "google.api_core.datetime_helpers.utcnow", return_value=utcnow
-        )
-        # ensure generator times out when awaiting past deadline
-        with pytest.raises(exceptions.RetryError):
-            infinite_gen = retry_(self._generator_mock, on_error)(sleep_time=60)
-            await infinite_gen.__anext__()
-        # ensure time between yields isn't counted
-        with utcnow_patcher as patched_utcnow:
-            generator = retry_(self._generator_mock)(sleep_time=0.05)
-            assert await generator.__anext__() == 0
-            patched_utcnow.return_value += datetime.timedelta(20)
-            assert await generator.__anext__() == 1
-        # ensure timeout budget is tracked
-        generator = retry_(self._generator_mock)(sleep_time=0.07)
-        assert await generator.__anext__() == 0
-        assert await generator.__anext__() == 1
-        with pytest.raises(exceptions.RetryError) as exc:
-            await generator.__anext__()
-            assert "Timeout of 0.2s exceeded" in str(exc.value)
-        # subsequent calls should also return a RetryError
-        with pytest.raises(exceptions.RetryError) as exc:
-            await generator.__anext__()
-            assert "Timeout of 0.2s exceeded" in str(exc.value)
-
     @pytest.mark.asyncio
     async def test___call___generator_await_cancel_retryable(self):
         """
@@ -794,3 +756,24 @@ async def __anext__(self):
         assert await retryable.__anext__() == 3
         with pytest.raises(StopAsyncIteration):
             await retryable.__anext__()
+
+    @pytest.mark.asyncio
+    async def test_iterate_stream_after_deadline(self):
+        """
+        Streaming retries should raise RetryError when calling next or send after deadline has passed
+        """
+        retry_ = retry_async.AsyncRetry(is_stream=True, deadline=0.01)
+        decorated = retry_(self._generator_mock)
+        generator = decorated(10)
+        starting_time_budget = generator.remaining_timeout_budget
+        assert starting_time_budget == 0.01
+        await generator.__anext__()
+        # ensure budget is used on each call
+        assert generator.remaining_timeout_budget < starting_time_budget
+        # simulate using up budget
+        generator.remaining_timeout_budget = 0
+        with pytest.raises(exceptions.RetryError):
+            await generator.__anext__()
+        with pytest.raises(exceptions.RetryError):
+            await generator.asend("test")
+

tests/unit/test_retry.py

@@ -304,7 +304,7 @@ def test_with_delay_partial_options(self):
         assert retry_ is not new_retry
         assert new_retry._initial == 1
         assert new_retry._maximum == 4
-        assert new_retry._multiplier == 3
+        assert n ew_retry._multiplier == 3
 
         new_retry = retry_.with_delay(multiplier=4)
         assert retry_ is not new_retry
@@ -768,3 +768,39 @@ def test___call___with_is_stream(self, sleep):
         gen = gen_retry_(wrapped)()
         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
+        """
+        from time import sleep
+        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)
+        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
+        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")
+