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 | pipe.execute() |
|
234 | pipe.execute() | |
235 |
|
235 | |||
236 | def get_mutex(self, key): |
|
236 | def get_mutex(self, key): | |
237 | u = redis_backend.u |
|
|||
238 | if self.distributed_lock: |
|
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 | log.debug('Trying to acquire Redis lock for key %s', lock_key) |
|
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 | else: |
|
248 | else: | |
243 | return None |
|
249 | return None | |
244 |
|
250 |
General Comments 0
You need to be logged in to leave comments.
Login now