spring-boot redis lua脚本实现滑动窗口限流

因为项目中没有集成redisson，但是又需要用到限流，所以简单的将redisson中限流的核心lua代码移植过来，并进行改造，因为公司版本的redis支持lua版本为5.1，针对于长字符串的数字，使用tonumber转换的时候会得到nil，而且还有各种奇怪的问题，可能是能力有限，所以对redisson的lua源码进行改造了一下

redisson源码：

设置限流器

//org.redisson.RedissonRateLimiter#trySetRateAsync
@Override
    public RFuture<Boolean> trySetRateAsync(RateType type, long rate, long rateInterval, RateIntervalUnit unit) {
        return commandExecutor.evalWriteNoRetryAsync(getRawName(), LongCodec.INSTANCE, RedisCommands.EVAL_BOOLEAN,
                "redis.call('hsetnx', KEYS[1], 'rate', ARGV[1]);"
              + "redis.call('hsetnx', KEYS[1], 'interval', ARGV[2]);"
              + "return redis.call('hsetnx', KEYS[1], 'type', ARGV[3]);",
                Collections.singletonList(getRawName()), rate, unit.toMillis(rateInterval), type.ordinal());
    }

请求限流

//org.redisson.RedissonRateLimiter#tryAcquireAsync(org.redisson.client.protocol.RedisCommand<T>, java.lang.Long)
private <T> RFuture<T> tryAcquireAsync(RedisCommand<T> command, Long value) {
        byte[] random = getServiceManager().generateIdArray();

        return commandExecutor.evalWriteAsync(getRawName(), LongCodec.INSTANCE, command,
                "local rate = redis.call('hget', KEYS[1], 'rate');"
              + "local interval = redis.call('hget', KEYS[1], 'interval');"
              + "local type = redis.call('hget', KEYS[1], 'type');"
              + "assert(rate ~= false and interval ~= false and type ~= false, 'RateLimiter is not initialized')"
              
              + "local valueName = KEYS[2];"
              + "local permitsName = KEYS[4];"
              + "if type == '1' then "
                  + "valueName = KEYS[3];"
                  + "permitsName = KEYS[5];"
              + "end;"

              + "assert(tonumber(rate) >= tonumber(ARGV[1]), 'Requested permits amount could not exceed defined rate'); "

              + "local currentValue = redis.call('get', valueName); "
              + "local res;"
              + "if currentValue ~= false then "
                     + "local expiredValues = redis.call('zrangebyscore', permitsName, 0, tonumber(ARGV[2]) - interval); "
                     + "local released = 0; "
                     + "for i, v in ipairs(expiredValues) do "
                          + "local random, permits = struct.unpack('Bc0I', v);"
                          + "released = released + permits;"
                     + "end; "

                     + "if released > 0 then "
                          + "redis.call('zremrangebyscore', permitsName, 0, tonumber(ARGV[2]) - interval); "
                          + "if tonumber(currentValue) + released > tonumber(rate) then "
                               + "currentValue = tonumber(rate) - redis.call('zcard', permitsName); "
                          + "else "
                               + "currentValue = tonumber(currentValue) + released; "
                          + "end; "
                          + "redis.call('set', valueName, currentValue);"
                     + "end;"

                     + "if tonumber(currentValue) < tonumber(ARGV[1]) then "
                         + "local firstValue = redis.call('zrange', permitsName, 0, 0, 'withscores'); "
                         + "res = 3 + interval - (tonumber(ARGV[2]) - tonumber(firstValue[2]));"
                     + "else "
                         + "redis.call('zadd', permitsName, ARGV[2], struct.pack('Bc0I', string.len(ARGV[3]), ARGV[3], ARGV[1])); "
                         + "redis.call('decrby', valueName, ARGV[1]); "
                         + "res = nil; "
                     + "end; "
              + "else "
                     + "redis.call('set', valueName, rate); "
                     + "redis.call('zadd', permitsName, ARGV[2], struct.pack('Bc0I', string.len(ARGV[3]), ARGV[3], ARGV[1])); "
                     + "redis.call('decrby', valueName, ARGV[1]); "
                     + "res = nil; "
              + "end;"

              + "local ttl = redis.call('pttl', KEYS[1]); "
              + "if ttl > 0 then "
                  + "redis.call('pexpire', valueName, ttl); "
                  + "redis.call('pexpire', permitsName, ttl); "
              + "end; "
              + "return res;",
                Arrays.asList(getRawName(), getValueName(), getClientValueName(), getPermitsName(), getClientPermitsName()),
                value, System.currentTimeMillis(), random);
    }

