Skip to content

Instantly share code, notes, and snippets.

@karanlyons karanlyons/locks.py
Created Dec 27, 2015

Embed
What would you like to do?
#!/usr/bin/env python2.7
# -*- coding: utf-8 -*-
from __future__ import division, absolute_import, print_function, unicode_literals
from contextlib import contextmanager
from os import getpid
from random import getrandbits
from django_redis import get_redis_connection
redis = get_redis_connection()
release_lock = redis.register_script('''
if redis.call('get', KEYS[1]) == ARGV[1] then
return redis.call('del', KEYS[1])
else
return 0
end
''')
class LockException(Exception):
pass
class Lock(object):
def __init__(self, key, duration=0, force=False, silent=False):
self.key = key
self.true_key = 'lock:{}'.format(self.key)
self.duration = duration
self.force = force
self.silent = silent
self.nonce = None
self.acquired = False
def acquire(self):
self.nonce = ':'.join([unicode(getpid()), unicode(getrandbits(32))])
self.acquired = bool(redis.set(self.true_key, self.nonce, nx=not self.force, ex=self.duration))
if not self.acquired:
self.nonce = None
if self.silent:
return False
else:
raise LockException('Could not acquire lock for {key}'.format(key=self.key))
else:
return True
def release(self):
if self.acquired and self.nonce:
result = release_lock(keys=(self.true_key,), args=(self.nonce,))
self.nonce = None
self.acquired = False
return bool(result)
else:
return None
class Locks(object):
def __init__(self, keys, duration=0, force=False):
self.keys = keys
self.duration = duration
self.force = force
self.acquired = False
self.locks = [Lock(key, self.duration, self.force) for key in self.keys]
def acquire(self):
try:
map(lambda l: l.acquire(), self.locks)
self.acquired = all(l.acquired for l in self.locks)
except LockException:
map(lambda l: l.release(), self.locks)
raise
def release(self):
if self.acquired:
result = reduce(lambda x, n: x + n, map(lambda l: l.release(), self.locks))
self.acquired = False
return result
else:
return None
@contextmanager
def lock(key, duration=0, force=False):
try:
l = Lock(key, duration, force)
l.acquire()
yield l
finally:
l.release()
@contextmanager
def locks(keys, duration=0, force=False):
try:
l = Locks(keys, duration, force)
l.acquire()
yield l
finally:
l.release()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.