lua脚本实现Redis令牌桶限流

背景

令牌桶限流是一种常见的流量控制算法，用于控制系统的请求处理速率，防止系统过载。在令牌桶限流算法中，可以将请求看作是令牌，而令牌桶则表示系统的处理能力。系统在处理请求时，首先需要从令牌桶中获取令牌，如果令牌桶中没有足够的令牌，就需要等待一定时间，直到令牌桶中有足够的令牌。
具体来说，令牌桶限流算法可以通过以下方式实现：
1.系统维护一个固定容量的令牌桶，每秒钟会向桶中添加一定数量的令牌，直到桶的容量达到上限。
2.每次请求来临时，需要先从令牌桶中获取令牌，如果桶中有足够的令牌，则请求被允许通过，并从桶中移除一个令牌；如果桶中的令牌数量不足，则请求被拒绝。具体来说，令牌桶算法会维护一个令牌桶，其中包含一定数量的令牌，每个令牌代表一个可以执行操作的许可。在每个时间段内，如果有令牌可用，就可以执行一个操作，并将令牌桶中的令牌数量减少一。如果没有令牌可用，就不能执行操作，需要等待一定时间，直到令牌桶中有足够的令牌。
3.由于令牌桶的容量是有限的，因此当桶中的令牌数量达到上限时，新的令牌会被丢弃，从而限制了请求的处理速率。
令牌桶限流算法可以在多种场景中进行流量控制，例如 Web 应用程序、消息队列、数据库等。在 Web 应用程序中，可以通过令牌桶限流算法控制 API 的访问速率，防止 API 被恶意攻击或者过载。在消息队列中，可以通过令牌桶限流算法控制消息的生产和消费速率，防止消息堆积和系统崩溃。在数据库中，可以通过令牌桶限流算法控制查询和写入操作的速率，防止数据库过载和响应时间过长。

lua脚本实现令牌桶算法

Lua 脚本可以用来实现 Redis 的令牌桶限流：
1.定义 Redis 数据结构
使用 Redis 的 Hash 数据结构存储当前令牌桶的状态。在 Hash 中，rate 表示速率（每秒生成的令牌数），capacity 表示桶的容量（最多可以同时存储的令牌数），tokens 表示当前桶中的令牌数量，timestamp 表示上次更新令牌数量的时间戳。示例代码：

HSET rdb:token_bucket rate 10 capacity 100 tokens 100 timestamp 0

2.编写 Lua 脚本
编写 Lua 脚本来实现限流逻辑。在脚本中，首先读取当前时间戳和桶的状态，计算出从上次更新时间戳到当前时间应该生成的令牌数量。然后，将当前桶中的令牌数量和应该生成的令牌数量相加，得到当前桶中的令牌数量。如果当前桶中的令牌数量超过了桶的容量，将其限制为桶的容量。
然后，判断当前桶中的令牌数量是否足够执行操作。如果令牌数量足够，将当前桶中的令牌数量减去操作所需的令牌数量，并更新桶的状态。如果令牌数量不足，则返回限流的错误信息。
示例代码：

-- 读取桶的状态
local rate = tonumber(redis.call('HGET', KEYS[1], 'rate'))
local capacity = tonumber(redis.call('HGET', KEYS[1], 'capacity'))
local tokens = tonumber(redis.call('HGET', KEYS[1], 'tokens'))
local timestamp = tonumber(redis.call('HGET', KEYS[1], 'timestamp'))

-- 计算应该生成的令牌数量
local now = redis.call('TIME')
local elapsed = now[1] - timestamp
local generated = math.floor(elapsed * rate)

-- 更新令牌数量并限制桶的容量
tokens = math.min(capacity, tokens + generated)

-- 执行操作
local required = tonumber(ARGV[1])
if tokens >= required then
    tokens = tokens - required
    redis.call('HSET', KEYS[1], 'tokens', tokens)
    redis.call('HSET', KEYS[1], 'timestamp', now[1])
    return 1
else
    return 0
end

