1__lazy_modules__ = frozenset(('asyncio',))
2from asyncutils import AsyncContextMixin, getcontext
3from asyncutils._internal.helpers import LoopMixinBase, fullname
4from asyncutils._internal.submodules import buckets_all as __all__
5from asyncio import Lock, sleep
6from sys import audit
7from time import monotonic
[docs]
8class TokenBucket:
9 __slots__ = '__c', '__lock', '__lu', '__rate', '__timer', '__tokens'
10 def __init__(self, rate, capacity, timer=monotonic): audit(fullname(self), rate, capacity); self.__c = self.__tokens = capacity; self.__rate, self.__lock, self.__lu, self.__timer = rate, Lock(), timer(), timer
[docs]
11 async def consume(self, tokens=None):
12 if tokens is None: tokens = float(getcontext().TOKEN_BUCKET_DEFAULT_CONSUME_TOKENS)
13 async with self.__lock:
14 e, self.__lu = (n := self.__timer())-self.__lu, n; self.__tokens = d = min(self.__c, self.__tokens+e*self.__rate)
15 if (d := d-tokens) >= 0: self.__tokens = d
16 else: await sleep(-d/self.__rate); self.__tokens = 0
17 @property
18 def capacity(self): return self.__c
[docs]
19class LeakyBucket(AsyncContextMixin, LoopMixinBase):
20 __slots__ = '__c', '__dr', '__efs', '__factor', '__hf', '__last', '__leak', '__lf', '__timer', '__tokens'
21 def __init__(self, capacity, leak, min_factor=None, max_factor=None, external_factor_settable=None, timer=monotonic): audit(fullname(self), capacity, leak); C = getcontext(); self.__leak, self.__last, self.__dr, self.__factor, self.__lf, self.__hf, self.__efs, self.__timer = leak, timer(), None, 1.0, C.LEAKY_BUCKET_DEFAULT_MIN_FACTOR if min_factor is None else min_factor, C.LEAKY_BUCKET_DEFAULT_MAX_FACTOR if max_factor is None else max_factor, C.LEAKY_BUCKET_DEFAULT_EXT_CAN_SET_FACTOR if external_factor_settable is None else external_factor_settable, timer; self.__c = self.__tokens = capacity
22 def _adjust_from_params(self, a, b, c, d, /):
23 if (f := self.__factor) < abs((n := min(f*b, self.__hf) if (r := self.__tokens/self.__c) <= a else max(f*d, self.__lf) if r >= c else min(f*1.05, 1) if f < 1 else f)-f)*100: self.__factor = n
24 def _add_tokens(self, amount):
25 if (s := self.__tokens+amount) <= self.__c: self.__tokens = s; return True
26 return False
27 def _set_tokens(self): e, self.__last = (c := self.__timer())-self.__last, c; self.__tokens = max(0, self.__tokens-e*self.__leak*self.__factor)
28 async def _drain(self):
29 while True: self._set_tokens(); self._adjust(); await sleep(min(1/(self.__leak*self.__factor), 1))
[docs]
30 async def acquire(self, amount=None): self._set_tokens(); return self._add_tokens(getcontext().LEAKY_BUCKET_DEFAULT_ACQUIRE_TOKENS if amount is None else amount)
[docs]
31 async def wait_for_tokens(self, amount=None):
32 c = getcontext()
33 if amount is None: amount = c.LEAKY_BUCKET_DEFAULT_WAIT_FOR_TOKENS_TOKENS
34 w, m = 0.0, c.LEAKY_BUCKET_WAIT_FOR_TOKENS_TICK
35 while True:
36 self._set_tokens(); self._adjust()
37 if self._add_tokens(amount): return w
38 await sleep(_ := min(m, (amount-self.__c+self.__tokens)/(self.__leak*self.__factor))); w += _
39 @property
40 def capacity(self): return self.__c
41 @property
42 def factor(self): return self.__factor
43 @factor.setter
44 def factor(self, v, /):
45 if self.__efs: self.__factor = max(self.__lf, min(self.__hf, v))
46 else: raise AttributeError(f'{fullname(self)}: cannot set factor because external_factor_settable=True not passed')
47 @factor.deleter
48 def factor(self): self.__factor = 1
49 def _adjust(self):
50 c = self.__c
51 for b, t in getcontext().LEAKY_BUCKET_ADJMAP:
52 if c >= b: return self._adjust_from_params(*t)
53 def __enter__(self):
54 if self.__dr is None: self.__dr = self.make(self._drain())
55 return self
[docs]
56 def __exit__(self, /, *a):
57 if d := self.__dr: self.__dr = None; d.cancel(a[1])