springboot整合redis及lua脚本实现接口限流

接口限流说明

接口限流是指在某些场景下，对某个接口的请求进行限制，以避免因请求过多而导致的系统负载过高、资源耗尽等问题。通常情况下，接口限流可以通过一定的算法来实现，比如令牌桶算法、漏桶算法、计数器算法等。这些算法可以根据接口的不同特点和业务需求，对请求进行限制和平滑处理，以达到系统资源的最优化利用。

令牌桶算法

令牌桶算法（Token Bucket Algorithm）：令牌桶算法可以通过限制请求的速率，来保护系统免受突发流量的冲击。该算法将请求和令牌都存放在一个桶中，每个请求需要从桶中取出一个令牌，如果桶中没有令牌，则请求将被拒绝。该算法能够平滑限制请求的速率，避免系统被突发流量打垮。

优点：

• 能够平滑限制请求的速率，适合对流量进行平滑的限制。
• 对于短时间内的流量突发，可以处理突发请求，保护系统不被打垮。

缺点：

• 实现相对较为复杂，需要维护令牌桶的状态。
• 无法应对突发的大量请求。

适用场景：

• 对于流量比较稳定的系统，需要对请求进行平滑限制的场景。
• 对于需要对请求进行按照一定速率限制的场景。

漏桶算法

漏桶算法（Leaky Bucket Algorithm）：漏桶算法与令牌桶算法相似，也是通过限制请求的速率来保护系统免受突发流量的冲击。该算法将请求放入一个漏桶中，每个请求都需要占据一定的空间，如果漏桶已满，则请求将被拒绝。该算法能够平滑限制请求的速率，但无法应对突发流量。

优点：

• 能够平滑限制请求的速率，适合对流量进行平滑的限制。
• 对于流量突发的情况，能够防止系统被过载。

缺点：

• 无法应对突发的大量请求。
• 实现相对较为复杂，需要维护漏桶的状态。

适用场景：

• 对于需要对请求进行按照一定速率限制的场景。

计数器算法

计数器算法（Counting Algorithm）：计数器算法是最简单的限流算法，通过对每个接口的请求数进行计数，并对其进行限制，来保护系统。该算法能够很好地限制请求的数量，但无法平滑限制请求的速率。

优点：

• 实现简单，易于实现。

缺点：

• 无法平滑限制请求的速率。
• 无法应对突发的大量请求。

适用场景：

• 对于需要对请求进行简单计数的场景。
• 对于不需要进行流量平滑限制的场景。

滑动窗口算法

滑动窗口算法（Sliding Window Algorithm）：滑动窗口算法可以通过限制请求的速率，来保护系统免受突发流量的冲击。该算法将请求按照时间顺序放入一个固定大小的窗口中，如果窗口已满，则新的请求将被拒绝。该算法能够平滑限制请求的速率，但无法应对突发流量。

优点：

• 能够平滑限制请求的速率，适合对流量进行平滑的限制。
• 对于短时间内的流量突发，可以处理突发请求，保护系统不被打垮。

缺点：

• 实现相对较为复杂，需要维护窗口的状态。
• 无法应对长时间的流量突发。

适用场景：

• 对于流量比较稳定的系统，需要对请求进行平滑限制的场景。
• 对于需要对请求进行按照一定速率限制的场景。

实现基于令牌桶+redis进行接口限流

这里我是基于令牌桶算法进行了变种，也就是针对不同的用户以及不同的方法在不同的时刻进行了限制，当然这个仅仅看个人业务

1️⃣：引入maven坐标

<!--lua脚本-->
<dependency>
  <groupId>org.luaj</groupId>
  <artifactId>luaj-jse</artifactId>
  <version>3.0.1</version>
</dependency>

<dependency>
  <groupId>org.springframework.boot</groupId>
  <artifactId>spring-boot-starter-data-redis</artifactId>
</dependency>

2️⃣：定义接口限制注解

package test.bo.work.redislimit.annotation;

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

