feat: retry generates backoff value after completing on_error callbacks

Author: daniel-sanche

google/api_core/retry/retry_base.py

@@ -25,7 +25,7 @@
 import time
 
 from enum import Enum
-from typing import Any, Callable, Optional, TYPE_CHECKING
+from typing import Any, Callable, Optional, Iterable, TYPE_CHECKING
 
 import requests.exceptions
 
@@ -174,7 +174,7 @@ def build_retry_error(
 def _retry_error_helper(
     exc: Exception,
     deadline: float | None,
-    next_sleep: float,
+    sleep_generator: Iterable[float],
     error_list: list[Exception],
     predicate_fn: Callable[[Exception], bool],
     on_error_fn: Callable[[Exception], None] | None,
@@ -183,7 +183,7 @@ def _retry_error_helper(
         tuple[Exception, Exception | None],
     ],
     original_timeout: float | None,
-):
+) -> float:
     """
     Shared logic for handling an error for all retry implementations
 
@@ -194,13 +194,15 @@ def _retry_error_helper(
     Args:
        - exc: the exception that was raised
        - deadline: the deadline for the retry, calculated as a diff from time.monotonic()
-       - next_sleep: the next sleep interval
+       - sleep_generator: iterable to draw the next backoff value from
        - error_list: the list of exceptions that have been raised so far
        - predicate_fn: takes `exc` and returns true if the operation should be retried
        - on_error_fn: callback to execute when a retryable error occurs
        - exc_factory_fn: callback used to build the exception to be raised on terminal failure
        - original_timeout_val: the original timeout value for the retry (in seconds),
            to be passed to the exception factory for building an error message
+    Returns:
+        - the next backoff value to use
     """
     error_list.append(exc)
     if not predicate_fn(exc):
@@ -212,6 +214,10 @@ def _retry_error_helper(
         raise final_exc from source_exc
     if on_error_fn is not None:
         on_error_fn(exc)
+    try:
+        next_sleep = next(sleep_generator)
+    except StopIteration:
+        raise ValueError("Sleep generator stopped yielding sleep values.") from exc
     if deadline is not None and time.monotonic() + next_sleep > deadline:
         final_exc, source_exc = exc_factory_fn(
             error_list,
@@ -222,6 +228,7 @@ def _retry_error_helper(
     _LOGGER.debug(
         "Retrying due to {}, sleeping {:.1f}s ...".format(error_list[-1], next_sleep)
     )
+    return next_sleep
 
 
 class _BaseRetry(object):

google/api_core/retry/retry_streaming.py

@@ -108,7 +108,7 @@ def retry_target_stream(
     )
     error_list: list[Exception] = []
 
-    for sleep in sleep_generator:
+    while True:
         # Start a new retry loop
         try:
             # Note: in the future, we can add a ResumptionStrategy object
@@ -121,20 +121,18 @@ def retry_target_stream(
         # This function explicitly must deal with broad exceptions.
         except Exception as exc:
             # defer to shared logic for handling errors
-            _retry_error_helper(
+            next_sleep = _retry_error_helper(
                 exc,
                 deadline,
-                sleep,
+                sleep_generator,
                 error_list,
                 predicate,
                 on_error,
                 exception_factory,
                 timeout,
             )
             # if exception not raised, sleep before next attempt
-            time.sleep(sleep)
-
-    raise ValueError("Sleep generator stopped yielding sleep values.")
+            time.sleep(next_sleep)
 
 
 class StreamingRetry(_BaseRetry):

google/api_core/retry/retry_streaming_async.py

@@ -111,7 +111,7 @@ async def retry_target_stream(
     error_list: list[Exception] = []
     target_is_generator: bool | None = None
 
-    for sleep in sleep_generator:
+    while True:
         # Start a new retry loop
         try:
             # Note: in the future, we can add a ResumptionStrategy object
@@ -174,22 +174,22 @@ async def retry_target_stream(
         # This function explicitly must deal with broad exceptions.
         except Exception as exc:
             # defer to shared logic for handling errors
-            _retry_error_helper(
+            next_sleep = _retry_error_helper(
                 exc,
                 deadline,
-                sleep,
+                sleep_generator,
                 error_list,
                 predicate,
                 on_error,
                 exception_factory,
                 timeout,
             )
             # if exception not raised, sleep before next attempt
-            await asyncio.sleep(sleep)
+            await asyncio.sleep(next_sleep)
+
         finally:
             if target_is_generator and target_iterator is not None:
                 await cast(AsyncGenerator["_Y", None], target_iterator).aclose()
-    raise ValueError("Sleep generator stopped yielding sleep values.")
 
 
 class AsyncStreamingRetry(_BaseRetry):

google/api_core/retry/retry_unary.py

@@ -139,7 +139,7 @@ def retry_target(
     deadline = time.monotonic() + timeout if timeout is not None else None
     error_list: list[Exception] = []
 
-    for sleep in sleep_generator:
+    while True:
         try:
             result = target()
             if inspect.isawaitable(result):
@@ -150,20 +150,18 @@ def retry_target(
         # This function explicitly must deal with broad exceptions.
         except Exception as exc:
             # defer to shared logic for handling errors
-            _retry_error_helper(
+            next_sleep = _retry_error_helper(
                 exc,
                 deadline,
-                sleep,
+                sleep_generator,
                 error_list,
                 predicate,
                 on_error,
                 exception_factory,
                 timeout,
             )
             # if exception not raised, sleep before next attempt
-            time.sleep(sleep)
-
-    raise ValueError("Sleep generator stopped yielding sleep values.")
+            time.sleep(next_sleep)
 
 
 class Retry(_BaseRetry):

google/api_core/retry/retry_unary_async.py

@@ -150,27 +150,25 @@ async def retry_target(
     deadline = time.monotonic() + timeout if timeout is not None else None
     error_list: list[Exception] = []
 
-    for sleep in sleep_generator:
+    while True:
         try:
             return await target()
         # pylint: disable=broad-except
         # This function explicitly must deal with broad exceptions.
         except Exception as exc:
             # defer to shared logic for handling errors
-            _retry_error_helper(
+            next_sleep = _retry_error_helper(
                 exc,
                 deadline,
-                sleep,
+                sleep_generator,
                 error_list,
                 predicate,
                 on_error,
                 exception_factory,
                 timeout,
             )
             # if exception not raised, sleep before next attempt
-            await asyncio.sleep(sleep)
-
-    raise ValueError("Sleep generator stopped yielding sleep values.")
+            await asyncio.sleep(next_sleep)
 
 
 class AsyncRetry(_BaseRetry):