feat(cluster): inject shard reconnection policy

Inject shard reconnection policy into cluster, session, connection and
host pool
Drop pending connections tracking logic, since policy does that.
Fix some tests that mocks Cluster, session, connection or host pool.

Author: dkropachev

cassandra/cluster.py

@@ -73,7 +73,8 @@
                                 ExponentialReconnectionPolicy, HostDistance,
                                 RetryPolicy, IdentityTranslator, NoSpeculativeExecutionPlan,
                                 NoSpeculativeExecutionPolicy, DefaultLoadBalancingPolicy,
-                                NeverRetryPolicy)
+                                NeverRetryPolicy, ShardReconnectionPolicy, NoDelayShardReconnectionPolicy,
+                                ShardReconnectionScheduler)
 from cassandra.pool import (Host, _ReconnectionHandler, _HostReconnectionHandler,
                             HostConnectionPool, HostConnection,
                             NoConnectionsAvailable)
@@ -757,6 +758,11 @@ def auth_provider(self, value):
 
         self._auth_provider = value
 
+    _shard_reconnection_policy: ShardReconnectionPolicy
+    @property
+    def shard_reconnection_policy(self) -> ShardReconnectionPolicy:
+        return self._shard_reconnection_policy
+
     _load_balancing_policy = None
     @property
     def load_balancing_policy(self):
@@ -1219,7 +1225,8 @@ def __init__(self,
                  shard_aware_options=None,
                  metadata_request_timeout=None,
                  column_encryption_policy=None,
-                 application_info:Optional[ApplicationInfoBase]=None
+                 application_info: Optional[ApplicationInfoBase] = None,
+                 shard_reconnection_policy: Optional[ShardReconnectionPolicy] = None,
                  ):
         """
         ``executor_threads`` defines the number of threads in a pool for handling asynchronous tasks such as
@@ -1325,6 +1332,13 @@ 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 not isinstance(shard_reconnection_policy, ShardReconnectionPolicy):
+                raise TypeError("shard_reconnection_policy should be an instance of class derived from ShardReconnectionPolicy")
+            self._shard_reconnection_policy = shard_reconnection_policy
+        else:
+            self._shard_reconnection_policy = NoDelayShardReconnectionPolicy()
+
         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 @@ def default_serial_consistency_level(self, cl):
     _metrics = None
     _request_init_callbacks = None
     _graph_paging_available = False
+    shard_reconnection_scheduler: ShardReconnectionScheduler
 
     def __init__(self, cluster, hosts, keyspace=None):
         self.cluster = cluster
@@ -2730,6 +2745,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()
@@ -4455,6 +4471,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/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/unit/test_host_connection_pool.py

@@ -26,7 +26,7 @@
 from cassandra.connection import Connection
 from cassandra.pool import HostConnection, HostConnectionPool
 from cassandra.pool import Host, NoConnectionsAvailable
-from cassandra.policies import HostDistance, SimpleConvictionPolicy
+from cassandra.policies import HostDistance, SimpleConvictionPolicy, _NoDelayShardReconnectionScheduler
 
 LOGGER = logging.getLogger(__name__)
 
@@ -41,6 +41,8 @@ def make_session(self):
         session.cluster.get_core_connections_per_host.return_value = 1
         session.cluster.get_max_requests_per_connection.return_value = 1
         session.cluster.get_max_connections_per_host.return_value = 1
+        session.shard_reconnection_scheduler = _NoDelayShardReconnectionScheduler(session)
+        session.is_shutdown = False
         return session
 
     def test_borrow_and_return(self):

tests/unit/test_shard_aware.py

@@ -11,17 +11,21 @@
 # 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 uuid
 
 try:
     import unittest2 as unittest
 except ImportError:
     import unittest  # noqa
 
+import time
 import logging
 from mock import MagicMock
 from concurrent.futures import ThreadPoolExecutor
 
-from cassandra.cluster import ShardAwareOptions
+from cassandra.cluster import ShardAwareOptions, _Scheduler
+from cassandra.policies import ConstantReconnectionPolicy, \
+    NoDelayShardReconnectionPolicy, NoConcurrentShardReconnectionPolicy, ShardReconnectionPolicyScope
 from cassandra.pool import HostConnection, HostDistance
 from cassandra.connection import ShardingInfo, DefaultEndPoint
 from cassandra.metadata import Murmur3Token
@@ -53,7 +57,15 @@ class OptionsHolder(object):
         self.assertEqual(shard_info.shard_id_from_token(Murmur3Token.from_key(b"e").value), 4)
         self.assertEqual(shard_info.shard_id_from_token(Murmur3Token.from_key(b"100000").value), 2)
 
-    def test_advanced_shard_aware_port(self):
+    def test_shard_aware_reconnection_policy_no_delay(self):
+        # with NoDelayReconnectionPolicy all the connections should be created right away
+        self._test_shard_aware_reconnection_policy(4, NoDelayShardReconnectionPolicy(), 4, 4)
+
+    def test_shard_aware_reconnection_policy_delay(self):
+        # with ConstantReconnectionPolicy first connection is created right away, others are delayed
+        self._test_shard_aware_reconnection_policy(4, NoConcurrentShardReconnectionPolicy(ShardReconnectionPolicyScope.Cluster, ConstantReconnectionPolicy(1)), 1, 4)
+
+    def _test_shard_aware_reconnection_policy(self, shard_count, shard_reconnection_policy, expected_count, expected_after):
         """
         Test that on given a `shard_aware_port` on the OPTIONS message (ShardInfo class)
         the next connections would be open using this port
