居于redis + lua脚本实现的滑动窗口

一、前言

我们常常使用滑动窗口实现限流操作，在单机时我们经常放在内存中实现，而在做全局接口限流时，我们除了可以通过查询接口调用记录外，还可以通过依赖redis实现的滑动窗口进行，比如限制1分钟可调用1000次，一小时可调用10000次。

二、滑动窗口的基本要素和操作

1、一个固定长度的循环队列

2、每个时间片的时长，可以是按秒、分、时。。。

3、每个时间窗口长度，由多个时间片组成一个时间窗口，也就是所需的一段时间

4、当前时间的所在时间片的索引

5、初始化循环队列的方法

6、选择当前时间所在时间片进行更新操作，增加调用次数

7、获取当前时间窗口的调用总数（1小时的调用总数）

8、获取当前时间片的调用总数（1分钟的调用总数）

9、处理时间窗口内可能存在的跳跃的时间片（更新时用到）

三、代码实现

1、定义基本要素：队列长度、每个时间片长度、窗口长度、当前时间所在时间片索引

/**
 * 每个时间片的时长 1min
 */
private static final int TIME_MILLIS_PER_SLICE = 60000;

/**
 * 窗口长度 多个时间片组成一个窗口,一个小时
 */
private static final int WINDOW_SIZE = 3600000;

/**
 * 队列总长度
 */
private static volatile int queueSize = 0;

/**
 * 最后记录的时间片索引
 */
private static volatile int slideIndex;

/**
 * list key
 */
private static final String REDIS_KEY_FOR_SLIDE_WINDOW = "SLIDE_WINDOW:INTERFACE_INVOKE_LIMIT";

/**
 * 最后一次记录时间片的索引，存储到redis，系统重启时可以拿到
 */
private static final String REDIS_KEY_FOR_SLIDE_INDEX = "SLIDE_WINDOW:SLIDE_INDEX";

1、定义初始化队列方法，在添加元素的时候用到

/**
 * 功能描述: 初始化队列
 * @author zcj
 * @date 2019/7/24
 */
private static void initQueue(RedisTemplate<String, String> redisTemplate) {

	//队列已初始化则直接返回
	if (queueSize != 0) {
		return;
	}

	queueSize = (WINDOW_SIZE / TIME_MILLIS_PER_SLICE) * 2 + 1;
	//启动时从redis获取最后记录的时间片
	String slideIndexStr = redisTemplate.opsForValue().get(REDIS_KEY_FOR_SLIDE_INDEX);
	if (StringUtils.isNotBlank(slideIndexStr)) {
		slideIndex = Integer.parseInt(slideIndexStr);
	}

        //判断队列是否已存在
	Long size = redisTemplate.opsForList().size(REDIS_KEY_FOR_SLIDE_WINDOW);
	if (size != null && size > 0) {
		return;
	}

	//队列未初始化，则初始化队列，设置队列长度为时间窗口的2被+1
	List<String> list = new ArrayList<>(queueSize);
	for (int i = 0; i < queueSize; i++) {
		list.add("0");
	}
	redisTemplate.opsForList().rightPushAll(REDIS_KEY_FOR_SLIDE_WINDOW, list);
}

2、接口调用时往当前时间片加1，因为这里的redisTemplate指定了value都是字符串，所以入参返回值都做了类型转换

增加调用数的方法：

/**
 * 功能描述: 往当前时间所在时间片增加1
 * @author zcj
 * @date 2019/7/25
 * @return 当前时间片的调用数
 */
public static Integer invokeIncr(RedisTemplate<String, String> redisTemplate) {

	//初始化队列
	initQueue(redisTemplate);

	//计算当前时间片的索引
	int currentIndex = (int) ((System.currentTimeMillis() / TIME_MILLIS_PER_SLICE) % queueSize);

	//通过lua脚本执行原子操作 如果当前时间片索引与旧索引一致，则该索引对应的值+1
        //传参，在lua脚本中通过KEYS[1]开始接收
	List<String> luaParams = new ArrayList<>();
	luaParams.add(REDIS_KEY_FOR_SLIDE_WINDOW);
	luaParams.add(REDIS_KEY_FOR_SLIDE_INDEX);

        //读取lua脚本，如果脚本比较短，可以直接通过redisScript.setScriptText()传入字符串就好
	DefaultRedisScript<String> redisScript = new DefaultRedisScript<>();
	redisScript.setResultType(String.class);
	redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("InterfaceInvokeLimit.lua")));
	String currentInvokeCount = redisTemplate.execute(redisScript, luaParams, String.valueOf(slideIndex),  String.valueOf(currentIndex), String.valueOf(queueSize));

	//把当前时间片的索引赋值给当前所属时间片属性
	slideIndex = currentIndex;
	return Integer.parseInt(currentInvokeCount);
}

