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

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

redisson源码：

设置限流器

//org.redisson.RedissonRateLimiter#trySetRateAsync
@Overridepublic 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]));";@Testpublic 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