redis + lua 实现接口限流

76 阅读1分钟

demo

@RestController
@RequestMapping("/api")
public class TestController {

    @Autowired
    private RedisTemplate redisTemplate;

    @RequestMapping("/test")
    public Result test(HttpServletRequest request) {

        // 取IP
        String ip = request.getRemoteHost();

        // 创建Lua脚本对象
        DefaultRedisScript script = new DefaultRedisScript();
        script.setResultType(String.class);
        script.setScriptText("local count = redis.call('GET', KEYS[1])\n" +
            "local ip = KEYS[1]\n" +
            "local maxCount = tonumber(KEYS[2])\n" +
            "local limitTime = tonumber(KEYS[3])\n" +
            "if (count)\n" +
            "then\n" +
            "    local num = tonumber(count)\n" +
            "    if (num < maxCount)\n" +
            "    then\n" +
            "        redis.call('SET', ip, num + 1)\n" +
            "        redis.call('EXPIRE', ip, limitTime)\n" +
            "        return '0'\n" +
            "    else\n" +
            "        return '1'\n" +
            "    end\n" +
            "else\n" +
            "    redis.call('SET', ip, 1)\n" +
            "    redis.call('EXPIRE', ip, limitTime)\n" +
            "    return '0'\n" +
            "end");

        // KEYS[0]=IP KEYS[1]=maxCount KEYS[3]=limitTime
        // 创建参数
        List<String> keys = new ArrayList<>();
        keys.add(ip);
        keys.add("3");
        keys.add("5");
        
        // 执行脚本
        Object result = redisTemplate.execute(script, keys, new String[]{});
        
        if ("0".equals(result.toString())) {
            return Result.success(ip);
        }
        return Result.fail(300, "请求频繁");
    }
}