令牌桶算法核心思想是：
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") } } }
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。