Merge pull request #298 from fruch/exponential_backoff

introducing `ExponentialBackoffRetryPolicy`

Author: Lorak-mmk

cassandra/cluster.py

@@ -5012,12 +5012,16 @@ def exception_from_response(response):
                 return response.to_exception()
             else:
                 return response
+        if len(retry_decision) == 2:
+            retry_type, consistency = retry_decision
+            delay = 0
+        elif len(retry_decision) == 3:
+            retry_type, consistency, delay = retry_decision
 
-        retry_type, consistency = retry_decision
         if retry_type in (RetryPolicy.RETRY, RetryPolicy.RETRY_NEXT_HOST):
             self._query_retries += 1
             reuse = retry_type == RetryPolicy.RETRY
-            self._retry(reuse, consistency, host)
+            self._retry(reuse, consistency, host, delay)
         elif retry_type is RetryPolicy.RETHROW:
             self._set_final_exception(exception_from_response(response))
         else:  # IGNORE
@@ -5027,7 +5031,7 @@ def exception_from_response(response):
 
         self._errors[host] = exception_from_response(response)
 
-    def _retry(self, reuse_connection, consistency_level, host):
+    def _retry(self, reuse_connection, consistency_level, host, delay):
         if self._final_exception:
             # the connection probably broke while we were waiting
             # to retry the operation
@@ -5039,7 +5043,7 @@ def _retry(self, reuse_connection, consistency_level, host):
             self.message.consistency_level = consistency_level
 
         # don't retry on the event loop thread
-        self.session.submit(self._retry_task, reuse_connection, host)
+        self.session.cluster.scheduler.schedule(delay, self._retry_task, reuse_connection, host)
 
     def _retry_task(self, reuse_connection, host):
         if self._final_exception:

cassandra/policies.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import random
 from itertools import islice, cycle, groupby, repeat
 import logging
 from random import randint, shuffle
