Move the sleep out of the locked portion.

This commit is contained in:
voussoir 2021-11-20 18:53:56 -08:00
parent 8f8a39eaca
commit f1aea14646
No known key found for this signature in database
GPG key ID: 5F7554F8C26DACCB

View file

@ -66,22 +66,22 @@ class Ratelimiter:
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)
if self.balance >= cost:
self.balance -= cost self.balance -= cost
successful = True
if self.balance >= 0:
success = True
sleep_needed = 0
elif self.mode == 'reject': elif self.mode == 'reject':
successful = False success = False
sleep_needed = 0
else: else:
deficit = cost - self.balance success = True
time_needed = deficit / self.gain_rate sleep_needed = abs(self.balance) / self.gain_rate
time.sleep(time_needed)
self.balance = 0
successful = True
self.last_operation = now self.last_operation = now
return successful return (success, sleep_needed)
def limit(self, cost=None): def limit(self, cost=None):
''' '''
@ -91,4 +91,9 @@ class Ratelimiter:
cost = self.operation_cost cost = self.operation_cost
with self.lock: with self.lock:
return self._limit(cost) (success, sleep_needed) = self._limit(cost)
if sleep_needed > 0:
time.sleep(sleep_needed)
return success