traffic_jam 1.0.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA1:
+  metadata.gz: d94e33ae19bb7a6f3780230e0236a937dde90cb5
+  data.tar.gz: aba261c52a890bf62392c0fda67742af28d9828c
+SHA512:
+  metadata.gz: aee5755e8a2585c309523cd1851ed2dbb81ae93e179e3605a0bdc945a02ca05ed3bc72e915bf2bfdb4574052cd3aaa3be6434a64e108186d8a1ec284d6c64e55
+  data.tar.gz: fafd3000ef48ddfde67bae79e072429f07cf8d11f088bcb418a754f3bfd5779cc6554fe4ebba6146cd9e98b2076b5618874d48ff28f7a4efb21679ddd27ae8b6

data/lib/traffic_jam.rb

@@ -0,0 +1,63 @@
+require 'ostruct'
+require 'digest/md5'
+require_relative 'traffic_jam/errors'
+require_relative 'traffic_jam/configuration'
+require_relative 'traffic_jam/limit'
+require_relative 'traffic_jam/limit_group'
+
+
+module TrafficJam
+  include Errors
+
+  @config = Configuration.new(
+    key_prefix: 'traffic_jam',
+    hash_length: 22
+  )
+
+  class << self
+    attr_reader :config
+
+    # Configure library in a block.
+    #
+    # @yield [TrafficJam::Configuration]
+    def configure
+      yield config
+    end
+
+    # Create limit with registed max/period.
+    #
+    # @param action [Symbol] registered action name
+    # @param value [String] limit target value
+    # @return [TrafficJam::Limit]
+    def limit(action, value)
+      limits = config.limits(action.to_sym)
+      TrafficJam::Limit.new(action, value, **limits)
+    end
+
+    # Reset all limits associated with the given action. If action is omitted or
+    # nil, this will reset all limits.
+    #
+    # @note Not recommended for use in production.
+    # @param action [Symbol] action to reset limits for
+    # @return [nil]
+    def reset_all(action: nil)
+      prefix =
+        if action.nil?
+          "#{config.key_prefix}:*"
+        else
+          "#{config.key_prefix}:#{action}:*"
+        end
+      config.redis.keys(prefix).each do |key|
+        config.redis.del(key)
+      end
+      nil
+    end
+
+    %w( exceeded? increment increment! decrement reset used remaining )
+      .each do |method|
+      define_method(method) do |action, value, *args|
+        limit(action, value).send(method, *args)
+      end
+    end
+  end
+end

data/lib/traffic_jam/configuration.rb

@@ -0,0 +1,63 @@
+module TrafficJam
+  # Configuration for TrafficJam library.
+  #
+  # @see TrafficJam#configure
+  class Configuration
+    OPTIONS = %i( key_prefix hash_length redis )
+
+    # @!attribute redis
+    #   @return [Redis] the connected Redis client the library uses
+    # @!attribute key_prefix
+    #   @return [String] the prefix of all limit keys in Redis
+    # @!attribute hash_length
+    #   @return [String] the number of characters to use from the Base64 encoded
+    #     hashes of the limit values
+    attr_accessor *OPTIONS
+
+    def initialize(options = {})
+      OPTIONS.each do |option|
+        self.send("#{option}=", options[option])
+      end
+    end
+
+    # Register a default cap and period with an action name. For use with
+    # {TrafficJam.limit}.
+    #
+    # @param action [Symbol] action name
+    # @param max [Integer] limit cap
+    # @param period [Fixnum] limit period in seconds
+    def register(action, max, period)
+      @limits ||= {}
+      @limits[action.to_sym] = { max: max, period: period }
+    end
+
+    # Get the limit cap registered to an action.
+    #
+    # @see #register
+    # @return [Integer] limit cap
+    def max(action)
+      limits(action)[:max]
+    end
+
+    # Get the limit period registered to an action.
+    #
+    # @see #register
+    # @return [Integer] limit period in seconds
+    def period(action)
+      limits(action)[:period]
+    end
+
+    # Get registered limit parameters for an action.
+    #
+    # @see #register
+    # @param action [Symbol] action name
+    # @return [Hash] max and period parameters in a hash
+    # @raise [TrafficJam::LimitNotFound] if action is not registered
+    def limits(action)
+      @limits ||= {}
+      limits = @limits[action.to_sym]
+      raise TrafficJam::LimitNotFound.new(action) if limits.nil?
+      limits
+    end
+  end
+end

