parrish/Congestion

module Congestion
  class RateLimiter
    attr_accessor :redis, :key, :options

    def initialize(redis, key, opts = { })
      self.redis = redis
      self.key = "#{ opts[:namespace] }:#{ key }"
      self.options = opts
      self.options[:interval] *= 1_000
      self.options[:min_delay] *= 1_000
      allowed?
    end

    def total_requests
      get_requests[1]
    end

    def first_request
      first = get_requests[2].first
      first ? first.to_i : nil
    end

    def last_request
      last = get_requests[3].first
      last ? last.to_i : nil
    end

    def allowed?
      add_request unless rejected?
      !rejected?
    end

    def rejected?
      too_many? || too_frequent?
    end

    def too_many?
      total_requests > options[:max_in_interval]
    end

    def too_frequent?
      last_request && time_since_last_request < options[:min_delay]
    end

    def backoff
      if too_many? && too_frequent?
        [quantity_backoff, frequency_backoff].max
      elsif too_many?
        quantity_backoff
      elsif too_frequent?
        frequency_backoff
      else
        0
      end
    end

    protected

    def current_time
      @current_time ||= (Time.now.utc.to_f * 1_000).round
    end

    def time_since_last_request
      current_time - last_request
    end

    def time_since_first_request
      current_time - first_request
    end

    def expired_at
      current_time - options[:interval]
    end

    def quantity_backoff
      millis = options[:interval] - time_since_first_request
      (millis / 1_000.0).ceil
    end

    def frequency_backoff
      millis = options[:min_delay] - time_since_last_request
      (millis / 1_000.0).ceil
    end

    def add_request
      unless options[:track_rejected]
        add_request = redis.multi do |t|
          t.zadd key, current_time, current_time # [0] - key added
          t.ttl key                              # [1] - key ttl
        end
        # TTL is -1 if not set, https://redis.io/commands/ttl
        if add_request[1] == -1
          # ensure we set the expire TTL on the request key
          # using the raw interval here after the 'get_requests' limit check
          # should be close enough to the actual interval folks desire
          redis.pexpire key, options[:interval]
        end
      end
    end

    def get_requests
      @requests ||= redis.multi do |t|
        t.zremrangebyscore key, 0, expired_at     # [0] - clear old requests
        t.zcount key, '-inf', '+inf'              # [1] - number of requests
        t.zrange key, 0, 0                        # [2] - first request
        t.zrange key, -1, -1                      # [3] - last request
        if options[:track_rejected]
          t.zadd(key, current_time, current_time) # [4] - Add the request if tracking rejected
        end
        # ensure the TTL is set after we've added the key if tracking rejected
        t.pexpire key, options[:interval]         # [5] - expire request key
      end
    end
  end
end