Lua脚本实现滑动窗口

一级目录

java 复制代码
import org.junit.jupiter.api.Test;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.util.*;

public class ApplicationTempTest {

    @Resource
    private StringRedisTemplate stringRedisTemplate;

    private String KEY_PREFIX  = "ni:ai:flow:limit:";
    private long WINDOW_MILLISECOND = 60 * 1000L;
    DefaultRedisScript<Long> script = new DefaultRedisScript<>();

    @PostConstruct
    public void scriptInit() {
        script.setScriptText(getFlowLimitLuaScript());
        script.setResultType(Long.class);
    }
        
    @Test
    public void testA() {
        String userId = "NEW_YY_1566";
        // 用户隔离级别
        // 设计一个滑动窗口算法。一分钟内，task1可以调用5次，task2可以调用10次
        Map<String, Integer> taskMap = new HashMap<>();
        taskMap.put("task1", 5);
        taskMap.put("task2", 10);

        taskMap.forEach((taskId, val) -> {
            boolean allow = allowRequest(taskId, val, WINDOW_MILLISECOND, userId);
            if (allow) {
                System.out.println(taskId + "可以继续执行！");
            } else {
                System.out.println(taskId + "不能继续执行！");
            }
        });
    }

    public boolean allowRequest(String key, long limit, Long windowMillisecond, String userId) {
        List<String> keys = Collections.singletonList(KEY_PREFIX + key);
        List<String> args = Arrays.asList(String.valueOf(limit), String.valueOf(windowMillisecond), String.valueOf(System.currentTimeMillis()), userId);
        Long result = stringRedisTemplate.execute(script, keys, args.toArray());
        // 判断脚本返回结果。限流0，允许1
        return result != null && result == 1;
    }

    public String getFlowLimitLuaScript() {
        return "local key = KEYS[1]\n" +
                "local maxRequests = tonumber(ARGC[1])\n" +
                "local windowMs = tonumber(ARGC[2])\n" +
                "local now = tonumber(ARGC[3])\n" +
                "local user_id = ARGC[4]\n" +
//                "-- 删除过期的记录（时间戳 < now - windows）\n" +
                "redis.call('ZREMRANGEBYSCORE', key, 0, now-windowMs)\n" +
//                "-- 获取当前窗口内的请求数\n" +
                "local count = redis.call('ZCARD', key)\n" +
//                "-- 判断是否超过限制（限流返回0，允许返回1）\n" +
                "if count >= maxRequests then\n" +
                "   return 0\n" +
                "else\n" +
//                "-- 添加当前请求\n" +
                "   redis.call('ZADD', key, now, user_id .. ':' .. now)\n" +
                "   return 1\n" +
                "end";
    }

}