Merge pull request #87 from fatkodima/connection-pool

Support ConnectionPool objects

Gemfile.lock

@@ -7,6 +7,7 @@ PATH
 GEM
   remote: https://rubygems.org/
   specs:
+    connection_pool (2.2.2)
     coveralls (0.8.22)
       json (>= 1.8, < 3)
       simplecov (~> 0.16.1)
@@ -45,6 +46,7 @@ PLATFORMS
   ruby
 
 DEPENDENCIES
+  connection_pool (~> 2.2)
   coveralls (~> 0.8)
   rake (~> 11.1, >= 11.1.2)
   redlock!

lib/redlock/client.rb

@@ -122,27 +122,38 @@ class RedisInstance
         end
       eos
 
+      module ConnectionPoolLike
+        def with
+          yield self
+        end
+      end
+
       def initialize(connection)
-        if connection.respond_to?(:client)
+        if connection.respond_to?(:with)
           @redis = connection
         else
-          @redis = Redis.new(connection)
+          if connection.respond_to?(:client)
+            @redis = connection
+          else
+            @redis = Redis.new(connection)
+          end
+          @redis.extend(ConnectionPoolLike)
         end
 
         load_scripts
       end
 
       def lock(resource, val, ttl, allow_new_lock)
         recover_from_script_flush do
-          @redis.evalsha @lock_script_sha, keys: [resource], argv: [val, ttl, allow_new_lock]
+          @redis.with { |conn| conn.evalsha @lock_script_sha, keys: [resource], argv: [val, ttl, allow_new_lock] }
         end
       rescue Redis::BaseConnectionError
         false
       end
 
       def unlock(resource, val)
         recover_from_script_flush do
-          @redis.evalsha @unlock_script_sha, keys: [resource], argv: [val]
+          @redis.with { |conn| conn.evalsha @unlock_script_sha, keys: [resource], argv: [val] }
         end
       rescue
         # Nothing to do, unlocking is just a best-effort attempt.
@@ -151,8 +162,8 @@ def unlock(resource, val)
       private
 
       def load_scripts
-        @unlock_script_sha = @redis.script(:load, UNLOCK_SCRIPT)
-        @lock_script_sha = @redis.script(:load, LOCK_SCRIPT)
+        @unlock_script_sha = @redis.with { |conn| conn.script(:load, UNLOCK_SCRIPT) }
+        @lock_script_sha = @redis.with { |conn| conn.script(:load, LOCK_SCRIPT) }
       end
 
       def recover_from_script_flush

redlock.gemspec

@@ -23,4 +23,5 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency "coveralls", "~> 0.8"
   spec.add_development_dependency 'rake', '~> 11.1', '>= 11.1.2'
   spec.add_development_dependency 'rspec', '~> 3', '>= 3.0.0'
+  spec.add_development_dependency 'connection_pool', '~> 2.2'
 end

spec/client_spec.rb

@@ -1,6 +1,7 @@
 require 'spec_helper'
 require 'securerandom'
 require 'redis'
+require 'connection_pool'
 
 RSpec.describe Redlock::Client do
   # It is recommended to have at least 3 servers in production
@@ -25,6 +26,15 @@
 
       expect(redlock_servers).to match_array([redis1_host, redis2_host])
     end
+
+    it 'accepts ConnectionPool objects' do
+      pool = ConnectionPool.new { Redis.new(url: "redis://#{redis1_host}:#{redis1_port}") }
+      redlock = Redlock::Client.new([pool])
+
+      lock_info = lock_manager.lock(resource_key, ttl)
+      expect(resource_key).to_not be_lockable(lock_manager, ttl)
+      lock_manager.unlock(lock_info)
+    end
   end
 
   describe 'lock' do
@@ -166,7 +176,7 @@
       it 'does not raise an error on connection issues' do
         # We re-route the lock manager to a (hopefully) non-existent Redis URL.
         redis_instance = lock_manager.instance_variable_get(:@servers).first
-        redis_instance.instance_variable_set(:@redis, Redis.new(url: 'redis://localhost:46864'))
+        redis_instance.instance_variable_set(:@redis, unreachable_redis)
 
         expect {
           expect(lock_manager.lock(resource_key, ttl)).to be_falsey
@@ -178,13 +188,22 @@
       it 'recovers from connection issues' do
         # Same as above.
         redis_instance = lock_manager.instance_variable_get(:@servers).first
-        redis_instance.instance_variable_set(:@redis, Redis.new(url: 'redis://localhost:46864'))
+        old_redis = redis_instance.instance_variable_get(:@redis)
+        redis_instance.instance_variable_set(:@redis, unreachable_redis)
         expect(lock_manager.lock(resource_key, ttl)).to be_falsey
-        redis_instance.instance_variable_set(:@redis, Redis.new(url: "redis://#{redis1_host}:#{redis1_port}"))
+        redis_instance.instance_variable_set(:@redis, old_redis)
         expect(lock_manager.lock(resource_key, ttl)).to be_truthy
       end
     end
 
+    def unreachable_redis
+      redis = Redis.new(url: 'redis://localhost:46864')
+      def redis.with
+        yield self
+      end
+      redis
+    end
+
     context 'when script cache has been flushed' do
       before(:each) do
         @manipulated_instance = lock_manager.instance_variable_get(:@servers).first