Actually rate limit when needed (#610)

Author: northeastprince

app/jobs/application_job.rb

@@ -1,3 +1,5 @@
 class ApplicationJob < ActiveJob::Base
   discard_on ActiveJob::DeserializationError # most likely a record that's not there anymore
+
+  include RateLimitable
 end

app/jobs/hackathons/website_archival_job.rb

@@ -1,5 +1,5 @@
 class Hackathons::WebsiteArchivalJob < ApplicationJob
-  limits_concurrency to: 15, duration: 1.minute, group: "Wayback Machine", key: "API"
+  rate_limit "Wayback Machine", to: 15, within: 1.minute
   queue_as :low
 
   def perform(hackathon)

app/models/concerns/rate_limitable.rb

@@ -0,0 +1,28 @@
+module RateLimitable
+  extend ActiveSupport::Concern
+
+  class Limit < StandardError
+    attr_reader :duration
+    def initialize(duration: nil)
+      @duration = duration
+      super
+    end
+  end
+
+  class_methods do
+    def rate_limit(key = name, to:, within:)
+      around_perform do |job, block|
+        unless Lock.acquire(key, limit: to, duration: within) { block }
+          raise Limit.new(duration:)
+        end
+      end
+    end
+  end
+
+  included do
+    rescue_from RateLimitable::Limit do |limit|
+      retry_job wait: limit.duration
+      logger.info "#{self.class.name} #{job_id} was rate limited for at least #{limit.duration.inspect}"
+    end
+  end
+end

app/models/hackathon/website/archivable.rb

@@ -46,7 +46,7 @@ def archive_with(job_id)
   end
 
   class FollowUpJob < ApplicationJob
-    limits_concurrency to: 15, duration: 1.minute, group: "Wayback Machine", key: "API"
+    rate_limit "Wayback Machine", to: 15, within: 1.minute
     queue_as :low
 
     def perform(hackathon, id)

app/models/lock.rb

@@ -0,0 +1,38 @@
+class Lock < ApplicationRecord
+  scope :expired, -> { where "expiration <= ?", Time.now }
+
+  class << self
+    def acquire(key, limit: 1, duration: nil, &block)
+      lock = nil
+      transaction do
+        lock = find_by(key:) || create!(key:, expiration: duration&.from_now)
+        if lock.capacity == limit
+          return false
+        end
+      end
+
+      lock.acquire
+
+      begin
+        yield block
+      ensure
+        lock.release
+      end
+    end
+  end
+
+  def acquire(capacity = 1)
+    increment!(:capacity, capacity)
+  end
+
+  def release(quantity = 1)
+    transaction do
+      reload
+      if capacity == 1
+        destroy!
+      else
+        decrement!(:capacity, quantity)
+      end
+    end
+  end
+end

config/initializers/action_mailer_concurrency.rb

@@ -1,4 +1,5 @@
 Rails.application.config.after_initialize do
   # AWS SES rate-limits to 14/second, so we'll set it to 10 to be safe
-  ActionMailer::MailDeliveryJob.limits_concurrency key: "mail", to: 10, duration: 1.second
+  ActionMailer::MailDeliveryJob.include(RateLimitable)
+    .rate_limit to: 10, within: 1.second
 end

config/schedule.yml

@@ -1,3 +1,7 @@
+lock_maintenance:
+  cron: "every second"
+  command: "Lock.expired.delete_all"
+
 digests_delivery:
   cron: "every tuesday at 10am on America/Los_Angeles"
   class: "Hackathons::DigestsDeliveryJob"

db/migrate/20250103164545_create_locks.rb

@@ -0,0 +1,11 @@
+class CreateLocks < ActiveRecord::Migration[8.1]
+  def change
+    create_table :locks do |t|
+      t.string :key, null: false, index: {unique: true}
+      t.integer :capacity, null: false, default: 0
+      t.datetime :expiration, index: true
+
+      t.timestamps
+    end
+  end
+end