From 28b4c8626bb30c09bfa517c2c91526cf117fe64d Mon Sep 17 00:00:00 2001 From: Ethan Dalool Date: Tue, 21 Dec 2021 17:02:08 -0800 Subject: [PATCH] Fix ratelimiter reject mode always subtracting balance. --- voussoirkit/ratelimiter.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/voussoirkit/ratelimiter.py b/voussoirkit/ratelimiter.py index 325631c..ba20da6 100644 --- a/voussoirkit/ratelimiter.py +++ b/voussoirkit/ratelimiter.py @@ -21,7 +21,8 @@ class Ratelimiter: Our spending balance per `period` seconds. 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: The default amount to remove from our balance after each operation. @@ -65,22 +66,21 @@ class Ratelimiter: time_diff = now - self.last_operation self.balance += time_diff * self.gain_rate self.balance = min(self.balance, self.allowance) + self.last_operation = now - self.balance -= cost - - if self.balance >= 0: - success = True - sleep_needed = 0 - - elif self.mode == 'reject': + if self.mode == 'reject' and self.balance < cost: success = False sleep_needed = 0 + return (success, sleep_needed) + self.balance -= cost + success = True + + if self.balance >= 0: + sleep_needed = 0 else: - success = True sleep_needed = abs(self.balance) / self.gain_rate - self.last_operation = now return (success, sleep_needed) def limit(self, cost=None):