Merge pull request #1112 from daveisfera/master

Add support for specifying the token when acquiring a lock

2 files changed, 41 insertions, 2 deletions

diff --git a/redis/lock.py b/redis/lock.py
index 2bb7794..c2e6887 100644
--- a/redis/lock.py
+++ b/redis/lock.py
@@ -149,7 +149,7 @@ class Lock(object):
     def __exit__(self, exc_type, exc_value, traceback):
         self.release()
 
-    def acquire(self, blocking=None, blocking_timeout=None):
+    def acquire(self, blocking=None, blocking_timeout=None, token=None):
         """
         Use Redis to hold a shared, distributed lock named ``name``.
         Returns True once the lock is acquired.
@@ -159,9 +159,13 @@ class Lock(object):
 
         ``blocking_timeout`` specifies the maximum number of seconds to
         wait trying to acquire the lock.
+
+        ``token`` specifies the value of the key used. Should have a unique
+        component to it when acquiring a lock with the same name.
         """
         sleep = self.sleep
-        token = uuid.uuid1().hex.encode()
+        if token is None:
+            token = uuid.uuid1().hex.encode()
         if blocking is None:
             blocking = self.blocking
         if blocking_timeout is None:
@@ -195,6 +199,13 @@ class Lock(object):
         """
         return self.redis.get(self.name) is not None
 
+    def owned(self):
+        """
+        Returns True if this key is locked by this lock, otherwise False.
+        """
+        return self.local.token is not None and \
+               self.redis.get(self.name) == self.local.token
+
     def release(self):
         "Releases the already acquired lock"
         expected_token = self.local.token
diff --git a/tests/test_lock.py b/tests/test_lock.py
index 945fe2f..42dadef 100644
--- a/tests/test_lock.py
+++ b/tests/test_lock.py
@@ -18,6 +18,16 @@ class TestLock(object):
         lock.release()
         assert r.get('foo') is None
 
+    def test_lock_token(self, r):
+        lock = self.get_lock(r, 'foo')
+        assert lock.acquire(blocking=False, token='test')
+        assert r.get('foo') == b'test'
+        assert lock.local.token == 'test'
+        assert r.ttl('foo') == -1
+        lock.release()
+        assert r.get('foo') is None
+        assert lock.local.token is None
+
     def test_locked(self, r):
         lock = self.get_lock(r, 'foo')
         assert lock.locked() is False
@@ -26,6 +36,24 @@ class TestLock(object):
         lock.release()
         assert lock.locked() is False
 
+    def test_owned(self, r):
+        lock = self.get_lock(r, 'foo')
+        assert lock.owned() is False
+        lock.acquire(blocking=False)
+        assert lock.owned() is True
+        lock.release()
+        assert lock.owned() is False
+
+        lock2 = self.get_lock(r, 'foo')
+        assert lock.owned() is False
+        assert lock2.owned() is False
+        lock2.acquire(blocking=False)
+        assert lock.owned() is False
+        assert lock2.owned() is True
+        lock2.release()
+        assert lock.owned() is False
+        assert lock2.owned() is False
+
     def test_competing_locks(self, r):
         lock1 = self.get_lock(r, 'foo')
         lock2 = self.get_lock(r, 'foo')