Merge 03021a9 into 365ff23

lib/traffic_jam.rb

@@ -4,6 +4,7 @@
 require_relative 'traffic_jam/configuration'
 require_relative 'traffic_jam/limit'
 require_relative 'traffic_jam/limit_group'
+require_relative 'traffic_jam/simple_limit'
 require_relative 'traffic_jam/rolling_limit'
 require_relative 'traffic_jam/lifetime_limit'
 

lib/traffic_jam/errors.rb

@@ -10,5 +10,8 @@ def initialize(limit)
         @limit = limit
       end
     end
+
+    class InvalidKeyError < StandardError; end
+    class UnknownReturnValue < StandardError; end
   end
 end

lib/traffic_jam/limit.rb

@@ -168,9 +168,13 @@ def key
           end
         hash = Digest::MD5.base64digest(converted_value.to_s)
         hash = hash[0...config.hash_length]
-        @key = "#{config.key_prefix}:#{action}:#{hash}"
+        @key = "#{key_prefix}:#{action}:#{hash}"
       end
       @key
     end
+
+    def key_prefix
+      config.key_prefix
+    end
   end
 end

lib/traffic_jam/scripts.rb

@@ -10,6 +10,8 @@ def self.load(name)
 
     INCREMENT_SCRIPT = load('increment')
     INCREMENT_SCRIPT_HASH = Digest::SHA1.hexdigest(INCREMENT_SCRIPT)
+    INCREMENT_SIMPLE = load('increment_simple')
+    INCREMENT_SIMPLE_HASH = Digest::SHA1.hexdigest(INCREMENT_SIMPLE)
     INCREMENT_ROLLING = load('increment_rolling')
     INCREMENT_ROLLING_HASH = Digest::SHA1.hexdigest(INCREMENT_ROLLING)
     INCRBY = load('incrby')

lib/traffic_jam/simple_limit.rb

@@ -0,0 +1,91 @@
+require_relative 'limit'
+require_relative 'scripts'
+
+module TrafficJam
+  # A SimpleLimit is a limit type that is more efficient for increments but does
+  # not support decrements or changing the max value without a complete reset.
+  # This means that if the period or max value for an action, value key changes,
+  # the used and remaining values cannot be preserved.
+  #
+  # This works by storing a key in Redis with a millisecond-precision expiry
+  # representing the time that the limit will be completely reset. Each
+  # increment operation converts the increment amount into the number of
+  # milliseconds to be added to the expiry.
+  #
+  # Example: Limit is 5 per 10 seconds.
+  #     An increment by 1 first sets the key to expire in 2s.
+  #     Another immediate increment by 4 sets the expiry to 10s.
+  #     Subsequent increments fail until clock time catches up to expiry
+  class SimpleLimit < Limit
+    # Increment the amount used by the given number. Does not perform increment
+    # if the operation would exceed the limit. Returns whether the operation was
+    # successful.
+    #
+    # @param amount [Integer] amount to increment by
+    # @param time [Time] time is ignored
+    # @return [Boolean] true if increment succeded and false if incrementing
+    #   would exceed the limit
+    def increment(amount = 1, time: Time.now)
+      return amount.zero? if max.zero?
+      raise ArgumentError.new("Amount must be positive") if amount < 0
+
+      if amount != amount.to_i
+        raise ArgumentError.new("Amount must be an integer")
+      end
+
+      incrby = (period * 1000 * amount / max).to_i
+      argv = [incrby, period * 1000]
+
+      result =
+        begin
+          redis.evalsha(
+            Scripts::INCREMENT_SIMPLE_HASH, keys: [key], argv: argv)
+        rescue Redis::CommandError
+          redis.eval(Scripts::INCREMENT_SIMPLE, keys: [key], argv: argv)
+        end
+
+      case result
+      when 0
+        return true
+      when -1
+        raise Errors::InvalidKeyError, "Redis key #{key} has no expire time set"
+      when -2
+        return false
+      else
+        raise Errors::UnknownReturnValue,
+              "Received unexpected return value #{result} from " \
+              "increment_simple eval"
+      end
+    end
+
+    # Decrement the amount used by the given number.
+    #
+    # @param amount [Integer] amount to decrement by
+    # @param time [Time] time is ignored
+    # @raise [NotImplementedError] decrement is not defined for SimpleLimit
+    def decrement(_amount = 1, time: Time.now)
+      raise NotImplementedError, "decrement is not defined for SimpleLimit"
+    end
+
+    # Return amount of limit used, taking time drift into account.
+    #
+    # @return [Integer] amount used
+    def used
+      return 0 if max.zero?
+
+      expiry = redis.pttl(key)
+      case expiry
+      when -1  # key exists but has no associated expire
+        raise Errors::InvalidKeyError, "Redis key #{key} has no expire time set"
+      when -2  # key does not exist
+        return 0
+      end
+
+      (max * expiry / (period * 1000.0)).ceil
+    end
+
+    def key_prefix
+      "#{config.key_prefix}:s"
+    end
+  end
+end

