rate_limiter.rb

# frozen_string_literal: true

class RateLimiter
  class Error < StandardError; end
  class UnknownIntervalError < RateLimiter::Error; end
  class WaitTimeError < RateLimiter::Error; end

  attr_reader :bucket_name, :interval, :maximum_wait_time, :rate, :redis

  def initialize(bucket_name, options = {})
    @bucket_name = bucket_name
    @interval = options[:interval]&.to_sym || :second
    @maximum_wait_time = options[:maximum_wait_time] || 5.minutes.in_seconds
    @rate = options[:rate].to_i
    @redis = options[:redis] || Redis.current
  end

  def within_limit(&block)
    return block.call if unlimited?
    increment_count!
    count <= rate ? with_execution_timeout(&block) : enqueue(&block)
  end

  private
  
  def count
    @count ||= redis.get(redis_key).to_i
  end

  def enqueue(&block)
    t = Time.zone.now.to_f
    sleep(t.ceil - t)

    within_limit(&block)
  end

  def increment_count!
    @count = redis.pipelined do
      redis.incr(redis_key)
      redis.expire(redis_key, ttl)
    end.first
  end

  def redis_key
    ["rate_limiter", bucket_name, timestamp_key].compact.join(":")
  end

  def timestamp_key
    case interval
    when :second then Time.zone.now.strftime("%M:%S")
    when :minute then Time.zone.now.strftime("%H:%M")
    when :hour then Time.zone.now.strftime("%F:%H")
    when :day then Time.zone.now.strftime("%F")
    else raise UnknownIntervalError, "Invalid interval: #{interval}"
    end
  end

  def ttl
    case interval
    when :second then 1
    when :minute then 60
    when :hour then 60 * 60
    when :day then 60 * 60 * 24
    else raise UnknownIntervalError, "Invalid interval: #{interval}"
    end
  end

  def unlimited?
    bucket_name.nil? || rate.to_i.zero?
  end

  def with_execution_timeout(&block)
    yield if maximum_wait_time.to_i.zero?
    error_message = "#{bucket_name} failed to execute within #{maximum_wait_time} seconds"
    Timeout.timeout(maximum_wait_time, WaitTimeError, error_message, &block)
  end
end