Implement LimitedConcurrencyShardConnectionBackoffPolicy

This policy is an implementation of `ShardConnectionBackoffPolicy`.
Its primary purpose is to prevent connection storms by imposing restrictions
on the number of concurrent pending connections per host and backoff
time between each connection attempt.

Author: dkropachev

cassandra/policies.py

@@ -14,15 +14,15 @@
 from __future__ import annotations
 
 import random
-
 from collections import namedtuple
+from functools import partial
 from itertools import islice, cycle, groupby, repeat
 import logging
 from random import randint, shuffle
 from threading import Lock
 import socket
 import warnings
-from typing import Callable, TYPE_CHECKING
+from typing import Callable, TYPE_CHECKING, Iterator, List, Tuple
 from abc import ABC, abstractmethod
 from cassandra import WriteType as WT
 
@@ -997,6 +997,236 @@ def shutdown(self):
             self.is_shutdown = True
 
 
+class ShardConnectionBackoffSchedule(ABC):
+    @abstractmethod
+    def new_schedule(self) -> Iterator[float]:
+        """
+        This should return a finite or infinite iterable of delays (each as a
+        floating point number of seconds).
+        Note that if the iterable is finite, schedule will be recreated right after iterable is exhausted.
+        """
+        raise NotImplementedError()
+
+
+class ConstantShardConnectionBackoffSchedule(ShardConnectionBackoffSchedule):
+    """
+    A :class:`.ShardConnectionBackoffSchedule` subclass which introduce a constant delay with jitter
+     between shard connections.
+    """
+
+    def __init__(self, delay, jitter: float = 0.0):
+        """
+        `delay` should be a floating point number of seconds to wait in-between
+        each connection attempt.
+
+        `jitter` is a random jitter in seconds.
+        """
+        if delay < 0:
+            raise ValueError("delay must not be negative")
+        if jitter < 0:
+            raise ValueError("jitter must not be negative")
+
+        self.delay = delay
+        self.jitter = jitter
+
+    def new_schedule(self):
+        if self.jitter == 0:
+            return repeat(self.delay)
+        def iterator():
+            while True:
+                yield self.delay + random.uniform(0.0, self.jitter)
+        return iterator()
+
+
+class LimitedConcurrencyShardConnectionBackoffPolicy(ShardConnectionBackoffPolicy):
+    """
+    A shard connection backoff policy that allows only `max_concurrent` concurrent connections per `host_id`.
+
+    For backoff calculation, it requires either a `cassandra.policies.ShardConnectionBackoffSchedule` or
+    a `cassandra.policies.ReconnectionPolicy`, as both expose the same API.
+
+    The backoff schedule is initiated when the first request for a given `host_id` is received.
+    If there are no remaining requests for that `host_id`, the schedule is reset.
+
+    This policy also prevents multiple pending or scheduled connections for the same (host, shard) pair;
+    any duplicate attempts to schedule a connection are silently ignored.
+    """
+    backoff_policy: ShardConnectionBackoffSchedule | ReconnectionPolicy
+
+    max_concurrent: int
+    """
+    Max concurrent connection creation requests per scope.
+    """
+
+    def __init__(
+            self,
+            backoff_policy: ShardConnectionBackoffSchedule | ReconnectionPolicy,
+            max_concurrent: int = 1,
+    ):
+        if not isinstance(backoff_policy, (ShardConnectionBackoffSchedule, ReconnectionPolicy)):
+            raise ValueError("backoff_policy must be a ShardConnectionBackoffSchedule or ReconnectionPolicy")
+        if max_concurrent < 1:
+            raise ValueError("max_concurrent must be a positive integer")
+        self.backoff_policy = backoff_policy
+        self.max_concurrent = max_concurrent
+
+    def new_connection_scheduler(self, scheduler: _Scheduler) -> ShardConnectionScheduler:
+        return _LimitedConcurrencyShardConnectionScheduler(scheduler, self.backoff_policy, self.max_concurrent)
+
+
+class _ScopeBucket:
+    """
+    Holds information for a shard connection backoff policy scope, schedules and executes requests to create connection.
+    """
+    session: _Scheduler
+    backoff_policy: ShardConnectionBackoffSchedule
+    lock: Lock
+    is_shutdown: bool = False
+
+    max_concurrent: int
+    """
+    Max concurrent connection creation requests in the scope.
+    """
+
+    currently_pending: int
+    """
+    Number of currently pending connections.
+    """
+
+    items: List[Callable[[], None]]
+    """
+    List of scheduled create connections requests.
+    """
+
+    def __init__(
+            self,
+            scheduler: _Scheduler,
+            backoff_policy: ShardConnectionBackoffSchedule,
+            max_concurrent: int,
+    ):
+        self.items = []
+        self.scheduler = scheduler
+        self.backoff_policy = backoff_policy
+        self.lock = Lock()
+        self.max_concurrent = max_concurrent
+        self.currently_pending = 0
+
+    def _get_delay(self, schedule: Iterator[float]) -> Tuple[Iterator[float], float]:
+        try:
+            return schedule, next(schedule)
+        except StopIteration:
+            # A bit of trickery to avoid having lock around self.schedule
+            schedule = self.backoff_policy.new_schedule()
+            delay = next(schedule)
+            self.schedule = schedule
+            return schedule, delay
+
+    def _run(self, schedule: Iterator[float]):
+        if self.is_shutdown:
+            return
+
+        with self.lock:
+            try:
+                request = self.items.pop(0)
+            except IndexError:
+                # Just in case
+                if self.currently_pending > 0:
+                    self.currently_pending -= 1
+                # When items are exhausted reset schedule to ensure that new items going to get another schedule
+                # It is important for exponential policy
+                return
+
+        try:
+            request()
+        finally:
+            schedule, delay = self._get_delay(schedule)
+            self.scheduler.schedule(delay, self._run, schedule)
+
+    def schedule_new_connection(self, cb: Callable[[], None]):
+        with self.lock:
+            if self.is_shutdown:
+                return
+            self.items.append(cb)
+            if self.currently_pending < self.max_concurrent:
+                self.currently_pending += 1
+                schedule = self.backoff_policy.new_schedule()
+                delay = next(schedule)
+                self.scheduler.schedule(delay, self._run, schedule)
+
+    def shutdown(self):
+        with self.lock:
+            self.is_shutdown = True
+
+
+class _LimitedConcurrencyShardConnectionScheduler(ShardConnectionScheduler):
+    """
+    A scheduler for ``cassandra.policies.LimitedConcurrencyShardConnectionPolicy``.
+
+    Limits concurrency for connection creation requests to ``max_concurrent`` per host_id.
+    """
+
+    already_scheduled: set[tuple[str, int]]
+    """
+    Set of (host_id, shard_id) of scheduled or pending requests.
+    """
+
+    per_host_scope: dict[str, _ScopeBucket]
+    """
+    Scopes storage, key is host_id, value is an instance that holds scope data.
+    """
+
+    backoff_policy: ShardConnectionBackoffSchedule
+    scheduler: _Scheduler
+    lock: Lock
+    is_shutdown: bool = False
+
+    max_concurrent: int
+    """
+    Max concurrent connection creation requests per host_id.
+    """
+
+    def __init__(
+            self,
+            scheduler: _Scheduler,
+            backoff_policy: ShardConnectionBackoffSchedule,
+            max_concurrent: int,
+    ):
+        self.already_scheduled = set()
+        self.per_host_scope = {}
+        self.backoff_policy = backoff_policy
+        self.max_concurrent = max_concurrent
+        self.scheduler = scheduler
+        self.lock = Lock()
+
+    def _execute(self, host_id: str, shard_id: int, method: Callable[[], 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))
+
+            scope_info = self.per_host_scope.get(host_id)
+            if not scope_info:
+                scope_info = _ScopeBucket(self.scheduler, self.backoff_policy, self.max_concurrent)
+                self.per_host_scope[host_id] = scope_info
+            scope_info.schedule_new_connection(partial(self._execute, host_id, shard_id, method))
+            return True
+
+    def shutdown(self):
+        with self.lock:
+            self.is_shutdown = True
+            for scope in self.per_host_scope.values():
+                scope.shutdown()
+
+
 class RetryPolicy(object):
     """
     A policy that describes whether to retry, rethrow, or ignore coordinator