1__lazy_modules__ = frozenset(('heapq',))
2import asyncutils as A, asyncio as I
3from asyncutils._internal.helpers import LoopMixinBase, fullname, subscriptable
4from asyncutils._internal.submodules import locks_all as __all__
5from _collections import defaultdict, deque
6from heapq import heappop, heappush
7from time import monotonic
[docs]
8class DynamicBoundedSemaphore(I.BoundedSemaphore):
9 def __init__(self, value=None): super().__init__(A.getcontext().DYNAMIC_BOUNDED_SEMAPHORE_DEFAULT_VALUE if value is None else value); self._waiters = deque()
10 @property
11 def bound(self): return self._bound_value
12 @bound.setter
13 def bound(self, value, /):
14 if value < 0: raise ValueError('asyncutils.locks.DynamicBoundedSemaphore: bound must be non-negative')
15 d, self._bound_value, f = value-self._bound_value, value, (W := self._waiters).popleft # ty: ignore[unresolved-attribute]
16 while d and W:
17 if not (w := f()).done(): w.set_result(None); d -= 1
18def d(m, /, _=__import__('functools').wraps):
19 async def w(self, *a, **k):
20 async with self._lock: self.update_tokens_lock_held(); m(self, *a, **k)
21 return _(m)(w)
[docs]
22class AdvancedRateLimit(LoopMixinBase, A.LockMixin):
23 __slots__ = '__fair', '__lock', '__lu', '__waiters', 'capacity', 'rate', 'tokens'
24 def __init__(self, rate, capacity=None, fair=True): super().__init__(); self.rate, self.__lock, self.__waiters, self.__fair, self.__lu = rate, I.Lock(), deque(), fair, monotonic(); self.tokens = self.capacity = capacity or rate
[docs]
25 async def acquire(self, tokens=None, timeout=None):
26 async with self.__lock:
27 self.update_tokens_lock_held()
28 if tokens > self.tokens: w = self.__waiters; (w.append if self.__fair else w.appendleft)((A.getcontext().ADVANCED_RATE_LIMIT_DEFAULT_TOKENS if tokens is None else tokens, F := self.loop.create_future()))
29 else: self.tokens -= tokens; return True
30 try: await I.wait_for(F, timeout); return True
31 except TimeoutError: return False
[docs]
32 @d
33 def release(self, tokens=None): self.tokens = min(self.tokens+(A.getcontext().ADVANCED_RATE_LIMIT_DEFAULT_TOKENS if tokens is None else tokens), self.capacity)
[docs]
34 @d
35 def set_rate(self, new): self.rate, self.capacity = new, max(self.capacity, new)
[docs]
36 def locked(self): return bool(self.__waiters)
[docs]
37 def update_tokens_lock_held(self):
38 if not (w := self.__waiters): return
39 e, p, self.__lu = (n := monotonic())-self.__lu, w.popleft, n; T = min(self.capacity, self.tokens+e*self.rate)
40 while (t := p())[0] <= T and w:
41 t, f = t; T -= t
42 if not f.done(): f.set_result(None)
43 self.tokens = T; w.appendleft(t)
[docs]
44class PrioritySemaphore(LoopMixinBase, A.LockMixin):
45 __slots__ = '__tiebreak', '__value', '__waiters'
46 def __init__(self, value=None): self.__value, self.__tiebreak, self.__waiters = A.getcontext().PRIORITY_SEMAPHORE_DEFAULT_VALUE if value is None else value, 0, []
[docs]
47 async def acquire(self, priority=0):
48 self.__value -= 1; self.__tiebreak += 1; w = self.__waiters
49 while self.__value < 0: heappush(w, (priority, self.__tiebreak, F := self.make_fut())); await F
50 return True
[docs]
51 def release(self, strict=True):
52 if w := self.__waiters: heappop(w)[-1].set_result(None)
53 elif strict: raise RuntimeError('asyncutils.locks.PrioritySemaphore: release called too many times')
54 self.__value += 1
[docs]
55 def locked(self): return self.__value < 0
[docs]
56 def reset(self):
57 for *_, e in self.__waiters: e.set_result(None)
58 self.__value, self.__tiebreak = 1, 0; self.__waiters.clear()
[docs]
59class KeyedCondition(LoopMixinBase, A.LockMixin):
60 __slots__ = '__lock', '__sw'
61 def __init__(self, lock=None): super().__init__(); self.__lock, self.__sw = lock or I.Lock(), defaultdict(set)
[docs]
62 async def acquire(self):
63 with A.ignore_noncritical:
64 if await self.__lock.acquire() != False: return True # noqa: E712
65 return False
[docs]
66 async def release(self):
67 if I.iscoroutine(r := self.__lock.release()): await r
[docs]
68 def locked(self): return self.__lock.locked()
[docs]
69 async def wait(self, key, timeout=None):
70 self.assert_locked(); (s := self.__sw[key]).add(F := self.make_fut())
71 try: await I.wait_for(F, timeout)
72 finally: s.discard(F)
[docs]
73 async def wait_for(self, key, pred, per_wait_timeout=None):
74 self.assert_locked(); f, g, h, F = (s := self.__sw[key]).add, s.discard, self.make_fut, None
75 try:
76 while not pred(): f(F := h()); await I.wait_for(F, per_wait_timeout); g(F)
77 finally: g(F)
[docs]
78 async def wait_all(self, timeout=None): self.assert_locked(); await I.wait_for(I.gather(*frozenset().union(*self.__sw.values()), return_exceptions=True), timeout)
[docs]
79 def assert_locked(self):
80 if not self.locked(): raise RuntimeError('asyncutils.locks.KeyedCondition: must acquire condition to notify')
[docs]
81 def notify(self, key, n=1, strict=False):
82 if n <= 0:
83 if strict: raise ValueError(f'{fullname(self)}: n must be positive')
84 return
85 self.assert_locked()
86 if (s := (S := self.__sw).pop(key, None)) is None:
87 if strict: raise ValueError(f'{fullname(self)}: no parties waiting for key {key!r}')
88 return
89 p = s.pop
90 while s:
91 if not (F := p()).done(): F.set_result(None); n -= 1
92 if n == 0: break
93 if s: S[key] = s
94 if strict and n > 0: raise ValueError(f'{fullname(self)}: not enough parties to notify')
[docs]
95 def notify_all(self, key=None):
96 self.assert_locked()
97 l = 0
98 for k in self.__sw if key is None else (key,):
99 if (s := self.__sw.pop(k, None)) is None: break
100 p = s.pop
101 while s:
102 if not (F := p()).done(): F.set_result(None)
103 l += len(s); s.clear()
104 return l
[docs]
105@subscriptable
106class MultiCountDownLatch:
107 __slots__ = '__cd', '__cts'
108 def __init__(self, counts): self.__cd, self.__cts = KeyedCondition(), {k: v for k, v in counts.items() if v > 0}
109 def _count_down_lock_held(self, key, strict):
110 if (c := (d := self.__cts).get(key)) is None:
111 if strict: raise KeyError(f'{fullname(self)}: cannot count down key {key!r} further')
112 return
113 if c > 1: d[key] = c-1
114 else: del d[key]
115 if c == 1: self.__cd.notify_all(key)
[docs]
116 async def count_down(self, key, strict=False):
117 async with self.__cd: self._count_down_lock_held(key, strict)
[docs]
118 async def count_down_all(self):
119 f = self._count_down_lock_held
120 async with self.__cd:
121 for key in self.__cts: f(key, True)
[docs]
122 async def wait(self, key, strict=False):
123 if key in self.__cts:
124 async with (C := self.__cd): await C.wait(key)
125 elif strict: raise KeyError(f'{fullname(self)}: no count for key {key!r}')
[docs]
126 async def wait_all(self, timeout=None):
127 async with (C := self.__cd): await C.wait_all(timeout)
128 @property
129 def broken(self): return not self.__cts
[docs]
130class RLock(A.LockWithOwnerMixin):
131 __slots__ = '__cnt', '__lock', '__owner'
132 def __init__(self, lock=None): self.__cnt, self.__owner, self.__lock = 0, None, lock or I.Lock()
[docs]
133 async def acquire(self):
134 async with self.__lock:
135 if self.is_owner: self.__cnt += 1; return True
136 while True:
137 async with self.__lock:
138 if self.__owner is None: self.__owner, self.__cnt = I.current_task(), 1; return True
[docs]
139 def _release(self):
140 if (c := self.__cnt) <= 0: raise RuntimeError(f'{fullname(self)}: release called too many times')
141 if c == 1: self.__owner = None
142 self.__cnt = c-1
[docs]
143 def locked(self): return self.__owner is not None
144 @property
145 def is_owner(self): return self.__owner is I.current_task()
[docs]
146class PriorityLock(LoopMixinBase, A.LockWithOwnerMixin):
147 __slots__ = '__owner', '__tiebreak', '__waiters'
148 def __init__(self): super().__init__(); self.__waiters, self.__tiebreak, self.__owner = [], 0, None
[docs]
149 async def acquire(self, priority=0, timeout=None):
150 heappush(self.__waiters, (priority, self.__tiebreak, F := self.make_fut())); self.__tiebreak += 1
151 try:
152 if len(self.__waiters) == 1 and self.__owner is None: F.set_result(True)
153 await I.wait_for(F, timeout); self.__owner = I.current_task(); return True
154 except TimeoutError: return False
155 finally:
156 if not F.done(): F.cancel()
[docs]
157 def _release(self, raise_=True):
158 self.__owner, w = None, self.__waiters
159 while w:
160 if not (F := heappop(w)[-1]).done(): return F.set_result(True)
161 if raise_: raise RuntimeError(f'{fullname(self)}: release called too many times')
[docs]
162 def locked(self): return self.__owner is not None
163 @property
164 def is_owner(self): return self.__owner is I.current_task()
[docs]
165class PriorityRLock(RLock):
166 __slots__ = ()
167 def __init__(self): super().__init__(PriorityLock())
168 @property
169 def owner(self): o = self.__lock._owner = self._owner; return o # ty: ignore[invalid-assignment]
170 @owner.setter
171 def owner(self, val, /): self._owner = self.__lock._owner = val # ty: ignore[invalid-assignment]
[docs]
172 async def acquire(self, priority=0, timeout=None):
173 if self.is_owner: self._count += 1; return True
174 if await self.__lock.acquire(priority, timeout): self._count = 1; return True # ty: ignore[too-many-positional-arguments]
175 return False
176del d