Merge branch 'exp-max-retry'

cassandra/policies.py

@@ -528,10 +528,14 @@ class ExponentialReconnectionPolicy(ReconnectionPolicy):
     a set maximum delay.
     """
 
-    def __init__(self, base_delay, max_delay):
+    def __init__(self, base_delay, max_delay, max_attempts=64):
         """
         `base_delay` and `max_delay` should be in floating point units of
         seconds.
+
+        `max_attempts` should be a total number of attempts to be made before
+        giving up, or :const:`None` to continue reconnection attempts forever.
+        The default is 64.
         """
         if base_delay < 0 or max_delay < 0:
             raise ValueError("Delays may not be negative")
@@ -539,11 +543,18 @@ class ExponentialReconnectionPolicy(ReconnectionPolicy):
         if max_delay < base_delay:
             raise ValueError("Max delay must be greater than base delay")
 
+        if max_attempts is not None and max_attempts < 0:
+            raise ValueError("max_attempts must not be negative")
+
         self.base_delay = base_delay
         self.max_delay = max_delay
+        self.max_attempts = max_attempts
 
     def new_schedule(self):
-        return (min(self.base_delay * (2 ** i), self.max_delay) for i in range(64))
+        i=0
+        while self.max_attempts == None or i < self.max_attempts:
+            yield min(self.base_delay * (2 ** i), self.max_delay)
+            i += 1
 
 
 class WriteType(object):

tests/unit/test_policies.py

@@ -814,9 +814,18 @@ class ExponentialReconnectionPolicyTest(unittest.TestCase):
         self.assertRaises(ValueError, ExponentialReconnectionPolicy, -1, 0)
         self.assertRaises(ValueError, ExponentialReconnectionPolicy, 0, -1)
         self.assertRaises(ValueError, ExponentialReconnectionPolicy, 9000, 1)
+        self.assertRaises(ValueError, ExponentialReconnectionPolicy, 1, 2,-1)
 
     def test_schedule(self):
-        policy = ExponentialReconnectionPolicy(base_delay=2, max_delay=100)
+        policy = ExponentialReconnectionPolicy(base_delay=2, max_delay=100, max_attempts=None)
+        i=0;
+        for delay in policy.new_schedule():
+            i += 1
+            if i > 10000:
+                break;
+        self.assertEqual(i, 10001)
+
+        policy = ExponentialReconnectionPolicy(base_delay=2, max_delay=100, max_attempts=64)
         schedule = list(policy.new_schedule())
         self.assertEqual(len(schedule), 64)
         for i, delay in enumerate(schedule):