1# ty: ignore[unresolved-attribute]
2__lazy_modules__ = frozenset(('functools',))
3from asyncutils.config import _randinst
4from asyncutils.constants import _NO_DEFAULT
5from asyncutils._internal import patch as P
6from asyncutils._internal.helpers import fullname
7from asyncutils._internal.submodules import altlocks_all as __all__
8from _collections import deque
9import asyncio as I, asyncutils as A
10from functools import wraps
11from itertools import count
12from sys import audit
13from time import monotonic
14from _warnings import warn
[docs]
15class Releasing:
16 __slots__ = '__lock',
17 def __init__(self, l, /): self.__lock = l
[docs]
18 async def __aenter__(self):
19 if not (l := self.__lock).locked(): raise RuntimeError('asyncutils.altlocks.Releasing: lock is not acquired')
20 if I.iscoroutine(r := l.release()): await r
[docs]
21 async def __aexit__(self, *_): await self.__lock.acquire()
22class Resource:
23 __inc_cnt = count(1).__next__; __slots__ = '__',
24 def __init__(self): self.__ = f'anonymous resource #{__class__.__inc_cnt()}'
25 def __repr__(self): return self.__
[docs]
26class ResourceGuard(A.AsyncContextMixin):
27 __slots__ = '__', '_t', '_u', 'action', 'guarded'
28 def __new__(cls, rsrc=_NO_DEFAULT, *, action='using', _=Resource):
29 if rsrc is _NO_DEFAULT: rsrc = _()
30 (s := object.__new__(cls)).__, s.action, s.guarded = rsrc, action, False; s._t = s._u = 0; return s
[docs]
31 def __enter__(self):
32 r = self.__; self._t += 1
33 if self.guarded: raise A.ResourceBusy(f'another task is already {self.action} resource: {r!r}')
34 self.guarded = True; self._u += 1
[docs]
35 def __exit__(self, /, *_, e='asyncutils.altlocks.ResourceGuard: __aexit__ called without prior __aenter__ call'):
36 if not self.guarded: raise RuntimeError(e)
37 self.guarded = False
[docs]
38 @A.dualcontextmanager(use_existing_executor=False, create_executor=False, strict=False)
39 def yields_resource(self, _=Resource):
40 if isinstance(r := self.__, _): raise TypeError('asyncutils.altlocks.ResourceGuard.yields_resource expected resource guard to have been instantiated with a resource')
41 with self: yield r
42 @property
43 def success_ratio(self): return u/self._t if (u := self._u) else 0.0
44 P.patch_method_signatures((__exit__, P.exit_sig), (yields_resource, ''), (__new__, "rsrc={0}, *, action='using'"))
[docs]
45class UniqueResourceGuard(ResourceGuard):
46 _cache = __import__('weakref').WeakValueDictionary(); __slots__ = '__weakref__',
47 def __init_subclass__(cls, /, **_): raise TypeError('cannot subclass asyncutils.altlocks.UniqueResourceGuard')
48 def __new__(cls, rsrc, **k):
49 if (r := (c := cls._cache).get(i := id(rsrc))) is None: audit('asyncutils.altlocks.UniqueResourceGuard', fullname(rsrc)); c[i] = r = super().__new__(cls, rsrc, **k)
50 elif k: warn('asyncutils.altlocks.UniqueResourceGuard: ignoring keyword arguments in favour of pre-existing guard', RuntimeWarning, 2)
51 return r
[docs]
52 @classmethod
53 def clear_cache(cls): audit('asyncutils.altlocks.UniqueResourceGuard.clear_cache'); cls._cache.clear()
[docs]
54class CircuitBreaker:
55 State = __import__('enum').IntEnum('State', ('CLOSED', 'HALF_OPEN', 'OPEN'), module=__name__)
56 __slots__ = '__exc', '__hoc', '__lock', '__max_hoc', '__mf', '__opened', '__reset', '__unlock', 'fails', 'name', 'state'; __inc_cnt = count(1).__next__
57 def __new__(cls, n, /, max_fails=None, reset=None, *, exc=Exception, max_half_open_calls=None, _='#%d'):
58 f = None
59 if callable(n) and (n := getattr(f := getattr(getattr(n, '__func__', n), '__wrapped__', n), '__qualname__', None)) is None is (n := getattr(f, '__name__', None)): n = _%cls.__inc_cnt()
60 audit('asyncutils.altlocks.CircuitBreaker', n, max_fails); s, C = super().__new__(cls), A.getcontext(); s.name, s.__mf, s.__reset, s.__exc, s.__opened, s.__max_hoc, s.__unlock, s.__lock, s.state = n, C.CIRCUIT_BREAKER_DEFAULT_MAX_FAILS if max_fails is None else max_fails, C.CIRCUIT_BREAKER_DEFAULT_RESET if reset is None else reset, exc, float('-inf'), C.CIRCUIT_BREAKER_DEFAULT_MAX_HALF_OPEN_CALLS if max_half_open_calls is None else max_half_open_calls, Releasing(l := I.Lock()), l, cls.State.CLOSED; s.fails = s.__hoc = 0; return s if f is None else s(f)
[docs]
61 def __call__(self, f, /, *, timer=monotonic, default=_NO_DEFAULT):
62 audit('asyncutils.altlocks.CircuitBreaker.__call__', self.name, fullname(f))
63 async def g(*a, **k):
64 C = self.State
65 async with self.__lock:
66 if (s := self.state) == C.OPEN:
67 if timer()-self.__opened > self.__reset: self.state, self.__hoc = C.HALF_OPEN, 0
68 else: raise A.CircuitOpen(f'asyncutils.altlocks.CircuitBreaker: circuit {self.name} is open')
69 elif s == C.HALF_OPEN:
70 if (c := self.__hoc) == (m := self.__max_hoc): raise A.CircuitHalfOpen(f'asyncutils.altlocks.CircuitBreaker: breaker {self.name} exceeded the maximum of {m} calls in the half-open state')
71 self.__hoc = c+1
72 try:
73 async with self.__unlock: r = await f(*a, **k)
74 if s == C.HALF_OPEN: self.__hoc = self.fails = 0; self.state = C.CLOSED
75 return r
76 except self.__exc:
77 if (x := self.fails+1) < self.__mf: self.fails = x
78 else: self.__opened, self.state, self.fails = timer(), C.OPEN, 0
79 if default is _NO_DEFAULT: raise
80 return default
81 except A.CRITICAL: raise A.Critical
82 except BaseException as e: raise A.CircuitBreakerError(f'asyncutils.altlocks.CircuitBreaker: unexpected {fullname(e)} in {fullname(f)} under breaker {self.name!r}') from e
83 return wraps(f)(g)
84 P.patch_classmethod_signatures((__new__, 'name, /, max_fails=None, reset=None, *, exc={}, max_half_open_calls=None'))
[docs]
85class StatefulBarrier(A.AwaitableMixin):
86 __slots__ = '__br', '__cd', '__cn', '__exc', '__ist', '__ps', '__st'
87 def __init__(self, parties, name='\b', init_state=(), max_state=None): self.__ps, self.__exc, self.__cn, self.__st, self.__cd, self.__ist, self.__br = parties, I.BrokenBarrierError(f'{fullname(self)} {name} is broken'), 0, deque(maxlen=max_state), I.Condition(), init_state, False
88 async def _wait(self, x, /):
89 self.raise_for_abort(); f = (S := self.__st).append
90 if (s := self.__ist) is not None:
91 async for i in A.iter_to_agen(s): f(i)
92 self.__ist = None
93 self.__cn, C = (c := self.__cn)+1, self.__cd
94 if x is not None: f(x)
95 if c == self.__ps-1: self.__cn = 0; C.notify_all(); self.__br = True
96 else:
97 w = C.wait
98 while not self.__br: await w()
99 return c, S.copy()
[docs]
100 async def wait(self, state=None, timeout=None):
101 try:
102 async with I.timeout(timeout), self.__cd: return await self._wait(state)
103 except TimeoutError: await self.abort(); raise
[docs]
104 async def abort(self):
105 async with (C := self.__cd):
106 if not self.__br: self.__br = True; C.notify_all()
[docs]
107 def raise_for_abort(self):
108 if self.__br: raise self.__exc
109 @property
110 def broken(self): return self.__br
111 @property
112 def remaining_parties(self): return self.__ps-self.__cn
113 @property
114 def parties(self): return self.__ps
115 @property
116 def n_waiting(self): return self.__cn
[docs]
117class DynamicThrottle:
118 __slots__ = '__fails', '__jitter', '__lb', '__lc', '__lf', '__lock', '__max', '__min', '__rate', '__rf', '__successes', '__timer', '__ub', '__uf', '__window'
119 def __init__(self, init_rate, min_rate=None, max_rate=None, window=None, *, ubound=None, lbound=None, ufactor=None, lfactor=None, jitter=None, timer=monotonic, rand=lambda j, u=_randinst.uniform: u(-j, j)): # noqa: PLR0913
120 C = A.getcontext()
121 if min_rate is None: min_rate = C.DYNAMIC_THROTTLE_DEFAULT_MIN_RATE
122 if max_rate is None: max_rate = C.DYNAMIC_THROTTLE_DEFAULT_MAX_RATE
123 if not 0 < min_rate <= init_rate <= max_rate: raise ValueError('asyncutils.altlocks.DynamicThrottle: inconsistent rates after applying bounds')
124 self.__min, self.__max, self.__window, self.__lock, self.__timer, self.__ub, self.__lb, self.__uf, self.__lf, self.jitter, self.__rf, self.__rate, self.__lc = min_rate, max_rate, C.DYNAMIC_THROTTLE_DEFAULT_WINDOW if window is None else window, I.Lock(), timer, C.DYNAMIC_THROTTLE_DEFAULT_UBOUND if ubound is None else ubound, C.DYNAMIC_THROTTLE_DEFAULT_LBOUND if lbound is None else lbound, C.DYNAMIC_THROTTLE_DEFAULT_UFACTOR if ufactor is None else ufactor, C.DYNAMIC_THROTTLE_DEFAULT_LFACTOR if lfactor is None else lfactor, C.DYNAMIC_THROTTLE_DEFAULT_JITTER if jitter is None else jitter, rand, init_rate, timer()-1.0/init_rate; self.reset()
125 @property
126 def rate(self): return self.__rate
127 @rate.setter
128 def rate(self, rate, /, _=0.01):
129 if abs(1-self.__rate/(rate := max(self.__min, min(self.__max, rate)))) > _: self.__rate = rate
130 @property
131 def jitter(self): return self.__jitter
132 @jitter.setter
133 def jitter(self, jitter, /): self.__jitter = max(0.0, float(jitter))
134 @property
135 def ctime(self): return self.__timer()
136 @property
137 def successes(self): return self.__successes
138 @property
139 def fails(self): return self.__fails
[docs]
140 async def __aenter__(self): await I.sleep((1.0/self.__rate-self.ctime+self.__lc)*(1.0+self.__rf(self.__jitter))); self.__lc = self.ctime
[docs]
141 async def __aexit__(self, e, /, *_):
142 async with self.__lock:
143 if (t := (s := self.__successes)+self.__fails) >= self.__window: self.rate *= self.__uf if (r := s/t) > self.__ub else self.__lf if r < self.__lb else 1.0; self.reset()
144 if e is None: self.__successes += 1
145 else: self.__fails += 1
[docs]
146 def reset(self): self.__successes = self.__fails = 0
147 P.patch_method_signatures((__aexit__, P.exit_sig))
148del Resource, _randinst, count, P