Add stop_max_estimate option

AUTHORS.rst

@@ -32,3 +32,4 @@ Patches and Suggestions
 - Jonathan Herriott
 - Job Evers
 - Cyrus Durgin
+- Smite Chow

README.rst

@@ -95,6 +95,13 @@ Most things don't like to be polled as fast as possible, so let's just wait 2 se
     def wait_2_s():
         print "Wait 2 second between retries"
 
+Also we can estimate next time retrying will cost time if expect or not.
+
+.. code-block:: python
+
+    @retry(stop_max_estimate=10000, wait_fixed=2000)
+    def stop_after_8_s():
+        print "Stopping after 8 seconds, because next time retrying will cost 2 second to wait"
 
 Some things perform best with a bit of randomness injected.
 

retrying.py

@@ -65,6 +65,7 @@ def __init__(self,
                  stop=None, wait=None,
                  stop_max_attempt_number=None,
                  stop_max_delay=None,
+                 stop_max_estimate=None,
                  wait_fixed=None,
                  wait_random_min=None, wait_random_max=None,
                  wait_incrementing_start=None, wait_incrementing_increment=None,
@@ -81,6 +82,7 @@ def __init__(self,
 
         self._stop_max_attempt_number = 5 if stop_max_attempt_number is None else stop_max_attempt_number
         self._stop_max_delay = 100 if stop_max_delay is None else stop_max_delay
+        self._stop_max_estimate = 100 if stop_max_estimate is None else stop_max_estimate
         self._wait_fixed = 1000 if wait_fixed is None else wait_fixed
         self._wait_random_min = 0 if wait_random_min is None else wait_random_min
         self._wait_random_max = 1000 if wait_random_max is None else wait_random_max
@@ -102,11 +104,14 @@ def __init__(self,
         if stop_max_delay is not None:
             stop_funcs.append(self.stop_after_delay)
 
+        if stop_max_estimate is not None:
+            stop_funcs.append(self.stop_after_estimate)
+
         if stop_func is not None:
             self.stop = stop_func
 
         elif stop is None:
-            self.stop = lambda attempts, delay: any(f(attempts, delay) for f in stop_funcs)
+            self.stop = lambda attempts, delay, estimate: any(f(attempts, delay, estimate) for f in stop_funcs)
 
         else:
             self.stop = getattr(self, stop)
@@ -142,7 +147,7 @@ def __init__(self,
             # this allows for providing a tuple of exception types that
             # should be allowed to retry on, and avoids having to create
             # a callback that does the same thing
-            if isinstance(retry_on_exception, (tuple)):
+            if isinstance(retry_on_exception, (tuple,)):
                 retry_on_exception = _retry_if_exception_of_type(
                     retry_on_exception)
             self._retry_on_exception = retry_on_exception
@@ -155,14 +160,21 @@ def __init__(self,
 
         self._wrap_exception = wrap_exception
 
-    def stop_after_attempt(self, previous_attempt_number, delay_since_first_attempt_ms):
+    def stop_after_attempt(self, previous_attempt_number, delay_since_first_attempt_ms,
+                           estimate_since_first_attempt_ms):
         """Stop after the previous attempt >= stop_max_attempt_number."""
         return previous_attempt_number >= self._stop_max_attempt_number
 
-    def stop_after_delay(self, previous_attempt_number, delay_since_first_attempt_ms):
+    def stop_after_delay(self, previous_attempt_number, delay_since_first_attempt_ms,
+                         estimate_since_first_attempt_ms):
         """Stop after the time from the first attempt >= stop_max_delay."""
         return delay_since_first_attempt_ms >= self._stop_max_delay
 
+    def stop_after_estimate(self, previous_attempt_number, delay_since_first_attempt_ms,
+                            estimate_since_first_attempt_ms):
+        """Stop after the time from the first attempt plus will sleep time>= stop_max_estimate."""
+        return estimate_since_first_attempt_ms >= self._stop_max_estimate
+
     @staticmethod
     def no_sleep(previous_attempt_number, delay_since_first_attempt_ms):
         """Don't sleep at all before retrying."""
@@ -234,17 +246,20 @@ def call(self, fn, *args, **kwargs):
                 self._after_attempts(attempt_number)
 
             delay_since_first_attempt_ms = int(round(time.time() * 1000)) - start_time
-            if self.stop(attempt_number, delay_since_first_attempt_ms):
+            sleep = self.wait(attempt_number, delay_since_first_attempt_ms)
+
+            if self._wait_jitter_max:
+                jitter = random.random() * self._wait_jitter_max
+                sleep = sleep + max(0, jitter)
+            estimate_since_first_attempt_ms = delay_since_first_attempt_ms + sleep
+
+            if self.stop(attempt_number, delay_since_first_attempt_ms, estimate_since_first_attempt_ms):
                 if not self._wrap_exception and attempt.has_exception:
                     # get() on an attempt with an exception should cause it to be raised, but raise just in case
                     raise attempt.get()
                 else:
                     raise RetryError(attempt)
             else:
-                sleep = self.wait(attempt_number, delay_since_first_attempt_ms)
-                if self._wait_jitter_max:
-                    jitter = random.random() * self._wait_jitter_max
-                    sleep = sleep + max(0, jitter)
                 time.sleep(sleep / 1000.0)
 
             attempt_number += 1

test_retrying.py

@@ -24,28 +24,34 @@ class TestStopConditions(unittest.TestCase):
 
     def test_never_stop(self):
         r = Retrying()
-        self.assertFalse(r.stop(3, 6546))
+        self.assertFalse(r.stop(3, 6546, 6546))
 
     def test_stop_after_attempt(self):
         r = Retrying(stop_max_attempt_number=3)
-        self.assertFalse(r.stop(2, 6546))
-        self.assertTrue(r.stop(3, 6546))
-        self.assertTrue(r.stop(4, 6546))
+        self.assertFalse(r.stop(2, 6546, 6546))
+        self.assertTrue(r.stop(3, 6546, 6546))
+        self.assertTrue(r.stop(4, 6546, 6546))
 
     def test_stop_after_delay(self):
         r = Retrying(stop_max_delay=1000)
-        self.assertFalse(r.stop(2, 999))
-        self.assertTrue(r.stop(2, 1000))
-        self.assertTrue(r.stop(2, 1001))
+        self.assertFalse(r.stop(2, 999, 999))
+        self.assertTrue(r.stop(2, 1000, 999))
+        self.assertTrue(r.stop(2, 1001, 999))
+
+    def test_stop_after_estimate(self):
+        r = Retrying(stop_max_estimate=1000)
+        self.assertFalse(r.stop(2, 999, 999))
+        self.assertTrue(r.stop(2, 999, 1000))
+        self.assertTrue(r.stop(2, 999, 1001))
 
     def test_legacy_explicit_stop_type(self):
         Retrying(stop="stop_after_attempt")
 
     def test_stop_func(self):
-        r = Retrying(stop_func=lambda attempt, delay: attempt == delay)
-        self.assertFalse(r.stop(1, 3))
-        self.assertFalse(r.stop(100, 99))
-        self.assertTrue(r.stop(101, 101))
+        r = Retrying(stop_func=lambda attempt, delay, estimate: attempt == delay)
+        self.assertFalse(r.stop(1, 3, 3))
+        self.assertFalse(r.stop(100, 99, 3))
+        self.assertTrue(r.stop(101, 101, 3))
 
 
 class TestWaitConditions(unittest.TestCase):