traffic_jam 1.0.1 → 1.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: dea90b9ea6d6cdd3f95e0c80541bf26fce1b3690
-  data.tar.gz: 6870c79d37669b7997aea051b45bffe041a176b8
+  metadata.gz: dc24485cfcc470e68b80dbb53f6b135ab9879bf0
+  data.tar.gz: ab8ea5c0a5e1075af9606fae0e4a0cc9bb1fc071
 SHA512:
-  metadata.gz: 79106d2df976e9826ecc5a1b8a7ae3882a1091a79c82f2193f0ea4b9a79bd83b0fa3c35e07e822a77609e25c7b2cb29505c204baae5b5248d5bd24bdc54e669d
-  data.tar.gz: 938bf7f0fcb98d9677015495399078358f50defc35962203f9de833d0cfea646ea0f86a10ac48ef96b6a28f70ba8d486d99f0c84175c662f1f4946d0a4026fc8
+  metadata.gz: f428099b9b945e9196a4e54146a79af29d8e8e8c430320f1e9237be2e96df3b9f39e29cad2b78c53694f9c76e85ee698f38af91b23ab715339ab9f578eaa4434
+  data.tar.gz: 7c9afc9302a854fa68c042c9329e73a6224858bec71e25fb0ff5ac2a72fdcd6ab7f5aeedba4af936242a20e5fb96bc4fc87dcacdf9a77b3291f582adb5ddd540

data/lib/traffic_jam/errors.rb

@@ -10,5 +10,8 @@ module TrafficJam
         @limit = limit
       end
     end
+
+    class InvalidKeyError < StandardError; end
+    class UnknownReturnValue < StandardError; end
   end
 end

data/lib/traffic_jam/lifetime_limit.rb

@@ -0,0 +1,53 @@
+module TrafficJam
+  # This class represents a lifetime limit on an action, value pair. For example, if
+  # limiting the amount of money a user can transfer, the action could be
+  # +:transfers+ and the value would be the user ID. The class exposes atomic
+  # increment operations and allows querying of the current amount used and
+  # amount remaining.
+  class LifetimeLimit < Limit
+    # Constructor takes an action name as a symbol, a maximum cap, and the
+    # period of limit. +max+ and +period+ are required keyword arguments.
+    #
+    # @param action [Symbol] action name
+    # @param value [String] limit target value
+    # @param max [Integer] required limit maximum
+    # @raise [ArgumentError] if max is nil
+    def initialize(action, value, max: nil)
+      super(action, value, max: max, period: -1)
+    end
+
+    # 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
+    # @return [Boolean] true if increment succeded and false if incrementing
+    #   would exceed the limit
+    def increment(amount = 1, time: Time.now)
+      raise ArgumentError, 'Amount must be an integer' if amount != amount.to_i
+      return amount <= 0 if max.zero?
+
+      !!run_script([amount.to_i, max])
+    end
+
+    # Return amount of limit used
+    #
+    # @return [Integer] amount used
+    def used
+      return 0 if max.zero?
+      amount = redis.get(key) || 0
+      [amount.to_i, max].min
+    end
+
+    private
+
+    def run_script(argv)
+      redis.evalsha(
+        Scripts::INCRBY_HASH, keys: [key], argv: argv
+      )
+    rescue Redis::CommandError => error
+      raise error if /ERR Error running script/ =~ error.message
+      redis.eval(Scripts::INCRBY, keys: [key], argv: argv)
+    end
+  end
+end

data/lib/traffic_jam/limit.rb

@@ -33,8 +33,8 @@ module TrafficJam
     # @param period [Integer] required limit period in seconds
     # @raise [ArgumentError] if max or period is nil
     def initialize(action, value, max: nil, period: nil)
-      raise ArgumentError('Max is required') if max.nil?
-      raise ArgumentError('Period is required') if period.nil?
+      raise ArgumentError.new('Max is required') if max.nil?
+      raise ArgumentError.new('Period is required') if period.nil?
       @action, @value, @max, @period = action, value, max, period
     end
 
@@ -71,14 +71,14 @@ module TrafficJam
         raise ArgumentError.new("Amount must be an integer")
       end
 
-      timestamp = (time.to_f * 1000).round
+      timestamp = (time.to_f * 1000).to_i
       argv = [timestamp, amount.to_i, max, period * 1000]
 
       result =
         begin
           redis.evalsha(
             Scripts::INCREMENT_SCRIPT_HASH, keys: [key], argv: argv)
-        rescue Redis::CommandError => e
+        rescue Redis::CommandError
           redis.eval(Scripts::INCREMENT_SCRIPT, keys: [key], argv: argv)
         end
 
@@ -127,9 +127,7 @@ module TrafficJam
     def used
       return 0 if max.zero?
 
-      obj = redis.hgetall(key)
-      timestamp = obj['timestamp']
-      amount = obj['amount']
+      timestamp, amount = redis.hmget(key, 'timestamp', 'amount')
       if timestamp && amount
         time_passed = Time.now.to_f - timestamp.to_i / 1000.0
         drift = max * time_passed / period
@@ -161,7 +159,7 @@ module TrafficJam
     end
 
     def key
-      if @key.nil?
+      if !defined?(@key) || @key.nil?
         converted_value =
           begin
             value.to_rate_limit_value
@@ -170,9 +168,13 @@ module TrafficJam
           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

data/lib/traffic_jam/rolling_limit.rb

@@ -0,0 +1,82 @@
+require_relative 'scripts'
+
+module TrafficJam
+  # This class represents a rolling limit on an action, value pair. For example,
+  # if limiting the amount of money a user can transfer in a week, the action
+  # could be +:transfers+ and the value would be the user ID. The class exposes
+  # atomic increment operations and allows querying of the current amount used
+  # and amount remaining.
+  #
+  # This class also handles 0 for period, where 0 is no period (each
+  # request is compared to the max).
+  #
+  # This class departs from the design of Limit by tracking a sum of the actions
+  # in a second, in a hash keyed by the timestamp. Therefore, this limit can put
+  # a lot of data size pressure on the Redis storage, so use it wisely.
+  class RollingLimit < Limit
+    # Constructor takes an action name as a symbol, a maximum cap, and the
+    # period of limit. +max+ and +period+ are required keyword arguments.
+    #
+    # @param action [Symbol] action name
+    # @param value [String] limit target value
+    # @param max [Integer] required limit maximum
+    # @param period [Integer] required limit period in seconds
+    # @raise [ArgumentError] if max or period is nil
+    def initialize(action, value, max: nil, period: nil)
+      super(action, value, max: max, period: period)
+    end
+
+    # Increment the amount used by the given number. Rolls back the increment
+    # if the operation exceeds the limit. Returns whether the operation was
+    # successful. Time of increment can be specified optionally with a keyword
+    # argument, which is not really useful since it be undone by used.
+    #
+    # @param amount [Integer] amount to increment by
+    # @param time [Time] time when increment occurs (ignored)
+    # @return [Boolean] true if increment succeded and false if incrementing
+    #   would exceed the limit
+    def increment(amount = 1, time: Time.now)
+      raise ArgumentError, 'Amount must be an integer' if amount != amount.to_i
+      return amount <= 0 if max.zero?
+      return amount <= max if period.zero?
+      return true if amount.zero?
+      return false if amount > max
+
+      !run_incr([time.to_i, amount.to_i, max, period]).nil?
+    end
+
+    # Return amount of limit used
+    #
+    # @return [Integer] amount used
+    def used
+      return 0 if max.zero? || period.zero?
+      [sum, max].min
+    end
+
+    private
+
+    def sum
+      run_sum([Time.now.to_i, period])
+    end
+
+    def clear_before
+      Time.now.to_i - period
+    end
+
+    def run_sum(argv)
+      redis.evalsha(Scripts::SUM_ROLLING_HASH, keys: [key], argv: argv)
+    rescue Redis::CommandError => error
+      raise error if /ERR Error running script/ =~ error.message
+      redis.eval(Scripts::SUM_ROLLING, keys: [key], argv: argv)
+    end
+
+    def run_incr(argv)
+      redis.evalsha(
+        Scripts::INCREMENT_ROLLING_HASH, keys: [key], argv: argv
+      )
+    rescue Redis::CommandError => error
+      raise error if /ERR Error running script/ =~ error.message
+      redis.eval(Scripts::INCREMENT_ROLLING, keys: [key], argv: argv)
+    end
+  end
+end

data/lib/traffic_jam/scripts.rb

@@ -10,5 +10,13 @@ module TrafficJam
 
     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')
+    INCRBY_HASH = Digest::SHA1.hexdigest(INCRBY)
+    SUM_ROLLING = load('sum_rolling')
+    SUM_ROLLING_HASH = Digest::SHA1.hexdigest(SUM_ROLLING)
   end
 end

data/lib/traffic_jam/simple_limit.rb

@@ -0,0 +1,94 @@
+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 true if amount == 0
+      return false if max == 0
+      raise ArgumentError.new("Amount must be positive") if amount < 0
+
+      if amount != amount.to_i
+        raise ArgumentError.new("Amount must be an integer")
+      end
+
+      return false if amount > max
+
+      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

