|
1
|
""" |
|
2
|
Distributed locking for critical code sections. |
|
3
|
|
|
4
|
Uses Redis SETNX with expiry for lock implementation. Each lock is stored |
|
5
|
as a Redis key with a TTL; SETNX guarantees only one holder at a time. |
|
6
|
""" |
|
7
|
|
|
8
|
from __future__ import annotations |
|
9
|
|
|
10
|
import logging |
|
11
|
import time |
|
12
|
import uuid |
|
13
|
from typing import Any |
|
14
|
|
|
15
|
logger = logging.getLogger(__name__) |
|
16
|
|
|
17
|
_LOCK_PREFIX = "navegador:lock:" |
|
18
|
|
|
19
|
|
|
20
|
class LockTimeout(Exception): |
|
21
|
"""Raised when a distributed lock cannot be acquired within the deadline.""" |
|
22
|
|
|
23
|
|
|
24
|
class DistributedLock: |
|
25
|
""" |
|
26
|
A distributed mutex backed by Redis SETNX. |
|
27
|
|
|
28
|
Usage (context manager):: |
|
29
|
|
|
30
|
lock = DistributedLock("redis://localhost:6379", "my-lock") |
|
31
|
with lock: |
|
32
|
# only one process runs this at a time |
|
33
|
... |
|
34
|
|
|
35
|
Usage (explicit):: |
|
36
|
|
|
37
|
lock = DistributedLock("redis://localhost:6379", "my-lock", timeout=10) |
|
38
|
if lock.acquire(): |
|
39
|
try: |
|
40
|
... |
|
41
|
finally: |
|
42
|
lock.release() |
|
43
|
|
|
44
|
Args: |
|
45
|
redis_url: Redis connection URL. |
|
46
|
name: Logical lock name (unique per resource). |
|
47
|
timeout: Lock expiry in seconds (default 30). Also used as the |
|
48
|
maximum time to wait when acquiring (via __enter__). |
|
49
|
retry_interval: Seconds to sleep between acquire retries (default 0.1). |
|
50
|
""" |
|
51
|
|
|
52
|
def __init__( |
|
53
|
self, |
|
54
|
redis_url: str, |
|
55
|
name: str, |
|
56
|
timeout: int = 30, |
|
57
|
retry_interval: float = 0.1, |
|
58
|
_redis_client: Any = None, |
|
59
|
) -> None: |
|
60
|
self._url = redis_url |
|
61
|
self._name = name |
|
62
|
self._timeout = timeout |
|
63
|
self._retry_interval = retry_interval |
|
64
|
self._token: str | None = None # unique value stored in Redis to own the lock |
|
65
|
self._redis: Any = _redis_client # injected in tests; lazily created otherwise |
|
66
|
|
|
67
|
# ── Internal ────────────────────────────────────────────────────────────── |
|
68
|
|
|
69
|
def _client(self) -> Any: |
|
70
|
if self._redis is None: |
|
71
|
try: |
|
72
|
import redis # type: ignore[import] |
|
73
|
except ImportError as exc: |
|
74
|
raise ImportError("Install redis: pip install redis") from exc |
|
75
|
self._redis = redis.from_url(self._url) |
|
76
|
return self._redis |
|
77
|
|
|
78
|
@property |
|
79
|
def _key(self) -> str: |
|
80
|
return f"{_LOCK_PREFIX}{self._name}" |
|
81
|
|
|
82
|
# ── Public API ──────────────────────────────────────────────────────────── |
|
83
|
|
|
84
|
def acquire(self, blocking: bool = False, deadline: float | None = None) -> bool: |
|
85
|
""" |
|
86
|
Try to acquire the lock. |
|
87
|
|
|
88
|
Args: |
|
89
|
blocking: If True, keep retrying until the lock is acquired or |
|
90
|
*deadline* is reached. |
|
91
|
deadline: Absolute time (``time.monotonic()``) after which |
|
92
|
acquisition fails. Only used when *blocking* is True. |
|
93
|
|
|
94
|
Returns: |
|
95
|
True if the lock was acquired, False otherwise. |
|
96
|
""" |
|
97
|
client = self._client() |
|
98
|
token = str(uuid.uuid4()) |
|
99
|
while True: |
|
100
|
acquired = client.set(self._key, token, nx=True, ex=self._timeout) |
|
101
|
if acquired: |
|
102
|
self._token = token |
|
103
|
logger.debug("Lock acquired: %s (%s)", self._name, token) |
|
104
|
return True |
|
105
|
if not blocking: |
|
106
|
return False |
|
107
|
if deadline is not None and time.monotonic() >= deadline: |
|
108
|
return False |
|
109
|
time.sleep(self._retry_interval) |
|
110
|
|
|
111
|
def release(self) -> None: |
|
112
|
"""Release the lock. No-op if this instance does not hold it.""" |
|
113
|
if self._token is None: |
|
114
|
return |
|
115
|
client = self._client() |
|
116
|
stored = client.get(self._key) |
|
117
|
# Decode bytes if necessary |
|
118
|
if isinstance(stored, bytes): |
|
119
|
stored = stored.decode() |
|
120
|
if stored == self._token: |
|
121
|
client.delete(self._key) |
|
122
|
logger.debug("Lock released: %s", self._name) |
|
123
|
self._token = None |
|
124
|
|
|
125
|
# ── Context manager ─────────────────────────────────────────────────────── |
|
126
|
|
|
127
|
def __enter__(self) -> "DistributedLock": |
|
128
|
deadline = time.monotonic() + self._timeout |
|
129
|
acquired = self.acquire(blocking=True, deadline=deadline) |
|
130
|
if not acquired: |
|
131
|
raise LockTimeout(f"Could not acquire lock '{self._name}' within {self._timeout}s") |
|
132
|
return self |
|
133
|
|
|
134
|
def __exit__(self, *_: object) -> None: |
|
135
|
self.release() |
|
136
|
|