From df3cb941f4e61968743fbf05c35783ae97d1e308 Mon Sep 17 00:00:00 2001
From: Jon Phillips <jphillips@totaldiscovery.com>
Date: Thu, 19 May 2016 10:38:48 -0700
Subject: [PATCH] Added extend semantics that fails when the lock isn't owned

---
 README.md             | 16 ++++++++++++++++
 lib/redlock/client.rb | 42 ++++++++++++++++++++++++++++++++++++++++++
 spec/client_spec.rb   | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 92 insertions(+)

diff --git a/README.md b/README.md
index 0b15539..55eaac0 100644
--- a/README.md
+++ b/README.md
@@ -95,6 +95,22 @@ rescue Redlock::LockError
 end
 ```
 
+To extend the life of the lock, provided that you didn't let it expire:
+
+```ruby
+begin
+  block_result = lock_manager.lock!("resource_key", 2000) do |lock_info|
+    # critical code
+    lock_manager.extend_life!(lock_info, 3000)
+    # more critical code
+  end
+rescue Redlock::LockError
+  # error handling
+end
+```
+There's also a non-bang version that returns true when the lock was
+extended
+
 ## Run tests
 
 Make sure you have at least 1 redis instances up.
diff --git a/lib/redlock/client.rb b/lib/redlock/client.rb
index 4aaac42..e4e7a38 100644
--- a/lib/redlock/client.rb
+++ b/lib/redlock/client.rb
@@ -52,6 +52,30 @@ def lock(resource, ttl, extend: nil, &block)
       end
     end
 
+    def extend_life(to_extend, ttl)
+      value = to_extend.fetch(:value)
+      resource = to_extend.fetch(:resource)
+
+
+      extended, time_elapsed = timed do
+        @servers.all? { |s| s.extend_life(resource, value, ttl) }
+      end
+
+      validity = ttl - time_elapsed - drift(ttl)
+
+      if extended
+        { validity: validity, resource: resource, value: value }
+      else
+        @servers.each { |s| s.unlock(resource, value) }
+        false
+      end
+    end
+
+    def extend_life!(to_extend, ttl)
+      new_lock_info = self.extend_life(to_extend, ttl)
+      raise LockError, 'failed to extend lock' unless new_lock_info
+    end
+
     # Unlocks a resource.
     # Params:
     # +lock_info+:: the lock that has been acquired when you locked the resource.
@@ -91,6 +115,16 @@ class RedisInstance
         end
       eos
 
+      EXTEND_LOCK_SCRIPT = <<-eos
+        if redis.call("get", KEYS[1]) == ARGV[1] then
+          redis.call("expire", KEYS[1], ARGV[2])
+          return 0
+        else
+          return 1
+        end
+      eos
+
+
       def initialize(connection)
         if connection.respond_to?(:client)
           @redis = connection
@@ -107,6 +141,13 @@ def lock(resource, val, ttl)
         end
       end
 
+      def extend_life(resource, val, ttl)
+        recover_from_script_flush do
+          rc = @redis.evalsha @extend_lock_script_sha, keys: [resource], argv: [val, ttl]
+          rc == 0
+        end
+      end
+
       def unlock(resource, val)
         recover_from_script_flush do
           @redis.evalsha @unlock_script_sha, keys: [resource], argv: [val]
@@ -120,6 +161,7 @@ def unlock(resource, val)
       def load_scripts
         @unlock_script_sha = @redis.script(:load, UNLOCK_SCRIPT)
         @lock_script_sha = @redis.script(:load, LOCK_SCRIPT)
+        @extend_lock_script_sha = @redis.script(:load, EXTEND_LOCK_SCRIPT)
       end
 
       def recover_from_script_flush
diff --git a/spec/client_spec.rb b/spec/client_spec.rb
index 7566b42..ad7c389 100644
--- a/spec/client_spec.rb
+++ b/spec/client_spec.rb
@@ -170,6 +170,40 @@
     end
   end
 
+  describe "extend" do
+    context 'when lock is available' do
+      before { @lock_info = lock_manager.lock(resource_key, ttl) }
+      after(:each) { lock_manager.unlock(@lock_info) if @lock_info }
+
+      it 'can extend its own lock' do
+        lock_info = lock_manager.extend_life(@lock_info, ttl)
+        expect(lock_info).to be_lock_info_for(resource_key)
+      end
+
+      it "can't extend a nonexistent lock" do
+        lock_manager.unlock(@lock_info)
+        lock_info = lock_manager.extend_life(@lock_info, ttl)
+        expect(lock_info).to eq(false)
+      end
+    end
+  end
+
+  describe "extend!" do
+    context 'when lock is available' do
+      before { @lock_info = lock_manager.lock(resource_key, ttl) }
+      after(:each) { lock_manager.unlock(@lock_info) if @lock_info }
+
+      it 'can extend its own lock' do
+        expect{ lock_manager.extend_life!(@lock_info, ttl) }.to_not raise_error
+      end
+
+      it "can't extend a nonexistent lock" do
+        lock_manager.unlock(@lock_info)
+        expect{ lock_manager.extend_life!(@lock_info, ttl) }.to raise_error(Redlock::LockError)
+      end
+    end
+  end
+
   describe 'lock!' do
     context 'when lock is available' do
       it 'locks' do