Add shard-aware reconnection policies with support for scheduling constraints

Introduce `ShardReconnectionPolicy` and its implementations:
- `NoDelayShardReconnectionPolicy`: avoids reconnection delay and ensures at most one reconnection per host+shard.
- `NoConcurrentShardReconnectionPolicy`: limits concurrent reconnections per scope (Cluster or Host) using a backoff policy.

This feature enables finer control over shard reconnection behavior, helping prevent reconnection storms.

Author: dkropachev

cassandra/cluster.py

@@ -72,7 +72,8 @@
                                 ExponentialReconnectionPolicy, HostDistance,
                                 RetryPolicy, IdentityTranslator, NoSpeculativeExecutionPlan,
                                 NoSpeculativeExecutionPolicy, DefaultLoadBalancingPolicy,
-                                NeverRetryPolicy)
+                                NeverRetryPolicy, ConstantReconnectionPolicy,
+                                ShardReconnectionPolicyScope, ShardReconnectionPolicy, NoConcurrentShardReconnectionPolicy)
 from cassandra.pool import (Host, _ReconnectionHandler, _HostReconnectionHandler,
                             HostConnectionPool, HostConnection,
                             NoConnectionsAvailable)
@@ -742,6 +743,19 @@ 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
+
     _load_balancing_policy = None
     @property
     def load_balancing_policy(self):