scripts/increment_simple.lua

@@ -0,0 +1,20 @@
+local arg_incrby = tonumber(ARGV[1])
+local arg_max = tonumber(ARGV[2])
+
+local old_value = redis.call("PTTL", KEYS[1])
+if old_value == -1 then -- key exists but has no associated expire
+   return -1 -- -1 signals key exists but has no associated expire
+elseif old_value == -2 then -- key does not exist
+   if arg_incrby > arg_max then
+      return -2 -- -2 signals increment exceeds max
+   end
+   redis.call("SET", KEYS[1], "", "PX", arg_incrby)
+else
+   local new_value = old_value + arg_incrby
+   if new_value > arg_max then
+      return -2 -- -2 signals increment exceeds max
+   end
+   redis.call("PEXPIRE", KEYS[1], new_value)
+end
+
+return 0 -- 0 signals success

spec/traffic_jam/simple_limit_spec.rb

@@ -0,0 +1,78 @@
+require_relative '../spec_helper'
+
+describe TrafficJam do
+  include RedisHelper
+
+  TrafficJam.configure do |config|
+    config.redis = RedisHelper.redis
+  end
+
+  let(:period) { 0.1 }
+  let(:limit) do
+    TrafficJam::SimpleLimit.new(:test, "user1", max: 3, period: period)
+  end
+
+  describe :increment do
+    after do
+      Spy.teardown
+    end
+
+    it "should be true when rate limit is not exceeded" do
+      assert limit.increment(1)
+    end
+
+    it "should be false when raise limit is exceeded" do
+      assert !limit.increment(4)
+      assert limit.increment(1)
+      assert limit.increment(2)
+      assert !limit.increment(1)
+    end
+
+    it "should raise an argument error if given a float" do
+      assert_raises(ArgumentError) do
+        limit.increment(1.5)
+      end
+    end
+
+    it "should be a no-op when limit would be exceeded" do
+      limit.increment(2)
+      assert !limit.increment(2)
+      assert limit.increment(1)
+    end
+
+    it "should be true when sufficient time passes" do
+      assert limit.increment(3)
+      sleep(period / 2)
+      assert limit.increment(1)
+      sleep(period * 2)
+      assert limit.increment(3)
+    end
+
+    describe "when max is zero" do
+      let(:limit) do
+        TrafficJam::SimpleLimit.new(:test, "user1", max: 0, period: period)
+      end
+
+      it "should be false for any positive amount" do
+        assert !limit.increment
+      end
+    end
+  end
+
+  describe :used do
+    it "should be 0 when there has been no incrementing" do
+      assert_equal 0, limit.used
+    end
+
+    it "should be the amount used" do
+      limit.increment(1)
+      assert_equal 1, limit.used
+    end
+
+    it "should decrease over time" do
+      limit.increment(2)
+      sleep(period / 2)
+      assert_equal 1, limit.used
+    end
+  end
+end

spec/traffic_jam_limit_spec.rb

@@ -36,6 +36,7 @@
     end
 
     it "should be false when raise limit is exceeded" do
+      assert !limit.increment(4)
       assert limit.increment(1)
       assert limit.increment(2)
       assert !limit.increment(1)