Class: RedisLocks::TokenBucket

Inherits:
Object
  • Object
show all
Defined in:
lib/redis_locks/token_bucket.rb

Constant Summary collapse

NAMESPACE = 
"token-bucket"
SCRIPT = 
<<-LUA
  local epoch = tonumber(ARGV[1])
  local rps = tonumber(ARGV[2])
  local burst = tonumber(ARGV[3])
  local key = KEYS[1]

  local token = 1.0 / rps
  local t = redis.call('get', key)
  if not t then
    t = epoch
  else
    t = tonumber(t)
  end

  if t < epoch then
    t = epoch
  elseif t > (epoch + (burst * token)) then
    return 0
  end

  redis.call('set', key, t + token)
  return 1
LUA
DIGEST = 
Digest::SHA1.hexdigest(SCRIPT)

Instance Method Summary collapse

Constructor Details

#initialize(key, period: 1, number: 1, redis: RedisLocks.redis) ⇒ TokenBucket

‘number` tokens are added to the bucket every `period` seconds (up to a max of `number` tokens being available). Each time a resource is used, a token is removed from the bucket; if no tokens are available, no resource may be used.



44
45
46
47
48
49
 
# File 'lib/redis_locks/token_bucket.rb', line 44

def initialize(key, period: 1, number: 1, redis: RedisLocks.redis)
  @key = "#{NAMESPACE}:#{key}".freeze
  @rps = number.to_f / period.to_i
  @burst = number.to_i
  @redis = Connections.ensure_pool(redis)
end

Instance Method Details

#take ⇒ Object 



51
52
53
54
55
56
57
58
59
60
61
62
63
 
# File 'lib/redis_locks/token_bucket.rb', line 51

def take
  took = @redis.with do |conn|
    RedisLocks.evalsha_or_eval(
      conn: conn,
      script: SCRIPT,
      digest: DIGEST,
      keys: [@key],
      args: [epoch_f(conn), @rps, @burst]
    )
  end

  took == 1
end

#take! ⇒ Object

Raises:



65
66
67
 
# File 'lib/redis_locks/token_bucket.rb', line 65

def take!
  raise RateLimitExceeded.new(@key, @rps) unless take
end