lua脚本：

--传入的key
local key = KEYS[1]
local slideKey = KEYS[2]

--传入的参数数组
local slideIndex = tonumber(ARGV[1])
local currentIndex = tonumber(ARGV[2])
local queueSize = tonumber(ARGV[3])

local indexValue = 0
local newValue = 0

--如果上一次记录的时间片与当前时间片相同
if(slideIndex == currentIndex)
then
    indexValue = redis.call("LINDEX", key, currentIndex)
    newValue = indexValue + 1
    redis.call("LSET", key, currentIndex, newValue)
else
    --如果上次记录的时间片与当前时间片不同，为当前时间片设置为1
    newValue = 1
    redis.call("LSET", key, currentIndex, newValue)

    --遍历设置跳跃时间片的值为0 index != currentIndex
    local index = (slideIndex + 1) % queueSize
    while(true)
    do
        -- 遍历到当前时间片即终止
        if(index == currentIndex)
        then
          break
        end

        redis.call("LSET", key, index, 0)
        index = (index + 1) % queueSize
    end
end

--记录最后的时间片
redis.call("SET", slideKey, currentIndex)

return tostring(newValue)

3、获取当前时间窗口内的接口调用总数的方法

/**
 * 功能描述: 返回当前时间窗口内调用总数
 * @param redisTemplate redisTemplate
 * @author zcj
 * @date 2019/7/25
 * @return 当前时间窗口的调用总数
 */
public static int getCurrentWindowSum(RedisTemplate<String,String> redisTemplate) {

	//计算队列长度，与初始化方法保持一致
	int queueSize = (WINDOW_SIZE / TIME_MILLIS_PER_SLICE) * 2 + 1;

	//计算当前时间窗口内包含的时间片索引
	int currentIndex = (int) ((System.currentTimeMillis() / TIME_MILLIS_PER_SLICE) % queueSize);
	List<Long> indexs = new ArrayList<>();
	for (int i = 0; i < WINDOW_SIZE / TIME_MILLIS_PER_SLICE; i++) {
		indexs.add((long)((currentIndex - i + queueSize) % queueSize));
	}

	//通过管道查询
	List<Object> results = redisTemplate.executePipelined((RedisCallback<Object>) redisConnection -> {
		indexs.forEach(aLong -> redisConnection.listCommands().lIndex(REDIS_KEY_FOR_SLIDE_WINDOW.getBytes(), aLong));
		return null;
	});

	//累加窗口内时间片的调用总和
	int invokeCount = 0;
	if (!CollectionUtils.isEmpty(results)) {
		for (Object result : results) {
			invokeCount += Integer.parseInt(result.toString());
		}
	}

	return invokeCount;
}

4、获取当前时间所在时间片的调用数方法

/**
 * @description 返回当前所在时间片的调用数量
 * @param redisTemplate redisTemplate
 * @author zcj
 * @date 2020/8/30 9:46
 * @return 当前时间所在时间片的调用次数
 */
public static int getCurrentSlideValue(RedisTemplate<String,String> redisTemplate) {
	int currentIndex = (int) ((System.currentTimeMillis() / TIME_MILLIS_PER_SLICE) % queueSize);
	String value = redisTemplate.opsForList().index(REDIS_KEY_FOR_SLIDE_WINDOW, currentIndex);
	return Integer.parseInt(value);
}

四、总结

总的来说，通过redis实现滑动窗口的原理并不难，主要的问题在lua脚本中对跳跃时间片的循环处理，处理不好会导致redis进入死循环，可以在redis中配置lua脚本执行的超时时间。