def __init__(self, fill_rate, queue=None, capacity=1):
self._bucket = TokenBucket(fill_rate, capacity)
self.queue = queue
if not self.queue:
self.queue = Queue()
class TokenBucketQueue(object):
"""Queue with rate limited get operations.
This uses the token bucket algorithm to rate limit the queue on get
operations.
:param fill_rate: The rate in tokens/second that the bucket will
be refilled.
:keyword capacity: Maximum number of tokens in the bucket.
Default is 1.
"""
RateLimitExceeded = RateLimitExceeded
def __init__(self, fill_rate, queue=None, capacity=1):
self._bucket = TokenBucket(fill_rate, capacity)
self.queue = queue
if not self.queue:
self.queue = Queue()
def put(self, item, block=True):
"""Put an item onto the queue."""
self.queue.put(item, block=block)
def put_nowait(self, item):
"""Put an item into the queue without blocking.
:raises Queue.Full: If a free slot is not immediately available.
"""
return self.put(item, block=False)
def get(self, block=True):
"""Remove and return an item from the queue.
:raises RateLimitExceeded: If a token could not be consumed from the
token bucket (consuming from the queue
too fast).
:raises Queue.Empty: If an item is not immediately available.
"""
get = block and self.queue.get or self.queue.get_nowait
if not block and not self.items:
raise Empty()
if not self._bucket.can_consume(1):
raise RateLimitExceeded()
return get()
def get_nowait(self):
"""Remove and return an item from the queue without blocking.
:raises RateLimitExceeded: If a token could not be consumed from the
token bucket (consuming from the queue
too fast).
:raises Queue.Empty: If an item is not immediately available.
"""
return self.get(block=False)
def qsize(self):
"""Returns the size of the queue."""
return self.queue.qsize()
def empty(self):
"""Returns :const:`True` if the queue is empty."""
return self.queue.empty()
def clear(self):
"""Delete all data in the queue."""
return self.items.clear()
def wait(self, block=False):
"""Wait until a token can be retrieved from the bucket and return
the next item."""
get = self.get
expected_time = self.expected_time
while 1:
remaining = expected_time()
if not remaining:
return get(block=block)
sleep(remaining)
def expected_time(self, tokens=1):
"""Returns the expected time in seconds of when a new token should be
available."""
if not self.items:
return 0
return self._bucket.expected_time(tokens)
@property
def items(self):
"""Underlying data. Do not modify."""
return self.queue.queue
