Spring boot 基于aop+redis+lua脚本自定义注解实现接口QPS限流

项目结构：

 

添加aop依赖：

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

1. 限流类型 LimitType

public enum LimitType {

    /**
     * 自定义key
     */
    CUSTOMER,

    /**
     * 根据请求者IP
     */
    IP;
}

2. 新增注解 Limit

@Target({ElementType.METHOD, ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Inherited
@Documented
public @interface Limit {
    /**
     * 限流key
     */
    String key() default "";

    /**
     * 用于描述接口功能
     */
    String name() default "";

    /**
     * 限流时间,单位秒
     */
    int time() default 60;

    /**
     * 限流次数
     */
    int count() default 100;

    /**
     * key的前缀名
     */
    String prefix() default "limit:";

    /**
     * 限流类型
     */
    LimitType limitType() default LimitType.IP;

}
​

3. 在resources资源文件夹下新建limit.lua文件

local key = KEYS[1]
local count = tonumber(ARGV[1])
local time = tonumber(ARGV[2])
local current = redis.call('get', key)
if current and tonumber(current) > count then
    return tonumber(current)
end
current = redis.call('incr', key)
if tonumber(current) == 1 then
    redis.call('expire', key, time)
end
return tonumber(current)

4. 新增 ScriptConfig 加载lua脚本

@Configuration
public class ScriptConfig {
    @Bean
    public DefaultRedisScript<Long> limitScript() {
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua/limit.lua")));
        redisScript.setResultType(Long.class);
        return redisScript;
    }
}

5. 新增 LimitAspect aop切面

@Slf4j
@Aspect
@Component
public class LimitAspect {

    @Autowired
    private RedisTemplate<String, Object> limitRedisTemplate;
    @Autowired
    private RedisScript<Long> limitScript;

    @Pointcut("@annotation(com.kingweb.common.limit.annotation.Limit)")
    public void pointcut() {}


    @Around("pointcut()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        HttpServletRequest request = ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getRequest();
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        Limit limitAnnotation = method.getAnnotation(Limit.class);
        LimitType limitType = limitAnnotation.limitType();
        String name = limitAnnotation.name();
        String key;
        int limitPeriod = limitAnnotation.time();
        int limitCount = limitAnnotation.count();
        switch (limitType) {
            case IP:
                key = IpUtils.getIpAddr(request);
                break;
            case CUSTOMER:
                key = limitAnnotation.key();
                break;
            default:
                key = StringUtils.upperCase(method.getName());
        }
        ImmutableList<String> keys = ImmutableList.of(StringUtils.join(limitAnnotation.prefix() + "_", key + "_" + request.getRequestedSessionId()));
        Number count = limitRedisTemplate.execute(limitScript, keys, limitCount, limitPeriod);
        log.info("第{}次访问key为 {}，描述为 [{}] 的接口", count, keys, name);
        if (count != null && count.intValue() <= limitCount) {
            return point.proceed();
        } else {
            throw new JeecgBootException("访问频繁，请稍后重试");
        }
    }
}

6. 测试是否可用 在接口添加注解

@Limit(time = 30, count = 5) // 30秒内最多访问5次