redis高阶第一篇：令牌桶算法限流

令牌桶算法核心思想是：

1、以恒定速率向桶中添加令牌

2、桶中令牌数有上限，为桶的容量，超过容量后令牌会直接丢弃

3、请求到来时会尝试从桶中取一个令牌，取成功则通过，服务端会处理此请求，否则拒绝。

token_bucket.lua，内容如下：

redis.replicate_commands()

local tokens_key = KEYS[1]
local timestamp_key = KEYS[2]

local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local now = tonumber(ARGV[3]) or redis.call('TIME')[1]
local requested = tonumber(ARGV[4])

local fill_time = capacity / rate
local ttl = math.floor(fill_time * 2)

local last_tokens = tonumber(redis.call("get", tokens_key)) or capacity
local last_refreshed = tonumber(redis.call("get", timestamp_key)) or 0

local delta = math.max(0, now-last_refreshed)
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
local allowed = filled_tokens >= requested
local new_tokens = allowed and filled_tokens - requested or filled_tokens

if ttl > 0 then
  redis.call("setex", tokens_key, ttl, new_tokens)
  redis.call("setex", timestamp_key, ttl, now)
end

return { allowed and 1 or 0, new_tokens }

上面lua脚本拷贝自spring-cloud-gateway-server-xxx.jar包中META-INF目录中scripts目录中的request_rate_limiter.lua文件。

<dependency>
    <groupId>org.springframework.cloud</groupId>
    <artifactId>spring-cloud-gateway-server</artifactId>
    <version>4.3.0</version>
</dependency>

在3.0.0版本中，还可以看到可以通过redis.log(redis.LOG_WARNING, "tokens_key " .. tokens_key)在redis中打印日志。

上面lua脚本的解释：

KEYS[1]是存放令牌数的key，KEYS[2]是存放请求时间的key

ARGV[1]是往桶中放令牌的速率，单位是个/秒，如值为5，表示每秒往桶中放5个令牌。ARGV[2]是桶的容量，ARGV[4]是当前请求要使用多少个令牌，一般都是1个。

allowed_num是请求是否可以被处理的标识，值为1表示请求可以被处理，值为0表示请求数已超过阈值，应该被拦截。

定义判断是否允许请求通过的方法，如果服务端是golang开发，则如下：

func getKeys(id, routeId string) []string {
    prefix := fmt.Sprintf("request_rate_limiter.{%v.%v}.", routeId, id)
    tokenKey := fmt.Sprintf("%vtokens", prefix)
    timestampKey := fmt.Sprintf("%vtimestamp", prefix)
    return []string{tokenKey, timestampKey}
}

func isAllowed(userId, routeId string, capacity, rate, tokens int) (bool, error) {
    result, err := rdb.Eval(ctx, script, getKeys(userId, routeId), rate, capacity, "", tokens).Result()
    if err != nil {
        return false, err
    }

    arr, _ := result.([]interface{})
    a, _ := arr[0].(int64)
    return a == 1, nil
}

这串代码是根据org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter中的getKeys()方法和isAllowed()方法翻译过来的。

调用上述函数代码示例：

func main() {
    for i := 0; i < 10000000; i++ {
        allowed, err := isAllowed("12345", "/feed/v1/channels", 10, 5, 1) // capacity=10, rate=5 tokens/sec, 取1个
        if err != nil {
            fmt.Println("Error:", err)
            return
        }
        if allowed {
            fmt.Println(time.Now().Format(time.DateTime), "Request allowed")
        } else {
            fmt.Println(time.Now().Format(time.DateTime), "Rate limit exceeded")
        }
    }
}