diff --git a/app/models/rate_limit.rb b/app/models/rate_limit.rb index 3236f1ac4..b3d0b8ab0 100644 --- a/app/models/rate_limit.rb +++ b/app/models/rate_limit.rb @@ -10,7 +10,7 @@ class RateLimit < ApplicationRecord # `cost` is the number of points the action costs. # `rate` is the number of points per second that are refilled. # `burst` is the maximum number of points that can be saved up. - def self.create_or_update!(action:, keys:, cost:, rate:, burst:) + def self.create_or_update!(action:, keys:, cost:, rate:, burst:, minimum_points: -30) # { key0: keys[0], ..., keyN: keys[N] } key_params = keys.map.with_index { |key, i| [:"key#{i}", key] }.to_h @@ -20,6 +20,12 @@ class RateLimit < ApplicationRecord # Do an upsert, creating a new rate limit object for each key that doesn't # already exist, and updating the limit for each limit that already exists. # + # If the current point count is negative, then we're limited. Penalize the + # caller 1 second (1 rate unit), up to a maximum penalty of 30 seconds (by default). + # + # Otherwise, if the point count is positive, then we're not limited. Update + # the point count and subtract the cost of the call. + # # https://www.postgresql.org/docs/current/sql-insert.html#SQL-ON-CONFLICT sql = <<~SQL INSERT INTO rate_limits (created_at, updated_at, action, key, points) @@ -30,9 +36,9 @@ class RateLimit < ApplicationRecord points = CASE WHEN rate_limits.points + :rate * EXTRACT(epoch FROM (:now - rate_limits.updated_at)) < 0 THEN - LEAST(:burst, rate_limits.points + :rate * EXTRACT(epoch FROM (:now - rate_limits.updated_at))) + GREATEST(:minimum_points, LEAST(:burst, rate_limits.points + :rate * EXTRACT(epoch FROM (:now - rate_limits.updated_at))) - :rate) ELSE - LEAST(:burst, rate_limits.points + :rate * EXTRACT(epoch FROM (:now - rate_limits.updated_at))) - :cost + GREATEST(:minimum_points, LEAST(:burst, rate_limits.points + :rate * EXTRACT(epoch FROM (:now - rate_limits.updated_at))) - :cost) END RETURNING * SQL @@ -44,6 +50,7 @@ class RateLimit < ApplicationRecord burst: burst, cost: cost, points: burst - cost, + minimum_points: minimum_points, **key_params } diff --git a/test/unit/rate_limit_test.rb b/test/unit/rate_limit_test.rb index d120a494a..6a8a47725 100644 --- a/test/unit/rate_limit_test.rb +++ b/test/unit/rate_limit_test.rb @@ -26,13 +26,13 @@ class RateLimitTest < ActiveSupport::TestCase assert_equal(9, limiter.rate_limits.first.points) end - should "be limited if the point count is negative" do + should "be limited and penalize the caller 1 second if the point count is negative" do freeze_time create(:rate_limit, action: "write", key: "users/1", points: -1) limiter = RateLimiter.new("write", ["users/1"], cost: 1) assert_equal(true, limiter.limited?) - assert_equal(-1, limiter.rate_limits.first.points) + assert_equal(-2, limiter.rate_limits.first.points) end should "not be limited if the point count was positive before the action" do @@ -50,17 +50,17 @@ class RateLimitTest < ActiveSupport::TestCase limiter = RateLimiter.new("write", ["users/1"], cost: 1, rate: 1, burst: 10) assert_equal(true, limiter.limited?) - assert_equal(-2, limiter.rate_limits.first.points) + assert_equal(-3, limiter.rate_limits.first.points) travel 1.second limiter = RateLimiter.new("write", ["users/1"], cost: 1, rate: 1, burst: 10) assert_equal(true, limiter.limited?) - assert_equal(-1, limiter.rate_limits.first.points) + assert_equal(-3, limiter.rate_limits.first.points) travel 5.second limiter = RateLimiter.new("write", ["users/1"], cost: 1, rate: 1, burst: 10) assert_equal(false, limiter.limited?) - assert_equal(3, limiter.rate_limits.first.points) + assert_equal(1, limiter.rate_limits.first.points) travel 60.second limiter = RateLimiter.new("write", ["users/1"], cost: 1, rate: 1, burst: 10)