@@ -1204,6 +1218,7 @@ def __init__(self,
                  shard_aware_options=None,
                  metadata_request_timeout=None,
                  column_encryption_policy=None,
+                 shard_reconnection_policy=None,
                  ):
         """
         ``executor_threads`` defines the number of threads in a pool for handling asynchronous tasks such as
@@ -1309,6 +1324,18 @@ def __init__(self,
         else:
             self._load_balancing_policy = default_lbp_factory()  # 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, ShardReconnectionPolicy):
+                raise TypeError("load_balancing_policy should be an instance of class derived from ReconnectionPolicy")
+            self.shard_reconnection_policy = shard_reconnection_policy
+        else:
+            # Limit to one pending reconnection per host with 200 milliseconds backoff at a time by default
+            self._shard_reconnection_policy = NoConcurrentShardReconnectionPolicy(
+                ShardReconnectionPolicyScope.Host,
+                ConstantReconnectionPolicy(0.2, 0))
+
         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 +2734,7 @@ def __init__(self, cluster, hosts, keyspace=None):
         self._protocol_version = self.cluster.protocol_version
 
         self.encoder = Encoder()
+        self.shard_reconnection_scheduler = cluster.shard_reconnection_policy.new_scheduler(self)
 
         # create connection pools in parallel
         self._initial_connect_futures = set()
@@ -4432,6 +4460,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,8 +12,12 @@
 # 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 enum import Enum
 from functools import lru_cache
 from itertools import islice, cycle, groupby, repeat
 import logging
@@ -778,6 +782,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 +876,172 @@ def _add_jitter(self, value):
         return min(max(self.base_delay, delay), self.max_delay)
 
 
+class _ShardReconnectionScheduler(object):
+    def schedule(self, host_id, shard_id, method, *args, **kwargs):
+        raise NotImplementedError()
+
+class ShardReconnectionPolicy(object):
+    """
+    Base class for shard reconnection policies.
+
+    On `new_scheduler` instantiate a scheduler that behaves according to the policy
+    """
+    def new_scheduler(self, session) -> _ShardReconnectionScheduler:
+        raise NotImplementedError()
+
+
+class NoDelayShardReconnectionPolicy(ShardReconnectionPolicy):
+    """
+    Shard reconnection policy that does not have delay.
+    Does not allow schedule multiple reconnections for same host+shard, silently ignores attempts to do that.
+
+    On `new_scheduler` instantiate a scheduler that behaves according to the policy
+    """
+    def new_scheduler(self, session) -> _ShardReconnectionScheduler:
+        return _NoDelayShardReconnectionScheduler(session)
+
+
+class _NoDelayShardReconnectionScheduler(_ShardReconnectionScheduler):
+    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 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 ShardReconnectionPolicyScope(Enum):
+    """
+    A scope for `ShardReconnectionPolicy`, in particular `NoConcurrentShardReconnectionPolicy`
+    """
+    Cluster = 0
+    Host = 1
+
+
+class NoConcurrentShardReconnectionPolicy(ShardReconnectionPolicy):
+    """
+    A shard reconnection policy that allows only one pending connection per scope, where scope could be `Host`, `Cluster`
+    For backoff it uses `ReconnectionPolicy`, when there is no more reconnections to scheduled backoff policy is reminded
+    For all scopes does not allow schedule multiple reconnections for same host+shard, it silently ignores attempts to do that.
+
+    On `new_scheduler` instantiate a scheduler that behaves according to the policy
+    """
+    def __init__(self, shard_reconnection_scope, reconnection_policy):
+        if not isinstance(shard_reconnection_scope, ShardReconnectionPolicyScope):
+            raise ValueError("shard_reconnection_scope must be a ShardReconnectionPolicyScope")
+        if not isinstance(reconnection_policy, ReconnectionPolicy):
+            raise ValueError("reconnection_policy must be a ReconnectionPolicy")
+        self.shard_reconnection_scope = shard_reconnection_scope
+        self.reconnection_policy = reconnection_policy
+
+    def new_scheduler(self, session) -> _ShardReconnectionScheduler:
+        return _NoConcurrentShardReconnectionScheduler(session, self.shard_reconnection_scope, self.reconnection_policy)
+
+
+class _ScopeBucket(object):
+    """
+    Holds information for a shard reconnection scope, schedules and executes reconnections.
+    """
+    def __init__(self, session, reconnection_policy):
+        self._items = []
+        self.session = session
+        self.reconnection_policy = reconnection_policy
+        self.lock = threading.Lock()
+        self.running = False
+        self.schedule = self.reconnection_policy.new_schedule()
+
+    def _get_delay(self):
+        if self.schedule is None:
+            self.schedule = self.reconnection_policy.new_schedule()
+        try:
+            return next(self.schedule)
+        except StopIteration:
+            self.schedule = self.reconnection_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
+                self.schedule = None
+                return
+
+        method, args, kwargs = item
+        try:
+            method(*args, **kwargs)
+        finally:
+            self._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()
+
+
+class _NoConcurrentShardReconnectionScheduler(_ShardReconnectionScheduler):
+    def __init__(self, session, shard_reconnection_scope, reconnection_policy):
+        self.already_scheduled = {}
+        self.scopes = {}
+        self.shard_reconnection_scope = shard_reconnection_scope
+        self.reconnection_policy = reconnection_policy
+        self.session = session
+        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):
+        if self.shard_reconnection_scope == ShardReconnectionPolicyScope.Cluster:
+            scope_hash = "global-cluster-scope"
+        else:
+            scope_hash = host_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_hash, 0)
+            if not scope_info:
+                scope_info = _ScopeBucket(self.session, self.reconnection_policy)
+                self.scopes[scope_hash] = scope_info
+            scope_info.add(self._execute, scheduled_key, method,*args, **kwargs)
+            return True
+
+
 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:

tests/integration/experiments/test_tablets.py

@@ -3,7 +3,8 @@
 import pytest
 
 from cassandra.cluster import Cluster
-from cassandra.policies import ConstantReconnectionPolicy, RoundRobinPolicy, TokenAwarePolicy
+from cassandra.policies import ConstantReconnectionPolicy, RoundRobinPolicy, TokenAwarePolicy, \
+    NoDelayShardReconnectionPolicy
 
 from tests.integration import PROTOCOL_VERSION, use_cluster
 from tests.unit.test_host_connection_pool import LOGGER
@@ -21,7 +22,9 @@ class TestTabletsIntegration:
     def setup_class(cls):
         cls.cluster = Cluster(contact_points=["127.0.0.1", "127.0.0.2", "127.0.0.3"], protocol_version=PROTOCOL_VERSION,
                               load_balancing_policy=TokenAwarePolicy(RoundRobinPolicy()),
-                              reconnection_policy=ConstantReconnectionPolicy(1))
+                              reconnection_policy=ConstantReconnectionPolicy(1),
+                              shard_reconnection_policy=NoDelayShardReconnectionPolicy(),
+                              )
         cls.session = cls.cluster.connect()
         cls.create_ks_and_cf(cls.session)
         cls.create_data(cls.session)