Class: RedisLocks::TokenBucket

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

Constant Summary collapse

NAMESPACE = 
"token-bucket"
SCRIPT = 
"  local epoch = tonumber(ARGV[1])\n  local rps = tonumber(ARGV[2])\n  local burst = tonumber(ARGV[3])\n  local key = KEYS[1]\n\n  local token = 1.0 / rps\n  local t = redis.call('get', key)\n  if not t then\n    t = epoch\n  else\n    t = tonumber(t)\n  end\n\n  if t < epoch then\n    t = epoch\n  elseif t > (epoch + (burst * token)) then\n    return 0\n  end\n\n  redis.call('set', key, t + token)\n  return 1\n"
DIGEST = 
Digest::SHA1.hexdigest(SCRIPT)

Instance Method Summary collapse

Constructor Details

#initialize(key, redis:, period: 1, number: 1) ⇒ 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, redis:, period: 1, number: 1)
  @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