add RETRY_NEXT_HOST decision, default to this on unavailable

aholmberg · aholmberg · commit 0b21341a2fca · 2016-04-11T15:49:00.000-05:00
PYTHON-285
diff --git a/cassandra/cluster.py b/cassandra/cluster.py
@@ -3058,15 +3058,17 @@ def _set_result(self, response):
                     return
 
                 retry_type, consistency = retry
-                if retry_type is RetryPolicy.RETRY:
+                if retry_type in (RetryPolicy.RETRY, RetryPolicy.RETRY_NEXT_HOST):
                     self._query_retries += 1
-                    self._retry(reuse_connection=True, consistency_level=consistency)
+                    reuse = retry_type  == RetryPolicy.RETRY
+                    self._retry(reuse_connection=reuse, consistency_level=consistency)
                 elif retry_type is RetryPolicy.RETHROW:
                     self._set_final_exception(response.to_exception())
                 else:  # IGNORE
                     if self._metrics is not None:
                         self._metrics.on_ignore()
                     self._set_final_result(None)
+                self._errors[self._current_host] = response.to_exception()
             elif isinstance(response, ConnectionException):
                 if self._metrics is not None:
                     self._metrics.on_connection_error()
diff --git a/cassandra/policies.py b/cassandra/policies.py
@@ -551,8 +551,8 @@ def __init__(self, base_delay, max_delay, max_attempts=64):
         self.max_attempts = max_attempts
 
     def new_schedule(self):
-        i=0
-        while self.max_attempts == None or i < self.max_attempts:
+        i = 0
+        while self.max_attempts is None or i < self.max_attempts:
             yield min(self.base_delay * (2 ** i), self.max_delay)
             i += 1
 
@@ -647,6 +647,12 @@ class or one of its subclasses.
     should be ignored but no more retries should be attempted.
     """
 
+    RETRY_NEXT_HOST = 3
+    """
+    This should be returned from the below methods if the operation
+    should be retried on another connection.
+    """
+
     def on_read_timeout(self, query, consistency, required_responses,
                         received_responses, data_retrieved, retry_num):
         """
