locks.py

#!/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()