用redis lua脚本实现时间窗分布式限流

需求背景：

限制某sql在30秒内最多只能执行3次

需求分析

微服务分布式部署，既然是分布式限流，首先自然就想到了结合redis的zset数据结构来实现。
分析对zset的操作，有几个步骤，首先，判断zset中符合rangeScore的元素个数是否已经达到阈值，如果未达到阈值，则add元素，并返回true。如果已达到阈值，则直接返回false。

代码实现

首先，我们需要根据需求编写一个lua脚本

redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, tonumber(ARGV[3]))
local res = 0
if(redis.call('ZCARD', KEYS[1]) < tonumber(ARGV[5])) then
    redis.call('ZADD', KEYS[1], tonumber(ARGV[2]), ARGV[1])
    res = 1
end
redis.call('EXPIRE', KEYS[1], tonumber(ARGV[4]))
return res

ARGV[1]: zset element
ARGV[2]: zset score(当前时间戳)
ARGV[3]: 30秒前的时间戳
ARGV[4]: zset key 过期时间30秒
ARGV[5]: 限流阈值

private final RedisTemplate<String, Object> redisTemplate;

public boolean execLuaScript(String luaStr, List<String> keys, List<Object> args){
	RedisScript<Boolean> redisScript = RedisScript.of(luaStr, Boolean.class)
	return redisTemplate.execute(redisScript, keys, args.toArray());
}

测试一下效果

@SpringBootTest
public class ApiApplicationTest {
    @Test
    public void test2() throws InterruptedException{
        String luaStr = "redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, tonumber(ARGV[3]))\n" +
                "local res = 0\n" +
                "if(redis.call('ZCARD', KEYS[1]) < tonumber(ARGV[5])) then\n" +
                "    redis.call('ZADD', KEYS[1], tonumber(ARGV[2]), ARGV[1])\n" +
                "    res = 1\n" +
                "end\n" +
                "redis.call('EXPIRE', KEYS[1], tonumber(ARGV[4]))\n" +
                "return res";
        for (int i = 0; i < 10; i++) {
            boolean res = execLuaScript(luaStr, Arrays.asList("aaaa"), Arrays.asList("ele"+i, System.currentTimeMillis(),System.currentTimeMillis()-30*1000, 30, 3));
            System.out.println(res);
            Thread.sleep(5000);
        }
    }
}

测试结果符合预期！