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

@@ -936,6 +936,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)