import time
class RedisLock(object):
"""
Distributed locking using Redis SETNX and GETSET.
Usage::
try:
with RedisLock('my_lock', expires=3600, timeout=0):
# Critical section
except RedisLockTimeout:
# Lock could not be acquired.
:param expires To keep crashed threads from hanging everything,
we consider locks expired after ``expires`` seconds.
This value should be at least 2X higher than it
takes the critical section to execute. A value of
0 means the lock will never expire.
Default: 3600 (1 hour)
:param timeout If we cannot obtain the lock, keep trying for up
to ``timeout`` seconds. A value of 0 will only
try once.
Default: 0
"""
def __init__(self, redis, key, expires=3600, timeout=0):
self.redis = redis
self.key = key
self.expires = expires
self.expires_at = -1
self.timeout = timeout
def __enter__(self):
timeout = self.timeout
while True:
expires = time.time() + self.expires + 1
if self.redis.setnx(self.key, expires):
# Got the lock. High five!
self.expires_at = expires
return
# A lock is set, get the current value to make sure it's not expired.
current_value = self.redis.get(self.key)
if (current_value and float(current_value) < time.time() and
self.redis.getset(self.key, expires) == current_value):
# The lock is expired and we nobody beat us to it.
self.expires_at = expires
return
if timeout > 0:
timeout -= 1
time.sleep(1)
else:
# Someone else has the lock.
raise RedisLockTimeout(self.key)
def __exit__(self, exc_type, exc_value, traceback):
# Theoretically, this is our lock. We need to delete it, but if we
# happen to do it after we expire then we might delete someone
# else's lock. To avoid this, we check that it is not expired yet.
# This is not perfect, but the chances of deleting the wrong key
# are low.
current_value = float(self.redis.get(self.key) or 0)
if time.time() < current_value:
self.redis.delete(self.key)
class RedisLockTimeout(BaseException):
pass
|
Loading...