Merge pull request #24 from leandromoreira/extend-rebased

Rebased extend

lib/redlock/client.rb

@@ -34,10 +34,11 @@ def initialize(servers = DEFAULT_REDIS_URLS, options = {})
     # 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 to be executed; after its execution, the lock (if successfully
     # acquired) is automatically unlocked.
-    def lock(resource, ttl, &block)
-      lock_info = try_lock_instances(resource, ttl)
+    def lock(resource, ttl, extend: nil, &block)
+      lock_info = try_lock_instances(resource, ttl, extend)
 
       if block_given?
         begin
@@ -60,15 +61,12 @@ def unlock(lock_info)
 
     # Locks a resource, executing the received block only after successfully acquiring the lock,
     # and returning its return value as a result.
-    # Params:
-    # +resource+:: the resource (or key) string to be locked.
-    # +ttl+:: the time-to-live in ms for the lock.
-    # +block+:: block to be executed after successful lock acquisition.
-    def lock!(resource, ttl)
+    # See Redlock::Client#lock for parameters.
+    def lock!(*args, **keyword_args)
       fail 'No block passed' unless block_given?
 
-      lock(resource, ttl) do |lock_info|
-        raise LockError, "Could not acquire lock #{resource}" unless lock_info
+      lock(*args, **keyword_args) do |lock_info|
+        raise LockError, 'failed to acquire lock' unless lock_info
         return yield
       end
     end
@@ -84,6 +82,15 @@ class RedisInstance
         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
+      # and @maltoe for https://github.com/leandromoreira/redlock-rb/pull/20#discussion_r38903633
+      LOCK_SCRIPT = <<-eos
+        if redis.call("exists", KEYS[1]) == 0 or redis.call("get", KEYS[1]) == ARGV[1] then
+          return redis.call("set", KEYS[1], ARGV[1], "PX", ARGV[2])
+        end
+      eos
+
       def initialize(connection)
         if connection.respond_to?(:client)
           @redis = connection
@@ -95,7 +102,7 @@ def initialize(connection)
       end
 
       def lock(resource, val, ttl)
-        @redis.set(resource, val, nx: true, px: ttl)
+        @redis.evalsha(@lock_script_sha, keys: [resource], argv: [val, ttl])
       end
 
       def unlock(resource, val)
@@ -108,12 +115,13 @@ def unlock(resource, val)
 
       def load_scripts
         @unlock_script_sha = @redis.script(:load, UNLOCK_SCRIPT)
+        @lock_script_sha = @redis.script(:load, LOCK_SCRIPT)
       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
@@ -123,8 +131,8 @@ def try_lock_instances(resource, ttl)
       false
     end
 
-    def lock_instances(resource, ttl)
-      value = SecureRandom.uuid
+    def lock_instances(resource, ttl, extend)
+      value = extend ? extend.fetch(:value) : SecureRandom.uuid
 
       locked, time_elapsed = timed do
         @servers.select { |s| s.lock(resource, value, ttl) }.size

lib/redlock/testing.rb

@@ -4,17 +4,17 @@ class Client
 
     alias_method :try_lock_instances_without_testing, :try_lock_instances
 
-    def try_lock_instances(resource, ttl)
+    def try_lock_instances(resource, ttl, extend)
       if @testing_mode == :bypass
         {
           validity: ttl,
           resource: resource,
-          value: SecureRandom.uuid
+          value: extend ? extend.fetch(:value) : SecureRandom.uuid
         }
       elsif @testing_mode == :fail
         false
       else
-        try_lock_instances_without_testing resource, ttl
+        try_lock_instances_without_testing resource, ttl, extend
       end
     end
 
@@ -25,12 +25,13 @@ def unlock(lock_info)
     end
 
     class RedisInstance
+      alias_method :load_scripts_without_testing, :load_scripts
+
       def load_scripts
-        begin
-          @unlock_script_sha = @redis.script(:load, UNLOCK_SCRIPT)
-        rescue Redis::CommandError
-          # ignore
-        end
+        load_scripts_without_testing
+      rescue Redis::CommandError
+        # FakeRedis doesn't have #script, but doesn't need it either.
+        raise unless defined?(::FakeRedis)
       end
     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 somebody else's lock" 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