|
- 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
|