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

Use sentinels for mode instead of strings.

Browse source
This commit is contained in:
voussoir 2022-01-19 18:37:49 -08:00
parent 0be376551f
commit 7fddee51df
No known key found for this signature in database
GPG key ID: 5F7554F8C26DACCB
1 changed files with 19 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
voussoirkit/ratelimiter.py
View file

@ -1,6 +1,11 @@
import threading import threading
import time import time
from voussoirkit import sentinel
REJECT = sentinel.Sentinel('reject mode')
SLEEP = sentinel.Sentinel('sleep mode')
class Ratelimiter: class Ratelimiter:
''' '''
The Ratelimiter class is used to limit how often you perform some other The Ratelimiter class is used to limit how often you perform some other
@ -19,7 +24,7 @@ class Ratelimiter:
self, self,
allowance, allowance,
*, *,
mode='sleep', mode=SLEEP,
operation_cost=1, operation_cost=1,
period=1, period=1,
starting_balance=None, starting_balance=None,
@ -29,11 +34,11 @@ class Ratelimiter:
Our spending balance per `period` seconds. Our spending balance per `period` seconds.
mode: mode:
'sleep': SLEEP:
If we do not have the balance for an operation, sleep until we If we do not have the balance for an operation, sleep until we
do. Then return True every time. do. Then return True every time.
'reject': REJECT:
If we do not have the balance for an operation, do nothing and If we do not have the balance for an operation, do nothing and
return False. Otherwise subtract the cost and return True. return False. Otherwise subtract the cost and return True.
@ -65,19 +70,25 @@ class Ratelimiter:
return x return x
if mode not in ('sleep', 'reject'):
raise ValueError(f'Invalid mode {repr(mode)}.')
self.allowance = positive(allowance, 'allowance') self.allowance = positive(allowance, 'allowance')
self.period = positive(period, 'period') self.period = positive(period, 'period')
self.operation_cost = positive(operation_cost, 'operation_cost') self.operation_cost = positive(operation_cost, 'operation_cost')
self.mode = mode
if starting_balance is None: if starting_balance is None:
self.balance = operation_cost self.balance = operation_cost
else: else:
self.balance = positive(starting_balance, 'starting_balance') self.balance = positive(starting_balance, 'starting_balance')
mode = {
'reject': REJECT,
'sleep': SLEEP,
}.get(mode, mode)
if mode is not SLEEP and mode is not REJECT:
raise ValueError(f'mode should be SLEEP or REJECT, not {repr(mode)}.')
self.mode = mode
self.lock = threading.Lock() self.lock = threading.Lock()
self.last_operation = time.monotonic() self.last_operation = time.monotonic()
@ -95,7 +106,7 @@ class Ratelimiter:
self.balance = min(self.balance, self.allowance) self.balance = min(self.balance, self.allowance)
self.last_operation = now self.last_operation = now
if self.mode == 'reject' and self.balance < cost: if self.mode is REJECT and self.balance < cost:
success = False success = False
sleep_needed = 0 sleep_needed = 0
return (success, sleep_needed) return (success, sleep_needed)