fixed send and asend retry logic

Author: daniel-sanche

google/api_core/retry_streaming.py

@@ -151,8 +151,8 @@ def send(self, value):
                 return self.active_target.send(value)
             except Exception as exc:
                 self._handle_exception(exc)
-            # if retryable exception was handled, try again with new active_target
-            return self.send(value)
+            # if exception was retryable, use new target for return value
+            return self.__next__()
         else:
             raise AttributeError(
                 "send() not implemented for {}".format(self.active_target)

google/api_core/retry_streaming_async.py

@@ -128,20 +128,13 @@ def _subtract_time_from_budget(self, start_timestamp):
                 datetime_helpers.utcnow() - start_timestamp
             ).total_seconds()
 
-    async def _iteration_helper(
-        self, iteration_fn: Callable[..., Awaitable], try_again_fn: Callable, *args
-    ):
+    async def _iteration_helper(self, iteration_routine: Awaitable):
         """
         Helper function for sharing logic between __anext__ and asend.
 
         Args:
-          - iteration_fn: The function to call to get the next value from the
-                iterator (e.g. __anext__ or asend)
-          - try_again_fn: The function to call after a retryable exception is
-                encountered, to get a value from the new active_target
-                (e.g. self.__anext__ or self.asend)
-          - *args: Any additional arguments to pass to iteration_fn and
-                try_again_fn (e.g. the value to send to asend)
+          - iteration_routine: The coroutine to await to get the next value
+                from the iterator (e.g. __anext__ or asend)
         """
         # check for expired timeouts before attempting to iterate
         if (
@@ -158,7 +151,7 @@ async def _iteration_helper(
             start_timestamp = datetime_helpers.utcnow()
             # grab the next value from the active_target
             next_val_routine = asyncio.wait_for(
-                iteration_fn(*args), self.remaining_timeout_budget
+                iteration_routine, self.remaining_timeout_budget
             )
             next_val = await next_val_routine
             # subtract the time spent waiting for the next value from the
@@ -168,16 +161,16 @@ async def _iteration_helper(
         except (Exception, asyncio.CancelledError) as exc:
             self._subtract_time_from_budget(start_timestamp)
             await self._handle_exception(exc)
-        # if retryable exception was handled, try again with new active_target
-        return await try_again_fn(*args)
+        # if retryable exception was handled, find the next value to return
+        return await self.__anext__()
 
     async def __anext__(self):
         """
         Implement the async iterator protocol.
         """
         await self._ensure_active_target()
         return await self._iteration_helper(
-            self.active_target.__anext__, self.__anext__
+            self.active_target.__anext__(),
         )
 
     async def aclose(self):
@@ -210,9 +203,7 @@ async def asend(self, value):
         """
         await self._ensure_active_target()
         if getattr(self.active_target, "asend", None):
-            return await self._iteration_helper(
-                self.active_target.asend, self.asend, value
-            )
+            return await self._iteration_helper(self.active_target.asend(value))
         else:
             raise AttributeError(
                 "asend() not implemented for {}".format(self.active_target)

tests/asyncio/test_retry_async.py

@@ -421,7 +421,12 @@ async def test___init___when_retry_is_executed(self, sleep, uniform):
         sleep.assert_any_call(retry_._initial)
 
     async def _generator_mock(
-        self, num=5, error_on=None, exceptions_seen=None, sleep_time=0
+        self,
+        num=5,
+        error_on=None,
+        exceptions_seen=None,
+        sleep_time=0,
+        ignore_sent=False,
     ):
         try:
             sent_in = None
@@ -431,6 +436,8 @@ async def _generator_mock(
                 if error_on and i == error_on:
                     raise ValueError("generator mock error")
                 sent_in = yield (sent_in if sent_in else i)
+                if ignore_sent:
+                    sent_in = None
         except (Exception, BaseException, GeneratorExit) as e:
             # keep track of exceptions seen by generator
             if exceptions_seen is not None:
@@ -607,6 +614,28 @@ async def test___call___with_generator_send(self, sleep):
         assert await generator.__anext__() == 4
         assert await generator.__anext__() == 5
 
+    @mock.patch("asyncio.sleep", autospec=True)
+    @pytest.mark.asyncio
+    async def test___call___generator_send_retry(self, sleep):
+        on_error = mock.Mock(return_value=None)
+        retry_ = retry_async.AsyncRetry(
+            on_error=on_error,
+            predicate=retry_async.if_exception_type(ValueError),
+            is_stream=True,
+            timeout=None,
+        )
+        generator = retry_(self._generator_mock)(error_on=3, ignore_sent=True)
+        with pytest.raises(TypeError) as exc_info:
+            await generator.asend("can not send to fresh generator")
+            assert exc_info.match("can't send non-None value")
+
+        # error thrown on 3
+        # generator should contain 0, 1, 2 looping
+        assert await generator.__anext__() == 0
+        unpacked = [await generator.asend(i) for i in range(10)]
+        assert unpacked == [1, 2, 0, 1, 2, 0, 1, 2, 0, 1]
+        assert on_error.call_count == 3
+
     @mock.patch("asyncio.sleep", autospec=True)
     @pytest.mark.asyncio
     async def test___call___with_generator_close(self, sleep):

tests/unit/test_retry.py

@@ -484,14 +484,21 @@ def test___init___when_retry_is_executed(self, sleep, uniform):
         sleep.assert_any_call(retry_._initial)
 
     def _generator_mock(
-        self, num=5, error_on=None, return_val=None, exceptions_seen=None
+        self,
+        num=5,
+        error_on=None,
+        return_val=None,
+        exceptions_seen=None,
+        ignore_sent=False,
     ):
         try:
             sent_in = None
             for i in range(num):
                 if error_on and i == error_on:
                     raise ValueError("generator mock error")
                 sent_in = yield (sent_in if sent_in else i)
+                if ignore_sent:
+                    sent_in = None
             return return_val
         except (Exception, BaseException, GeneratorExit) as e:
             # keep track of exceptions seen by generator
@@ -610,14 +617,16 @@ def test___call___with_generator_send_retry(self, sleep):
             is_stream=True,
             timeout=None,
         )
-        result = retry_(self._generator_mock)(error_on=3)
+        result = retry_(self._generator_mock)(error_on=3, ignore_sent=True)
         with pytest.raises(TypeError) as exc_info:
             result.send("can not send to fresh generator")
             assert exc_info.match("can't send non-None value")
+        # initiate iteration with None
+        assert result.send(None) == 0
         # error thrown on 3
         # generator should contain 0, 1, 2 looping
-        unpacked = [result.send(None) for i in range(10)]
-        assert unpacked == [0, 1, 2, 0, 1, 2, 0, 1, 2, 0]
+        unpacked = [result.send(i) for i in range(10)]
+        assert unpacked == [1, 2, 0, 1, 2, 0, 1, 2, 0, 1]
         assert on_error.call_count == 3
 
     @mock.patch("time.sleep", autospec=True)