Merge e09f327 into 75bc721

lib/redlock/client.rb

@@ -38,8 +38,9 @@ def testing=(mode)
     # Params:
     # +resource+:: the resource (or key) string to be locked.
     # +ttl+:: The time-to-live in ms for the lock.
+    # +extend+: A lock ("lock_info") to extend.
     # +block+:: an optional block that automatically unlocks the lock.
-    def lock(resource, ttl, &block)
+    def lock(resource, ttl, extend: nil, &block)
       if @testing_mode == :bypass
         lock_info = {
           validity: ttl,
@@ -49,7 +50,7 @@ def lock(resource, ttl, &block)
       elsif @testing_mode == :fail
         lock_info = false
       else
-        lock_info = try_lock_instances(resource, ttl)
+        lock_info = try_lock_instances(resource, ttl, extend)
       end
 
       if block_given?
@@ -83,6 +84,17 @@ class RedisInstance
           return 0
         end
       eos
+      # thanks to https://github.com/sbertrang/redis-distlock/blob/master/lib/Redis/DistLock.pm
+      # also https://github.com/sbertrang/redis-distlock/issues/2 which proposes the value-checking
+      EXTEND_SCRIPT = <<-eos
+        if redis.call( "get", KEYS[1] ) == ARGV[1] then
+          if redis.call( "set", KEYS[1], ARGV[1], "XX", "PX", ARGV[2] ) then
+            return "OK"
+          end
+        else
+          return redis.call( "set", KEYS[1], ARGV[1], "NX", "PX", ARGV[2] )
+        end
+      eos
 
       def initialize(connection)
         if connection.respond_to?(:client)
@@ -92,10 +104,15 @@ def initialize(connection)
         end
 
         @unlock_script_sha = @redis.script(:load, UNLOCK_SCRIPT)
+        @extend_script_sha = @redis.script(:load, EXTEND_SCRIPT)
       end
 
-      def lock(resource, val, ttl)
-        @redis.set(resource, val, nx: true, px: ttl)
+      def lock(resource, val, ttl, extend)
+        if extend
+          @redis.evalsha(@extend_script_sha, keys: [resource], argv: [extend[:value], ttl])
+        else
+          @redis.set(resource, val, nx: true, px: ttl)
+        end
       end
 
       def unlock(resource, val)
@@ -105,9 +122,9 @@ def unlock(resource, val)
       end
     end
 
-    def try_lock_instances(resource, ttl)
+    def try_lock_instances(resource, ttl, extend)
       @retry_count.times do
-        lock_info = lock_instances(resource, ttl)
+        lock_info = lock_instances(resource, ttl, extend)
         return lock_info if lock_info
 
         # Wait a random delay before retrying
@@ -117,19 +134,20 @@ def try_lock_instances(resource, ttl)
       false
     end
 
-    def lock_instances(resource, ttl)
+    def lock_instances(resource, ttl, extend)
       value = SecureRandom.uuid
 
       locked, time_elapsed = timed do
-        @servers.select { |s| s.lock(resource, value, ttl) }.size
+        @servers.select { |s| s.lock(resource, value, ttl, extend) }.size
       end
 
       validity = ttl - time_elapsed - drift(ttl)
+      used_value = extend ? extend[:value] : value
 
       if locked >= @quorum && validity >= 0
-        { validity: validity, resource: resource, value: value }
+        { validity: validity, resource: resource, value: used_value }
       else
-        @servers.each { |s| s.unlock(resource, value) }
+        @servers.each { |s| s.unlock(resource, used_value) }
         false
       end
     end

spec/client_spec.rb

@@ -35,6 +35,25 @@
 
         expect(@lock_info).to be_lock_info_for(resource_key)
       end
+
+      it 'can extend its own lock' do
+        my_lock_info = lock_manager.lock(resource_key, ttl)
+        @lock_info = lock_manager.lock(resource_key, ttl, extend: my_lock_info)
+        expect(@lock_info).to be_lock_info_for(resource_key)
+        expect(@lock_info[:value]).to eq(my_lock_info[:value])
+      end
+
+      it "sets the given value when trying to extend a non-existent lock" do
+        @lock_info = lock_manager.lock(resource_key, ttl, extend: {value: 'hello world'})
+        expect(@lock_info).to be_lock_info_for(resource_key)
+        expect(@lock_info[:value]).to eq('hello world') # really we should test what's in redis
+      end
+
+      it "doesn't extend lock by default" do
+        @lock_info = lock_manager.lock(resource_key, ttl)
+        second_attempt = lock_manager.lock(resource_key, ttl)
+        expect(second_attempt).to eq(false)
+      end
     end
 
     context 'when lock is not available' do
@@ -46,6 +65,12 @@
 
         expect(lock_info).to eql(false)
       end
+
+      it "can't extend somebody else's lock" do
+        yet_another_lock_info = @another_lock_info.merge value: 'gibberish'
+        lock_info = lock_manager.lock(resource_key, ttl, extend: yet_another_lock_info)
+        expect(lock_info).to eql(false)
+      end
     end
 
     describe 'block syntax' do