Introduce NoDelayShardConnectionBackoffPolicy

This policy is implementation of ShardConnectionBackoffPolicy.
It implements same behavior that driver currently has:
1. No delay between creating shard connections
2. It avoids creating multiple connections to same host_id, shard_id

Author: dkropachev

cassandra/policies.py

@@ -934,6 +934,69 @@ def new_connection_scheduler(self, scheduler: _Scheduler) -> ShardConnectionSche
         raise NotImplementedError()
 
 
+class NoDelayShardConnectionBackoffPolicy(ShardConnectionBackoffPolicy):
+    """
+    A shard connection backoff policy with no delay between attempts.
+    Ensures that at most one pending request connection per (host, shard) pair.
+    If connection attempts for the same (host, shard) it is silently dropped.
+    """
+
+    def new_connection_scheduler(self, scheduler: _Scheduler) -> ShardConnectionScheduler:
+        return _NoDelayShardConnectionBackoffScheduler(scheduler)
+
+
+class _NoDelayShardConnectionBackoffScheduler(ShardConnectionScheduler):
+    """
+    A scheduler for ``cassandra.policies.NoDelayShardConnectionBackoffPolicy``.
+
+    A shard connection backoff policy with no delay between attempts.
+    Ensures that at most one pending request connection per (host, shard) pair.
+    If connection attempts for the same (host, shard) it is silently dropped.
+    """
+
+    scheduler: _Scheduler
+    already_scheduled: set[tuple[str, int]]
+    lock: Lock
+    is_shutdown: bool = False
+
+    def __init__(self, scheduler: _Scheduler):
+        self.scheduler = scheduler
+        self.already_scheduled = set()
+        self.lock = Lock()
+
+    def _execute(
+            self,
+            host_id: str,
+            shard_id: int,
+            method: Callable[[], None],
+    ) -> None:
+        if self.is_shutdown:
+            return
+        try:
+            method()
+        finally:
+            with self.lock:
+                self.already_scheduled.remove((host_id, shard_id))
+
+    def schedule(
+            self,
+            host_id: str,
+            shard_id: int,
+            method: Callable[[], None],
+    ) -> bool:
+        with self.lock:
+            if self.is_shutdown or (host_id, shard_id) in self.already_scheduled:
+                return False
+            self.already_scheduled.add((host_id, shard_id))
+
+        self.scheduler.schedule(0, self._execute, host_id, shard_id, method)
+        return True
+
+    def shutdown(self):
+        with self.lock:
+            self.is_shutdown = True
+
+
 class RetryPolicy(object):
     """
     A policy that describes whether to retry, rethrow, or ignore coordinator

tests/unit/test_policies.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import unittest
+from functools import partial
 
 from itertools import islice, cycle
 from unittest.mock import Mock, patch, call
@@ -26,13 +27,15 @@
 from cassandra import ConsistencyLevel
 from cassandra.cluster import Cluster, ControlConnection
 from cassandra.metadata import Metadata
-from cassandra.policies import (RackAwareRoundRobinPolicy, RoundRobinPolicy, WhiteListRoundRobinPolicy, DCAwareRoundRobinPolicy,
+from cassandra.policies import (RackAwareRoundRobinPolicy, RoundRobinPolicy, WhiteListRoundRobinPolicy,
+                                DCAwareRoundRobinPolicy,
                                 TokenAwarePolicy, SimpleConvictionPolicy,
                                 HostDistance, ExponentialReconnectionPolicy,
                                 RetryPolicy, WriteType,
                                 DowngradingConsistencyRetryPolicy, ConstantReconnectionPolicy,
                                 LoadBalancingPolicy, ConvictionPolicy, ReconnectionPolicy, FallthroughRetryPolicy,
-                                IdentityTranslator, EC2MultiRegionTranslator, HostFilterPolicy, ExponentialBackoffRetryPolicy)
+                                IdentityTranslator, EC2MultiRegionTranslator, HostFilterPolicy,
+                                ExponentialBackoffRetryPolicy, _NoDelayShardConnectionBackoffScheduler)
 from cassandra.connection import DefaultEndPoint, UnixSocketEndPoint
 from cassandra.pool import Host
 from cassandra.query import Statement
@@ -1579,3 +1582,79 @@ def test_create_whitelist(self):
         # Only the filtered replicas should be allowed
         self.assertEqual(set(query_plan), {Host(DefaultEndPoint("127.0.0.1"), SimpleConvictionPolicy),
                                            Host(DefaultEndPoint("127.0.0.4"), SimpleConvictionPolicy)})
+
+
+class MockScheduler:
+    def __init__(self):
+        self.requests = []
+
+    def schedule(self, delay, fn, *args, **kwargs):
+        self.requests.append((delay, fn, args, kwargs))
+
+    def execute(self):
+        for delay, fn, args, kwargs in self.requests:
+            fn(*args, **kwargs)
+        self.requests = []
+
+
+class NoDelayShardConnectionBackoffSchedulerTests(unittest.TestCase):
+    def test_schedule_executes_method_immediately(self):
+        method = Mock()
+        scheduler = MockScheduler()
+        policy = _NoDelayShardConnectionBackoffScheduler(scheduler)
+
+        self.assertTrue(policy.schedule('host1', 0, partial(method, 1, 2, key='val')))
+
+        self.assertEqual(scheduler.requests[0][0], 0)
+        scheduler.execute()
+
+        method.assert_called_once_with(1, 2, key='val')
+
+    def test_schedule_skips_if_host_shard_already_scheduled(self):
+        method = Mock()
+        scheduler = MockScheduler()
+        policy = _NoDelayShardConnectionBackoffScheduler(scheduler)
+
+        self.assertTrue(policy.schedule('host1', 0, method))
+        self.assertFalse(policy.schedule('host1', 0, method))
+
+        self.assertEqual(len(scheduler.requests), 1)
+        scheduler.execute()
+        method.assert_called_once()
+
+    def test_schedule_does_not_skip_if_shard_is_different(self):
+        method = Mock()
+        scheduler = MockScheduler()
+        policy = _NoDelayShardConnectionBackoffScheduler(scheduler)
+
+        self.assertTrue(policy.schedule('host1', 0, method))
+        self.assertTrue(policy.schedule('host1', 1, method))
+
+        self.assertEqual(len(scheduler.requests), 2)
+        scheduler.execute()
+
+        self.assertEqual(method.call_count, 2)
+
+    def test_already_scheduled_resets_after_execution(self):
+        method = Mock()
+        scheduler = MockScheduler()
+        policy = _NoDelayShardConnectionBackoffScheduler(scheduler)
+        self.assertTrue(policy.schedule('host1', 0, method))
+
+        scheduler.execute()
+
+        self.assertTrue(policy.schedule('host1', 0, method))
+
+        scheduler.execute()
+
+        self.assertEqual(method.call_count, 2)
+
+    def test_schedule_skips_if_shutdown(self):
+        method = Mock()
+        scheduler = MockScheduler()
+        policy = _NoDelayShardConnectionBackoffScheduler(scheduler)
+        policy.shutdown()
+
+        policy.schedule('host1', 0, method)
+
+        self.assertEqual(len(scheduler.requests), 0)

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(len(session.hosts) * (shard_count - 1), found_related_calls)
+        finally:
+            if session:
+                session.cluster.scheduler.shutdown()
+                session.cluster.executor.shutdown(wait=True)