1# ruff: noqa: PLR6301
2import asyncio as I, asyncutils as A, asyncutils._internal.log as L
3from asyncutils.constants import _NO_DEFAULT
4from asyncutils._internal.helpers import check_methods, fullname, get_loop_and_set
5from asyncutils._internal.submodules import locksmiths_all as __all__
6from enum import IntEnum as E
7from sys import audit
8ForceResult, RecognitionResult = E('ForceResult', 'UNFORCEABLE NO_CURRENT_TASK OWNER_COMPLETED ALREADY_BEING_FORCED FAILURE RELEASED_WITH_FALSE SUCCESS RELEASED', module=__name__), E('RecognitionResult', 'FAILED_PRELIM FAILED_ACK ALREADY_RECOGNIZED SUCCESS', module=__name__)
9succeeded = frozenset((ForceResult.SUCCESS, ForceResult.RELEASED, RecognitionResult.ALREADY_RECOGNIZED, RecognitionResult.SUCCESS)).__contains__
[docs]
10class LocksmithBase:
11 __slots__ = '__lock', '__loop', '__recognized'; handlers = {} # noqa: RUF012
[docs]
12 @classmethod
13 def register_handler(cls, h, /, *, shadow=True):
14 def register(t, H=cls.handlers, h=h):
15 if not isinstance(t, type): raise TypeError('asyncutils.locksmiths.LocksmithBase: non-type cannot be registered')
16 if shadow: H[t] = h
17 elif h is not (h := H.setdefault(t, h)): raise KeyError('asyncutils.locksmiths.LocksmithBase: handler for type already registered', t, h)
18 return t
19 return register
20 @property
21 def currently_recognized(self): return frozenset(self.__recognized)
22 def __init__(self, loop=None, lcls=I.Lock): self.__recognized, self.__loop, self.__lock = __import__('_weakrefset').WeakSet((l := lcls(),)), loop or get_loop_and_set(), l
[docs]
23 async def recognize_lock(self, l, /):
24 if not self.preliminary_check_lock(l): return RecognitionResult.FAILED_PRELIM
25 async with self.__lock:
26 if l in (r := self.__recognized): return RecognitionResult.ALREADY_RECOGNIZED
27 if callable(f := getattr(l, 'acknowledge_locksmith_lock_held', None)):
28 try: return bool((await f) if I.iscoroutine(f := f(self)) else f)
29 except A.CRITICAL: raise A.Critical
30 except: return RecognitionResult.FAILED_ACK # noqa: E722
31 r.add(l); return RecognitionResult.SUCCESS
[docs]
32 async def force(self, l, /, info=_NO_DEFAULT, *, purge_waiters=True):
33 audit('asyncutils.locksmiths.LocksmithBase.force', id(self), id(l))
34 async with self.__lock:
35 if not self.can_force_lock_held(l): return ForceResult.UNFORCEABLE
36 if info is _NO_DEFAULT: info = await self.get_info(l)
37 try:
38 if I.iscoroutine(r := l.release()): r = await r
39 except A.CRITICAL: raise A.Critical
40 except: return await self._force_except(l, info) # noqa: E722
41 else: return await self.release_returned_false(l) if r is False else ForceResult.RELEASED
42 finally:
43 if purge_waiters: await self.purge_waiters(l)
44 async def _force_except(self, l, i, /):
45 if self.find_owner(l) is (o := I.current_task(self.__loop)) and (r := await self._force_is_owner(l, i, o)): return r
46 try:
47 if callable(f := self.handlers.get(type(l))) and I.iscoroutine(r := f(l)): await r
48 except A.CRITICAL: raise A.Critical
49 return ForceResult.SUCCESS
50 async def _force_is_owner(self, l, i, o, /):
51 if o is None: return await self.throw_fallback(l)
52 if (c := o.get_coro()) is None: return await self.eager_fallback(l)
53 E = A.LockForceRequest(self, (F := self.__loop.create_future()).set_result, l, i) # ty: ignore[invalid-argument-type]
54 try: c.throw(E)
55 except A.CRITICAL as e: return self.task_raised_critical(l, e)
56 except A.LockForceRequest as e:
57 if (r := e.requester) is not self: await self.lock_busy(l, r, {})
58 elif e is E: await self.task_propagated_request(l)
59 else: return await self.already_forcing(l)
60 except BaseException as e: await self.task_raised_other(l, e) # noqa: BLE001
61 else: await self.answer_received(l, await F)
[docs]
62 async def purge_waiters(self, l, /):
63 if w := getattr(l, '_waiters', None): await A.safe_cancel_batch(w, disembowel=True)
[docs]
64 async def host(self, t, l, /, *, timeout1=_NO_DEFAULT, timeout2=_NO_DEFAULT, timeout3=_NO_DEFAULT):
65 await I.wait(f := tuple(map(self.wrap_task, (self.force(l, purge_waiters=False), l.acquire()))), return_when='FIRST_COMPLETED'); f, a, T = *f, A.getcontext().LOCKSMITH_BASE_DEFAULT_TIMEOUTS
66 if await I.wait_for(f, T[0] if timeout1 is _NO_DEFAULT else timeout1): await a
67 else:
68 try: await I.wait_for(a, T[1] if timeout2 is _NO_DEFAULT else timeout2)
69 except TimeoutError: raise TimeoutError(f'{fullname(self)}.host: failed to acquire lock {l!r} within {timeout2} seconds') from None
70 self.patch_owner(t := self.wrap_task(t), l); return await I.wait_for(self._wait_on(t, l), T[2] if timeout3 is _NO_DEFAULT else timeout3)
[docs]
71 async def get_info(self, l, /): return f'potential deadlock situation involving {fullname(l)} at {id(l):#x}'
[docs]
72 async def lock_busy(self, l, r, _, /): await A.transient_block(self.__loop, L.info, 'lock busy: %r; requesters: %r, %r', l, self, r)
[docs]
73 async def task_propagated_request(self, l, /): await A.transient_block(self.__loop, L.warning, '%s.force: running task did not handle request to release %s at %#x properly', fullname(self), fullname(l), id(l))
[docs]
74 async def answer_received(self, l, a, /): await A.transient_block(self.__loop, L.info, '%r received answer %r from %r', self, a, l)
[docs]
75 async def throw_fallback(self, _, /): return ForceResult.NO_CURRENT_TASK
[docs]
76 async def eager_fallback(self, _, /): return ForceResult.OWNER_COMPLETED
[docs]
77 async def release_returned_false(self, _, /): return ForceResult.RELEASED_WITH_FALSE
[docs]
78 async def already_forcing(self, _, /): return ForceResult.ALREADY_BEING_FORCED
[docs]
79 async def _wait_on(self, t, l, /):
80 try: return await t
81 finally:
82 if l.locked() and I.iscoroutine(a := l.release()): await a
[docs]
83 async def task_raised_other(self, l, e, /):
84 if not isinstance(e, RuntimeError): await A.transient_block(self.__loop, L.error, 'error encountered in attempt to force %s at %#x', fullname(l), id(l), exc_info=e)
[docs]
85 def wrap_task(self, a, /): return self.__loop.create_task(A.wrap_in_coro(a))
[docs]
86 def patch_owner(self, t, l, /):
87 if hasattr(l, '_owner'): l._owner = t
[docs]
88 def find_owner(self, l, /): return getattr(l, '_owner', None)
[docs]
89 def preliminary_check_lock(self, l, /): return check_methods(l, 'acquire', 'release', 'locked')
[docs]
90 def task_raised_critical(self, _, e, /): raise A.Critical(e) from None
[docs]
91 def can_force_lock_held(self, l, /): return l in self.__recognized and l.locked()
92del E