Add ability to inspect upcoming sleep in stop funcs, and add `stop_…

…before_delay` (#423)

* Add upcoming_sleep to retry_state, and add stop_before_delay stop.

* Add unit test to cover stop_before_delay.

* Changelog.

* Update docs for stop_before_delay.

* More docs for the two stop_x_delay functions.

* Add test to ensure it acts the same as stop_after_delay when upcoming sleep is 0.

* Linter fixups.

---------

Co-authored-by: Julien Danjou <[email protected]>

doc/source/index.rst

@@ -124,6 +124,16 @@ retrying stuff.
         print("Stopping after 10 seconds")
         raise Exception
 
+If you're on a tight deadline, and exceeding your delay time isn't ok, 
+then you can give up on retries one attempt before you would exceed the delay. 
+
+.. testcode::
+
+    @retry(stop=stop_before_delay(10))
+    def stop_before_10_s():
+        print("Stopping 1 attempt before 10 seconds")
+        raise Exception
+
 You can combine several stop conditions by using the `|` operator:
 
 .. testcode::

releasenotes/notes/add-stop-before-delay-a775f88ac872c923.yaml

@@ -0,0 +1,7 @@
+---
+features:
+  - |
+    Added a new stop function: stop_before_delay, which will stop execution
+    if the next sleep time would cause overall delay to exceed the specified delay. 
+    Useful for use cases where you have some upper bound on retry times that you must
+    not exceed, so returning before that timeout is preferable than returning after that timeout.

tenacity/__init__.py

@@ -50,6 +50,7 @@
 # Import all built-in stop strategies for easier usage.
 from .stop import stop_after_attempt  # noqa
 from .stop import stop_after_delay  # noqa
+from .stop import stop_before_delay  # noqa
 from .stop import stop_all  # noqa
 from .stop import stop_any  # noqa
 from .stop import stop_never  # noqa
@@ -316,6 +317,13 @@ def iter(self, retry_state: "RetryCallState") -> t.Union[DoAttempt, DoSleep, t.A
         if self.after is not None:
             self.after(retry_state)
 
+        if self.wait:
+            sleep = self.wait(retry_state)
+        else:
+            sleep = 0.0
+
+        retry_state.upcoming_sleep = sleep
+
         self.statistics["delay_since_first_attempt"] = retry_state.seconds_since_start
         if self.stop(retry_state):
             if self.retry_error_callback:
@@ -325,10 +333,6 @@ def iter(self, retry_state: "RetryCallState") -> t.Union[DoAttempt, DoSleep, t.A
                 raise retry_exc.reraise()
             raise retry_exc from fut.exception()
 
-        if self.wait:
-            sleep = self.wait(retry_state)
-        else:
-            sleep = 0.0
         retry_state.next_action = RetryAction(sleep)
         retry_state.idle_for += sleep
         self.statistics["idle_for"] += sleep
@@ -451,6 +455,8 @@ def __init__(
         self.idle_for: float = 0.0
         #: Next action as decided by the retry manager
         self.next_action: t.Optional[RetryAction] = None
+        #: Next sleep time as decided by the retry manager.
+        self.upcoming_sleep: float = 0.0
 
     @property
     def seconds_since_start(self) -> t.Optional[float]:
@@ -568,6 +574,7 @@ def wrap(f: WrappedFn) -> WrappedFn:
     "sleep_using_event",
     "stop_after_attempt",
     "stop_after_delay",
+    "stop_before_delay",
     "stop_all",
     "stop_any",
     "stop_never",

tenacity/stop.py

@@ -92,7 +92,14 @@ def __call__(self, retry_state: "RetryCallState") -> bool:
 
 
 class stop_after_delay(stop_base):
-    """Stop when the time from the first attempt >= limit."""
+    """
+    Stop when the time from the first attempt >= limit.
+
+    Note: `max_delay` will be exceeded, so when used with a `wait`, the actual total delay will be greater
+    than `max_delay` by some of the final sleep period before `max_delay` is exceeded.
+
+    If you need stricter timing with waits, consider `stop_before_delay` instead.
+    """
 
     def __init__(self, max_delay: _utils.time_unit_type) -> None:
         self.max_delay = _utils.to_seconds(max_delay)
@@ -101,3 +108,20 @@ def __call__(self, retry_state: "RetryCallState") -> bool:
         if retry_state.seconds_since_start is None:
             raise RuntimeError("__call__() called but seconds_since_start is not set")
         return retry_state.seconds_since_start >= self.max_delay
+
+
+class stop_before_delay(stop_base):
+    """
+    Stop right before the next attempt would take place after the time from the first attempt >= limit.
+
+    Most useful when you are using with a `wait` function like wait_random_exponential, but need to make
+    sure that the max_delay is not exceeded.
+    """
+
+    def __init__(self, max_delay: _utils.time_unit_type) -> None:
+        self.max_delay = _utils.to_seconds(max_delay)
+
+    def __call__(self, retry_state: "RetryCallState") -> bool:
+        if retry_state.seconds_since_start is None:
+            raise RuntimeError("__call__() called but seconds_since_start is not set")
+        return retry_state.seconds_since_start + retry_state.upcoming_sleep >= self.max_delay

tests/test_tenacity.py

@@ -51,7 +51,7 @@ def _set_delay_since_start(retry_state, delay):
     assert retry_state.seconds_since_start == delay
 
 
-def make_retry_state(previous_attempt_number, delay_since_first_attempt, last_result=None):
+def make_retry_state(previous_attempt_number, delay_since_first_attempt, last_result=None, upcoming_sleep=0):
     """Construct RetryCallState for given attempt number & delay.
 
     Only used in testing and thus is extra careful about timestamp arithmetics.
@@ -70,6 +70,9 @@ def make_retry_state(previous_attempt_number, delay_since_first_attempt, last_re
         retry_state.outcome = last_result
     else:
         retry_state.set_result(None)
+
+    retry_state.upcoming_sleep = upcoming_sleep
+
     _set_delay_since_start(retry_state, delay_since_first_attempt)
     return retry_state
 
@@ -163,6 +166,19 @@ def test_stop_after_delay(self):
                 self.assertTrue(r.stop(make_retry_state(2, 1)))
                 self.assertTrue(r.stop(make_retry_state(2, 1.001)))
 
+    def test_stop_before_delay(self):
+        for delay in (1, datetime.timedelta(seconds=1)):
+            with self.subTest():
+                r = Retrying(stop=tenacity.stop_before_delay(delay))
+                self.assertFalse(r.stop(make_retry_state(2, 0.999, upcoming_sleep=0.0001)))
+                self.assertTrue(r.stop(make_retry_state(2, 1, upcoming_sleep=0.001)))
+                self.assertTrue(r.stop(make_retry_state(2, 1, upcoming_sleep=1)))
+
+                # It should act the same as stop_after_delay if upcoming sleep is 0
+                self.assertFalse(r.stop(make_retry_state(2, 0.999, upcoming_sleep=0)))
+                self.assertTrue(r.stop(make_retry_state(2, 1, upcoming_sleep=0)))
+                self.assertTrue(r.stop(make_retry_state(2, 1.001, upcoming_sleep=0)))
+
     def test_legacy_explicit_stop_type(self):
         Retrying(stop="stop_after_attempt")