voussoir/voussoirkit
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Add a lock to make Ratelimiter thread-safe.

Browse source
This commit is contained in:
Ethan Dalool 2020-02-17 23:39:43 -08:00
parent 9b8b0b6339
commit 941b9b5350
1 changed files with 13 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
voussoirkit/ratelimiter.py
View file

@ -1,6 +1,6 @@
import threading
import time import time
class Ratelimiter: class Ratelimiter:
def __init__(self, allowance, period=1, operation_cost=1, mode='sleep'): def __init__(self, allowance, period=1, operation_cost=1, mode='sleep'):
''' '''
@ -30,6 +30,7 @@ class Ratelimiter:
self.period = period self.period = period
self.operation_cost = operation_cost self.operation_cost = operation_cost
self.mode = mode self.mode = mode
self.lock = threading.Lock()
self.last_operation = time.time() self.last_operation = time.time()
self.balance = 0 self.balance = 0
@ -38,13 +39,7 @@ class Ratelimiter:
def gain_rate(self): def gain_rate(self):
return self.allowance / self.period return self.allowance / self.period
def limit(self, cost=None): def _limit(self, cost):
'''
See the main class docstring for info about cost and mode behavior.
'''
if cost is None:
cost = self.operation_cost
time_diff = time.time() - self.last_operation time_diff = time.time() - self.last_operation
self.balance += time_diff * self.gain_rate self.balance += time_diff * self.gain_rate
self.balance = min(self.balance, self.allowance) self.balance = min(self.balance, self.allowance)
@ -65,3 +60,13 @@ class Ratelimiter:
self.last_operation = time.time() self.last_operation = time.time()
return successful 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)