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 = redis
end

Instance Method Details

#take ⇒ Object 



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

def take
  epoch_i, microseconds = @redis.time
  epoch_f = epoch_i + (microseconds.to_f/1_000_000)
  took = RedisLocks.evalsha_or_eval(
    redis: @redis,
    script: SCRIPT,
    digest: DIGEST,
    keys: [@key],
    args: [epoch_f, @rps, @burst]
  )
  took == 1
end

#take! ⇒ Object

Raises:



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

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