redis实现限流

在面试中金时问到了，卡壳了。

限流算法有

固定窗口、滑动窗口、漏桶算法、令牌桶算法

漏桶算法：

把请求到来想象成往桶里加水，把处理请求想象成桶往外滴水，滴水的速率是固定的。如果桶满了，那么再进来的请求就会拒绝。

令牌桶算法：

以固定频率往桶里放置令牌，当桶满的时候，多出来的令牌被丢弃。每来一个请求，在处理前，先从桶里取一个令牌，如果能取到，就处理请求，如果取不到，则表示达到限流条件了，请求被拒绝。

guava实现了令牌桶算法。核心类是RateLimiter。

spring-cloud-gateway-server.jar有一个内嵌的lua脚本(META-INF/scripts/request_rate_limiter.lua文件)，实现了令牌桶算法。

redis.replicate_commands()

local tokens_key = KEYS[1]
local timestamp_key = KEYS[2]
--redis.log(redis.LOG_WARNING, "tokens_key " .. tokens_key)

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

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

--redis.log(redis.LOG_WARNING, "rate " .. ARGV[1])
--redis.log(redis.LOG_WARNING, "capacity " .. ARGV[2])
--redis.log(redis.LOG_WARNING, "now " .. now)
--redis.log(redis.LOG_WARNING, "requested " .. ARGV[4])
--redis.log(redis.LOG_WARNING, "filltime " .. fill_time)
--redis.log(redis.LOG_WARNING, "ttl " .. ttl)

local last_tokens = tonumber(redis.call("get", tokens_key))
if last_tokens == nil then
  last_tokens = capacity
end
--redis.log(redis.LOG_WARNING, "last_tokens " .. last_tokens)

local last_refreshed = tonumber(redis.call("get", timestamp_key))
if last_refreshed == nil then
  last_refreshed = 0
end
--redis.log(redis.LOG_WARNING, "last_refreshed " .. last_refreshed)

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 = filled_tokens
local allowed_num = 0
if allowed then
  new_tokens = filled_tokens - requested
  allowed_num = 1
end

--redis.log(redis.LOG_WARNING, "delta " .. delta)
--redis.log(redis.LOG_WARNING, "filled_tokens " .. filled_tokens)
--redis.log(redis.LOG_WARNING, "allowed_num " .. allowed_num)
--redis.log(redis.LOG_WARNING, "new_tokens " .. new_tokens)

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

-- return { allowed_num, new_tokens, capacity, filled_tokens, requested, new_tokens }
return { allowed_num, new_tokens }

上面lua脚本的解释：

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

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

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

 

假如编程语言是go，则代码是：

	var script = redis.NewScript("xxx")
	keys := []string{"kou_tokens_key", "kou_timestamp_key"}
	values := []interface{}{1, 5, 5}
	if res, err := script.Run(ctx, rdb, keys, values...).Int64Slice(); err != nil {
		zap.L().Panic("SearchRiskAggregation run lua script error!!!", zap.Error(err))
	} else {
		zap.L().Info("SearchRiskAggregation!!!", zap.Any("res", res),
			zap.String("res type", reflect.TypeOf(res).String()))
	}

使用go-redis依赖，先调用redis包的NewScript函数生成Script指针，然后调用Script指针的Run方法生成Cmd指针，然后根据返回值的不同类型(lua脚本最终return的内容)调用Cmd指针不同的方法获取返回结果。