Initial implementation

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.

Author: errm

.github/workflows/main.yml

@@ -16,12 +16,24 @@ 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
+        ports:
+          - 6379:6379
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
     - name: Set up Ruby
       uses: ruby/setup-ruby@v1
       with:
         ruby-version: ${{ matrix.ruby }}
         bundler-cache: true
     - name: Run the default task
       run: bundle exec rake
+      env:
+        MILLRACE_REDIS_URL: redis://localhost:6379

Gemfile

@@ -11,3 +11,6 @@ gem "rubocop", "~> 1.21"
 gem "rubocop-performance"
 gem "rubocop-rails"
 gem "rubocop-rspec"
+
+gem "hiredis-client"
+gem "redis"

Gemfile.lock

@@ -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

lib/millrace.rb

@@ -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

lib/millrace/rate_limit.rb

@@ -0,0 +1,72 @@
+require "digest"
+require "prorate"
+
+module Millrace
+  class RateLimit
+    def initialize(name:, rate:, window:, penalty: 0, redis_config: nil)
+      @name         = name
+      @rate         = rate
+      @window       = window
+      @penalty      = penalty
+      @redis_config = redis_config
+    end
+
+    attr_reader :name, :rate, :window
+
+    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_config
+        @redis_config || { url: ENV.fetch("MILLRACE_REDIS_URL", nil) }.compact
+      end
+
+      def redis
+        Thread.current["millrace_#{name}_redis"] ||= Redis.new(redis_config)
+      end
+  end
+end

lib/millrace/rate_limited.rb

@@ -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

millrace.gemspec

@@ -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

spec/rate_limit_spec.rb

@@ -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