wip

Author: ceache

kazoo/client.py

@@ -125,6 +125,7 @@ def __init__(
         ca=None,
         use_ssl=False,
         verify_certs=True,
+        max_async_requests=0,
         **kwargs,
     ):
         """Create a :class:`KazooClient` instance. All time arguments
@@ -246,6 +247,18 @@ def __init__(
         self.keyfile = keyfile
         self.keyfile_password = keyfile_password
         self.ca = ca
+        if max_async_requests > 0:
+            self.logger.info(
+                "Zookeeper client rate-limited to %d concurent async requests",
+                max_async_requests,
+            )
+            self.rate_limiting_sem = self.handler.semaphore_impl(
+                max_async_requests
+            )
+
+        else:
+            self.rate_limiting_sem = None
+
         # Curator like simplified state tracking, and listeners for
         # state transitions
         self._state = KeeperState.CLOSED
@@ -640,6 +653,10 @@ def _call(self, request, async_object):
             async_object.set_exception(SessionExpiredError())
             return False
 
+        if self.rate_limiting_sem:
+            self.rate_limiting_sem.acquire(blocking=True)
+            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()

kazoo/handlers/eventlet.py

@@ -10,6 +10,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
 import kazoo.python2atexit as python2atexit
@@ -80,6 +81,7 @@ class SequentialEventletHandler(object):
     name = "sequential_eventlet_handler"
     queue_impl = green_queue.LightQueue
     queue_empty = green_queue.Empty
+    semaphore_impl = green_semaphore.BoundedSemaphore
 
     def __init__(self):
         """Create a :class:`SequentialEventletHandler` instance"""

kazoo/handlers/gevent.py

@@ -13,10 +13,7 @@
 
 from kazoo.handlers.utils import selector_select
 
-try:
-    from gevent.lock import Semaphore, RLock
-except ImportError:
-    from gevent.coros import Semaphore, RLock
+from gevent.lock import BoundedSemaphore as Semaphore, RLock as RLock
 
 from kazoo.handlers import utils
 from kazoo import python2atexit
@@ -55,6 +52,7 @@ class SequentialGeventHandler(object):
     queue_impl = gevent.queue.Queue
     queue_empty = gevent.queue.Empty
     sleep_func = staticmethod(gevent.sleep)
+    semaphore_impl = Semaphore
 
     def __init__(self):
         """Create a :class:`SequentialGeventHandler` instance"""

kazoo/handlers/threading.py

@@ -100,6 +100,7 @@ class SequentialThreadingHandler(object):
     sleep_func = staticmethod(time.sleep)
     queue_impl = Queue.Queue
     queue_empty = Queue.Empty
+    semaphore_impl = threading.BoundedSemaphore
 
     def __init__(self):
         """Create a :class:`SequentialThreadingHandler` instance"""