diff --git a/voussoirkit/ratelimiter.py b/voussoirkit/ratelimiter.py index 02c2ad2..26629af 100644 --- a/voussoirkit/ratelimiter.py +++ b/voussoirkit/ratelimiter.py @@ -1,6 +1,6 @@ +import threading import time - class Ratelimiter: def __init__(self, allowance, period=1, operation_cost=1, mode='sleep'): ''' @@ -30,6 +30,7 @@ class Ratelimiter: self.period = period self.operation_cost = operation_cost self.mode = mode + self.lock = threading.Lock() self.last_operation = time.time() self.balance = 0 @@ -38,13 +39,7 @@ class Ratelimiter: def gain_rate(self): return self.allowance / self.period - def limit(self, cost=None): - ''' - See the main class docstring for info about cost and mode behavior. - ''' - if cost is None: - cost = self.operation_cost - + def _limit(self, cost): time_diff = time.time() - self.last_operation self.balance += time_diff * self.gain_rate self.balance = min(self.balance, self.allowance) @@ -65,3 +60,13 @@ class Ratelimiter: self.last_operation = time.time() return successful + + def limit(self, cost=None): + ''' + See the main class docstring for info about cost and mode behavior. + ''' + if cost is None: + cost = self.operation_cost + + with self.lock: + return self._limit(cost)