/**
 * @author xiaobo
 * @date 2023/3/13
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {

    // 限制名称
    String key();
    // 次数
    int limit();
    // 秒
    int seconds();

}

3️⃣：lua脚本实现

local current = tonumber(redis.call('get', KEYS[1]) or '0')
-- 判断是否还有令牌可用
if current + 1 > tonumber(ARGV[1]) then
    return 0
else
    redis.call('incrby', KEYS[1], 1)
    redis.call('expire', KEYS[1], ARGV[2])
    return 1
end

4️⃣：引入lua脚本

package test.bo.work.redislimit;

import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.nio.file.Files;

/**
 * @author xiaobo
 * @date 2023/3/13
 */
@Component
public class RedisLuaScripts {
    @Value("classpath:/lua/ratelimit.lua")
    private Resource rateLimitScriptResource;

    public RedisScript<Long> getRateLimitScript() throws IOException {
        String script = new String(Files.readAllBytes(rateLimitScriptResource.getFile().toPath()));
        return new DefaultRedisScript<>(script, Long.class);
    }
}

5️⃣： 定义一个接口限制的AOP

package test.bo.work.redislimit.assept;

import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import test.bo.work.config.exception.BusinessException;
import test.bo.work.redislimit.RedisLuaScripts;
import test.bo.work.redislimit.annotation.RateLimit;

import javax.servlet.http.HttpServletRequest;
import java.time.LocalDate;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

/**
 * @author xiaobo
 * @date 2023/3/13
 */
@Aspect
@Component
public class RateLimitAspect {

    @Autowired
    private StringRedisTemplate redisTemplate;

    @Autowired
    private RedisLuaScripts redisLuaScripts;

    @Around("@annotation(rateLimit)")
    public Object checkRateLimit(ProceedingJoinPoint joinPoint, RateLimit rateLimit) throws Throwable {

        // 获取请求头中的AREA_TOKEN
        // 获取RequestAttributes
        RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
        // 从获取RequestAttributes中获取HttpServletRequest的信息
        HttpServletRequest request = (HttpServletRequest) requestAttributes
                .resolveReference(RequestAttributes.REFERENCE_REQUEST);
        String areaToken = request.getHeader("AREA_TOKEN");

        // 获取注解上的参数
        String key = rateLimit.key();
        int limit = rateLimit.limit();
        int seconds = rateLimit.seconds();

        String redisKey = key + areaToken + "-" + LocalDate.now();

        RedisScript<Long> script = redisLuaScripts.getRateLimitScript();
        List<String> keys = Collections.singletonList(redisKey);
        List<String> args = Arrays.asList(String.valueOf(limit), String.valueOf(seconds));
        Long result = redisTemplate.execute(script, keys, args.toArray());

        if (result != null && result == 1) {
            return joinPoint.proceed();
        } else {
            throw new BusinessException(500,"Rate limit exceeded for " + key);
        }
    }

}

🔚：接口实现

package test.bo.work.redislimit.controller;

import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RestController;
import test.bo.work.entity.OpResult;
import test.bo.work.redislimit.annotation.RateLimit;

/**
 * @author xiaobo
 * @date 2023/3/13
 */
@RestController
@Slf4j
public class RateLimitController {

    @PostMapping("rateLimit")
    @RateLimit(key = "rateLimit", limit = 10, seconds = 1)
    public OpResult testRateLimit() {
        return OpResult.Ok("success");
    }
}

说明

基于令牌桶算法的变种在正常情况下可以很好地限制接口的访问速率，但在短时间内突然出现大量请求的情况下，该算法可能会出现较大的误差。原因是在突发请求的情况下，桶内已经有很多令牌，但这些令牌并不能很快地被消耗，导致一些请求得到了允许，而另一些请求被拒绝。而漏桶算法则不会出现这个问题，因为它是基于固定速率漏水的方式进行限流，无论突发请求多少，都不会超出限制速率。

⚠️：经过jmeter测试确实出现了这样的情况

如果是需要对大量用户进行限流，建议使用更高效的限流算法，比如漏桶算法，或基于漏桶算法的Token Bucket算法