Introduce #renew for ExclusiveLease

1 files changed, 18 insertions, 1 deletions

diff --git a/lib/gitlab/exclusive_lease.rb b/lib/gitlab/exclusive_lease.rb
index 62ddd45785d..a0f46594eb1 100644
--- a/lib/gitlab/exclusive_lease.rb
+++ b/lib/gitlab/exclusive_lease.rb
@@ -10,13 +10,21 @@ module Gitlab
   # ExclusiveLease.
   #
   class ExclusiveLease
-    LUA_CANCEL_SCRIPT = <<-EOS.freeze
+    LUA_CANCEL_SCRIPT = <<~EOS.freeze
       local key, uuid = KEYS[1], ARGV[1]
       if redis.call("get", key) == uuid then
         redis.call("del", key)
       end
     EOS
 
+    LUA_RENEW_SCRIPT = <<~EOS.freeze
+      local key, uuid, ttl = KEYS[1], ARGV[1], ARGV[2]
+      if redis.call("get", key) == uuid then
+        redis.call("expire", key, ttl)
+        return uuid
+      end
+    EOS
+
     def self.cancel(key, uuid)
       Gitlab::Redis.with do |redis|
         redis.eval(LUA_CANCEL_SCRIPT, keys: [redis_key(key)], argv: [uuid])
@@ -42,6 +50,15 @@ module Gitlab
       end
     end
 
+    # Try to renew an existing lease. Return lease UUID on success,
+    # false if the lease is taken by a different UUID or inexistent.
+    def renew
+      Gitlab::Redis.with do |redis|
+        result = redis.eval(LUA_RENEW_SCRIPT, keys: [@redis_key], argv: [@uuid, @timeout])
+        result == @uuid
+      end
+    end
+
     # Returns true if the key for this lease is set.
     def exists?
       Gitlab::Redis.with do |redis|