@@ -71,17 +83,25 @@ def __init__(self, is_ssl=False, *args, **kwargs):
                     self.cluster.ssl_options = None
                 self.cluster.shard_aware_options = ShardAwareOptions()
                 self.cluster.executor = ThreadPoolExecutor(max_workers=2)
+                self._executor_submit_original = self.cluster.executor.submit
+                self.cluster.executor.submit = self._executor_submit
+                self.cluster.scheduler = _Scheduler(self.cluster.executor)
                 self.cluster.signal_connection_failure = lambda *args, **kwargs: False
                 self.cluster.connection_factory = self.mock_connection_factory
                 self.connection_counter = 0
+                self.shard_reconnection_scheduler = shard_reconnection_policy.new_scheduler(self)
                 self.futures = []
 
             def submit(self, fn, *args, **kwargs):
+                if self.is_shutdown:
+                    return None
+                return self.cluster.executor.submit(fn, *args, **kwargs)
+
+            def _executor_submit(self, fn, *args, **kwargs):
                 logging.info("Scheduling %s with args: %s, kwargs: %s", fn, args, kwargs)
-                if not self.is_shutdown:
-                    f = self.cluster.executor.submit(fn, *args, **kwargs)
-                    self.futures += [f]
-                    return f
+                f = self._executor_submit_original(fn, *args, **kwargs)
+                self.futures += [f]
+                return f
 
             def mock_connection_factory(self, *args, **kwargs):
                 connection = MagicMock()
@@ -90,26 +110,50 @@ def mock_connection_factory(self, *args, **kwargs):
                 connection.is_closed = False
                 connection.orphaned_threshold_reached = False
                 connection.endpoint = args[0] 
-                sharding_info = ShardingInfo(shard_id=1, shards_count=4, partitioner="", sharding_algorithm="", sharding_ignore_msb=0, shard_aware_port=19042, shard_aware_port_ssl=19045)
+                sharding_info = ShardingInfo(shard_id=1, shards_count=shard_count, partitioner="", sharding_algorithm="", sharding_ignore_msb=0, shard_aware_port=19042, shard_aware_port_ssl=19045)
                 connection.features = ProtocolFeatures(shard_id=kwargs.get('shard_id', self.connection_counter), sharding_info=sharding_info)
                 self.connection_counter += 1
 
                 return connection
 
         host = MagicMock()
+        host.host_id = uuid.uuid4()
         host.endpoint = DefaultEndPoint("1.2.3.4")
+        session = None
+        reconnection_policy = None
+        if isinstance(shard_reconnection_policy, NoConcurrentShardReconnectionPolicy):
+            reconnection_policy = shard_reconnection_policy.reconnection_policy
+        try:
+            for port, is_ssl in [(19042, False), (19045, True)]:
+                session = MockSession(is_ssl=is_ssl)
+                pool = HostConnection(host=host, host_distance=HostDistance.REMOTE, session=session)
+                for f in session.futures:
+                    f.result()
+                assert len(pool._connections) == expected_count
+                for shard_id, connection in pool._connections.items():
+                    assert connection.features.shard_id == shard_id
+                    if shard_id == 0:
+                        assert connection.endpoint == DefaultEndPoint("1.2.3.4")
+                    else:
+                        assert connection.endpoint == DefaultEndPoint("1.2.3.4", port=port)
 
-        for port, is_ssl in [(19042, False), (19045, True)]:
-            session = MockSession(is_ssl=is_ssl)
-            pool = HostConnection(host=host, host_distance=HostDistance.REMOTE, session=session)
-            for f in session.futures:
-                f.result()
-            assert len(pool._connections) == 4
-            for shard_id, connection in pool._connections.items():
-                assert connection.features.shard_id == shard_id
-                if shard_id == 0:
-                    assert connection.endpoint == DefaultEndPoint("1.2.3.4")
-                else:
-                    assert connection.endpoint == DefaultEndPoint("1.2.3.4", port=port)
+                sleep_time = 0
+                if reconnection_policy:
+                    # Check that connections to shards are being established according to the policy
+                    # Calculate total time it will need to establish all connections
+                    # Sleep half of the time and check that connections are not there yet
+                    # Sleep rest of the time + 1 second and check that all connections has been established
+                    schedule = reconnection_policy.new_schedule()
+                    for _ in range(shard_count):
+                        sleep_time += next(schedule)
+                    if sleep_time > 0:
+                        time.sleep(sleep_time/2)
+                        # Check that connection are not being established quicker than expected
+                        assert len(pool._connections) < expected_after
+                        time.sleep(sleep_time/2 + 1)
 
-        session.cluster.executor.shutdown(wait=True)
+                assert len(pool._connections) == expected_after
+        finally:
+            if session:
+                session.cluster.scheduler.shutdown()
+                session.cluster.executor.shutdown(wait=True)