diff --git a/kazoo/client.py b/kazoo/client.py index 3f2c3b943..aba99aa0d 100644 --- a/kazoo/client.py +++ b/kazoo/client.py @@ -175,6 +175,7 @@ def __init__( use_ssl: bool = False, verify_certs: bool = True, check_hostname: bool = False, + concurrent_request_limit: int = 0, ) -> None: ... @@ -206,6 +207,7 @@ def __init__( use_ssl: bool = False, verify_certs: bool = True, check_hostname: bool = False, + concurrent_request_limit: int = 0, **kwargs: Unpack[LegacyRetryParams], ) -> None: ... @@ -231,6 +233,7 @@ def __init__( use_ssl: bool = False, verify_certs: bool = True, check_hostname: bool = False, + concurrent_request_limit: int = 0, **kwargs: Unpack[LegacyRetryParams], ) -> None: """Create a :class:`KazooClient` instance. All time arguments @@ -366,6 +369,18 @@ def __init__( self.keyfile = keyfile self.keyfile_password = keyfile_password self.ca = ca + if concurrent_request_limit > 0: + self.logger.info( + "Zookeeper client rate-limited to %d concurrent requests", + concurrent_request_limit, + ) + self.rate_limiting_sem = self.handler.semaphore_impl( + concurrent_request_limit + ) + + else: + self.rate_limiting_sem = None + # Curator like simplified state tracking, and listeners for # state transitions self._state: KeeperState = KeeperState.CLOSED @@ -800,6 +815,16 @@ def _call( async_object.set_exception(SessionExpiredError()) return False + if self.rate_limiting_sem: + if not self.rate_limiting_sem.acquire(blocking=False): + self.logger.info( + "Limiting concurrent requests. Waiting for completion." + ) + # Actually block on the sempahore here + self.rate_limiting_sem.acquire(blocking=True) + # Register the release of the semaphore on async request completion + async_object.rawlink(lambda _res: self.rate_limiting_sem.release()) + self._queue.append((request, async_object)) # wake the connection, guarding against a race with close() diff --git a/kazoo/handlers/eventlet.py b/kazoo/handlers/eventlet.py index 1c60545f5..24bc8b912 100644 --- a/kazoo/handlers/eventlet.py +++ b/kazoo/handlers/eventlet.py @@ -14,6 +14,7 @@ from eventlet.green import threading as green_threading from eventlet.green import selectors as green_selectors from eventlet import queue as green_queue +from eventlet import semaphore as green_semaphore from kazoo.handlers import utils from kazoo.handlers.utils import selector_select @@ -98,6 +99,7 @@ class SequentialEventletHandler: name = "sequential_eventlet_handler" queue_impl = green_queue.LightQueue queue_empty = green_queue.Empty + semaphore_impl = green_semaphore.BoundedSemaphore def __init__(self) -> None: """Create a :class:`SequentialEventletHandler` instance""" diff --git a/kazoo/handlers/gevent.py b/kazoo/handlers/gevent.py index d8e1a8381..7fa20628b 100644 --- a/kazoo/handlers/gevent.py +++ b/kazoo/handlers/gevent.py @@ -22,7 +22,11 @@ import gevent.thread # type: ignore[import] import gevent.selectors # type: ignore[import] -from gevent.lock import Semaphore, RLock # type: ignore[import] +from gevent.lock import ( # type: ignore[import] + BoundedSemaphore, + RLock, + Semaphore, +) if TYPE_CHECKING: @@ -71,6 +75,7 @@ class SequentialGeventHandler: queue_impl = gevent.queue.Queue queue_empty = gevent.queue.Empty sleep_func = staticmethod(gevent.sleep) + semaphore_impl = BoundedSemaphore def __init__(self) -> None: """Create a :class:`SequentialGeventHandler` instance""" diff --git a/kazoo/handlers/threading.py b/kazoo/handlers/threading.py index 1c7a20c62..c4a4db464 100644 --- a/kazoo/handlers/threading.py +++ b/kazoo/handlers/threading.py @@ -107,6 +107,7 @@ class SequentialThreadingHandler(IHandler): sleep_func = staticmethod(time.sleep) queue_impl = queue.Queue queue_empty = queue.Empty + semaphore_impl = threading.BoundedSemaphore def __init__(self) -> None: """Create a :class:`SequentialThreadingHandler` instance"""