Add better support for throttling exceptions

opnieuw/exceptions.py

@@ -3,9 +3,25 @@
 #
 # Licensed under the 3-clause BSD license, see the LICENSE file in the repository root.
 
+from typing import Union
+
 
 class RetryException(Exception):
     """
     Defines a custom RetryException that can be raised for specific errors we
     want to retry on.
     """
+
+
+class BackoffAndRetryException(Exception):
+    """
+    A custom exception that can be raised for specific errors when a fixed
+    wait time before retrying is in order. This can be particularly useful
+    when requests are throttled and the response includes a Retry-After header.
+
+    This will reset the counters for `max_calls_total`
+    and `retry_window_after_first_call_in_seconds`
+    """
+
+    def __init__(self, seconds: Union[int, float]):
+        self.seconds = seconds

opnieuw/retries.py

@@ -26,6 +26,8 @@
 )
 
 from .clock import Clock, MonotonicClock
+from .exceptions import BackoffAndRetryException
+
 
 logger = logging.getLogger(__name__)
 
@@ -162,6 +164,10 @@ def retry(
      - `namespace` - A name with which the wait behavior can be controlled
        using the `opnieuw.test_util.retry_immediately` contextmanager.
 
+    Note: if the decorated function raises a `BackoffAndRetryException`, the
+    counters for `max_calls_total` and `retry_window_after_first_call_in_seconds`
+    will reset.
+
     This function will:
 
      - Calculate how to fit `max_calls_total` executions of function in the
@@ -202,24 +208,43 @@ def foo() -> None:
     Opnieuw is based on a retry algorithm off of:
         https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/
     """
+
+    def get_state_generator():
+        """ Create a generator object that produces DoCall and DoWait objects. """
+        return iter(
+            __retry_state_namespaces[namespace](
+                MonotonicClock(),
+                max_calls_total=max_calls_total,
+                retry_window_after_first_call_in_seconds=retry_window_after_first_call_in_seconds,
+            )
+        )
+
     def decorator(f: F) -> F:
         @functools.wraps(f)
         def wrapper(*args: Any, **kwargs: Any) -> Any:
 
             last_exception = None
 
-            retry_state = __retry_state_namespaces[namespace](
-                MonotonicClock(),
-                max_calls_total=max_calls_total,
-                retry_window_after_first_call_in_seconds=retry_window_after_first_call_in_seconds,
-            )
+            retry_state = get_state_generator()
 
-            for retry_action in retry_state:
+            while True:
+                try:
+                    retry_action = next(retry_state)
+                except StopIteration:
+                    break
 
                 if isinstance(retry_action, DoCall):
                     try:
                         return f(*args, **kwargs)
 
+                    except BackoffAndRetryException as e:
+                        last_exception = e
+                        logger.debug(
+                            f"Encountered fixed backoff, sleeping for {e.seconds}"
+                        )
+                        time.sleep(e.seconds)
+                        retry_state = get_state_generator()
+
                     except retry_on_exceptions as e:
                         last_exception = e
 
@@ -267,24 +292,42 @@ def retry_async(
     retry_window_after_first_call_in_seconds: int = 60,
     namespace: Optional[str] = None,
 ) -> Callable[[AF], AF]:
+    def get_state_generator():
+        """ Create a generator object that produces DoCall and DoWait objects. """
+        return iter(
+            __retry_state_namespaces[namespace](
+                MonotonicClock(),
+                max_calls_total=max_calls_total,
+                retry_window_after_first_call_in_seconds=retry_window_after_first_call_in_seconds,
+            )
+        )
+
     def decorator(f: AF) -> AF:
         @functools.wraps(f)
         async def wrapper(*args: Any, **kwargs: Any) -> Any:
 
             last_exception = None
 
-            retry_state = __retry_state_namespaces[namespace](
-                MonotonicClock(),
-                max_calls_total=max_calls_total,
-                retry_window_after_first_call_in_seconds=retry_window_after_first_call_in_seconds,
-            )
+            retry_state = get_state_generator()
 
-            for retry_action in retry_state:
+            while True:
+                try:
+                    retry_action = next(retry_state)
+                except StopIteration:
+                    break
 
                 if isinstance(retry_action, DoCall):
                     try:
                         return await f(*args, **kwargs)
 
+                    except BackoffAndRetryException as e:
+                        last_exception = e
+                        logger.debug(
+                            f"Encountered fixed backoff, sleeping for {e.seconds}"
+                        )
+                        await asyncio.sleep(e.seconds)
+                        retry_state = get_state_generator()
+
                     except retry_on_exceptions as e:
                         last_exception = e
 

tests/test_opnieuw.py

@@ -9,6 +9,11 @@
 from opnieuw.clock import TestClock, MonotonicClock
 from opnieuw.retries import RetryState, DoCall, retry
 from opnieuw.test_util import retry_immediately
+from opnieuw.exceptions import BackoffAndRetryException
+
+
+class CustomBackoffThrottleError(BackoffAndRetryException):
+    pass
 
 
 class TestRetryState(unittest.TestCase):
@@ -110,6 +115,32 @@ def test_mixed_states(self) -> None:
         with retry_immediately("bar_retry"):
             self.test_retry_with_waits()
 
+    @retry(
+        retry_on_exceptions=ValueError,
+        max_calls_total=4,
+        retry_window_after_first_call_in_seconds=10,
+    )
+    def namespaced_retry_fixed_backoff(self) -> None:
+        self.counter += 1
+        if self.counter == 4:
+            raise CustomBackoffThrottleError(seconds=1)
+        raise ValueError
+
+    def test_fixed_backoff_reset_max_calls(self):
+        self.counter = 0
+        start = time.monotonic()
+
+        with retry_immediately():
+            self.assertRaises(ValueError, self.namespaced_retry_fixed_backoff)
+
+        # expect to sleep at least 1 second for the BackoffAndRetry
+        end = time.monotonic()
+        runtime_seconds = end - start
+        self.assertGreater(runtime_seconds, 1)
+
+        # we exhaust max_calls_total twice
+        assert self.counter == 8
+
 
 if __name__ == "__main__":
     unittest.main()