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

Fix ratelimiter reject mode always subtracting balance.

Browse source
This commit is contained in:
voussoir 2021-12-21 17:02:08 -08:00
parent 0a0bbae88c
commit 28b4c8626b
No known key found for this signature in database
GPG key ID: 5F7554F8C26DACCB
1 changed files with 10 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
voussoirkit/ratelimiter.py
View file

@ -21,7 +21,8 @@ class Ratelimiter:
Our spending balance per `period` seconds. Our spending balance per `period` seconds.
period: period:
The number of seconds over which we can perform `allowance` operations. The number of seconds over which we can perform `allowance`
operations.
operation_cost: operation_cost:
The default amount to remove from our balance after each operation. The default amount to remove from our balance after each operation.
@ -65,22 +66,21 @@ class Ratelimiter:
time_diff = now - self.last_operation time_diff = now - 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)
self.last_operation = now
self.balance -= cost if self.mode == 'reject' and self.balance < cost:
if self.balance >= 0:
success = True
sleep_needed = 0
elif self.mode == 'reject':
success = False success = False
sleep_needed = 0 sleep_needed = 0
return (success, sleep_needed)
self.balance -= cost
success = True
if self.balance >= 0:
sleep_needed = 0
else: else:
success = True
sleep_needed = abs(self.balance) / self.gain_rate sleep_needed = abs(self.balance) / self.gain_rate
self.last_operation = now
return (success, sleep_needed) return (success, sleep_needed)
def limit(self, cost=None): def limit(self, cost=None):