Class: Prorate::Throttle

Inherits:
Object
  • Object
show all
Defined in:
lib/prorate/throttle.rb

Constant Summary collapse

CURRENT_SCRIPT_HASH =
get_script_hash

Class Method Summary collapse

Instance Method Summary collapse

Constructor Details

#initializeThrottle

Returns a new instance of Throttle.



28
29
30
31
32
33
34
# File 'lib/prorate/throttle.rb', line 28

def initialize(*)
  super
  @discriminators = [name.to_s]
  self.redis = NullPool.new(redis) unless redis.respond_to?(:with)
  raise MisconfiguredThrottle if ((period <= 0) || (limit <= 0))
  @leak_rate = limit.to_f / period # tokens per second;
end

Class Method Details

.get_script_hashObject



20
21
22
23
24
# File 'lib/prorate/throttle.rb', line 20

def self.get_script_hash
  script_filepath = File.join(__dir__,"rate_limit.lua")
  script = File.read(script_filepath)
  Digest::SHA1.hexdigest(script)
end

Instance Method Details

#<<(discriminator) ⇒ Object



36
37
38
# File 'lib/prorate/throttle.rb', line 36

def <<(discriminator)
  @discriminators << discriminator
end

#run_lua_throttler(redis:, identifier:, bucket_capacity:, leak_rate:, block_for:) ⇒ Object



56
57
58
59
60
61
62
63
64
65
66
67
68
69
# File 'lib/prorate/throttle.rb', line 56

def run_lua_throttler(redis: , identifier: , bucket_capacity: , leak_rate: , block_for: )
  redis.evalsha(CURRENT_SCRIPT_HASH, [], [identifier, bucket_capacity, leak_rate, block_for])
rescue Redis::CommandError => e
  if e.message.include? "NOSCRIPT"
    # The Redis server has never seen this script before. Needs to run only once in the entire lifetime of the Redis server (unless the script changes)
    script_filepath = File.join(__dir__,"rate_limit.lua")
    script = File.read(script_filepath)
    raise ScriptHashMismatch if Digest::SHA1.hexdigest(script) != CURRENT_SCRIPT_HASH
    redis.script(:load, script)
    redis.evalsha(CURRENT_SCRIPT_HASH, [], [identifier, bucket_capacity, leak_rate, block_for])
  else
    raise e
  end
end

#throttle!Object



40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
# File 'lib/prorate/throttle.rb', line 40

def throttle!
  discriminator = Digest::SHA1.hexdigest(Marshal.dump(@discriminators))
  identifier = [name, discriminator].join(':')
  
  redis.with do |r|
    logger.info { "Applying throttle counter %s" % name }
    remaining_block_time, bucket_level = run_lua_throttler(redis: r, identifier: identifier, bucket_capacity: limit, leak_rate: @leak_rate, block_for: block_for)

    if remaining_block_time > 0
      logger.warn { "Throttle %s exceeded limit of %d in %d seconds and is blocked for the next %d seconds" % [name, limit, period, remaining_block_time] }
      raise Throttled.new(remaining_block_time)
    end
    available_calls = limit - bucket_level
  end
end