使用Redis调用Lua脚本的方式对SpringBoot接口进行限流

前言

在日常开发中，限流功能时常被使用，用于对某些接口进行限流熔断，譬如限制单位时间内接口访问次数；或者按照某种规则进行限流，如限制ip的单位时间访问次数等，在SpringBoot中，常用的限流方式有：Lua脚本限流、阿里巴巴开源的限流器Sentinel等等，本篇主要介绍使用Lua脚本进行接口限流。

一、步骤

1、自定义限流注解 Limit.java，用于标注在需要限流的接口上

@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
public @interface Limit {

    /**
     * 资源key，用于自定义限流标识时使用，当limitType为key或key+ip时，key不能为空
     */
    String key() default "";

    /**
     * redis中key的前缀
     */
    String prefix() default "limit_";

    /**
     * 时间，以毫秒为限流单位
     */
    int period();

    /**
     * 限制访问次数
     */
    int count();

    /**
     * 类型
     */
    LimitType limitType();

}

2、编写限流类型枚举类 LimitType.java

public enum LimitType {

    /**
     * 自定义key限流，即限制所有用户在单位时间内对同一限流标识标识的接口的访问次数
     */
    KEY,

    /**
     * 根据请求者ip限流，即限制同一ip在单位时间内对限流类型为IP类型接口的访问次数
     */
    IP,

    /**
     * key+ip限流，即自定义一个限流标识，限制同一ip在单位时间内对同一限流标识标识的接口的访问次数
     */
    KEYIP,

    /**
     * 根据方法名限流，即限制所有用户在单位时间内对同一接口的访问次数
     */
    METHOD,

    /**
     * 根据方法名+ip限流,即限制同一ip在单位时间内对同一接口的访问次数
     */
    METHODIP;

}

3、编写限流具体实现类 LimitAspect.java，通过AOP方式进行限流

@Aspect
@Order(1)
@Component
public class LimitAspect {

    private final RedisTemplate<Object, Object> redisTemplate;

    public LimitAspect(RedisTemplate<Object, Object> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    @Pointcut("@annotation(com.tsing.bootadmin.common.annotations.Limit)")//这里是限流注解所在的路径名
    public void pointcut() {
    }

    @Around("pointcut()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        Limit limitAnnotation = method.getAnnotation(Limit.class);
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
        String key;
        switch (limitAnnotation.limitType()) {
            case KEY:
                key = limitAnnotation.key();
                break;
            case IP:
                key = ServletUtil.getClientIP(request);
                break;
            case KEYIP:
                key = limitAnnotation.key() + "_" + ServletUtil.getClientIP(request);
                break;
            case METHOD:
                key = method.getName();
                break;
            default:
                key = method.getName() + "_" + ServletUtil.getClientIP(request);
        }
        List<Object> keys = Collections.singletonList(limitAnnotation.prefix() + key);
        RedisScript<Boolean> redisScript = new DefaultRedisScript<>(buildLuaScript(), Boolean.class);
        if (redisTemplate.execute(redisScript, keys, limitAnnotation.count(), limitAnnotation.period())) {
            return joinPoint.proceed();
        } else {
            throw new BusinessException(CommonException.Proxy.SERVER_IS_TOO_BUSY);//此处为达到限流次数的逻辑处理，我这里是直接返回异常给前端，可根据实际情况做调整
        }
    }

    /**
     * lua限流脚本
     */
    private String buildLuaScript() {
        return "local num = redis.call('incr',KEYS[1])\n" +
                "if tonumber(num) == 1 then\n" +
                "redis.call('pexpire',KEYS[1],ARGV[2])\n" +
                "return true\n" +
                "elseif tonumber(num) > tonumber(ARGV[1]) then\n" +
                "return false\n" +
                "else\n" +
                "return true\n" +
                "end";
    }

}

4、Controller限流测试，我以登录接口为例，实现每个ip在一秒内只能访问一次登录接口

@PostMapping("/login")
@ApiOperation(value = "用户登录")
@Limit(period = 1000, count = 1, limitType = LimitType.METHODIP)
public ResultData<String> login(@Validated @RequestBody UserLoginReqVo reqVo, HttpServletRequest request) {
	return ResultUtil.success(userService.login(reqVo, request));
}

5、Swagger限流测试

正常访问：

快速连续点击两次：

总结

如果这篇博客对你有帮助的话，记得给我点个赞，你的鼓励是对我最大的支持！谢谢｡◕‿◕｡