Bottleneck - simple rate limiter for Rails

Some time ago I need to implement rate limiter in ruby, to use in my Ruby on Rails applications.

I splitted all code for two parts:

Rails gem called Bottleneck.

And Rails Application to demonstrate gem usage.

So, let’s move deeper and check what throttling exactly mean.

Throttling is a process that is used to control the usage of APIs by consumers during a given period. You can define throttling at the application level and API level. Throttling limit is considered as cumulative at API level.

Administrators and publishers of API manager can use throttling to limit the number of API requests per day/week/month. For example, you can limit the number of total API requests as 10000/day.

When a throttle limit is crossed, the server sends 429 message as HTTP status to the user with message content as “too many requests”.

Great, in my case, I just need to implemept two things:

• Maximum request count per IP address.

• Time limit for reset all counters.

Great. Let’s move one.

Step one - storage.

It can be Rails cache store, Database, NoSQL storage, Key-Value storage etc

I decided to use Redis here - extremely fast in-memory key-value data structure store with built-in key expiration functionality. It can be used as a database, cache or message broker for example.

P.S. One interesting and great thing - Rails 5.2 will be shipped with built-in Redis Cache Store.

Step two - our Bottleneck module.

What interesting here?

We read our limits from bottleneck.yml file:

limits:
  time_period_seconds: 3600
  max_requests_count: 100

And redis.yml file:

host: '127.0.0.1'
port: 6379

Firstly we need to init our Redis::Namespace object with Redis.new instance. We use :bottleneck namespace for our limiter.

Second thing here - load_config method uses a filename and load config file from Rails root path, or from current directory.

And third step - check method, which receive current request’s IP address. It’s a starting point for main logic of our gem.

require "bottleneck/version"
require "bottleneck/core"
require "yaml"
require "redis"
require "redis-namespace"

module Bottleneck
  class << self
    def check(ip)
      Core.new(ip).run
    end

    def storage
      init_storage
    end

    def redis_conn
      redis_conf = load_config("redis.yml")
      Redis.new(host: redis_conf["host"], port: redis_conf["port"])
    end

    def init_storage
      Redis::Namespace.new(:bottleneck, redis: redis_conn)
    end

    def config
      load_config("bottleneck.yml")
    end

    private

    def load_config(file)
      root = (defined?(Rails) && Rails.respond_to?(:root) && Rails.root) || Dir.pwd
      path = "#{root}/config/#{file}"
      raise "No #{file} file found in your config directory!" unless File.exist?(path)
      YAML.load_file(path)
    end
  end
end

Step three - Constants class.

Simple step with constants class. Just success status for our response, expiration status and fixed message.

module Bottleneck
  class Constants
    SUCCESS_STATUS = 200
    EXPIRED_STATUS = 429
    OK_MESSAGE = "ok".freeze
  end
end

Step four - main Core class.

Very simple logic, and all power of Redis EXPIRE and TTL commands usage.

Let’s check this class step-by-step:

• We initialize class with IP address from request. Init our storage and get limits from config

• run method creates uniq key, based on IP address, and default response.

• then we checking our storage for existing key, and setting it, if key is nil.

• also we set expiration time in seconds from our config file, using expire method

• next we check if our requests count in our limit, and if no - we set status and message in our result hash, with seconds from period method, which uses ttl method for remaining number of seconds.

• otherwise, we increment our key in Redis with incr method.

require "bottleneck/constants"

module Bottleneck
  class Core
    def initialize(ip)
      @ip = ip.to_s
      @storage = Bottleneck.storage
      @limits = Bottleneck.config["limits"]
    end

    def run
      client_ip = @ip
      key = "request_count:#{client_ip}"
      result = { status: Constants::SUCCESS_STATUS, message: Constants::OK_MESSAGE }
      requests_count = @storage.get(key)
      unless requests_count
        @storage.set(key, 0)
        @storage.expire(key, @limits["time_period_seconds"])
      end
      if requests_count.to_i >= @limits["max_requests_count"]
        result[:status] = Constants::EXPIRED_STATUS
        result[:message] = message(period(key))
      else
        @storage.incr(key)
      end
      result
    end

    private

    def period(key)
      @storage.ttl(key)
    end

    def message(secs)
      "Rate limit exceeded. Try again in #{secs} seconds"
    end
  end
end

That’s it! :)