data/lib/traffic_jam/errors.rb

@@ -0,0 +1,14 @@
+module TrafficJam
+  module Errors
+    class LimitNotFound < StandardError; end
+
+    class LimitExceededError < StandardError
+      attr_accessor :limit
+
+      def initialize(limit)
+        super("Rate limit exceeded: #{limit.action}")
+        @limit = limit
+      end
+    end
+  end
+end

data/lib/traffic_jam/limit.rb

@@ -0,0 +1,178 @@
+require_relative 'scripts'
+
+module TrafficJam
+  # This class represents a rate limit on an action, value pair. For example, if
+  # rate limiting the number of requests per IP address, the action could be
+  # +:requests+ and the value would be the IP address. The class exposes atomic
+  # increment operations and allows querying of the current amount used and
+  # amount remaining.
+  class Limit
+    # @!attribute [r] action
+    #   @return [Symbol] the name of the action being rate limited.
+    # @!attribute [r] value
+    #   @return [String] the target of the limit. The value should be a string
+    #     or convertible to a distinct string when +to_s+ is called. If you
+    #     would like to use objects that can be converted to a unique string,
+    #     like a database-mapped object with an ID, you can implement
+    #     +to_rate_limit_value+ on the object, which returns a deterministic
+    #     string unique to that object.
+    # @!attribute [r] max
+    #   @return [Integer] the integral cap of the limit amount.
+    # @!attribute [r] period
+    #   @return [Integer] the duration of the limit in seconds. Regardless of
+    #     the current amount used, after the period passes, the amount used will
+    #     be 0.
+    attr_reader :action, :max, :period, :value
+
+    # 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)
+      raise ArgumentError('Max is required') if max.nil?
+      raise ArgumentError('Period is required') if period.nil?
+      @action, @value, @max, @period = action, value, max, period
+    end
+
+    # Return whether incrementing by the given amount would exceed limit. Does
+    # not change amount used.
+    #
+    # @param amount [Integer]
+    # @return [Boolean]
+    def exceeded?(amount = 1)
+      used + amount > max
+    end
+
+    # Return itself if incrementing by the given amount would exceed limit,
+    # otherwise nil. Does not change amount used.
+    #
+    # @return [TrafficJam::Limit, nil]
+    def limit_exceeded(amount = 1)
+      self if exceeded?(amount)
+    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. Time of increment can be specified optionally with a keyword
+    # argument, which is useful for rolling back with a decrement.
+    #
+    # @param amount [Integer] amount to increment by
+    # @param time [Time] time when increment occurs
+    # @return [Boolean] true if increment succeded and false if incrementing
+    #   would exceed the limit
+    def increment(amount = 1, time: Time.now)
+      return amount <= 0 if max.zero?
+
+      if amount != amount.to_i
+        raise ArgumentError.new("Amount must be an integer")
+      end
+
+      timestamp = (time.to_f * 1000).round
+      argv = [timestamp, amount.to_i, max, period * 1000]
+
+      result =
+        begin
+          redis.evalsha(
+            Scripts::INCREMENT_SCRIPT_HASH, keys: [key], argv: argv)
+        rescue Redis::CommandError => e
+          redis.eval(Scripts::INCREMENT_SCRIPT, keys: [key], argv: argv)
+        end
+
+      !!result
+    end
+
+    # Increment the amount used by the given number. Does not perform increment
+    # if the operation would exceed the limit. Raises an exception if the
+    # operation is unsuccessful. Time of# increment can be specified optionally
+    # with a keyword argument, which is useful for rolling back with a
+    # decrement.
+    #
+    # @param amount [Integer] amount to increment by
+    # @param time [Time] time when increment occurs
+    # @return [nil]
+    # @raise [TrafficJam::LimitExceededError] if incrementing would exceed the
+    #   limit
+    def increment!(amount = 1, time: Time.now)
+      if !increment(amount, time: time)
+        raise TrafficJam::LimitExceededError.new(self)
+      end
+    end
+
+    # Decrement the amount used by the given number. Time of decrement can be
+    # specified optionally with a keyword argument, which is useful for rolling
+    # back an increment operation at a certain time.
+    #
+    # @param amount [Integer] amount to increment by
+    # @param time [Time] time when increment occurs
+    # @return [true]
+    def decrement(amount = 1, time: Time.now)
+      increment(-amount, time: time)
+    end
+
+    # Reset amount used to 0.
+    #
+    # @return [nil]
+    def reset
+      redis.del(key)
+      nil
+    end
+
+    # Return amount of limit used, taking time drift into account.
+    #
+    # @return [Integer] amount used
+    def used
+      return 0 if max.zero?
+
+      obj = redis.hgetall(key)
+      timestamp = obj['timestamp']
+      amount = obj['amount']
+      if timestamp && amount
+        time_passed = Time.now.to_f - timestamp.to_i / 1000.0
+        drift = max * time_passed / period
+        last_amount = [amount.to_f, max].min
+        [(last_amount - drift).ceil, 0].max
+      else
+        0
+      end
+    end
+
+    # Return amount of limit remaining, taking time drift into account.
+    #
+    # @return [Integer] amount remaining
+    def remaining
+      max - used
+    end
+
+    def flatten
+      [self]
+    end
+
+    private
+    def config
+      TrafficJam.config
+    end
+
+    def redis
+      config.redis
+    end
+
+    def key
+      if @key.nil?
+        converted_value =
+          begin
+            value.to_rate_limit_value
+          rescue NoMethodError
+            value
+          end
+        hash = Digest::MD5.base64digest(converted_value.to_s)
+        hash = hash[0...config.hash_length]
+        @key = "#{config.key_prefix}:#{action}:#{hash}"
+      end
+      @key
+    end
+  end
+end

