##// END OF EJS Templates
redis: added better distributed lock implementation.
super-admin -
r4705:dce8d447 stable
parent child Browse files
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
@@ -286,11 +286,17 b' class BaseRedisBackend(redis_backend.Red'
286 286 pipe.execute()
287 287
288 288 def get_mutex(self, key):
289 u = redis_backend.u
290 289 if self.distributed_lock:
291 lock_key = u('_lock_{0}').format(key)
290 import redis_lock
291 lock_key = redis_backend.u('_lock_{0}').format(key)
292 292 log.debug('Trying to acquire Redis lock for key %s', lock_key)
293 return self.client.lock(lock_key, self.lock_timeout, self.lock_sleep)
293 lock = redis_lock.Lock(
294 redis_client=self.client,
295 name=lock_key,
296 expire=self.lock_timeout,
297 auto_renewal=False,
298 )
299 return lock
294 300 else:
295 301 return None
296 302
General Comments 0
You need to be logged in to leave comments. Login now