1"""Token bucket implementation for rate limiting.""" 2from __future__ import absolute_import, unicode_literals 3 4from collections import deque 5 6from kombu.five import monotonic 7 8__all__ = ('TokenBucket',) 9 10 11class TokenBucket(object): 12 """Token Bucket Algorithm. 13 14 See Also: 15 https://en.wikipedia.org/wiki/Token_Bucket 16 17 Most of this code was stolen from an entry in the ASPN Python Cookbook: 18 https://code.activestate.com/recipes/511490/ 19 20 Warning: 21 Thread Safety: This implementation is not thread safe. 22 Access to a `TokenBucket` instance should occur within the critical 23 section of any multithreaded code. 24 """ 25 26 #: The rate in tokens/second that the bucket will be refilled. 27 fill_rate = None 28 29 #: Maximum number of tokens in the bucket. 30 capacity = 1 31 32 #: Timestamp of the last time a token was taken out of the bucket. 33 timestamp = None 34 35 def __init__(self, fill_rate, capacity=1): 36 self.capacity = float(capacity) 37 self._tokens = capacity 38 self.fill_rate = float(fill_rate) 39 self.timestamp = monotonic() 40 self.contents = deque() 41 42 def add(self, item): 43 self.contents.append(item) 44 45 def pop(self): 46 return self.contents.popleft() 47 48 def clear_pending(self): 49 self.contents.clear() 50 51 def can_consume(self, tokens=1): 52 """Check if one or more tokens can be consumed. 53 54 Returns: 55 bool: true if the number of tokens can be consumed 56 from the bucket. If they can be consumed, a call will also 57 consume the requested number of tokens from the bucket. 58 Calls will only consume `tokens` (the number requested) 59 or zero tokens -- it will never consume a partial number 60 of tokens. 61 """ 62 if tokens <= self._get_tokens(): 63 self._tokens -= tokens 64 return True 65 return False 66 67 def expected_time(self, tokens=1): 68 """Return estimated time of token availability. 69 70 Returns: 71 float: the time in seconds. 72 """ 73 _tokens = self._get_tokens() 74 tokens = max(tokens, _tokens) 75 return (tokens - _tokens) / self.fill_rate 76 77 def _get_tokens(self): 78 if self._tokens < self.capacity: 79 now = monotonic() 80 delta = self.fill_rate * (now - self.timestamp) 81 self._tokens = min(self.capacity, self._tokens + delta) 82 self.timestamp = now 83 return self._tokens 84