Show More
@@ -0,0 +1,389 b'' | |||
|
1 | import sys | |
|
2 | import threading | |
|
3 | import weakref | |
|
4 | from base64 import b64encode | |
|
5 | from logging import getLogger | |
|
6 | from os import urandom | |
|
7 | ||
|
8 | from redis import StrictRedis | |
|
9 | ||
|
10 | __version__ = '3.7.0' | |
|
11 | ||
|
12 | loggers = { | |
|
13 | k: getLogger(".".join((__name__, k))) | |
|
14 | for k in [ | |
|
15 | "acquire", | |
|
16 | "refresh.thread.start", | |
|
17 | "refresh.thread.stop", | |
|
18 | "refresh.thread.exit", | |
|
19 | "refresh.start", | |
|
20 | "refresh.shutdown", | |
|
21 | "refresh.exit", | |
|
22 | "release", | |
|
23 | ] | |
|
24 | } | |
|
25 | ||
|
26 | PY3 = sys.version_info[0] == 3 | |
|
27 | ||
|
28 | if PY3: | |
|
29 | text_type = str | |
|
30 | binary_type = bytes | |
|
31 | else: | |
|
32 | text_type = unicode # noqa | |
|
33 | binary_type = str | |
|
34 | ||
|
35 | ||
|
36 | # Check if the id match. If not, return an error code. | |
|
37 | UNLOCK_SCRIPT = b""" | |
|
38 | if redis.call("get", KEYS[1]) ~= ARGV[1] then | |
|
39 | return 1 | |
|
40 | else | |
|
41 | redis.call("del", KEYS[2]) | |
|
42 | redis.call("lpush", KEYS[2], 1) | |
|
43 | redis.call("pexpire", KEYS[2], ARGV[2]) | |
|
44 | redis.call("del", KEYS[1]) | |
|
45 | return 0 | |
|
46 | end | |
|
47 | """ | |
|
48 | ||
|
49 | # Covers both cases when key doesn't exist and doesn't equal to lock's id | |
|
50 | EXTEND_SCRIPT = b""" | |
|
51 | if redis.call("get", KEYS[1]) ~= ARGV[1] then | |
|
52 | return 1 | |
|
53 | elseif redis.call("ttl", KEYS[1]) < 0 then | |
|
54 | return 2 | |
|
55 | else | |
|
56 | redis.call("expire", KEYS[1], ARGV[2]) | |
|
57 | return 0 | |
|
58 | end | |
|
59 | """ | |
|
60 | ||
|
61 | RESET_SCRIPT = b""" | |
|
62 | redis.call('del', KEYS[2]) | |
|
63 | redis.call('lpush', KEYS[2], 1) | |
|
64 | redis.call('pexpire', KEYS[2], ARGV[2]) | |
|
65 | return redis.call('del', KEYS[1]) | |
|
66 | """ | |
|
67 | ||
|
68 | RESET_ALL_SCRIPT = b""" | |
|
69 | local locks = redis.call('keys', 'lock:*') | |
|
70 | local signal | |
|
71 | for _, lock in pairs(locks) do | |
|
72 | signal = 'lock-signal:' .. string.sub(lock, 6) | |
|
73 | redis.call('del', signal) | |
|
74 | redis.call('lpush', signal, 1) | |
|
75 | redis.call('expire', signal, 1) | |
|
76 | redis.call('del', lock) | |
|
77 | end | |
|
78 | return #locks | |
|
79 | """ | |
|
80 | ||
|
81 | ||
|
82 | class AlreadyAcquired(RuntimeError): | |
|
83 | pass | |
|
84 | ||
|
85 | ||
|
86 | class NotAcquired(RuntimeError): | |
|
87 | pass | |
|
88 | ||
|
89 | ||
|
90 | class AlreadyStarted(RuntimeError): | |
|
91 | pass | |
|
92 | ||
|
93 | ||
|
94 | class TimeoutNotUsable(RuntimeError): | |
|
95 | pass | |
|
96 | ||
|
97 | ||
|
98 | class InvalidTimeout(RuntimeError): | |
|
99 | pass | |
|
100 | ||
|
101 | ||
|
102 | class TimeoutTooLarge(RuntimeError): | |
|
103 | pass | |
|
104 | ||
|
105 | ||
|
106 | class NotExpirable(RuntimeError): | |
|
107 | pass | |
|
108 | ||
|
109 | ||
|
110 | class Lock(object): | |
|
111 | """ | |
|
112 | A Lock context manager implemented via redis SETNX/BLPOP. | |
|
113 | """ | |
|
114 | unlock_script = None | |
|
115 | extend_script = None | |
|
116 | reset_script = None | |
|
117 | reset_all_script = None | |
|
118 | ||
|
119 | def __init__(self, redis_client, name, expire=None, id=None, auto_renewal=False, strict=True, signal_expire=1000): | |
|
120 | """ | |
|
121 | :param redis_client: | |
|
122 | An instance of :class:`~StrictRedis`. | |
|
123 | :param name: | |
|
124 | The name (redis key) the lock should have. | |
|
125 | :param expire: | |
|
126 | The lock expiry time in seconds. If left at the default (None) | |
|
127 | the lock will not expire. | |
|
128 | :param id: | |
|
129 | The ID (redis value) the lock should have. A random value is | |
|
130 | generated when left at the default. | |
|
131 | ||
|
132 | Note that if you specify this then the lock is marked as "held". Acquires | |
|
133 | won't be possible. | |
|
134 | :param auto_renewal: | |
|
135 | If set to ``True``, Lock will automatically renew the lock so that it | |
|
136 | doesn't expire for as long as the lock is held (acquire() called | |
|
137 | or running in a context manager). | |
|
138 | ||
|
139 | Implementation note: Renewal will happen using a daemon thread with | |
|
140 | an interval of ``expire*2/3``. If wishing to use a different renewal | |
|
141 | time, subclass Lock, call ``super().__init__()`` then set | |
|
142 | ``self._lock_renewal_interval`` to your desired interval. | |
|
143 | :param strict: | |
|
144 | If set ``True`` then the ``redis_client`` needs to be an instance of ``redis.StrictRedis``. | |
|
145 | :param signal_expire: | |
|
146 | Advanced option to override signal list expiration in milliseconds. Increase it for very slow clients. Default: ``1000``. | |
|
147 | """ | |
|
148 | if strict and not isinstance(redis_client, StrictRedis): | |
|
149 | raise ValueError("redis_client must be instance of StrictRedis. " | |
|
150 | "Use strict=False if you know what you're doing.") | |
|
151 | if auto_renewal and expire is None: | |
|
152 | raise ValueError("Expire may not be None when auto_renewal is set") | |
|
153 | ||
|
154 | self._client = redis_client | |
|
155 | ||
|
156 | if expire: | |
|
157 | expire = int(expire) | |
|
158 | if expire < 0: | |
|
159 | raise ValueError("A negative expire is not acceptable.") | |
|
160 | else: | |
|
161 | expire = None | |
|
162 | self._expire = expire | |
|
163 | ||
|
164 | self._signal_expire = signal_expire | |
|
165 | if id is None: | |
|
166 | self._id = b64encode(urandom(18)).decode('ascii') | |
|
167 | elif isinstance(id, binary_type): | |
|
168 | try: | |
|
169 | self._id = id.decode('ascii') | |
|
170 | except UnicodeDecodeError: | |
|
171 | self._id = b64encode(id).decode('ascii') | |
|
172 | elif isinstance(id, text_type): | |
|
173 | self._id = id | |
|
174 | else: | |
|
175 | raise TypeError("Incorrect type for `id`. Must be bytes/str not %s." % type(id)) | |
|
176 | self._name = 'lock:' + name | |
|
177 | self._signal = 'lock-signal:' + name | |
|
178 | self._lock_renewal_interval = (float(expire) * 2 / 3 | |
|
179 | if auto_renewal | |
|
180 | else None) | |
|
181 | self._lock_renewal_thread = None | |
|
182 | ||
|
183 | self.register_scripts(redis_client) | |
|
184 | ||
|
185 | @classmethod | |
|
186 | def register_scripts(cls, redis_client): | |
|
187 | global reset_all_script | |
|
188 | if reset_all_script is None: | |
|
189 | reset_all_script = redis_client.register_script(RESET_ALL_SCRIPT) | |
|
190 | cls.unlock_script = redis_client.register_script(UNLOCK_SCRIPT) | |
|
191 | cls.extend_script = redis_client.register_script(EXTEND_SCRIPT) | |
|
192 | cls.reset_script = redis_client.register_script(RESET_SCRIPT) | |
|
193 | cls.reset_all_script = redis_client.register_script(RESET_ALL_SCRIPT) | |
|
194 | ||
|
195 | @property | |
|
196 | def _held(self): | |
|
197 | return self.id == self.get_owner_id() | |
|
198 | ||
|
199 | def reset(self): | |
|
200 | """ | |
|
201 | Forcibly deletes the lock. Use this with care. | |
|
202 | """ | |
|
203 | self.reset_script(client=self._client, keys=(self._name, self._signal), args=(self.id, self._signal_expire)) | |
|
204 | ||
|
205 | @property | |
|
206 | def id(self): | |
|
207 | return self._id | |
|
208 | ||
|
209 | def get_owner_id(self): | |
|
210 | owner_id = self._client.get(self._name) | |
|
211 | if isinstance(owner_id, binary_type): | |
|
212 | owner_id = owner_id.decode('ascii', 'replace') | |
|
213 | return owner_id | |
|
214 | ||
|
215 | def acquire(self, blocking=True, timeout=None): | |
|
216 | """ | |
|
217 | :param blocking: | |
|
218 | Boolean value specifying whether lock should be blocking or not. | |
|
219 | :param timeout: | |
|
220 | An integer value specifying the maximum number of seconds to block. | |
|
221 | """ | |
|
222 | logger = loggers["acquire"] | |
|
223 | ||
|
224 | logger.debug("Getting %r ...", self._name) | |
|
225 | ||
|
226 | if self._held: | |
|
227 | raise AlreadyAcquired("Already acquired from this Lock instance.") | |
|
228 | ||
|
229 | if not blocking and timeout is not None: | |
|
230 | raise TimeoutNotUsable("Timeout cannot be used if blocking=False") | |
|
231 | ||
|
232 | if timeout: | |
|
233 | timeout = int(timeout) | |
|
234 | if timeout < 0: | |
|
235 | raise InvalidTimeout("Timeout (%d) cannot be less than or equal to 0" % timeout) | |
|
236 | ||
|
237 | if self._expire and not self._lock_renewal_interval and timeout > self._expire: | |
|
238 | raise TimeoutTooLarge("Timeout (%d) cannot be greater than expire (%d)" % (timeout, self._expire)) | |
|
239 | ||
|
240 | busy = True | |
|
241 | blpop_timeout = timeout or self._expire or 0 | |
|
242 | timed_out = False | |
|
243 | while busy: | |
|
244 | busy = not self._client.set(self._name, self._id, nx=True, ex=self._expire) | |
|
245 | if busy: | |
|
246 | if timed_out: | |
|
247 | return False | |
|
248 | elif blocking: | |
|
249 | timed_out = not self._client.blpop(self._signal, blpop_timeout) and timeout | |
|
250 | else: | |
|
251 | logger.warning("Failed to get %r.", self._name) | |
|
252 | return False | |
|
253 | ||
|
254 | logger.info("Got lock for %r.", self._name) | |
|
255 | if self._lock_renewal_interval is not None: | |
|
256 | self._start_lock_renewer() | |
|
257 | return True | |
|
258 | ||
|
259 | def extend(self, expire=None): | |
|
260 | """Extends expiration time of the lock. | |
|
261 | ||
|
262 | :param expire: | |
|
263 | New expiration time. If ``None`` - `expire` provided during | |
|
264 | lock initialization will be taken. | |
|
265 | """ | |
|
266 | if expire: | |
|
267 | expire = int(expire) | |
|
268 | if expire < 0: | |
|
269 | raise ValueError("A negative expire is not acceptable.") | |
|
270 | elif self._expire is not None: | |
|
271 | expire = self._expire | |
|
272 | else: | |
|
273 | raise TypeError( | |
|
274 | "To extend a lock 'expire' must be provided as an " | |
|
275 | "argument to extend() method or at initialization time." | |
|
276 | ) | |
|
277 | ||
|
278 | error = self.extend_script(client=self._client, keys=(self._name, self._signal), args=(self._id, expire)) | |
|
279 | if error == 1: | |
|
280 | raise NotAcquired("Lock %s is not acquired or it already expired." % self._name) | |
|
281 | elif error == 2: | |
|
282 | raise NotExpirable("Lock %s has no assigned expiration time" % self._name) | |
|
283 | elif error: | |
|
284 | raise RuntimeError("Unsupported error code %s from EXTEND script" % error) | |
|
285 | ||
|
286 | @staticmethod | |
|
287 | def _lock_renewer(lockref, interval, stop): | |
|
288 | """ | |
|
289 | Renew the lock key in redis every `interval` seconds for as long | |
|
290 | as `self._lock_renewal_thread.should_exit` is False. | |
|
291 | """ | |
|
292 | while not stop.wait(timeout=interval): | |
|
293 | loggers["refresh.thread.start"].debug("Refreshing lock") | |
|
294 | lock = lockref() | |
|
295 | if lock is None: | |
|
296 | loggers["refresh.thread.stop"].debug( | |
|
297 | "The lock no longer exists, stopping lock refreshing" | |
|
298 | ) | |
|
299 | break | |
|
300 | lock.extend(expire=lock._expire) | |
|
301 | del lock | |
|
302 | loggers["refresh.thread.exit"].debug("Exit requested, stopping lock refreshing") | |
|
303 | ||
|
304 | def _start_lock_renewer(self): | |
|
305 | """ | |
|
306 | Starts the lock refresher thread. | |
|
307 | """ | |
|
308 | if self._lock_renewal_thread is not None: | |
|
309 | raise AlreadyStarted("Lock refresh thread already started") | |
|
310 | ||
|
311 | loggers["refresh.start"].debug( | |
|
312 | "Starting thread to refresh lock every %s seconds", | |
|
313 | self._lock_renewal_interval | |
|
314 | ) | |
|
315 | self._lock_renewal_stop = threading.Event() | |
|
316 | self._lock_renewal_thread = threading.Thread( | |
|
317 | group=None, | |
|
318 | target=self._lock_renewer, | |
|
319 | kwargs={'lockref': weakref.ref(self), | |
|
320 | 'interval': self._lock_renewal_interval, | |
|
321 | 'stop': self._lock_renewal_stop} | |
|
322 | ) | |
|
323 | self._lock_renewal_thread.setDaemon(True) | |
|
324 | self._lock_renewal_thread.start() | |
|
325 | ||
|
326 | def _stop_lock_renewer(self): | |
|
327 | """ | |
|
328 | Stop the lock renewer. | |
|
329 | ||
|
330 | This signals the renewal thread and waits for its exit. | |
|
331 | """ | |
|
332 | if self._lock_renewal_thread is None or not self._lock_renewal_thread.is_alive(): | |
|
333 | return | |
|
334 | loggers["refresh.shutdown"].debug("Signalling the lock refresher to stop") | |
|
335 | self._lock_renewal_stop.set() | |
|
336 | self._lock_renewal_thread.join() | |
|
337 | self._lock_renewal_thread = None | |
|
338 | loggers["refresh.exit"].debug("Lock refresher has stopped") | |
|
339 | ||
|
340 | def __enter__(self): | |
|
341 | acquired = self.acquire(blocking=True) | |
|
342 | assert acquired, "Lock wasn't acquired, but blocking=True" | |
|
343 | return self | |
|
344 | ||
|
345 | def __exit__(self, exc_type=None, exc_value=None, traceback=None): | |
|
346 | self.release() | |
|
347 | ||
|
348 | def release(self): | |
|
349 | """Releases the lock, that was acquired with the same object. | |
|
350 | ||
|
351 | .. note:: | |
|
352 | ||
|
353 | If you want to release a lock that you acquired in a different place you have two choices: | |
|
354 | ||
|
355 | * Use ``Lock("name", id=id_from_other_place).release()`` | |
|
356 | * Use ``Lock("name").reset()`` | |
|
357 | """ | |
|
358 | if self._lock_renewal_thread is not None: | |
|
359 | self._stop_lock_renewer() | |
|
360 | loggers["release"].debug("Releasing %r.", self._name) | |
|
361 | error = self.unlock_script(client=self._client, keys=(self._name, self._signal), args=(self._id, self._signal_expire)) | |
|
362 | if error == 1: | |
|
363 | raise NotAcquired("Lock %s is not acquired or it already expired." % self._name) | |
|
364 | elif error: | |
|
365 | raise RuntimeError("Unsupported error code %s from EXTEND script." % error) | |
|
366 | ||
|
367 | def locked(self): | |
|
368 | """ | |
|
369 | Return true if the lock is acquired. | |
|
370 | ||
|
371 | Checks that lock with same name already exists. This method returns true, even if | |
|
372 | lock have another id. | |
|
373 | """ | |
|
374 | return self._client.exists(self._name) == 1 | |
|
375 | ||
|
376 | ||
|
377 | reset_all_script = None | |
|
378 | ||
|
379 | ||
|
380 | def reset_all(redis_client): | |
|
381 | """ | |
|
382 | Forcibly deletes all locks if its remains (like a crash reason). Use this with care. | |
|
383 | ||
|
384 | :param redis_client: | |
|
385 | An instance of :class:`~StrictRedis`. | |
|
386 | """ | |
|
387 | Lock.register_scripts(redis_client) | |
|
388 | ||
|
389 | reset_all_script(client=redis_client) # noqa |
@@ -234,11 +234,17 b' class BaseRedisBackend(redis_backend.Red' | |||
|
234 | 234 | pipe.execute() |
|
235 | 235 | |
|
236 | 236 | def get_mutex(self, key): |
|
237 | u = redis_backend.u | |
|
238 | 237 | if self.distributed_lock: |
|
239 | lock_key = u('_lock_{0}').format(key) | |
|
238 | import redis_lock | |
|
239 | lock_key = redis_backend.u('_lock_{0}').format(key) | |
|
240 | 240 | log.debug('Trying to acquire Redis lock for key %s', lock_key) |
|
241 | return self.client.lock(lock_key, self.lock_timeout, self.lock_sleep) | |
|
241 | lock = redis_lock.Lock( | |
|
242 | redis_client=self.client, | |
|
243 | name=lock_key, | |
|
244 | expire=self.lock_timeout, | |
|
245 | auto_renewal=False, | |
|
246 | ) | |
|
247 | return lock | |
|
242 | 248 | else: |
|
243 | 249 | return None |
|
244 | 250 |
General Comments 0
You need to be logged in to leave comments.
Login now