github.com/godaddy-x/freego@v1.0.156/cache/limiter/rate_limiter_redis.go (about)

     1  package rate
     2  
     3  import (
     4  	"github.com/garyburd/redigo/redis"
     5  	"github.com/godaddy-x/freego/cache"
     6  	"github.com/godaddy-x/freego/utils"
     7  	"github.com/godaddy-x/freego/zlog"
     8  	"time"
     9  )
    10  
    11  const (
    12  	limiterKey = "redis:limiter:"
    13  )
    14  
    15  var limiterScript = redis.NewScript(1, `
    16  	-- KEYS = [resource]
    17  	local cache_key = redis.pcall('HMGET', KEYS[1], 'last_request_time', 'surplus_token')
    18  	local last_request_time = cache_key[1] -- 上次请求时间
    19  	local surplus_token = tonumber(cache_key[2]) -- 剩余的令牌数
    20  	local bucket_token = tonumber(ARGV[1]) -- 令牌桶最大数
    21  	local token_rate = tonumber(ARGV[2]) -- 令牌数生成速率/秒
    22  	local now_request_time = tonumber(ARGV[3]) -- 当前请求时间/毫秒
    23  	local token_ms_rate = token_rate/1000 -- 每毫秒生产令牌速率
    24  	local past_time = 0 -- 两次请求时间差
    25  	if surplus_token == nil then
    26  		surplus_token = bucket_token -- 填充剩余令牌数最大值
    27  		last_request_time = now_request_time -- 填充上次请求时间
    28  	else
    29  		past_time = now_request_time - last_request_time -- 填充两次请求时间差
    30  		if past_time <= 0 then
    31  			past_time = 0 -- 防止多台服务器出现时间差小于0
    32  		end
    33  		local add_token = math.floor(past_time * token_ms_rate)  -- 通过时间差生成令牌数,向下取整
    34  		surplus_token = math.min((surplus_token + add_token), bucket_token) -- 剩余令牌数+生成令牌数 <= 令牌桶最大数
    35  	end
    36  	local status = 0 -- 判定状态 0.拒绝 1.通过
    37  	if surplus_token > 0 then
    38  		surplus_token = surplus_token - 1 -- 通过则剩余令牌数-1
    39  		last_request_time = last_request_time + past_time -- 刷新最后请求时间
    40  		redis.call('HMSET', KEYS[1], 'last_request_time', last_request_time, 'surplus_token', surplus_token) -- 更新剩余令牌数和最后请求时间
    41  		status = 1
    42  	end
    43  	if surplus_token < 0 then
    44  		redis.call('PEXPIRE', KEYS[1], 3000) -- 设置超时重置数据
    45  	end
    46  	return status
    47  `)
    48  
    49  type RedisRateLimiter struct {
    50  	option Option
    51  }
    52  
    53  func (self *RedisRateLimiter) key(resource string) string {
    54  	return utils.AddStr(limiterKey, resource)
    55  }
    56  
    57  func (self *RedisRateLimiter) Allow(resource string) bool {
    58  	client, err := cache.NewRedis()
    59  	if err != nil {
    60  		zlog.Error("redis rate limiter get client failed", 0, zlog.AddError(err))
    61  		return false
    62  	}
    63  	rds := client.Pool.Get()
    64  	defer client.Close(rds)
    65  	res, err := limiterScript.Do(rds, self.key(resource), self.option.Bucket, self.option.Limit, time.Now().UnixNano()/1e6)
    66  	if err != nil {
    67  		zlog.Error("redis rate limiter client do lua script failed", 0, zlog.AddError(err))
    68  		return false
    69  	}
    70  	if v, b := res.(int64); b && v == 1 {
    71  		return true
    72  	}
    73  	return false
    74  }