1

Author: dkropachev

cassandra/cluster.py

@@ -20,6 +20,7 @@
 
 import atexit
 import datetime
+import threading
 from binascii import hexlify
 from collections import defaultdict
 from collections.abc import Mapping
@@ -72,7 +73,8 @@
                                 ExponentialReconnectionPolicy, HostDistance,
                                 RetryPolicy, IdentityTranslator, NoSpeculativeExecutionPlan,
                                 NoSpeculativeExecutionPolicy, DefaultLoadBalancingPolicy,
-                                NeverRetryPolicy)
+                                NeverRetryPolicy, ConstantReconnectionPolicy, ReconnectionPolicy,
+                                ShardReconnectionScope, ShardReconnectionScopeHost, NoDelayReconnectionPolicy)
 from cassandra.pool import (Host, _ReconnectionHandler, _HostReconnectionHandler,
                             HostConnectionPool, HostConnection,
                             NoConnectionsAvailable)
@@ -742,6 +744,32 @@ def auth_provider(self, value):
 
         self._auth_provider = value
 
+    _shard_reconnection_policy = None
+    @property
+    def shard_reconnection_policy(self):
+        return self._shard_reconnection_policy
+
+    @shard_reconnection_policy.setter
+    def shard_reconnection_policy(self, srp):
+        if self._config_mode == _ConfigMode.PROFILES:
+            raise ValueError(
+                "Cannot set Cluster.shard_reconnection_policy while using Configuration Profiles. Set this in a profile instead.")
+        self._shard_reconnection_policy = srp
+        self._config_mode = _ConfigMode.LEGACY
+
+    _shard_reconnection_scope = None
+    @property
+    def shard_reconnection_scope(self):
+        return self._shard_reconnection_scope
+
+    @shard_reconnection_scope.setter
+    def shard_reconnection_scope(self, scope):
+        if self._config_mode == _ConfigMode.PROFILES:
+            raise ValueError(
+                "Cannot set Cluster.shard_reconnection_scope while using Configuration Profiles. Set this in a profile instead.")
+        self._shard_reconnection_scope = scope
+        self._config_mode = _ConfigMode.LEGACY
+
     _load_balancing_policy = None
     @property
     def load_balancing_policy(self):
@@ -1204,6 +1232,8 @@ def __init__(self,
                  shard_aware_options=None,
                  metadata_request_timeout=None,
                  column_encryption_policy=None,
+                 shard_reconnection_policy=None,
+                 shard_reconnection_scope=None,
                  ):
         """
         ``executor_threads`` defines the number of threads in a pool for handling asynchronous tasks such as
@@ -1309,6 +1339,24 @@ def __init__(self,
         else:
             self._load_balancing_policy = default_lbp_factory()  # set internal attribute to avoid committing to legacy config mode
 
+        if shard_reconnection_scope is not None:
+            if isinstance(shard_reconnection_scope, type):
+                raise TypeError("shard_reconnection_scope should not be a class, it should be an instance of that class")
+            if not isinstance(shard_reconnection_policy, ShardReconnectionScope):
+                raise TypeError("load_balancing_policy should be an instance of class derived from ReconnectionPolicy")
+            self.shard_reconnection_scope = shard_reconnection_scope
+        else:
+            self._shard_reconnection_scope = ShardReconnectionScopeHost()  # set internal attribute to avoid committing to legacy config mode
+
+        if shard_reconnection_policy is not None:
+            if isinstance(shard_reconnection_policy, type):
+                raise TypeError("load_balancing_policy should not be a class, it should be an instance of that class")
+            if not isinstance(shard_reconnection_policy, ReconnectionPolicy):
+                raise TypeError("load_balancing_policy should be an instance of class derived from ReconnectionPolicy")
+            self.shard_reconnection_policy = shard_reconnection_policy
+        else:
+            self._shard_reconnection_policy = ConstantReconnectionPolicy(2, 0)  # set internal attribute to avoid committing to legacy config mode
+
         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")
@@ -2707,6 +2755,11 @@ def __init__(self, cluster, hosts, keyspace=None):
         self._protocol_version = self.cluster.protocol_version
 
         self.encoder = Encoder()
+        if isinstance(cluster.shard_reconnection_policy, NoDelayReconnectionPolicy):
+            self.shard_reconnection_scheduler = NoDelayShardReconnectionScheduler(self)
+        else:
+            self.shard_reconnection_scheduler = ShardReconnectionScheduler(
+                self, cluster, cluster.shard_reconnection_scope, cluster.shard_reconnection_policy)
 
         # create connection pools in parallel
         self._initial_connect_futures = set()
@@ -3546,6 +3599,105 @@ class UserTypeDoesNotExist(Exception):
     pass
 
 
+class ScopeBucket(object):
+    def __init__(self, session, shard_reconnection_policy):
+        self._items = []
+        self.last_run = None
+        self.session = weakref.proxy(session)
+        self.policy = shard_reconnection_policy
+        self.lock = threading.Lock()
+        self.running = False
+        self.schedule = self.policy.new_schedule()
+
+    def add(self, method, *args, **kwargs):
+        if self.session.is_shutdown:
+           return
+
+        with self.lock:
+            self._items.append([method, args, kwargs])
+            if not self.running:
+                if not self.session.is_shutdown:
+                    self.session.submit(self.run)
+
+    def get_delay(self):
+        try:
+            return next(self.schedule)
+        except StopIteration:
+            self.schedule = self.policy.new_schedule()
+            return next(self.schedule)
+
+    def run(self):
+        if self.session.is_shutdown:
+           return
+
+        with self.lock:
+            try:
+                item = self._items.pop()
+                self.running = True
+            except IndexError:
+                self.running = False
+                return
+
+        method, args, kwargs = item
+        try:
+            method(*args, **kwargs)
+        finally:
+            if not self.session.is_shutdown:
+                delay = self.get_delay()
+                if delay:
+                    self.session.cluster.scheduler.schedule(delay, self.run)
+                else:
+                    self.session.submit(self.run)
+
+
+class ShardReconnectionSchedulerBase(object):
+    def schedule(self, host_id, shard_id, method, *args, **kwargs):
+        raise NotImplementedError()
+
+
+class NoDelayShardReconnectionScheduler(ShardReconnectionSchedulerBase):
+    def __init__(self, session):
+        self.session = weakref.proxy(session)
+
+    def schedule(self, host_id, shard_id, method, *args, **kwargs):
+        if not self.session.is_shutdown:
+            self.session.submit(method, *args, **kwargs)
+
+
+class ShardReconnectionScheduler(ShardReconnectionSchedulerBase):
+    def __init__(self, session, cluster, shard_reconnection_scope, shard_reconnection_policy):
+        self.already_scheduled = {}
+        self.scopes = {}
+        self.session = weakref.proxy(session)
+        self.cluster = weakref.proxy(cluster)
+        self.shard_reconnection_scope = shard_reconnection_scope
+        self.shard_reconnection_policy = shard_reconnection_policy
+        self.lock = threading.Lock()
+
+    def _execute(self, scheduled_key, method, *args, **kwargs):
+        try:
+            method(*args, **kwargs)
+        finally:
+            with self.lock:
+                self.already_scheduled[scheduled_key] = False
+
+    def schedule(self, host_id, shard_id, method, *args, **kwargs):
+        scope_id = self.shard_reconnection_scope.get_hash(self.cluster, host_id, shard_id)
+        scheduled_key = f'{host_id}-{shard_id}'
+
+        with self.lock:
+            if self.already_scheduled.get(scheduled_key):
+                return False
+            self.already_scheduled[scheduled_key] = True
+
+            scope_info = self.scopes.get(scope_id, 0)
+            if not scope_info:
+                scope_info = ScopeBucket(self.session, self.shard_reconnection_policy)
+                self.scopes[scope_id] = scope_info
+            scope_info.add(self._execute, scheduled_key, method,*args, **kwargs)
+            return True
+
+
 class _ControlReconnectionHandler(_ReconnectionHandler):
     """
     Internal

