Merge 7d1156f into aa8bd7a

.ruby-version

@@ -1 +1 @@
-ruby-2.1.2
+ruby-2.3.1

lib/ratelimit.rb

@@ -11,6 +11,7 @@ class Ratelimit
   # @option options [Integer] :bucket_interval (5) How many seconds each bucket represents
   # @option options [Integer] :bucket_expiry (@bucket_span) How long we keep data in each bucket before it is auto expired. Cannot be larger than the bucket_span.
   # @option options [Redis]   :redis (nil) Redis client if you need to customize connection options
+  # @option options [Lambda]  :checkout_redis_with (nil) Lambda that yields to a passed block with a Redis instance, for use with a Redis connection pool
   #
   # @return [RateLimit] RateLimit instance
   #
@@ -30,6 +31,7 @@ def initialize(key, options = {})
       raise ArgumentError.new("Cannot have less than 3 buckets")
     end
     @redis = options[:redis]
+    @checkout_redis_with = options[:checkout_redis_with]
   end
 
   # Add to the counter for a given subject.
@@ -41,12 +43,14 @@ def initialize(key, options = {})
   def add(subject, count = 1)
     bucket = get_bucket
     subject = "#{@key}:#{subject}"
-    redis.pipelined do
-      redis.hincrby(subject, bucket, count)
-      redis.hdel(subject, (bucket + 1) % @bucket_count)
-      redis.hdel(subject, (bucket + 2) % @bucket_count)
-      redis.expire(subject, @bucket_expiry)
-    end.first
+    use_redis do |r|
+      r.multi do
+        r.hincrby(subject, bucket, count)
+        r.hdel(subject, (bucket + 1) % @bucket_count)
+        r.hdel(subject, (bucket + 2) % @bucket_count)
+        r.expire(subject, @bucket_expiry)
+      end.first
+    end
   end
 
   # Returns the count for a given subject and interval
@@ -62,7 +66,10 @@ def count(subject, interval)
     keys = (0..count - 1).map do |i|
       (bucket - i) % @bucket_count
     end
-    return redis.hmget(subject, *keys).inject(0) {|a, i| a + i.to_i}
+
+    return use_redis do |r|
+      r.hmget(subject, *keys).inject(0) { |a, i| a + i.to_i }
+    end
   end
 
   # Check if the rate limit has been exceeded.
@@ -113,7 +120,19 @@ def get_bucket(time = Time.now.to_i)
     ((time % @bucket_span) / @bucket_interval).floor
   end
 
-  def redis
-    @redis ||= Redis::Namespace.new(:ratelimit, :redis => @redis || Redis.new)
+  def use_redis
+    if @checkout_redis_with
+      @checkout_redis_with.call { |redis| yield(namespaced_redis_instance(redis)) }
+    else
+      yield single_redis_instance
+    end
+  end
+
+  def single_redis_instance
+    @redis ||= namespaced_redis_instance(@redis || Redis.new)
+  end
+
+  def namespaced_redis_instance(redis)
+    Redis::Namespace.new(:ratelimit, redis: redis)
   end
 end

spec/ratelimit_spec.rb

@@ -1,10 +1,9 @@
 require 'spec_helper'
 
 describe Ratelimit do
-
   before do
     @r = Ratelimit.new("key")
-    @r.send(:redis).flushdb
+    @r.send(:use_redis) { |r| r.flushdb }
   end
 
   it "should set_bucket_expiry to the bucket_span if not defined" do
@@ -106,4 +105,20 @@
     expect(@r.count('value1', 10)).to eql(1)
   end
 
+  context "using the checkout_redis_with option" do
+    let(:redis) { Redis.new }
+    let(:redis_checkout_lambda) do
+      lambda do |&block|
+        block.call(redis)
+      end
+    end
+
+    let(:key) { SecureRandom.hex(6) }
+    subject { Ratelimit.new(key, checkout_redis_with: redis_checkout_lambda) }
+
+    it "adds correctly" do
+      subject.add("value1", 3)
+      expect(subject.count("value1", 1)).to eq(3)
+    end
+  end
 end