forked from Gitlink/forgeplus
48 lines
1.0 KiB
Ruby
48 lines
1.0 KiB
Ruby
class TokenBucketRateLimiter
|
||
def initialize(user_id,rate, capacity)
|
||
@rate = rate # 令牌发放速率（每秒发放的令牌数）
|
||
@capacity = capacity # 令牌桶容量
|
||
@redis = $redis_cache
|
||
@key = "#{user_id}:TokenBucket"
|
||
end
|
||
|
||
def refill
|
||
now = Time.now.to_f * 1000
|
||
tokens = $redis_cache.hget(@key, 'tokens')
|
||
timestamp = $redis_cache.hget(@key, 'timestamp')
|
||
|
||
if tokens.nil?
|
||
tokens = @capacity
|
||
timestamp = now
|
||
else
|
||
tokens = [@capacity, tokens.to_i + (now - timestamp.to_f) * @rate / 1000].min
|
||
timestamp = now
|
||
end
|
||
|
||
$redis_cache.hset(@key, 'tokens', tokens)
|
||
$redis_cache.hset(@key, 'timestamp', timestamp)
|
||
end
|
||
|
||
def allow_request
|
||
refill
|
||
tokens = @redis.hget(@key, 'tokens')
|
||
if !tokens.nil? && tokens.to_i >= 1
|
||
@redis.hset(@key, 'tokens', tokens.to_i - 1)
|
||
return true
|
||
end
|
||
false
|
||
end
|
||
end
|
||
|
||
=begin
|
||
#测试通过
|
||
limiter = TokenBucketRateLimiter.new(user_id,10, 40) # 设置令牌发放速率为10，令牌桶容量为40
|
||
30.times do
|
||
if limiter.allow_request
|
||
puts "Allow"
|
||
else
|
||
puts "Reject"
|
||
end
|
||
end
|
||
|
||
=end |