Skip to content

Commit

Permalink
Actually rate limit when needed
Browse files Browse the repository at this point in the history
  • Loading branch information
northeastprince committed Jan 3, 2025
1 parent af8c8ee commit b86151b
Show file tree
Hide file tree
Showing 9 changed files with 98 additions and 4 deletions.
2 changes: 2 additions & 0 deletions app/jobs/application_job.rb
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
class ApplicationJob < ActiveJob::Base
discard_on ActiveJob::DeserializationError # most likely a record that's not there anymore

include RateLimitable
end
2 changes: 1 addition & 1 deletion app/jobs/hackathons/website_archival_job.rb
Original file line number Diff line number Diff line change
@@ -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)
Expand Down
28 changes: 28 additions & 0 deletions app/models/concerns/rate_limitable.rb
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion app/models/hackathon/website/archivable.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
38 changes: 38 additions & 0 deletions app/models/lock.rb
Original file line number Diff line number Diff line change
@@ -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
3 changes: 2 additions & 1 deletion config/initializers/action_mailer_concurrency.rb
Original file line number Diff line number Diff line change
@@ -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
4 changes: 4 additions & 0 deletions config/schedule.yml
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
11 changes: 11 additions & 0 deletions db/migrate/20250103164545_create_locks.rb
Original file line number Diff line number Diff line change
@@ -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
12 changes: 11 additions & 1 deletion db/schema.rb

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit b86151b

Please sign in to comment.