improve ExponentialReconnectionPolicy,now can custom max attempts time

root · vipjml · commit bcfa8123c165 · 2016-01-07T10:34:14.000+08:00
diff --git a/cassandra/policies.py b/cassandra/policies.py
@@ -528,22 +528,33 @@ 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")
 
         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):
diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py
@@ -814,9 +814,18 @@ def test_bad_vals(self):
         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):