data/lib/traffic_jam/limit_group.rb

@@ -0,0 +1,129 @@
+module TrafficJam
+  # A limit group is a way of enforcing a cap over a set of limits with the
+  # guarantee that either all limits will be incremented or none. This is useful
+  # if you must check multiple limits before allowing an action to be taken.
+  # Limit groups can contain other limit groups.
+  class LimitGroup
+    attr_reader :limits
+
+    # Creates a limit group from a collection of limits or other limit groups.
+    #
+    # @param limits [Array<TrafficJam::Limit>] either an array or splat of
+    #   limits or other limit groups
+    # @param ignore_nil_values [Boolean] silently drop limits with a nil value
+    def initialize(*limits, ignore_nil_values: false)
+      @limits = limits.flatten
+      @ignore_nil_values = ignore_nil_values
+      if @ignore_nil_values
+        @limits.reject! do |limit|
+          limit.respond_to?(:value) && limit.value.nil?
+        end
+      end
+    end
+
+    # Add a limit to the group.
+    #
+    # @param limit [TrafficJam::Limit, TrafficJam::LimitGroup]
+    def <<(limit)
+      if !(@ignore_nil_values && limit.value.nil?)
+        limits << limit
+      end
+    end
+
+    # Attempt to increment the limits by the given amount. Does not increment
+    # if incrementing would exceed any limit.
+    #
+    # @param amount [Integer] amount to increment by
+    # @param time [Time] optional time of increment
+    # @return [Boolean] whether increment operation was successful
+    def increment(amount = 1, time: Time.now)
+      exceeded_index = limits.find_index do |limit|
+        !limit.increment(amount, time: time)
+      end
+      if exceeded_index
+        limits[0...exceeded_index].each do |limit|
+          limit.decrement(amount, time: time)
+        end
+      end
+      exceeded_index.nil?
+    end
+
+    # Increment the limits by the given amount. Raises an error and does not
+    # increment if doing so would exceed any limit.
+    #
+    # @param amount [Integer] amount to increment by
+    # @param time [Time] optional time of increment
+    # @return [nil]
+    # @raise [TrafficJam::LimitExceededError] if increment would exceed any
+    #   limits
+    def increment!(amount = 1, time: Time.now)
+      exception = nil
+      exceeded_index = limits.find_index do |limit|
+        begin
+          limit.increment!(amount, time: time)
+        rescue TrafficJam::LimitExceededError => e
+          exception = e
+          true
+        end
+      end
+      if exceeded_index
+        limits[0...exceeded_index].each do |limit|
+          limit.decrement(amount, time: time)
+        end
+        raise exception
+      end
+    end
+
+    # Decrement the limits by the given amount.
+    #
+    # @param amount [Integer] amount to decrement by
+    # @param time [Time] optional time of decrement
+    # @return [true]
+    def decrement(amount = 1, time: Time.now)
+      limits.all? { |limit| limit.decrement(amount, time: time) }
+    end
+
+    # Return whether incrementing by the given amount would exceed any limit.
+    # Does not change amount used.
+    #
+    # @param amount [Integer]
+    # @return [Boolean] whether any limit would be exceeded
+    def exceeded?(amount = 1)
+      limits.any? { |limit| limit.exceeded?(amount) }
+    end
+
+    # Return the first limit to be exceeded if incrementing by the given amount,
+    # or nil otherwise. Does not change amount used for any limit.
+    #
+    # @param amount [Integer]
+    # @return [TrafficJam::Limit, nil]
+    def limit_exceeded(amount = 1)
+      limits.each do |limit|
+        limit_exceeded = limit.limit_exceeded(amount)
+        return limit_exceeded if limit_exceeded
+      end
+      nil
+    end
+
+    # Resets all limits to 0.
+    def reset
+      limits.each(&:reset)
+      nil
+    end
+
+    # Return minimum amount remaining of any limit.
+    #
+    # @return [Integer] amount remaining in limit group
+    def remaining
+      limits.map(&:remaining).min
+    end
+
+    # Return flattened list of limit. Will return list limits even if this group
+    # contains nested limit groups.
+    #
+    # @return [Array<TrafficJam::Limit>] list of limits
+    def flatten
+      limits.map(&:flatten).flatten
+    end
+  end
+end