data/lib/traffic_jam.rb

@@ -4,7 +4,9 @@ require_relative 'traffic_jam/errors'
 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'
 
 module TrafficJam
   include Errors

data/scripts/incrby.lua

@@ -0,0 +1,18 @@
+local arg_amount = tonumber(ARGV[1])
+local arg_max = tonumber(ARGV[2])
+
+local old_amount = tonumber(redis.call("GET", KEYS[1]))
+local new_amount
+
+if not old_amount then
+  new_amount = arg_amount
+else
+  new_amount = old_amount + arg_amount
+
+  if new_amount > arg_max then
+    return false
+  end
+end
+
+redis.call("INCRBY", KEYS[1], arg_amount)
+return true

data/scripts/increment_rolling.lua

@@ -0,0 +1,44 @@
+-- gets all fields from a hash as a dictionary
+local hgetall = function (key)
+  local bulk = redis.call('HGETALL', key)
+	local result = {}
+	local nextkey
+	for i, v in ipairs(bulk) do
+		if i % 2 == 1 then
+			nextkey = v
+		else
+			result[nextkey] = v
+		end
+	end
+	return result
+end
+
+local arg_timestamp = tonumber(ARGV[1])
+local arg_amount = tonumber(ARGV[2])
+local arg_max = tonumber(ARGV[3])
+local arg_period = tonumber(ARGV[4])
+
+local sum = arg_amount
+
+local clear_before = arg_timestamp - arg_period
+local mytable = hgetall(KEYS[1])
+if mytable ~= false then
+  -- print key -> value for mytable
+  print('mytable:')
+  for key, val in pairs(mytable) do
+    print('    ' .. key .. ' -> ' .. val)
+    if tonumber(key) < clear_before then
+      redis.call('HDEL', KEYS[1], key)
+    else
+      sum = sum + tonumber(val)
+    end
+  end
+end
+
+if sum > arg_max then
+  return false
+end
+
+redis.call("HINCRBY", KEYS[1], arg_timestamp, arg_amount)
+redis.call("EXPIRE", KEYS[1], arg_period)
+return true

data/scripts/increment_simple.lua

@@ -0,0 +1,17 @@
+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
+   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

data/scripts/sum_rolling.lua

@@ -0,0 +1,32 @@
+-- gets all fields from a hash as a dictionary
+local hgetall = function (key)
+  local bulk = redis.call('HGETALL', key)
+	local result = {}
+	local nextkey
+	for i, v in ipairs(bulk) do
+		if i % 2 == 1 then
+			nextkey = v
+		else
+			result[nextkey] = v
+		end
+	end
+	return result
+end
+
+local arg_timestamp = tonumber(ARGV[1])
+local arg_period = tonumber(ARGV[2])
+local sum = 0
+
+local clear_before = arg_timestamp - arg_period
+local mytable = hgetall(KEYS[1])
+if mytable ~= false then
+  for key, val in pairs(mytable) do
+    if tonumber(key) < clear_before then
+      redis.call('HDEL', KEYS[1], key)
+    else
+      sum = sum + tonumber(val)
+    end
+  end
+end
+
+return sum

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: traffic_jam
 version: !ruby/object:Gem::Version
-  version: 1.0.1
+  version: 1.1.0
 platform: ruby
 authors:
 - Jim Posen
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2015-08-13 00:00:00.000000000 Z
+date: 2018-01-09 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: redis
@@ -47,10 +47,17 @@ files:
 - lib/traffic_jam.rb
 - lib/traffic_jam/configuration.rb
 - lib/traffic_jam/errors.rb
+- lib/traffic_jam/lifetime_limit.rb
 - lib/traffic_jam/limit.rb
 - lib/traffic_jam/limit_group.rb
+- lib/traffic_jam/rolling_limit.rb
 - lib/traffic_jam/scripts.rb
+- lib/traffic_jam/simple_limit.rb
+- scripts/incrby.lua
 - scripts/increment.lua
+- scripts/increment_rolling.lua
+- scripts/increment_simple.lua
+- scripts/sum_rolling.lua
 homepage: https://github.com/coinbase/traffic_jam
 licenses:
 - MIT
@@ -71,9 +78,8 @@ required_rubygems_version: !ruby/object:Gem::Requirement
       version: '0'
 requirements: []
 rubyforge_project: 
-rubygems_version: 2.2.2
+rubygems_version: 2.6.13
 signing_key: 
 specification_version: 4
 summary: Library for time-based rate limiting
 test_files: []
-has_rdoc: