Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Initial implementation
Browse files Browse the repository at this point in the history
Impliments a leaky bucket rate limiter, that unlike
prorate continues to to count requests against the
limit even when the rate limiter is in the blocking
state.

This means that the client has to slow down, or
they will remain blocked indefinately.

Optionally a penalty can be added, that adds additonal tokens
to the bucket at the point that the limit is breached, to
futher ensure that the block lasts longer for clients that
are only marginly breaching the rate limit.
errm committed May 31, 2024

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
1 parent d17e46b commit a240b11
Showing 8 changed files with 194 additions and 2 deletions.
9 changes: 9 additions & 0 deletions .github/workflows/main.yml
Original file line number Diff line number Diff line change
@@ -16,6 +16,14 @@ jobs:
ruby:
- '3.3.1'

services:
redis:
image: redis:7.2.5
options: >-
--health-cmd "redis-cli ping"
--health-interval 10s
--health-timeout 5s
--health-retries 5
steps:
- uses: actions/checkout@v3
- name: Set up Ruby
@@ -25,3 +33,4 @@ jobs:
bundler-cache: true
- name: Run the default task
run: bundle exec rake
env: MILLRACE_REDIS_URL: redis://redis:6379
3 changes: 3 additions & 0 deletions Gemfile
Original file line number Diff line number Diff line change
@@ -11,3 +11,6 @@ gem "rubocop", "~> 1.21"
gem "rubocop-performance"
gem "rubocop-rails"
gem "rubocop-rspec"

gem "hiredis-client"
gem "redis"
11 changes: 11 additions & 0 deletions Gemfile.lock
Original file line number Diff line number Diff line change
@@ -2,6 +2,7 @@ PATH
remote: .
specs:
millrace (0.1.0)
prorate

GEM
remote: https://rubygems.org/
@@ -23,6 +24,8 @@ GEM
connection_pool (2.4.1)
diff-lcs (1.5.1)
drb (2.2.1)
hiredis-client (0.22.2)
redis-client (= 0.22.2)
i18n (1.14.5)
concurrent-ruby (~> 1.0)
json (2.7.2)
@@ -33,10 +36,16 @@ GEM
parser (3.3.1.0)
ast (~> 2.4.1)
racc
prorate (0.7.3)
redis (>= 2)
racc (1.7.3)
rack (3.0.11)
rainbow (3.1.1)
rake (13.2.1)
redis (5.2.0)
redis-client (>= 0.22.0)
redis-client (0.22.2)
connection_pool
regexp_parser (2.9.0)
rexml (3.2.8)
strscan (>= 3.0.9)
@@ -96,8 +105,10 @@ PLATFORMS
ruby

DEPENDENCIES
hiredis-client
millrace!
rake (~> 13.0)
redis
rspec (~> 3.0)
rubocop (~> 1.21)
rubocop-performance
4 changes: 2 additions & 2 deletions lib/millrace.rb
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
# frozen_string_literal: true

require_relative "millrace/version"
require_relative "millrace/rate_limited"
require_relative "millrace/rate_limit"

module Millrace
class Error < StandardError; end
# Your code goes here...
end
68 changes: 68 additions & 0 deletions lib/millrace/rate_limit.rb
Original file line number Diff line number Diff line change
@@ -0,0 +1,68 @@
require "digest"
require "prorate"

module Millrace
class RateLimit
def initialize(name:, rate:, window:, penalty: 0, redis_url: ENV.fetch("MILLRACE_REDIS_URL", {}))
@name = name
@rate = rate
@window = window
@penalty = penalty
@redis_url = redis_url
end

attr_reader :name, :rate, :window, :redis_url

def before(controller)
bucket = get_bucket(controller.request.remote_ip)
level = record_request(bucket)

return unless level > threshold

if level - 1 < threshold
level = bucket.fillup(penalty).level
end

raise RateLimited.new(limit_name: name, retry_after: retry_after(level))
end

private

def retry_after(level)
((level - threshold) / rate).to_i
end

def record_request(bucket)
bucket.fillup(1).level
end

def get_bucket(ip)
Prorate::LeakyBucket.new(
redis: redis,
redis_key_prefix: key(ip),
leak_rate: rate,
bucket_capacity: capacity,
)
end

def key(ip)
"millrace.#{name}.#{Digest::SHA1.hexdigest(ip)}"
end

def capacity
(threshold * 2) + penalty
end

def threshold
window * rate
end

def penalty
@penalty * rate
end

def redis
Thread.current["millrace_#{name}_redis"] ||= Redis.new(redis_url)
end
end
end
10 changes: 10 additions & 0 deletions lib/millrace/rate_limited.rb
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
module Millrace
class RateLimited < StandardError
def initialize(limit_name:, retry_after:)
@limit_name = limit_name
@retry_after = retry_after
end

attr_reader :limit_name, :retry_after
end
end
2 changes: 2 additions & 0 deletions millrace.gemspec
Original file line number Diff line number Diff line change
@@ -38,4 +38,6 @@ Gem::Specification.new do |spec|
# For more information and examples about making a new gem, check out our
# guide at: https://bundler.io/guides/creating_gem.html
spec.metadata["rubygems_mfa_required"] = "true"

spec.add_dependency "prorate"
end
89 changes: 89 additions & 0 deletions spec/rate_limit_spec.rb
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
# frozen_string_literal: true

RSpec.describe Millrace::RateLimit do
let(:subject) do
described_class.new(
name: "test",
rate: 10,
window: 2,
penalty: penalty,
)
end

let(:penalty) { 1 }

let(:controller) do
double(:controller, request: double(:request, remote_ip: to_s))
end

describe "#before" do
it "rate limits" do
# Fill the bucket
20.times { subject.before(controller) }

# hit the threshold and get a penalty
expect { subject.before(controller) }.to raise_error Millrace::RateLimited

sleep 1
# Still blocked for the penalty duration
expect { subject.before(controller) }.to raise_error Millrace::RateLimited

# Not blocked after the penalty duration is over
sleep 1
subject.before(controller)
end

it "returns an exeption with the correct name" do
# Fill the bucket
20.times { subject.before(controller) }

# hit the threshold and get an error
expect { subject.before(controller) }.to raise_error do |exception|
expect(exception.limit_name).to eq "test"
end
end

it "returns an exeption with the correct retry time" do
# Fill the bucket
20.times { subject.before(controller) }

# hit the threshold and get an error
expect { subject.before(controller) }.to raise_error do |exception|
expect(exception.retry_after).to eq 1
end
end

context "a longer penalty" do
let(:penalty) { 10 }

it "returns an exeption with the correct retry time" do
# Fill the bucket
20.times { subject.before(controller) }

# hit the threshold and get an error
expect { subject.before(controller) }.to raise_error do |exception|
expect(exception.retry_after).to eq 10
end
end
end

context "additional requests" do
let(:penalty) { 0 }

it "returns an exeption with the correct retry time" do
# Fill the bucket
40.times do
subject.before(controller)
# Keep making requests even though we are rate limited
rescue Millrace::RateLimited
nil
end

# hit the threshold and get an error
expect { subject.before(controller) }.to raise_error do |exception|
expect(exception.retry_after).to eq 2
end
end
end
end
end

0 comments on commit a240b11

Please sign in to comment.