From ff120df78ccd85d6e2e2938ee02d1eb831676724 Mon Sep 17 00:00:00 2001
From: Ihor Kalnytskyi <ihor@kalnytskyi.com>
Date: Tue, 24 Jul 2018 19:14:50 +0300
Subject: Add `.reacquire()` method to Lock

`Lock` class provides a method called `.extend()` to manage a TTL of the
acquired lock. However, the method allows only to extend a timeout of
existing lock by N seconds, there's no way you can reset a TTL to the
timeout value you passed to this lock. There could be multiple use cases
for such behaviour. For instance, one may want to use a lock to
implement active/passive behaviour where only one process owns a lock
and resets its TTL all over again until it dies. This commit adds a new
method called `.reacquire()` to reset a TTL of the acquired lock back to
the passed timeout value.
---
 redis/lock.py | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

(limited to 'redis/lock.py')

diff --git a/redis/lock.py b/redis/lock.py
index 5524d7a..2bb7794 100644
--- a/redis/lock.py
+++ b/redis/lock.py
@@ -16,6 +16,7 @@ class Lock(object):
 
     lua_release = None
     lua_extend = None
+    lua_reacquire = None
 
     # KEYS[1] - lock name
     # ARGS[1] - token
@@ -49,6 +50,19 @@ class Lock(object):
         return 1
     """
 
+    # KEYS[1] - lock name
+    # ARGS[1] - token
+    # ARGS[2] - milliseconds
+    # return 1 if the locks time was reacquired, otherwise 0
+    LUA_REACQUIRE_SCRIPT = """
+        local token = redis.call('get', KEYS[1])
+        if not token or token ~= ARGV[1] then
+            return 0
+        end
+        redis.call('pexpire', KEYS[1], ARGV[2])
+        return 1
+    """
+
     def __init__(self, redis, name, timeout=None, sleep=0.1,
                  blocking=True, blocking_timeout=None, thread_local=True):
         """
@@ -121,6 +135,9 @@ class Lock(object):
             cls.lua_release = client.register_script(cls.LUA_RELEASE_SCRIPT)
         if cls.lua_extend is None:
             cls.lua_extend = client.register_script(cls.LUA_EXTEND_SCRIPT)
+        if cls.lua_reacquire is None:
+            cls.lua_reacquire = \
+                client.register_script(cls.LUA_REACQUIRE_SCRIPT)
 
     def __enter__(self):
         # force blocking, as otherwise the user would have to check whether
@@ -214,3 +231,22 @@ class Lock(object):
             raise LockNotOwnedError("Cannot extend a lock that's"
                                     " no longer owned")
         return True
+
+    def reacquire(self):
+        """
+        Resets a TTL of an already acquired lock back to a timeout value.
+        """
+        if self.local.token is None:
+            raise LockError("Cannot reacquire an unlocked lock")
+        if self.timeout is None:
+            raise LockError("Cannot reacquire a lock with no timeout")
+        return self.do_reacquire()
+
+    def do_reacquire(self):
+        timeout = int(self.timeout * 1000)
+        if not bool(self.lua_reacquire(keys=[self.name],
+                                       args=[self.local.token, timeout],
+                                       client=self.redis)):
+            raise LockNotOwnedError("Cannot reacquire a lock that's"
+                                    " no longer owned")
+        return True
-- 
cgit v1.2.1