Integrate ShardConnectionBackoffPolicy

Add code that integrates ShardConnectionBackoffPolicy into:
1. Cluster
2. Session
3. HostConnection

Main idea is to put ShardConnectionBackoffPolicy in control of
shard connection creation proccess.
Removing duplicate logic from HostConnection that tracks pending
connection creation requests.

Author: dkropachev

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,
                             HostConnection,
                             NoConnectionsAvailable)
@@ -754,6 +755,11 @@ def auth_provider(self, value):
 
         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):
@@ -1216,7 +1222,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_connection_backoff_policy: Optional[ShardConnectionBackoffPolicy] = None,
                  ):
         """
         ``executor_threads`` defines the number of threads in a pool for handling asynchronous tasks such as
@@ -1322,6 +1329,13 @@ def __init__(self,
         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")
@@ -2659,6 +2673,7 @@ def default_serial_consistency_level(self, cl):
     _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
@@ -2673,6 +2688,7 @@ def __init__(self, cluster, hosts, keyspace=None):
         self._protocol_version = self.cluster.protocol_version
 
         self.encoder = Encoder()
+        self.shard_connection_backoff_scheduler = cluster.shard_connection_backoff_policy.new_connection_scheduler(self.cluster.scheduler)
 
         # create connection pools in parallel
         self._initial_connect_futures = set()
@@ -3281,6 +3297,7 @@ def shutdown(self):
             else:
                 self.is_shutdown = True
 
+        self.shard_connection_backoff_scheduler.shutdown()
         # PYTHON-673. If shutdown was called shortly after session init, avoid
         # a race by cancelling any initial connection attempts haven't started,
         # then blocking on any that have.

cassandra/pool.py

@@ -16,7 +16,7 @@
 Connection pooling and host management.
 """
 from concurrent.futures import Future
-from functools import total_ordering
+from functools import total_ordering, partial
 import logging
 import socket
 import time