3.在应用程序中调用 Lua 脚本
在应用程序中，使用 Redis 的 EVAL 命令来调用 Lua 脚本。示例代码：

@Component
public class TokenBucketLimiter {

    @Autowired
    private RedisTemplate<String, String> redisTemplate;

    public boolean tryAcquire(String key, int tokens) {
        List<String> keys = Arrays.asList(key);
        List<String> args = Arrays.asList(Integer.toString(tokens));
        Long result = redisTemplate.execute(new DefaultRedisScript<>(
                "local rate = tonumber(redis.call('HGET', KEYS[1], 'rate')) " +
                        "local capacity = tonumber(redis.call('HGET', KEYS[1], 'capacity')) " +
                        "local tokens = tonumber(redis.call('HGET', KEYS[1], 'tokens')) " +
                        "local timestamp = tonumber(redis.call('HGET', KEYS[1], 'timestamp')) " +
                        "local now = redis.call('TIME') " +
                        "local elapsed = now[1] - timestamp " +
                        "local generated = math.floor(elapsed * rate) " +
                        "tokens = math.min(capacity, tokens + generated) " +
                        "if tokens >= tonumber(ARGV[1]) then " +
                        "    tokens = tokens - tonumber(ARGV[1]) " +
                        "    redis.call('HSET', KEYS[1], 'tokens', tokens) " +
                        "    redis.call('HSET', KEYS[1], 'timestamp', now[1]) " +
                        "    return 1 " +
                        "else " +
                        "    return 0 " +
                        "end",
                Long.class), keys, args);
        return result != null && result == 1L;
    }
}

或者如下脚本：

-- 返回码 1:通过限流 0:不通过
-- rate ARGV[1] 每秒填充速率
-- now  ARGV[2] 当前时间
-- capacity ARGV[3] 令牌桶最大数量
-- request ARGV[4] 需要令牌数量
local SUCCESS = "1"
local FAIL = "0"
local rate = tonumber(ARGV[1]) -- replenishRate 令令牌桶填充平均速率
local capacity = tonumber(ARGV[2]) -- burstCapacity 令牌桶上限
local now = tonumber(ARGV[3]) -- 机器传入的当前时间 秒
local requested = tonumber(ARGV[4]) -- 消耗令牌数量，默认取1

local fill_time = capacity/rate   -- 计算令牌桶填充满令牌需要多久时间
local ttl = math.floor(fill_time*2) -- *2 保证时间充足

local result = SUCCESS;




-- ttl 防止小于0
if ttl < 1 then
    ttl = 10
end


-- 1、获取桶内令牌剩余数量
local last_tokens = tonumber(redis.call("get", KEYS[1]))
-- 获得令牌桶剩余令牌数
if last_tokens == nil then -- 第一次时，没有数值，所以桶时满的
    last_tokens = capacity
end

-- 2、获取上次更新时间
local last_refreshed = tonumber(redis.call("get", KEYS[2]))
-- 令牌桶最后填充令牌时间
if last_refreshed == nil then
    last_refreshed = 0
end

-- 3、本次验证和上次更新时间的间隔
local delta = math.max(0, now-last_refreshed)
-- 填充令牌，计算新的令牌桶剩余令牌数 填充不超过令牌桶令牌上限。
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))

-- 4、判断令牌数量是否足够
local allowed = filled_tokens >= requested
local new_tokens = filled_tokens
local allowed_num = "0"
if allowed then
    -- 若成功，令牌桶剩余令牌数(new_tokens) 减消耗令牌数( requested )，并设置获取成功( allowed_num = 1 ) 。
    new_tokens = filled_tokens - requested
    allowed_num = SUCCESS
end

-- 5、设置令牌桶剩余令牌数( new_tokens ) ，令牌桶最后填充令牌时间(now) ttl是超时时间
redis.call("setex", KEYS[1], ttl, new_tokens)
redis.call("setex", KEYS[2], ttl, now)

if not allowed then
    return FAIL
end

return SUCCESS