cassandra/policies.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import random
+import threading
+import time
+import weakref
 
 from collections import namedtuple
 from functools import lru_cache
@@ -778,6 +781,14 @@ def new_schedule(self):
         raise NotImplementedError()
 
 
+class NoDelayReconnectionPolicy(ReconnectionPolicy):
+    """
+    A :class:`.ReconnectionPolicy` subclass which does not sleep.
+    """
+    def new_schedule(self):
+        return repeat(0)
+
+
 class ConstantReconnectionPolicy(ReconnectionPolicy):
     """
     A :class:`.ReconnectionPolicy` subclass which sleeps for a fixed delay
@@ -864,6 +875,26 @@ def _add_jitter(self, value):
         return min(max(self.base_delay, delay), self.max_delay)
 
 
+class ShardReconnectionScope(object):
+    def get_hash(self, cluster, host_id, shard_id):
+        raise NotImplementedError()
+
+
+class ShardReconnectionScopeCluster(ShardReconnectionScope):
+    def get_hash(self, cluster, host_id, shard_id):
+        return id(cluster)
+
+
+class ShardReconnectionScopeHost(ShardReconnectionScope):
+    def get_hash(self, cluster, host_id, shard_id):
+        return hash(host_id)
+
+
+class ShardReconnectionScopeShard(ShardReconnectionScope):
+    def get_hash(self, cluster, host_id, shard_id):
+        return hash((host_id, shard_id))
+
+
 class RetryPolicy(object):
     """
     A policy that describes whether to retry, rethrow, or ignore coordinator

cassandra/pool.py

@@ -610,7 +610,8 @@ def _replace(self, connection):
                     del self._connections[connection.features.shard_id]
                 if self.host.sharding_info and not self._session.cluster.shard_aware_options.disable:
                     self._connecting.add(connection.features.shard_id)
-                    self._session.submit(self._open_connection_to_missing_shard, connection.features.shard_id)
+                    self._session.shard_reconnection_scheduler.schedule(
+                        self.host.host_id, connection.features.shard_id, self._open_connection_to_missing_shard, connection.features.shard_id)
                 else:
                     connection = self._session.cluster.connection_factory(self.host.endpoint,
                                                                           on_orphaned_stream_released=self.on_orphaned_stream_released)