改造源码：

private static final String LIMIT_KEY_PREFIX = "api:camera:offline:limit:";
    private static final String ACQUIRE_LUA = "local rate = redis.call('hget', KEYS[1], 'rate');"
            + "local interval = redis.call('hget', KEYS[1], 'interval') * 1000;"
            + "local type = redis.call('hget', KEYS[1], 'type');"
            + "assert(rate ~= false and interval ~= false and type ~= false, 'RateLimiter is not initialized')"

            + "local valueName = KEYS[2];"
            + "local permitsName = KEYS[4];"
            + "if type == '1' then "
            + "valueName = KEYS[3];"
            + "permitsName = KEYS[5];"
            + "end;"

            + "assert(tonumber(rate) >= tonumber(ARGV[1]), 'Requested permits amount could not exceed defined rate'); "

            + "local currentValue = redis.call('get', valueName); "
            + "local res;"
            + "local time = redis.call('TIME');"
            + "local milliseconds = time[1] * 1000 + math.floor(time[2] / 1000);"

//            + "return interval * 1000;";

            + "if currentValue ~= false then "
//            + "return KEYS[1]; end;";

            + "local expiredValues = redis.call('zrangebyscore', permitsName, 0, milliseconds - interval); "
            + "local released = 0; "
            + "for i, v in ipairs(expiredValues) do "
            + "local random, permits = struct.unpack('Bc0I', v);"
            + "released = released + permits;"
            + "end; "

            + "if released > 0 then "
            + "redis.call('zremrangebyscore', permitsName, 0, milliseconds - interval); "
            + "if tonumber(currentValue) + released > tonumber(rate) then "
            + "currentValue = tonumber(rate) - redis.call('zcard', permitsName); "
            + "else "
            + "currentValue = tonumber(currentValue) + released; "
            + "end; "
            + "redis.call('set', valueName, currentValue);"
            + "end;"

            + "if tonumber(currentValue) < tonumber(ARGV[1]) then "
            + "local firstValue = redis.call('zrange', permitsName, 0, 0, 'withscores'); "
            + "res = 3 + interval - (milliseconds - tonumber(firstValue[2]));"
            + "else "
            + "redis.call('zadd', permitsName, milliseconds, struct.pack('Bc0I', string.len(ARGV[2]), ARGV[2], ARGV[1])); "
            + "redis.call('decrby', valueName, ARGV[1]); "
            + "res = nil; "
            + "end; "
            + "else "
            + "redis.call('set', valueName, rate); "
            + "redis.call('zadd', permitsName, milliseconds, struct.pack('Bc0I', string.len(ARGV[2]), ARGV[2], ARGV[1])); "
            + "redis.call('decrby', valueName, ARGV[1]); "
            + "res = nil; "
            + "end;"

            + "local ttl = redis.call('pttl', KEYS[1]); "
            + "if ttl > 0 then "
            + "redis.call('pexpire', valueName, ttl); "
            + "redis.call('pexpire', permitsName, ttl); "
            + "end; "
            + "return res;";
    private static final String CREATE_LIMIT_LUA = "redis.call('hsetnx', KEYS[1], 'rate', ARGV[1]);"
            + "redis.call('hsetnx', KEYS[1], 'interval', ARGV[2]);"
            + "return tostring(redis.call('hsetnx', KEYS[1], 'type', ARGV[3]));";

    @Test
    public void testRedis() {
        String key = "randdodd";
        String redisKey = LIMIT_KEY_PREFIX+ key;
        String rateKey = "rate";
        RedisTemplate redisTemplate = getRedisTemplate();
//        long secondsMillis = TimeUnit.SECONDS.toMillis(windowSize);
        long windowSize = 60L;
        Integer limit = 10;
        try {
            Boolean existList = redisTemplate.opsForHash().hasKey(LIMIT_KEY_PREFIX + key, rateKey);
            if (Boolean.FALSE.equals(existList)) {
                DefaultRedisScript<String> redisScript = new DefaultRedisScript<>(CREATE_LIMIT_LUA, String.class);
//                redisScript.setScriptText(CREATE_LIMIT_LUA);
                redisTemplate.execute(redisScript, Collections.singletonList(redisKey), limit, windowSize, "0");
                System.out.println("创建限流成功");
            }
            DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>(ACQUIRE_LUA, Long.class);
//            redisScript.setScriptText(ACQUIRE_LUA);
            Integer time = Math.toIntExact(System.currentTimeMillis() % 1000_000);
            for (int i = 0; i < 20; i++) {
                Object result = redisTemplate.execute(redisScript, Arrays.asList(getRawName(key), getValueName(key), getClientValueName(key, limit),
                        getPermitsName(key), getClientPermitsName(key, limit)), 1, IdUtil.fastSimpleUUID());
                System.out.println("获取令牌结果={}"+result);
                System.out.println(Objects.isNull(result));
            }

        } catch (Exception e) {
            e.printStackTrace();
        }

    }

    private RedisTemplate getRedisTemplate() {
        RedisStandaloneConfiguration redisStandaloneConfiguration = new RedisStandaloneConfiguration();
        redisStandaloneConfiguration.setHostName("123.1.1.1");
        redisStandaloneConfiguration.setPassword("2345262345234");
        redisStandaloneConfiguration.setPort(6379);

        GenericObjectPoolConfig genericObjectPoolConfig = new GenericObjectPoolConfig();
        genericObjectPoolConfig.setMaxTotal(10);
        genericObjectPoolConfig.setMaxIdle(10);
        genericObjectPoolConfig.setMaxWait(Duration.ofSeconds(10));
        genericObjectPoolConfig.setMinIdle(10);
        genericObjectPoolConfig.setTestOnBorrow(false);
        LettuceClientConfiguration configuration = LettucePoolingClientConfiguration.builder().poolConfig(genericObjectPoolConfig).build();

        RedisTemplate<Object, Object> redisTemplate = new RedisTemplate<>();
        LettuceConnectionFactory connectionFactory = new LettuceConnectionFactory(redisStandaloneConfiguration, configuration);
        connectionFactory.afterPropertiesSet();
        redisTemplate.setConnectionFactory(connectionFactory);
        Jackson2JsonRedisSerializer jackson2JsonRedisSerializer = new Jackson2JsonRedisSerializer(Object.class);
        ObjectMapper objectMapper = new ObjectMapper();
        objectMapper.setVisibility(PropertyAccessor.ALL, JsonAutoDetect.Visibility.ANY);
        objectMapper.enableDefaultTyping(ObjectMapper.DefaultTyping.NON_FINAL);
        objectMapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
        jackson2JsonRedisSerializer.setObjectMapper(objectMapper);

        redisTemplate.setValueSerializer(jackson2JsonRedisSerializer);
        redisTemplate.setKeySerializer(new StringRedisSerializer());
        redisTemplate.setHashKeySerializer(new StringRedisSerializer());
        redisTemplate.afterPropertiesSet();
        return redisTemplate;
    }

    private String getClientPermitsName(String redisKey, int limit) {
        return suffixName(getPermitsName(redisKey), redisKey+":"+limit);
    }

    private String getPermitsName(String redisKey) {
        return suffixName(getRawName(redisKey), "permits");
    }

    private String getClientValueName(String key, int limit) {
        return suffixName(getValueName(key), key+":"+limit);
    }

    private String suffixName(String name, String suffix) {
        if (name.contains("{")) {
            return name + ":" + suffix;
        }
        return "{" + name + "}:" + suffix;
    }

    private String getValueName(String key) {
        return suffixName(getRawName(key), "value");
    }

    private String getRawName(String key) {
        return LIMIT_KEY_PREFIX+key;
    }

关于redisson限流源码解读：Redisson分布式限流器RRateLimiter原理解析 · Issue #13 · oneone1995/blog · GitHub