redis_token_bucket 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA1:
+  metadata.gz: e1991459cf61f197e4b366ade53dd41ecbaf26bd
+  data.tar.gz: 4931d16f9aa7fb622bf167fdbc74d4781bbdc548
+SHA512:
+  metadata.gz: 3aa2fdf2aca0e7c374cd6845dece306ef0de74651f50b17328269d0abb5a38f9d742ae32ed9fb25119bd5922e166d00aceb42ce2574dae1da02df712d31e581a
+  data.tar.gz: 7f99aef70badd3708e784d240a81c1fa0f1c25a543f6d7fb66d40f7c5e7eb3d6cb4c7e64d2985aeb90aab834ea652ed8f940109fd4f39214e303d33e01de7ccc

data/.gitignore

@@ -0,0 +1,9 @@
+/.bundle/
+/.yardoc
+/Gemfile.lock
+/_yardoc/
+/coverage/
+/doc/
+/pkg/
+/spec/reports/
+/tmp/

data/.rspec

@@ -0,0 +1,2 @@
+--format documentation
+--color

data/.travis.yml

@@ -0,0 +1,12 @@
+language: ruby
+rvm:
+  - 2.2
+before_install:
+  - gem update bundler
+script: bundle exec rspec spec
+services:
+  - redis-server
+addons:
+  apt:
+    packages:
+      - redis-server

data/Gemfile

@@ -0,0 +1,7 @@
+source 'https://rubygems.org'
+
+# Specify your gem's dependencies in redis_token_bucket.gemspec
+gemspec
+
+# needed by "demo.rb"
+gem "concurrent-ruby"

data/LICENSE.txt

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2016 Kristian Hanekamp
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

data/README.md

