class Rookout::Augs::AugRateLimiter

Attributes

windows[R]

Public Class Methods

new(quota, window_size, active_limit) click to toggle source
# File lib/rookout/augs/aug_rate_limiter.rb, line 7
def initialize quota, window_size, active_limit
  @quota = quota
  @has_quota = !@quota.nil? && (@quota > 0)
  @window_size = window_size
  @active_weight = quota / active_limit

  @mutex = Mutex.new
  @active_count = 0
  @windows = {}
end

Public Instance Methods

with_limit(start_time = nil) { || ... } click to toggle source
# File lib/rookout/augs/aug_rate_limiter.rb, line 18
def with_limit start_time = nil
  active = false
  start_time ||= Time.now

  # If quota, verify it
  if @has_quota
    # Get current time
    now_ns = Utils.time_to_nanoseconds start_time

    # Calculate window keys
    current_window_key, prev_window_key = timestamp_to_window_keys now_ns

    @mutex.synchronize do
      # Clean old windows
      cleanup now_ns
      # Increase active count
      @active_count += 1
      active = true

      # If exceeding quota
      if current_usage(now_ns, current_window_key, prev_window_key) > @quota
        warning = Processor::RookError.new Exceptions::RookRuleRateLimited.new
        UserWarnings.notify_warning warning
        return
      end
    end
  end

  begin
    yield
  ensure
    if @has_quota
      @mutex.synchronize { record_usage current_window_key, Utils.time_to_nanoseconds(Time.now) - now_ns }
    end
  end
ensure
  # Reduce active count
  @mutex.synchronize { @active_count -= 1 } if active
end

Private Instance Methods

cleanup(now_ns) click to toggle source
# File lib/rookout/augs/aug_rate_limiter.rb, line 93
def cleanup now_ns
  @windows.reject! { |key, _| key < (now_ns - @window_size * 5) } if @windows.length > 10
end
current_usage(now_ns, current_window_key, prev_window_key) click to toggle source
# File lib/rookout/augs/aug_rate_limiter.rb, line 68
def current_usage now_ns, current_window_key, prev_window_key
  # Get usage information
  current_window_usage = @windows[current_window_key]
  if current_window_usage.nil?
    # Create new window
    current_window_usage = @windows[current_window_key] = 0
  end
  prev_window_usage = @windows[prev_window_key] || 0

  # Previous window weight
  prev_weight = 1 - (now_ns - current_window_key) / @window_size.to_f

  # Final weighted usage
  (prev_window_usage * prev_weight) + (@active_count * @active_weight) + current_window_usage
end
record_usage(current_window_key, duration) click to toggle source
# File lib/rookout/augs/aug_rate_limiter.rb, line 84
def record_usage current_window_key, duration
  # windows might be cleared while aug is running (unlikely)
  prev_value = @windows[current_window_key]
  return if prev_value.nil?

  # Add value to quota
  @windows[current_window_key] += [duration, 5].max.to_f
end
timestamp_to_window_keys(now_ns) click to toggle source
# File lib/rookout/augs/aug_rate_limiter.rb, line 62
def timestamp_to_window_keys now_ns
  current_window_key = (now_ns / @window_size) * @window_size
  prev_window_key = current_window_key - @window_size
  [current_window_key, prev_window_key]
end