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