@@ -0,0 +1,149 @@
+# RedisTokenBucket
+
+A [Token Bucket](https://en.wikipedia.org/wiki/Token_bucket) rate limiting implementation in Ruby using a Redis backend.
+
+Features:
+* Lightweight and efficient
+  * Uses a single Redis key per bucket
+  * Buckets are automatically created when first used
+  * Buckets are automatically removed when no longer used
+* Fast and concurrency safe
+  * Each operation uses just a single network roundtrip to Redis
+  * Charging tokens is done with all-or-nothing semantics
+* Computed continuously
+  * Token values (rate, size, current level, cost) use floating point numbers
+  * Bucket level is computed with microsecond precision
+* Powerful and flexible
+  * Ability to charge multiple buckets with arbitrary token amounts at once
+  * Ability to "reserve" tokens and to create "token debt"
+
+Redis version 3.2 or newer is needed.
+
+## Installation
+
+Add this line to your application's Gemfile:
+
+```ruby
+gem 'redis_token_bucket'
+```
+
+## Usage
+
+Basic rate limiting:
+
+```ruby
+require 'redis'
+require 'redis_token_bucket'
+
+# create connection to redis server
+# details see: https://github.com/redis/redis-rb/
+redis = Redis.new
+
+# create a limiter instance which uses the redis connection
+limiter = RedisTokenBucket.limiter(redis)
+
+# define the bucket
+bucket = {
+  key: "RedisKeyForMyBucket",
+  rate: 100,
+  size: 1000,
+}
+
+# charge 10 tokens to the bucket
+success, level = limiter.charge(bucket, 10)
+
+# check if charging was successful
+if success
+  # rate limiter permits request
+  call_my_business_logic
+else
+  # rate limiter denies request
+  raise "Rate Limit exceeded. Increase you calm!"
+end
+
+# print the resulting level of tokens in the bucket
+puts "The current level of tokens in my bucket: #{level}"
+
+```
+
+Reading the current level of tokens of a bucket:
+
+```ruby
+puts "Current level of tokens: #{limiter.read_level(bucket)}"
+```
+
+Charging multiple buckets at once:
+
+```ruby
+long_bucket = {
+  key: "RedisKeyForLongBucket",
+  rate: 100,
+  size: 10000
+}
+
+short_bucket = {
+  key: "RedisKeyForShortBucket",
+  rate: 1000,
+  size: 3000
+}
+
+success, levels = limiter.batch_charge(
+  [long_bucket, 1],
+  [short_bucket, 1]
+)
+
+puts "The current level of tokens in bucket short: #{levels[short_bucket[:key]]}"
+puts "The current level of tokens in bucket long: #{levels[long_bucket[:key]]}"
+
+if success
+  # rate limiter permits request (all buckets were charged)
+  call_my_business_logic
+else
+  # rate limiter denies request (none of the buckets was charged)
+  raise "Rate Limit exceeded. Increase you calm!"
+end
+```
+
+Reading the current level of tokens from multiple buckets:
+
+```ruby
+levels = limiter.read_levels(short_bucket, long_bucket)
+
+puts "The current level of tokens in bucket short: #{levels[short_bucket[:key]]}"
+puts "The current level of tokens in bucket long: #{levels[long_bucket[:key]]}"
+```
+
+Advanced: Bucket with Reserved Tokens
+
+```ruby
+# this reserves the last 10 tokens,
+# i.e. charging will fail if it would result in less than 10 tokens
+
+RedisTokenBucket.charge(bucket, 1, {limit: 10})
+
+# also possible with batch_charge
+RedisTokenBucket.batch_charge(
+  [short_bucket, 1, {limit: 10}],
+  [long_bucket, 2, {limit: 5}],
+)
+```
+
+Advanced: Bucket with Token Debt
+
+```ruby
+# this allows up to 10 "negative" tokens
+# i.e. charging will only fail if it would result in less than -10 tokens
+RedisTokenBucket.charge(bucket, 1, {limit: -10})
+```
+
+## Development
+
+After checking out the repo, run `bundle` to install dependencies.
+
+Use `bundle exec rspec` to run tests.
+
+Use `bundle exec ruby demo.rb` to run a demo.
+
+## Contributors
+
+Original author: Kristian Hanekamp

data/Rakefile

@@ -0,0 +1 @@
+require "bundler/gem_tasks"

data/bin/console

@@ -0,0 +1,14 @@
+#!/usr/bin/env ruby
+
+require "bundler/setup"
+require "redis_token_bucket"
+
+# You can add fixtures and/or initialization code here to make experimenting
+# with your gem easier. You can also use a different console, if you like.
+
+# (If you use this, don't forget to add pry to your Gemfile!)
+# require "pry"
+# Pry.start
+
+require "irb"
+IRB.start

data/bin/setup

@@ -0,0 +1,7 @@
+#!/bin/bash
+set -euo pipefail
+IFS=$'\n\t'
+
+bundle install
+
+# Do any other automated setup that you need to do here

data/demo.rb

@@ -0,0 +1,109 @@
+require "redis_token_bucket"
+require "concurrent"
+require "redis"
+require "securerandom"
+
+puts <<-EOS
+  The script attempts to continuously make requests against the rate limiter.
+  Each second, the number of requests accepted and denied by the rate limiter is printed.
+
+  You should see the following pattern:
+  * briefly, a burst of requests is accepted
+  * then the 'short' buckets starts limiting to 1000 requests per second
+  * after a few seconds, the 'long' buckets starts limiting to 100 requests per second
+
+EOS
+
+def random_key
+  "RedisTokenBucket:demo:#{SecureRandom.hex}"
+end
+
+buckets = {
+  long: {
+    key: random_key,
+    rate: 100,
+    size: 10000
+  },
+  short: {
+    key: random_key,
+    rate: 1000,
+    size: 3000
+  }
+}
+
+consumed = Concurrent::Atom.new(0)
+
+rejected = {}
+buckets.each { |name, _| rejected[name] = Concurrent::Atom.new(0) }
+
+def increase(atom)
+  atom.swap { |before| before + 1 }
+end
+
+def reset(atom)
+  last = nil
+
+  atom.swap do |before|
+    last = before
+
+    0
+  end
+
+  last
+end
+
+NUM_FORKS = 1
+NUM_THREADS_PER_FORK = 10
+
+child_processes = NUM_FORKS.times.map do
+  Process.fork do
+
+    # each fork has an independent output thread
+    output = Thread.new do
+      last_output = 0
+      while true
+        now = Time.now.to_i
+
+        if last_output < now
+          denied_stats = rejected.map { |name, atom| "#{name}: #{reset(atom)}" }
+          puts "Accepted: #{reset(consumed)} / Denied: #{denied_stats}"
+
+          last_output = now
+        end
+
+        sleep 0.001
+      end
+    end
+
+    # and a number of worker threads, charging tokens
+    workers = NUM_THREADS_PER_FORK.times.map do |i|
+      Thread.new do
+        begin
+          limiter = RedisTokenBucket.limiter(Redis.new)
+
+          while true
+            success, levels = limiter.batch_charge([buckets[:short], 1], [buckets[:long], 1])
+
+            if success
+              increase(consumed)
+            else
+              levels.map do |key, level|
+                name = buckets.keys.detect { |n| buckets[n][:key] == key }
+                increase(rejected[name]) if level < 1
+              end
+            end
+          end
+        rescue Exception => e
+          puts "ERROR"
+          puts e
+          puts e.backtrace
+        end
+      end
+    end
+
+    output.join
+    workers.join
+  end
+end
+
+child_processes.map { |pid| Process.waitpid(pid) }

data/lib/redis_token_bucket.rb

@@ -0,0 +1,8 @@
+require "redis_token_bucket/version"
+require "redis_token_bucket/limiter"
+
+module RedisTokenBucket
+  def self.limiter(redis)
+    Limiter.new(redis)
+  end
+end

data/lib/redis_token_bucket/limiter.lua

@@ -0,0 +1,73 @@
+redis.replicate_commands()
+local injected_time = tonumber(ARGV[1])
+local redis_time = redis.call('time')
+local local_time = redis_time[1] + 0.000001 * redis_time[2]
+
+local now = injected_time or local_time
+
+local current_bucket_levels = {}
+local new_bucket_levels = {}
+local timeouts = {}
+local exceeded = false
+
+for key_index, key in ipairs(KEYS) do
+  local arg_index = key_index * 4 - 2
+  local rate = tonumber(ARGV[arg_index])
+  local size = tonumber(ARGV[arg_index + 1])
+  local amount = tonumber(ARGV[arg_index + 2])
+
+  local bucket = redis.call('hmget', key, 'time', 'level')
+  local last_time = tonumber(bucket[1]) or now
+  local before_level = tonumber(bucket[2]) or size
+
+  local elapsed = math.max(0, now - last_time)
+  local gained = rate * elapsed
+
+  local current_level = math.min(size, before_level + gained)
+
+  current_bucket_levels[key_index] = current_level
+
+  if amount > 0 then
+    local limit = tonumber(ARGV[arg_index + 3]) or 0
+
+    local new_level = current_level - amount
+    new_bucket_levels[key_index] = new_level
+
+    local seconds_to_full = (size - new_level) / rate
+    timeouts[key_index] = seconds_to_full
+
+    if new_level < limit then
+      exceeded = true
+    end
+  end
+end
+
+local levels_to_report
+local charged
+
+if exceeded or #new_bucket_levels == 0 then
+  levels_to_report = current_bucket_levels
+  charged = 0
+else
+  levels_to_report = new_bucket_levels
+  charged = 1
+
+  for key_index, key in ipairs(KEYS) do
+    local new_level = new_bucket_levels[key_index]
+    local timeout = timeouts[key_index]
+
+    redis.call('hmset', key,
+      'time', string.format("%.16g", now),
+      'level', string.format("%.16g", new_level)
+      )
+
+    redis.call('expire', key, math.ceil(timeout))
+  end
+end
+
+local formatted_levels = {}
+for index, value in ipairs(levels_to_report) do
+  formatted_levels[index] = string.format("%.16g", value)
+end
+
+return {charged, formatted_levels}

data/lib/redis_token_bucket/limiter.rb

@@ -0,0 +1,106 @@
+module RedisTokenBucket
+class Limiter
+  def initialize(redis, clock = nil)
+    @redis = redis
+    @clock = clock
+  end
+
+  # charges `amount` tokens to the specified `bucket`.
+  #
+  # charging only happens if the bucket has sufficient tokens.
+  # the level of "sufficient tokens" can be adjusted by passing in option[:limit]
+  #
+  # returns a tuple (= Array with two elements) containing
+  # `success:boolean` and `level:Numeric`
+  def charge(bucket, amount, options = nil)
+    success, levels = batch_charge([bucket, amount, options])
+
+    return success, levels[bucket[:key]]
+  end
+
+  # performs several bucket charge operations in batch.
+  #
+  # each operation is passed in as an Array, containing the parameters
+  # for `batch`.
+  #
+  # charging only happens if all buckets have sufficient tokens.
+  # the charges are done transactionally, so either all buckets are charged or none.
+  #
+  # returns a tuple (= Array with two elements) containing
+  # `success:boolean` and `levels:Hash<String, Numeric>`
+  # where `levels` is a hash from bucket keys to bucket levels.
+  def batch_charge(*charges)
+    charges.each do |(bucket, amount, options)|
+      unless amount > 0
+        message = "tried to charge #{amount}, needs to be Numeric and > 0"
+        raise ArgumentError, message
+      end
+    end
+
+    run_script(charges)
+  end
+
+  # returns the current level of tokens in the specified `bucket`.
+  def read_level(bucket)
+    read_levels(bucket)[bucket[:key]]
+  end
+
+  # reports the current level of tokens for each of the specified `buckets`.
+  # returns the levels as a Hash from bucket keys to bucket levels.
+  def read_levels(*buckets)
+    _, levels = run_script(buckets.map { |bucket| [bucket, 0] })
+
+    levels
+  end
+
+  private
+
+  def run_script(charges)
+    props = charges.map(&method(:props_for_charge)).flatten
+    time = @clock.call if @clock
+
+    argv = [time] + props
+    keys = charges.map { |(bucket, _, _)| bucket[:key] }
+
+    success, levels = eval_script(:keys => keys, :argv => argv)
+
+    levels_as_hash = {}
+    levels.each_with_index do |level, index|
+      levels_as_hash[keys[index]] = level.to_f
+    end
+
+    [success > 0, levels_as_hash]
+  end
+
+  def props_for_charge(charge)
+    bucket, amount, options = charge
+
+    [bucket[:rate], bucket[:size], amount, options ? options[:limit] : nil]
+  end
+
+  def eval_script(options)
+    retries = 0
+
+    begin
+      @redis.evalsha(script_sha, options)
+    rescue Redis::CommandError => e
+      if retries > 0
+        raise
+      end
+
+      @@script_sha = nil
+
+      retries = 1
+      retry
+    end
+  end
+
+  def script_sha
+    @@script_sha ||= @redis.script(:load, script_code)
+  end
+
+  def script_code
+    @@script ||= File.read(File.expand_path("../limiter.lua", __FILE__))
+  end
+end
+end

data/lib/redis_token_bucket/version.rb

@@ -0,0 +1,3 @@
+module RedisTokenBucket
+  VERSION = "0.1.0"
+end

data/redis_token_bucket.gemspec

@@ -0,0 +1,26 @@
+# coding: utf-8
+lib = File.expand_path('../lib', __FILE__)
+$LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
+require 'redis_token_bucket/version'
+
+Gem::Specification.new do |spec|
+  spec.name          = "redis_token_bucket"
+  spec.version       = RedisTokenBucket::VERSION
+  spec.authors       = ["Kristian Hanekamp"]
+  spec.email         = ["kris.hanekamp@gmail.com"]
+
+  spec.summary       = %q{Token Bucket Rate Limiting using Redis}
+  spec.homepage      = "https://github.com/krishan/redis_token_bucket"
+  spec.license       = "MIT"
+
+  spec.files         = `git ls-files -z`.split("\x0").reject { |f| f.match(%r{^(test|spec|features)/}) }
+  spec.bindir        = "exe"
+  spec.executables   = spec.files.grep(%r{^exe/}) { |f| File.basename(f) }
+  spec.require_paths = ["lib"]
+
+  spec.add_dependency "redis", "~> 3.0"
+
+  spec.add_development_dependency "bundler", "~> 1.9"
+  spec.add_development_dependency "rake", "~> 10.0"
+  spec.add_development_dependency "rspec", "~> 3.5"
+end

metadata

@@ -0,0 +1,115 @@
+--- !ruby/object:Gem::Specification
+name: redis_token_bucket
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Kristian Hanekamp
+autorequire: 
+bindir: exe
+cert_chain: []
+date: 2016-12-02 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: bundler
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.9'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.9'
+- !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'
+- !ruby/object:Gem::Dependency
+  name: rspec
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.5'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.5'
+description: 
+email:
+- kris.hanekamp@gmail.com
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- ".gitignore"
+- ".rspec"
+- ".travis.yml"
+- Gemfile
+- LICENSE.txt
+- README.md
+- Rakefile
+- bin/console
+- bin/setup
+- demo.rb
+- lib/redis_token_bucket.rb
+- lib/redis_token_bucket/limiter.lua
+- lib/redis_token_bucket/limiter.rb
+- lib/redis_token_bucket/version.rb
+- redis_token_bucket.gemspec
+homepage: https://github.com/krishan/redis_token_bucket
+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: Token Bucket Rate Limiting using Redis
+test_files: []