基于Redis的限流

以下代码演示同一个IP在30秒内只能获取一次验证码：

限流类型枚举

public enum LimitType {
    //自定义限流
    CUSTOMER,
    //手机号限流
    MOBILE,
    //IP限流
    IP;
}

定义限流注解

/**
 * 控制一个接口在指定时间内被访问次数限制
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface Limit {

    // 资源名称，用于描述接口功能
    String name() default "";

    // 资源 key
    String key() default "";

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

    // 时间的，单位秒
    int period();

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

    // 限制类型
    LimitType limitType() default LimitType.CUSTOMER;
}

使用Spring的Aspect拦截具有Limit注解的接口

@Slf4j
@Aspect
@Component
public class LimitAspect {

    private final RedisTemplate<String, Serializable> limitRedisTemplate;

    @Autowired
    public LimitAspect(RedisTemplate<String, Serializable> limitRedisTemplate) {
        this.limitRedisTemplate = limitRedisTemplate;
    }

    @Pointcut("@annotation(com.dpqwl.xunmishijie.common.annotation.Limit)")
    public void pointcut() {
        // 拦截所有具有limit注解的bean
    }

    @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;
        String ip = IPUtil.getIpAddr(request);
        int limitPeriod = limitAnnotation.period();
        int limitCount = limitAnnotation.count();
        //判断限流的类型，可以是IP，手机号，自定义
        switch (limitType) {
            case IP:
                key = ip;
                break;
            case CUSTOMER:
                key = limitAnnotation.key();
                break;
            case MOBILE:
                key=request.getParameter("phoneNum");;
                break;
            default:
                key = StringUtils.upperCase(method.getName());
        }
        ImmutableList<String> keys = ImmutableList.of(StringUtils.join(limitAnnotation.prefix() + "_", key));
        String luaScript = buildLuaScript();
        RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
        Number count = limitRedisTemplate.execute(redisScript, keys, limitCount, limitPeriod);
        log.info("IP:{} 第 {} 次访问key为 {}，描述为 [{}] 的接口", ip, count, keys, name);
        if (count != null && count.intValue() <= limitCount) {
            return point.proceed();
        } else {
            throw new LimitAccessException("访问频繁,稍后再试");
        }
    }

    /**
     * 限流脚本
     * 调用的时候不超过阈值，则直接返回并执行计算器自加。
     *
     * @return lua脚本
     */
    private String buildLuaScript() {
        // \n表示回车换行 ARGV[1]表示参数1：限流次数，ARGV[2]表示参数2：过期时间
        return "local c" +
                "\nc = redis.call('get',KEYS[1])" +
                "\nif c and tonumber(c) > tonumber(ARGV[1]) then" +
                "\nreturn c;" +
                "\nend" +
                "\nc = redis.call('incr',KEYS[1])" +
                "\nif tonumber(c) == 1 then" +
                "\nredis.call('expire',KEYS[1],ARGV[2])" +
                "\nend" +
                "\nreturn c;";
    }

}

以上的 buildLuaScript 方法，是核心部分，取出对应key的值和限流次数比较（如果key不存在执行redis.call(‘incr’,KEYS[1])会自动设置key并初始化为0再加1）：

• 如果当前次数大于限流次数，直接返回当前次数。
• 如果当前次数小于限流次数，将当前key值加1，如果当前值为1，表示该时间段内第一次访问，则设置key的过期时间为限流策略的时间值，当key过期之后开始重新计数。

在接口上使用Limit注解

@GetMapping("/getCaptcha")
@Limit(key = "captcha", period = 30, count = 1, name = "获取验证码", prefix = "limit", limitType = LimitType.IP)
public Response getCaptcha(@NotBlank String mobile){
    //30秒内同一个IP地址只能调用一次该接口
    return new Response().code("000").data(null).message("验证码获取成功!");
}

以上限流算法为固定窗口限流算法，其他限流算法可参考：常见限流算法