@@ -401,7 +401,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
@@ -417,7 +416,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:
@@ -482,25 +480,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_connection_backoff_scheduler.schedule(
+                        self.host.host_id, shard_id, partial(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_connection_backoff_scheduler.schedule(
+                    self.host.host_id, shard_id, partial(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()),
@@ -610,8 +608,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_connection_backoff_scheduler.schedule(
+                        self.host.host_id, connection.features.shard_id, partial(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)
@@ -636,9 +634,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()
@@ -848,7 +843,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):
         """
@@ -861,10 +855,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_connection_backoff_scheduler.schedule(
+                    self.host.host_id, shard_id, partial(self._open_connection_to_missing_shard, shard_id))
 
         trash_conns = None
         with self._lock:

tests/unit/test_host_connection_pool.py

@@ -22,15 +22,23 @@
 from threading import Thread, Event, Lock
 from unittest.mock import Mock, NonCallableMagicMock, MagicMock
 
-from cassandra.cluster import Session, ShardAwareOptions
+from cassandra.cluster import Session, ShardAwareOptions, _Scheduler
 from cassandra.connection import Connection
 from cassandra.pool import HostConnection
 from cassandra.pool import Host, NoConnectionsAvailable
-from cassandra.policies import HostDistance, SimpleConvictionPolicy
+from cassandra.policies import HostDistance, SimpleConvictionPolicy, _NoDelayShardConnectionBackoffScheduler
 
 LOGGER = logging.getLogger(__name__)
 
 
+class FakeScheduler(_Scheduler):
+    def __init__(self):
+        super(FakeScheduler, self).__init__(ThreadPoolExecutor())
+
+    def schedule(self, delay, fn, *args, **kwargs):
+        super().schedule(0, fn, *args, **kwargs)
+
+
 class _PoolTests(unittest.TestCase):
     __test__ = False
     PoolImpl = None
@@ -40,6 +48,9 @@ def make_session(self):
         session = NonCallableMagicMock(spec=Session, keyspace='foobarkeyspace')
         session.cluster.get_core_connections_per_host.return_value = 1
         session.cluster.get_max_connections_per_host.return_value = 1
+        session.shard_connection_backoff_scheduler = _NoDelayShardConnectionBackoffScheduler(FakeScheduler())
+        session.shard_connection_backoff_scheduler.schedule = Mock(wraps=session.shard_connection_backoff_scheduler.schedule)
+        session.is_shutdown = False
         return session
 
     def test_borrow_and_return(self):
@@ -173,9 +184,9 @@ def test_return_defunct_connection_on_down_host(self):
         if self.PoolImpl is HostConnection:
             # on shard aware implementation we use submit function regardless
             self.assertTrue(host.signal_connection_failure.call_args)
-            self.assertTrue(session.submit.called)
+            self.assertTrue(session.shard_connection_backoff_scheduler.schedule.called)
         else:
-            self.assertFalse(session.submit.called)
+            self.assertFalse(session.shard_connection_backoff_scheduler.schedule.called)
             self.assertTrue(session.cluster.signal_connection_failure.call_args)
         self.assertTrue(pool.is_shutdown)
 

tests/unit/test_shard_aware.py

@@ -11,6 +11,10 @@
 # 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
+from unittest.mock import Mock
+
+from cassandra.policies import NoDelayShardConnectionBackoffPolicy, _NoDelayShardConnectionBackoffScheduler
 
 try:
     import unittest2 as unittest
@@ -21,7 +25,7 @@
 from mock import MagicMock
 from concurrent.futures import ThreadPoolExecutor
 
-from cassandra.cluster import ShardAwareOptions
+from cassandra.cluster import ShardAwareOptions, _Scheduler
 from cassandra.pool import HostConnection, HostDistance
 from cassandra.connection import ShardingInfo, DefaultEndPoint
 from cassandra.metadata import Murmur3Token
@@ -53,11 +57,18 @@ 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, NoDelayShardConnectionBackoffPolicy(), 4)
+
+    def _test_shard_aware_reconnection_policy(self, shard_count, shard_connection_backoff_policy, expected_connections):
         """
         Test that on given a `shard_aware_port` on the OPTIONS message (ShardInfo class)
-        the next connections would be open using this port
+        It checks that:
+        1. Next connections are opened using this port
+        2. Connection creation pase matches `shard_connection_backoff_policy`
         """
+
         class MockSession(MagicMock):
             is_shutdown = False
             keyspace = "ks1"
@@ -71,45 +82,82 @@ 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)
+
+                # Collect scheduled calls and execute them right away
+                self.scheduler_calls = []
+                original_schedule = self.cluster.scheduler.schedule
+
+                def new_schedule(delay, fn, *args, **kwargs):
+                    self.scheduler_calls.append((delay, fn, args, kwargs))
+                    return original_schedule(0, fn, *args, **kwargs)
+
+                self.cluster.scheduler.schedule = Mock(side_effect=new_schedule)
                 self.cluster.signal_connection_failure = lambda *args, **kwargs: False
                 self.cluster.connection_factory = self.mock_connection_factory
                 self.connection_counter = 0
+                self.shard_connection_backoff_scheduler = shard_connection_backoff_policy.new_connection_scheduler(
+                    self.cluster.scheduler)
                 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()
                 connection.is_shutdown = False
                 connection.is_defunct = False
                 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)
-                connection.features = ProtocolFeatures(shard_id=kwargs.get('shard_id', self.connection_counter), sharding_info=sharding_info)
+                connection.endpoint = args[0]
+                sharding_info = None
+                if shard_count:
+                    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
+        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_connections
+                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)
-
-        session.cluster.executor.shutdown(wait=True)
+                sleep_time = 0
+                found_related_calls = 0
+                for delay, fn, args, kwargs in session.scheduler_calls:
+                    if fn.__self__.__class__ is _NoDelayShardConnectionBackoffScheduler:
+                        found_related_calls += 1
+                        self.assertEqual(delay, sleep_time)
+                self.assertLessEqual(shard_count - 1, found_related_calls)
+        finally:
+            if session:
+                session.cluster.scheduler.shutdown()
+                session.cluster.executor.shutdown(wait=True)