data/lib/traffic_jam/scripts.rb

@@ -0,0 +1,14 @@
+require 'digest/sha1'
+
+module TrafficJam
+  module Scripts
+    def self.load(name)
+      scripts_dir = File.join(File.dirname(__FILE__), '..', '..', 'scripts')
+      File.read(File.join(scripts_dir, "#{name}.lua"))
+    end
+    private_class_method :load
+
+    INCREMENT_SCRIPT = load('increment')
+    INCREMENT_SCRIPT_HASH = Digest::SHA1.hexdigest(INCREMENT_SCRIPT)
+  end
+end

data/scripts/increment.lua

@@ -0,0 +1,46 @@
+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 old_timestamp = redis.call("HGET", KEYS[1], "timestamp")
+
+local new_amount
+local new_timestamp
+
+if not old_timestamp
+then
+   new_amount = arg_amount
+   new_timestamp = arg_timestamp
+else
+   local time_diff = arg_timestamp - tonumber(old_timestamp)
+   local drift_amount = time_diff * arg_max / arg_period
+   if time_diff < 0
+   then
+      local incr_amount = arg_amount + drift_amount
+      if incr_amount <= 0
+      then
+         return true
+      end
+      local old_amount = tonumber(redis.call("HGET", KEYS[1], "amount"))
+      old_amount = math.min(old_amount, arg_max)
+      new_amount = old_amount + incr_amount
+      new_timestamp = old_timestamp
+   else
+      local old_amount = tonumber(redis.call("HGET", KEYS[1], "amount"))
+      old_amount = math.min(old_amount, arg_max)
+      local current_amount = math.max(old_amount - drift_amount, 0)
+      new_amount = current_amount + arg_amount
+      new_timestamp = arg_timestamp
+   end
+end
+
+if new_amount > arg_max
+then
+   return false
+end
+
+redis.call("HSET", KEYS[1], "amount", new_amount)
+redis.call("HSET", KEYS[1], "timestamp", new_timestamp)
+redis.call("EXPIRE", KEYS[1], arg_period)
+return true

metadata

@@ -0,0 +1,79 @@
+--- !ruby/object:Gem::Specification
+name: traffic_jam
+version: !ruby/object:Gem::Version
+  version: 1.0.0
+platform: ruby
+authors:
+- Jim Posen
+autorequire: 
+bindir: bin
+cert_chain: []
+date: 2015-05-05 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: redis
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.0'
+- !ruby/object:Gem::Dependency
+  name: rake
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '10.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '10.0'
+description: Library for Redis-backed time-based rate limiting
+email: jimpo@coinbase.com
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- lib/traffic_jam.rb
+- lib/traffic_jam/configuration.rb
+- lib/traffic_jam/errors.rb
+- lib/traffic_jam/limit.rb
+- lib/traffic_jam/limit_group.rb
+- lib/traffic_jam/scripts.rb
+- scripts/increment.lua
+homepage: ''
+licenses:
+- MIT
+metadata: {}
+post_install_message: 
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubyforge_project: 
+rubygems_version: 2.2.2
+signing_key: 
+specification_version: 4
+summary: Library for time-based rate limiting
+test_files: []
+has_rdoc: