Add shard-aware reconnection policies with support for scheduling constraints

cassandra/cluster.py

@@ -73,7 +73,8 @@
                                 ExponentialReconnectionPolicy, HostDistance,
                                 RetryPolicy, IdentityTranslator, NoSpeculativeExecutionPlan,
                                 NoSpeculativeExecutionPolicy, DefaultLoadBalancingPolicy,
-                                NeverRetryPolicy)
+                                NeverRetryPolicy, ShardConnectionBackoffPolicy, NoDelayShardConnectionBackoffPolicy,
+                                ShardConnectionScheduler)
 from cassandra.pool import (Host, _ReconnectionHandler, _HostReconnectionHandler,
                             HostConnectionPool, HostConnection,
                             NoConnectionsAvailable)
@@ -757,6 +758,11 @@
 
         self._auth_provider = value
 
+    _shard_connection_backoff_policy: ShardConnectionBackoffPolicy
+    @property
+    def shard_connection_backoff_policy(self) -> ShardConnectionBackoffPolicy:
+        return self._shard_connection_backoff_policy
+
     _load_balancing_policy = None
     @property
     def load_balancing_policy(self):
@@ -1219,7 +1225,8 @@
                  shard_aware_options=None,
                  metadata_request_timeout=None,
                  column_encryption_policy=None,
-                 application_info:Optional[ApplicationInfoBase]=None
+                 application_info: Optional[ApplicationInfoBase] = None,
+                 shard_connection_backoff_policy: Optional[ShardConnectionBackoffPolicy] = None,
                  ):
         """
         ``executor_threads`` defines the number of threads in a pool for handling asynchronous tasks such as
@@ -1325,6 +1332,13 @@
         else:
             self._load_balancing_policy = default_lbp_factory()  # set internal attribute to avoid committing to legacy config mode
 
+        if shard_connection_backoff_policy is not None:
+            if not isinstance(shard_connection_backoff_policy, ShardConnectionBackoffPolicy):
+                raise TypeError("shard_connection_backoff_policy should be an instance of class derived from ShardConnectionBackoffPolicy")
+            self._shard_connection_backoff_policy = shard_connection_backoff_policy
+        else:
+            self._shard_connection_backoff_policy = NoDelayShardConnectionBackoffPolicy()
+
         if reconnection_policy is not None:
             if isinstance(reconnection_policy, type):
                 raise TypeError("reconnection_policy should not be a class, it should be an instance of that class")
@@ -2716,6 +2730,7 @@
     _metrics = None
     _request_init_callbacks = None
     _graph_paging_available = False
+    shard_connection_backoff_scheduler: ShardConnectionScheduler
 
     def __init__(self, cluster, hosts, keyspace=None):
         self.cluster = cluster
@@ -2730,6 +2745,7 @@
         self._protocol_version = self.cluster.protocol_version
 
         self.encoder = Encoder()
+        self.shard_connection_backoff_scheduler = cluster.shard_connection_backoff_policy.new_scheduler(self)
 
         # create connection pools in parallel
         self._initial_connect_futures = set()
@@ -4489,7 +4505,7 @@
                         self._scheduled_tasks.discard(task)
                         fn, args, kwargs = task
                         kwargs = dict(kwargs)
                         future = self._executor.submit(fn, *args, **kwargs)
                         future.add_done_callback(self._log_if_failed)
                     else:
                         self._queue.put_nowait((run_at, i, task))