51 lines
1.4 KiB
Python
51 lines
1.4 KiB
Python
from __future__ import annotations
|
|
|
|
import threading
|
|
from dataclasses import dataclass
|
|
|
|
|
|
@dataclass
|
|
class InboxRunLease:
|
|
inbox_id: str
|
|
_locks: list[threading.Lock]
|
|
_released: bool = False
|
|
|
|
def release(self) -> None:
|
|
if not self._released:
|
|
self._released = True
|
|
for lock in reversed(self._locks):
|
|
lock.release()
|
|
|
|
def __enter__(self) -> "InboxRunLease":
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb) -> None:
|
|
self.release()
|
|
|
|
|
|
class InboxRunLocks:
|
|
def __init__(self) -> None:
|
|
self._guard = threading.Lock()
|
|
self._global_lock = threading.Lock()
|
|
self._locks: dict[str, threading.Lock] = {}
|
|
|
|
def acquire(self, inbox_id: str, *, blocking: bool = False) -> InboxRunLease | None:
|
|
if not self._global_lock.acquire(blocking=blocking):
|
|
return None
|
|
with self._guard:
|
|
lock = self._locks.setdefault(inbox_id, threading.Lock())
|
|
if not lock.acquire(blocking=blocking):
|
|
self._global_lock.release()
|
|
return None
|
|
return InboxRunLease(inbox_id=inbox_id, _locks=[self._global_lock, lock])
|
|
|
|
def active(self, inbox_id: str) -> bool:
|
|
lease = self.acquire(inbox_id, blocking=False)
|
|
if not lease:
|
|
return True
|
|
lease.release()
|
|
return False
|
|
|
|
|
|
inbox_run_locks = InboxRunLocks()
|