@@ -1019,6 +1019,57 @@ def on_unavailable(self, query, consistency, required_replicas, alive_replicas,
             return self._pick_consistency(alive_replicas)
 
 
+class ExponentialBackoffRetryPolicy(RetryPolicy):
+    """
+    A policy that do retries with exponential backoff
+    """
+
+    def __init__(self, max_num_retries: float, min_interval: float = 0.1, max_interval: float = 10.0,
+                 *args, **kwargs):
+        """
+        `max_num_retries` counts how many times the operation would be retried,
+        `min_interval` is the initial time in seconds to wait before first retry
+        `max_interval` is the maximum time to wait between retries
+        """
+        self.min_interval = min_interval
+        self.max_num_retries = max_num_retries
+        self.max_interval = max_interval
+        super(ExponentialBackoffRetryPolicy).__init__(*args, **kwargs)
+
+    def _calculate_backoff(self, attempt: int):
+        delay = min(self.max_interval, self.min_interval * 2 ** attempt)
+        # add some jitter
+        delay += random.random() * self.min_interval - (self.min_interval / 2)
+        return delay
+
+    def on_read_timeout(self, query, consistency, required_responses,
+                        received_responses, data_retrieved, retry_num):
+        if retry_num < self.max_num_retries and received_responses >= required_responses and not data_retrieved:
+            return self.RETRY, consistency, self._calculate_backoff(retry_num)
+        else:
+            return self.RETHROW, None, None
+
+    def on_write_timeout(self, query, consistency, write_type,
+                         required_responses, received_responses, retry_num):
+        if retry_num < self.max_num_retries and write_type == WriteType.BATCH_LOG:
+            return self.RETRY, consistency, self._calculate_backoff(retry_num)
+        else:
+            return self.RETHROW, None, None
+
+    def on_unavailable(self, query, consistency, required_replicas,
+                       alive_replicas, retry_num):
+        if retry_num < self.max_num_retries:
+            return self.RETRY_NEXT_HOST, None, self._calculate_backoff(retry_num)
+        else:
+            return self.RETHROW, None, None
+
+    def on_request_error(self, query, consistency, error, retry_num):
+        if retry_num < self.max_num_retries:
+            return self.RETRY_NEXT_HOST, None, self._calculate_backoff(retry_num)
+        else:
+            return self.RETHROW, None, None
+
+
 class AddressTranslator(object):
     """
     Interface for translating cluster-defined endpoints.

tests/integration/standard/test_policies.py

@@ -16,7 +16,7 @@
 
 from cassandra.cluster import ExecutionProfile, EXEC_PROFILE_DEFAULT
 from cassandra.policies import HostFilterPolicy, RoundRobinPolicy,  SimpleConvictionPolicy, \
-    WhiteListRoundRobinPolicy
+    WhiteListRoundRobinPolicy, ExponentialBackoffRetryPolicy
 from cassandra.pool import Host
 from cassandra.connection import DefaultEndPoint
 
@@ -90,3 +90,20 @@ def test_only_connects_to_subset(self):
             queried_hosts.update(response.response_future.attempted_hosts)
         queried_hosts = set(host.address for host in queried_hosts)
         self.assertEqual(queried_hosts, only_connect_hosts)
+
+
+class ExponentialRetryPolicyTests(unittest.TestCase):
+
+    def setUp(self):
+        self.cluster = TestCluster(default_retry_policy=ExponentialBackoffRetryPolicy(max_num_retries=3))
+        self.session = self.cluster.connect()
+
+    def tearDown(self):
+        self.cluster.shutdown()
+
+    def test_exponential_retries(self):
+        self.session.execute(
+            """
+            CREATE KEYSPACE preparedtests
+            WITH replication = {'class': 'SimpleStrategy', 'replication_factor': '1'}
+            """)

tests/unit/test_policies.py

@@ -32,7 +32,7 @@
                                 RetryPolicy, WriteType,
                                 DowngradingConsistencyRetryPolicy, ConstantReconnectionPolicy,
                                 LoadBalancingPolicy, ConvictionPolicy, ReconnectionPolicy, FallthroughRetryPolicy,
-                                IdentityTranslator, EC2MultiRegionTranslator, HostFilterPolicy)
+                                IdentityTranslator, EC2MultiRegionTranslator, HostFilterPolicy, ExponentialBackoffRetryPolicy)
 from cassandra.pool import Host
 from cassandra.connection import DefaultEndPoint, UnixSocketEndPoint
 from cassandra.query import Statement
@@ -1247,6 +1247,24 @@ def test_unavailable(self):
         self.assertEqual(consistency, ConsistencyLevel.ONE)
 
 
+class ExponentialRetryPolicyTest(unittest.TestCase):
+    def test_calculate_backoff(self):
+        policy = ExponentialBackoffRetryPolicy(max_num_retries=2)
+
+        cases = [
+            (0, 0.1),
+            (1, 2 * 0.1),
+            (2, (2 * 2) * 0.1),
+            (3, (2 * 2 * 2) * 0.1),
+        ]
+
+        for attempts, delay in cases:
+            for i in range(100):
+                d = policy._calculate_backoff(attempts)
+                assert d > delay - (0.1 / 2), f"d={d} attempts={attempts}, delay={delay}"
+                assert d < delay + (0.1 / 2), f"d={d} attempts={attempts}, delay={delay}"
+
+
 class WhiteListRoundRobinPolicyTest(unittest.TestCase):
 
     def test_hosts_with_hostname(self):

tests/unit/test_response_future.py

@@ -30,7 +30,7 @@
                                 RESULT_KIND_ROWS, RESULT_KIND_SET_KEYSPACE,
                                 RESULT_KIND_SCHEMA_CHANGE, RESULT_KIND_PREPARED,
                                 ProtocolHandler)
-from cassandra.policies import RetryPolicy
+from cassandra.policies import RetryPolicy, ExponentialBackoffRetryPolicy
 from cassandra.pool import NoConnectionsAvailable
 from cassandra.query import SimpleStatement
 
@@ -265,7 +265,7 @@ def test_retry_policy_says_retry(self):
         host = Mock()
         rf._set_result(host, None, None, result)
 
-        session.submit.assert_called_once_with(rf._retry_task, True, host)
+        rf.session.cluster.scheduler.schedule.assert_called_once_with(ANY, rf._retry_task, True, host)
         self.assertEqual(1, rf._query_retries)
 
         connection = Mock(spec=Connection)
@@ -300,7 +300,7 @@ def test_retry_with_different_host(self):
         host = Mock()
         rf._set_result(host, None, None, result)
 
-        session.submit.assert_called_once_with(rf._retry_task, False, host)
+        rf.session.cluster.scheduler.schedule.assert_called_once_with(ANY, rf._retry_task, False, host)
         # query_retries does get incremented for Overloaded/Bootstrapping errors (since 3.18)
         self.assertEqual(1, rf._query_retries)
 
@@ -332,7 +332,8 @@ def test_all_retries_fail(self):
         rf._set_result(host, None, None, result)
 
         # simulate the executor running this
-        session.submit.assert_called_once_with(rf._retry_task, False, host)
+        rf.session.cluster.scheduler.schedule.assert_called_once_with(ANY, rf._retry_task, False, host)
+
         rf._retry_task(False, host)
 
         # it should try with a different host
@@ -342,11 +343,34 @@ def test_all_retries_fail(self):
         rf._set_result(host, None, None, result)
 
         # simulate the executor running this
-        session.submit.assert_called_with(rf._retry_task, False, host)
+        rf.session.cluster.scheduler.schedule.assert_called_with(ANY, rf._retry_task, False, host)
         rf._retry_task(False, host)
 
         self.assertRaises(NoHostAvailable, rf.result)
 
+    def test_exponential_retry_policy_fail(self):
+        session = self.make_session()
+        pool = session._pools.get.return_value
+        connection = Mock(spec=Connection)
+        pool.borrow_connection.return_value = (connection, 1)
+
+        query = SimpleStatement("SELECT * FROM foo")
+        message = QueryMessage(query=query, consistency_level=ConsistencyLevel.ONE)
+        rf = ResponseFuture(session, message, query, 1, retry_policy=ExponentialBackoffRetryPolicy(2))
+        rf.send_request()
+        rf.session._pools.get.assert_called_once_with('ip1')
+
+        result = Mock(spec=IsBootstrappingErrorMessage, info={})
+        host = Mock()
+        rf._set_result(host, None, None, result)
+
+        # simulate the executor running this
+        rf.session.cluster.scheduler.schedule.assert_called_once_with(ANY, rf._retry_task, False, host)
+
+        delay = rf.session.cluster.scheduler.schedule.mock_calls[-1][1][0]
+        assert delay > 0.05
+        rf._retry_task(False, host)
+
     def test_all_pools_shutdown(self):
         session = self.make_basic_session()
         session.cluster._default_load_balancing_policy.make_query_plan.return_value = ['ip1', 'ip2']