@@ -674,11 +680,11 @@ def on_read_timeout(self, query, consistency, required_responses,
         a sufficient number of replicas responded (with data digests).
         """
         if retry_num != 0:
-            return (self.RETHROW, None)
+            return self.RETHROW, None
         elif received_responses >= required_responses and not data_retrieved:
-            return (self.RETRY, consistency)
+            return self.RETRY, consistency
         else:
-            return (self.RETHROW, None)
+            return self.RETHROW, None
 
     def on_write_timeout(self, query, consistency, write_type,
                          required_responses, received_responses, retry_num):
@@ -707,11 +713,11 @@ def on_write_timeout(self, query, consistency, write_type,
         :attr:`~.WriteType.BATCH_LOG`.
         """
         if retry_num != 0:
-            return (self.RETHROW, None)
+            return self.RETHROW, None
         elif write_type == WriteType.BATCH_LOG:
-            return (self.RETRY, consistency)
+            return self.RETRY, consistency
         else:
-            return (self.RETHROW, None)
+            return self.RETHROW, None
 
     def on_unavailable(self, query, consistency, required_replicas, alive_replicas, retry_num):
         """
@@ -736,7 +742,7 @@ def on_unavailable(self, query, consistency, required_replicas, alive_replicas,
 
         By default, no retries will be attempted and the error will be re-raised.
         """
-        return (self.RETHROW, None)
+        return (self.RETRY_NEXT_HOST, consistency) if retry_num == 0 else (self.RETHROW, None)
 
 
 class FallthroughRetryPolicy(RetryPolicy):
@@ -746,13 +752,13 @@ class FallthroughRetryPolicy(RetryPolicy):
     """
 
     def on_read_timeout(self, *args, **kwargs):
-        return (self.RETHROW, None)
+        return self.RETHROW, None
 
     def on_write_timeout(self, *args, **kwargs):
-        return (self.RETHROW, None)
+        return self.RETHROW, None
 
     def on_unavailable(self, *args, **kwargs):
-        return (self.RETHROW, None)
+        return self.RETHROW, None
 
 
 class DowngradingConsistencyRetryPolicy(RetryPolicy):
@@ -804,46 +810,46 @@ class DowngradingConsistencyRetryPolicy(RetryPolicy):
     """
     def _pick_consistency(self, num_responses):
         if num_responses >= 3:
-            return (self.RETRY, ConsistencyLevel.THREE)
+            return self.RETRY, ConsistencyLevel.THREE
         elif num_responses >= 2:
-            return (self.RETRY, ConsistencyLevel.TWO)
+            return self.RETRY, ConsistencyLevel.TWO
         elif num_responses >= 1:
-            return (self.RETRY, ConsistencyLevel.ONE)
+            return self.RETRY, ConsistencyLevel.ONE
         else:
-            return (self.RETHROW, None)
+            return self.RETHROW, None
 
     def on_read_timeout(self, query, consistency, required_responses,
                         received_responses, data_retrieved, retry_num):
         if retry_num != 0:
-            return (self.RETHROW, None)
+            return self.RETHROW, None
         elif received_responses < required_responses:
             return self._pick_consistency(received_responses)
         elif not data_retrieved:
-            return (self.RETRY, consistency)
+            return self.RETRY, consistency
         else:
-            return (self.RETHROW, None)
+            return self.RETHROW, None
 
     def on_write_timeout(self, query, consistency, write_type,
                          required_responses, received_responses, retry_num):
         if retry_num != 0:
-            return (self.RETHROW, None)
+            return self.RETHROW, None
 
         if write_type in (WriteType.SIMPLE, WriteType.BATCH, WriteType.COUNTER):
             if received_responses > 0:
                 # persisted on at least one replica
-                return (self.IGNORE, None)
+                return self.IGNORE, None
             else:
-                return (self.RETHROW, None)
+                return self.RETHROW, None
         elif write_type == WriteType.UNLOGGED_BATCH:
             return self._pick_consistency(received_responses)
         elif write_type == WriteType.BATCH_LOG:
-            return (self.RETRY, consistency)
+            return self.RETRY, consistency
 
-        return (self.RETHROW, None)
+        return self.RETHROW, None
 
     def on_unavailable(self, query, consistency, required_replicas, alive_replicas, retry_num):
         if retry_num != 0:
-            return (self.RETHROW, None)
+            return self.RETHROW, None
         else:
             return self._pick_consistency(alive_replicas)
 
diff --git a/tests/integration/standard/test_metrics.py b/tests/integration/standard/test_metrics.py
@@ -145,13 +145,13 @@ def test_unavailable(self):
             query = SimpleStatement("INSERT INTO test (k, v) VALUES (2, 2)", consistency_level=ConsistencyLevel.ALL)
             with self.assertRaises(Unavailable):
                 self.session.execute(query)
-            self.assertEqual(1, self.cluster.metrics.stats.unavailables)
+            self.assertEqual(2, self.cluster.metrics.stats.unavailables)
 
             # Test write
             query = SimpleStatement("SELECT * FROM test", consistency_level=ConsistencyLevel.ALL)
             with self.assertRaises(Unavailable):
                 self.session.execute(query, timeout=None)
-            self.assertEqual(2, self.cluster.metrics.stats.unavailables)
+            self.assertEqual(4, self.cluster.metrics.stats.unavailables)
         finally:
             get_node(1).start(wait_other_notice=True, wait_for_binary_proto=True)
             # Give some time for the cluster to come back up, for the next test
diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py
@@ -916,14 +916,14 @@ def test_unavailable(self):
         retry, consistency = policy.on_unavailable(
             query=None, consistency=ONE,
             required_replicas=1, alive_replicas=2, retry_num=0)
-        self.assertEqual(retry, RetryPolicy.RETHROW)
-        self.assertEqual(consistency, None)
+        self.assertEqual(retry, RetryPolicy.RETRY_NEXT_HOST)
+        self.assertEqual(consistency, ONE)
 
         retry, consistency = policy.on_unavailable(
             query=None, consistency=ONE,
             required_replicas=10000, alive_replicas=1, retry_num=0)
-        self.assertEqual(retry, RetryPolicy.RETHROW)
-        self.assertEqual(consistency, None)
+        self.assertEqual(retry, RetryPolicy.RETRY_NEXT_HOST)
+        self.assertEqual(consistency, ONE)
 
 
 class FallthroughRetryPolicyTest(unittest.TestCase):
