mirror of
https://gitlink.org.cn/Gitlink/forgeplus.git
synced 2026-07-13 21:44:32 +08:00
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 |