From 43c3fa4455c5a4d943b0ef9db68b111f52a1177a Mon Sep 17 00:00:00 2001
From: Gabriel Mazetto <brodock@gmail.com>
Date: Fri, 23 Jun 2017 07:36:19 +0200
Subject: [PATCH] Introduce #renew for ExclusiveLease

---
 lib/gitlab/exclusive_lease.rb           | 19 ++++++++++++++++++-
 spec/lib/gitlab/exclusive_lease_spec.rb | 13 +++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

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|
diff --git a/spec/lib/gitlab/exclusive_lease_spec.rb b/spec/lib/gitlab/exclusive_lease_spec.rb
index a366d68a146..81bbd70ffb8 100644
--- a/spec/lib/gitlab/exclusive_lease_spec.rb
+++ b/spec/lib/gitlab/exclusive_lease_spec.rb
@@ -19,6 +19,19 @@ describe Gitlab::ExclusiveLease, type: :redis do
     end
   end
 
+  describe '#renew' do
+    it 'returns true when we have the existing lease' do
+      lease = described_class.new(unique_key, timeout: 3600)
+      expect(lease.try_obtain).to be_present
+      expect(lease.renew).to be_truthy
+    end
+
+    it 'returns false when we dont have a lease' do
+      lease = described_class.new(unique_key, timeout: 3600)
+      expect(lease.renew).to be_falsey
+    end
+  end
+
   describe '#exists?' do
     it 'returns true for an existing lease' do
       lease = described_class.new(unique_key, timeout: 3600)