【redis】redis滑动时间窗口算法实现限流

滑动时间窗口实现限流

依赖

首先创建一个Springboot项目

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-data-redis</artifactId>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-web</artifactId>
        </dependency>
        <dependency>
            <groupId>org.projectlombok</groupId>
            <artifactId>lombok</artifactId>
            <optional>true</optional>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-test</artifactId>
            <scope>test</scope>
        </dependency>
        <dependency>
            <groupId>cn.hutool</groupId>
            <artifactId>hutool-all</artifactId>
            <version>5.5.4</version>
        </dependency>

配置

server:
  port: 7777
spring:
  redis:
    host: 127.0.0.1
    port: 6379
    password: 你的redis密码，一般默认为空

实现

@Service
@Slf4j
public class RedisSlidingWindowDemo {
    @Resource
    RedisTemplate<Object, Object> redisTemplate;
    @Resource
    StringRedisTemplate stringRedisTemplate;
    String key = "redis_limiter";
    // 窗口大小, 单位：毫秒
    long windowTime = 1000;
    // 每 窗口大小时间 最多 多少个请求
    int limitCount = 5;
    long ttl = 10000;

    public void req() {
        // 当前时间
        long currentTime = System.currentTimeMillis();
        // 窗口起始时间
        long windowStartMs = currentTime - windowTime;
        ZSetOperations<Object, Object> zSetOperations = redisTemplate.opsForZSet();
        // 清除窗口过期成员
        zSetOperations.removeRangeByScore(key, 0, windowStartMs);
        // 添加当前时间 score=当前时间 value=当前时间+随机数，防止并发时重复
        zSetOperations.add(key, currentTime, currentTime + RandomUtil.randomInt());

    }

    public boolean canReq() {
        long currentTime = System.currentTimeMillis();
        int count = redisTemplate.opsForZSet().rangeByScore(key, currentTime - windowTime, currentTime).size();
        log.info("当前线程进入判断能否请求，当前时间={}，窗口={}-{}，数量={}", currentTime, (currentTime - windowTime), currentTime, count);
        if (count < limitCount) {
            req();
            return true;
        } else {
            return false;
        }
    }

    public boolean canReqByLua() {
        DefaultRedisScript<Long> script = new DefaultRedisScript<>();
        script.setScriptText(lua());
        script.setResultType(Long.class);
        long currentTime = System.currentTimeMillis();
        Long execute = stringRedisTemplate.execute(script, Collections.singletonList(key), String.valueOf(currentTime),
                String.valueOf(ttl), String.valueOf(windowTime), String.valueOf(limitCount),
                String.valueOf(currentTime + RandomUtil.randomInt()));
        boolean result = execute != 0;
        log.info("{}线程进入判断能否请求，当前时间={}，窗口={}-{}，数量={}，result={}", Thread.currentThread().getName(), currentTime,
                (currentTime - windowTime), currentTime, execute, result);
        return result;
    }

    public String lua() {
        return "local key = KEYS[1]\n" +
                "local currentTime = tonumber(ARGV[1])\n" +
                "local ttl = tonumber(ARGV[2])\n" +
                "local windowTime = tonumber(ARGV[3]) --\n" +
                "local limitCount = tonumber(ARGV[4])\n" +
                "local value = tonumber(ARGV[5])\n" +
                "redis.call('zremrangebyscore', key, 0, currentTime - windowTime)\n" +
                "local currentNum = tonumber(redis.call('zcard', key))\n" +
                "local next = currentNum + 1\n" +
                "if next > limitCount then\n" +
                "return 0;\n" +
                "else\n" +
                "redis.call(\"zadd\", key, currentTime, value)\n" +
                "redis.call(\"expire\", key, ttl)\n" +
                "return next\n" +
                "end";
    }
}

测试

@Slf4j
@RunWith(SpringJUnit4ClassRunner.class)
@SpringBootTest(classes = SpringBootDemoApplication.class)
public class RedisSlidingWindowDemoTest {

    @Resource
    RedisSlidingWindowDemo redisSlidingWindowDemo;

    @Test
    public void req() {
        Integer threadNum = 10;
        CountDownLatch downLatch = new CountDownLatch(threadNum);
        for (int i = 0; i < threadNum; i++) {
            new Thread(() -> {
                for (int j = 0; j < 20; j++) {
                    boolean access = redisSlidingWindowDemo.canReqByLua();
                    try {
                        TimeUnit.MILLISECONDS.sleep(100);
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                }
                downLatch.countDown();
            }, "t" + i).start();
        }
        try {
            downLatch.await();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }
}