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,141 @@ class UserTypeDoesNotExist(Exception):
     pass
 
 
+class ScopeBucket(object):
+    def __init__(self, session, shard_reconnection_policy):
+        self._items = []
+        self.last_run = None
+        self.session = 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):
+        with self.lock:
+            self._items.append([method, args, kwargs])
+            if not self.running:
+                self.running = True
+                self._schedule()
+
+    def _get_delay(self):
+        try:
+            return next(self.schedule)
+        except StopIteration:
+            self.schedule = self.policy.new_schedule()
+            return next(self.schedule)
+
+    def _schedule(self):
+        if self.session.is_shutdown:
+            return
+        delay = self._get_delay()
+        if delay:
+            self.session.cluster.scheduler.schedule(delay, self.run)
+        else:
+            self.session.submit(self.run)
+
+    def run(self):
+        if self.session.is_shutdown:
+            return
+
+        with self.lock:
+            try:
+                item = self._items.pop()
+            except IndexError:
+                self.running = False
+                return
+
+        method, args, kwargs = item
+        try:
+            method(*args, **kwargs)
+        finally:
+            self._schedule()
+
+
+class ShardReconnectionSchedulerBase(object):
+    def schedule(self, host_id, shard_id, method, *args, **kwargs):
+        raise NotImplementedError()
+
+    def forced_schedule(self, host_id, shard_id, method, *args, **kwargs):
+        raise NotImplementedError()
+
+
+class NoDelayShardReconnectionScheduler(ShardReconnectionSchedulerBase):
+    def __init__(self, session):
+        self.session = weakref.proxy(session)
+        self.already_scheduled = {}
+
+    def _execute(self, scheduled_key, method, *args, **kwargs):
+        try:
+            method(*args, **kwargs)
+        finally:
+            self.already_scheduled[scheduled_key] = False
+
+    def forced_schedule(self, host_id, shard_id, method, *args, **kwargs):
+        scheduled_key = f'{host_id}-{shard_id}'
+        self.already_scheduled[scheduled_key] = True
+
+        if not self.session.is_shutdown:
+            self.session.submit(self._execute, scheduled_key, method, *args, **kwargs)
+
+    def schedule(self, host_id, shard_id, method, *args, **kwargs):
+        scheduled_key = f'{host_id}-{shard_id}'
+        if self.already_scheduled.get(scheduled_key):
+            return
+
+        self.already_scheduled[scheduled_key] = True
+        if not self.session.is_shutdown:
+            self.session.submit(self._execute, scheduled_key, 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 forced_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:
+            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
+
+    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
@@ -4432,6 +4620,9 @@ def shutdown(self):
         self._queue.put_nowait((0, 0, None))
         self.join()
 
+    def empty(self):
+        return len(self._scheduled_tasks) == 0 and self._queue.empty()
+
     def schedule(self, delay, fn, *args, **kwargs):
         self._insert_task(delay, (fn, args, tuple(kwargs.items())))
 

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

@@ -402,7 +402,6 @@ def __init__(self, host, host_distance, session):
         # this is used in conjunction with the connection streams. Not using the connection lock because the connection can be replaced in the lifetime of the pool.
         self._stream_available_condition = Condition(Lock())
         self._is_replacing = False
-        self._connecting = set()
         self._connections = {}
         self._pending_connections = []
         # A pool of additional connections which are not used but affect how Scylla
@@ -418,7 +417,6 @@ def __init__(self, host, host_distance, session):
         # and are waiting until all requests time out or complete
         # so that we can dispose of them.
         self._trash = set()
-        self._shard_connections_futures = []
         self.advanced_shardaware_block_until = 0
 
         if host_distance == HostDistance.IGNORED:
@@ -483,25 +481,25 @@ def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table
                     self.host,
                     routing_key
                 )
-                if conn.orphaned_threshold_reached and shard_id not in self._connecting:
+                if conn.orphaned_threshold_reached:
                     # The connection has met its orphaned stream ID limit
                     # and needs to be replaced. Start opening a connection
                     # to the same shard and replace when it is opened.
-                    self._connecting.add(shard_id)
-                    self._session.submit(self._open_connection_to_missing_shard, shard_id)
+                    self._session.shard_reconnection_scheduler.schedule(
+                        self.host.host_id, shard_id, self._open_connection_to_missing_shard, shard_id)
                     log.debug(
-                        "Connection to shard_id=%i reached orphaned stream limit, replacing on host %s (%s/%i)",
+                        "Scheduling Connection to shard_id=%i reached orphaned stream limit, replacing on host %s (%s/%i)",
                         shard_id,
                         self.host,
                         len(self._connections.keys()),
                         self.host.sharding_info.shards_count
                     )
-            elif shard_id not in self._connecting:
+            else:
                 # rate controlled optimistic attempt to connect to a missing shard
-                self._connecting.add(shard_id)
-                self._session.submit(self._open_connection_to_missing_shard, shard_id)
+                self._session.shard_reconnection_scheduler.schedule(
+                    self.host.host_id, shard_id, self._open_connection_to_missing_shard, shard_id)
                 log.debug(
-                    "Trying to connect to missing shard_id=%i on host %s (%s/%i)",
+                    "Scheduling connection to missing shard_id=%i on host %s (%s/%i)",
                     shard_id,
                     self.host,
                     len(self._connections.keys()),
@@ -609,8 +607,8 @@ def _replace(self, connection):
                 if connection.features.shard_id in self._connections.keys():
                     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)
@@ -635,9 +633,6 @@ def shutdown(self):
             with self._stream_available_condition:
                 self._stream_available_condition.notify_all()
 
-            for future in self._shard_connections_futures:
-                future.cancel()
-
             connections_to_close = self._connections.copy()
             pending_connections_to_close = self._pending_connections.copy()
             self._connections.clear()
@@ -843,7 +838,6 @@ def _open_connection_to_missing_shard(self, shard_id):
                     self._excess_connections.add(conn)
             if close_connection:
                 conn.close()
-        self._connecting.discard(shard_id)
 
     def _open_connections_for_all_shards(self, skip_shard_id=None):
         """
@@ -856,10 +850,8 @@ def _open_connections_for_all_shards(self, skip_shard_id=None):
             for shard_id in range(self.host.sharding_info.shards_count):
                 if skip_shard_id is not None and skip_shard_id == shard_id:
                     continue
-                future = self._session.submit(self._open_connection_to_missing_shard, shard_id)
-                if isinstance(future, Future):
-                    self._connecting.add(shard_id)
-                    self._shard_connections_futures.append(future)
+                self._session.shard_reconnection_scheduler.schedule(
+                    self.host.host_id, shard_id, self._open_connection_to_missing_shard, shard_id)
 
         trash_conns = None
         with self._lock: