注解+redis做限流，诶嘛真香

前言

做项目时提供对外接口，为了防止频繁调用而引起的服务器崩溃，我使用注解+redis做了一个限流的功能，使用之后只能说真香。

功能需求

• 能通过请求ip地址限流，比如该接口1分钟只能调用10次
• 能通过请求参数限制流量：比如当user=zhangsan，我们限制其1分钟只能调用20次
• 当基础注解不能满足时，还能自定义限流策略和实时修改策略。
• 更高级功能：限流控制台，可以实时查看限流情况和限流策略（暂未实现）。

功能实现

1.定义注解

ip限流注解

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

    String[] value() default "*";

    int limit() default 0;

    int second() default 0;

    String returnValue() default "请求频繁，请稍后访问！";

    Class<? extends CustomReturnInvocation> CustomReturnObject() default CustomReturnInvocation.class;

}
复制代码

String[] value(): 需要被限制的ip地址，比如：127.0.0.1 。该值是一个数组可以看出，支持配置多个ip地址。

int limit(): 被限制数。 和second()的值搭配使用，表示多长时间内限制多少次数。

int second(): 时常，单位是秒。和 limit() 的值搭配使用，表示多长时间内限制多少次数。

returnValue(): 返回值提醒。比如限流提醒：请求频繁，请稍后访问！

CustomReturnObject(): 自定义返回值，比如当想返回一个对象时，就需要用到该配置。

参数限流注解

参数限流和ip注解限流大同小异，不同的只是限流的参数不同。分开的原因一是分开使用逻辑更清晰，二是支持多配置限流，比如ip分流和参数限流都可作用于一个方法，方案采用 "熔断器"（Circuit Breaker）模式，当一个策略被满足时，即时触发限流。

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

    String[] value() default "";
  
    int limit() default 0;
 
    int second() default 0;

    String returnValue() default "请求频繁，请稍后访问！";

    Class<? extends CustomReturnInvocation> CustomReturnObject() default CustomReturnInvocation.class;
}
复制代码

配置参数都几乎相同。

String[] value(): 被限制的参数设置，支持 * 号配置 和 参数值的正则匹配。比如：

name=* ， name参数的值都将被限流限制。

user=n({a-z}+)e ，当参数user的值满足正则表达式的值时，将会被限流。比如：user = name , user = nsime 。

user=zhangsan ， 当参数user的值等于 zhangsan 时，将被限流，其他的都不会被限流

自定义限流策略

自定义限流可以实现更加灵活的限流的策略和实时更新限流策略，实现策略的动态增加和删除。

@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface CustomCurrentLimit {
    
    Class<? extends CustomCurrentLimitInjection> value() default CustomCurrentLimitInjection.class;

    Class<? extends CustomReturnInvocation> CustomReturnObject() default CustomReturnInvocation.class;
}
复制代码

public class CustomCurrentLimitElement implements Serializable {

    private final int limit;

    private final int second;

    private final String[] limitObject;

    private final CurrentLimitStrategyTypeEnum getStrategyType;

    private final String returnValue;
}
复制代码

CustomCurrentLimit ：

Class<? extends CustomCurrentLimitInjection> value() 可以传入自定义限量策略组。

Class<? extends CustomReturnInvocation> CustomReturnObject() 自定义返回值。

CustomCurrentLimitElement ： 该类的属性和注解配置的含义差不多。只是需要用CurrentLimitStrategyTypeEnum 指定限流策略的类型。

2.请求拦截

这里我是用的AOP的方式，在项目中我实现了一个AOP的统一拦截框架，用于实现一个请求的统一处理（当然也可以使用 HandlerInterceptor 拦截请求实现改功能）。两者主要逻辑基本一致。

IpCurrentLimit ipCurrentLimit = context.getAnnotationOnMethod(IpCurrentLimit.class);
ParameterCurrentLimit parameterAnnotation = context.getAnnotationOnMethod(ParameterCurrentLimit.class);
CustomCurrentLimit customCurrentLimit = context.getAnnotationOnMethod(CustomCurrentLimit.class);
String requestURI = context.getRequest().getRequestURI();
boolean isLimit = false;
try {
    if (ipCurrentLimit != null) {
        CurrentLimitExecute.registerStrategy(requestURI, ipCurrentLimit);
        invocation = ipCurrentLimit.CustomReturnObject();
        isLimit = true;
    }
    if (parameterAnnotation != null) {
        CurrentLimitExecute.registerStrategy(requestURI, parameterAnnotation);
        invocation = parameterAnnotation.CustomReturnObject();
        isLimit = true;
    }
    if (customCurrentLimit != null) {
        Class<? extends CustomCurrentLimitInjection> value = customCurrentLimit.value();
        CustomCurrentLimitElement[] element = value.newInstance().registerStrategy(requestURI);
        if (element != null) {
            invocation = customCurrentLimit.CustomReturnObject();
            CurrentLimitExecute.registerStrategy(requestURI, element);
            isLimit = true;
        }
    }
    if (isLimit) {
        try {
            CurrentLimitExecute.accessible(context.getRequest());
        } catch (CurrentLimitException e) {
            Object returnValue = e.getStrategy().getMessage();
            if (!invocation.getName().equals(CustomCurrentLimitInjection.class.getName())) {
                returnValue = invocation.newInstance().invoke(requestURI);
            }
            throw new DplAspectException("", returnValue);
        }
    }
} catch (InstantiationException | IllegalAccessException e) {
    throw new DplAspectException("", "服务器限流策略执行错误");
}
复制代码

该方法的作用是：在AOP中对请求进行限流判断和处理。具体的实现方式是，首先获取方法上标注的IpCurrentLimit、ParameterCurrentLimit和CustomCurrentLimit注解，这些注解分别表示对IP、请求参数和自定义限流策略进行限制。然后根据这些注解和请求的URI注册限流策略，在访问方法前对请求进行访问控制，如果超过了限制就会抛出CurrentLimitException异常。如果抛出了异常，就会根据不同情况返回不同的信息。

其中CurrentLimitExecute是一个限流策略注册和访问控制的工具类，CustomCurrentLimitInjection是一个抽象类，表示自定义的限流策略实现，它需要实现invoke方法来返回限流时的响应信息。invocation是一个CustomCurrentLimitInjection对象，表示当前使用的自定义限流策略。如果当前限流异常是由自定义限流策略触发的，就会调用当前限流策略的invoke方法来返回响应信息。

3.限流解析

首先我定了一个 private static final Map<String, List<CurrentLimitStrategy>> map = new HashMap<>(); 定一个Map用于缓存请求方法和该方法的限流策略组的绑定关系。核心执行逻辑如下：

String requestURI = request.getRequestURI();
String[] parameterKey = strategy.limitParameterKey(requestURI);
if (parameterKey != null) {
    CurrentLimitStrategy[] strategyArr = strategy.getParameterStrategy(requestURI);
    List<String> limitValue = new ArrayList<>();
    for (String key : parameterKey) {
        limitValue.add("pm:" + key + "=" + request.getParameter(key));
    }
    getCurrentLimitServices().accessible(requestURI, limitValue.toArray(new String[0]), strategyArr);
}

if (strategy.limitIp(requestURI)) {
    CurrentLimitStrategy[] strategyArr = strategy.getIpStrategy(requestURI);
    getCurrentLimitServices().accessible(requestURI, new String[]{"ip:" + getIpAdder(request)}, strategyArr);
}
复制代码

该方法大概是：从缓存中获取保存的参数限流策略和Ip限流策略，然后获取调用 getCurrentLimitServices().accessible(requestURI, limitValue.toArray(new String[0]), strategyArr);

进行限流判断。

parameterKey：是一种自定义拼装的Key,将参数值进行拼装：

limitValue.add("pm:" + key + "=" + request.getParameter(key));

4.核心判断逻辑

for (String key : limitKey) {
    String finalKey = key;
    // 1. 多种策略时，当满足一种策略后，就不再像后匹配。所以在多个策略的时候，需要对策略排序的功能。 采用优先匹配策略
    List<CurrentLimitStrategy> currentLimitStrategies = Arrays.stream(strategy).
        filter(cell -> cell.match(finalKey)).collect(Collectors.toList());

    for (CurrentLimitStrategy currentLimitStrategy : currentLimitStrategies) {
        // 2
        if (currentLimitStrategy.limit() == -1 || currentLimitStrategy.second() == -1) {
            continue;
        }
        // 3
        if (currentLimitStrategy.limit() == 0 || currentLimitStrategy.second() == 0) {
            throw new CurrentLimitException("超出请求限制", currentLimitStrategy);
        }

        key = url + key;
        // 4
        Integer maxLimit = redisTemplate().opsForValue().get(key);
        // 5
        if (maxLimit == null) {
            redisTemplate.opsForValue().set(key, 1, currentLimitStrategy.second(), TimeUnit.SECONDS);
        } else if (maxLimit < currentLimitStrategy.limit()) { // 6
            redisTemplate.opsForValue().set(key, maxLimit + 1, currentLimitStrategy.second(), TimeUnit.SECONDS);
        } else { // 7
            log.info("参数{}执行的受限制策略：{}", String.join(",", limitKey), currentLimitStrategy);
            throw new CurrentLimitException("超出请求限制", currentLimitStrategy);
        }
    }
}
复制代码

方法解析：

1.循环找出满足条件的策略组。然后循环匹配，一旦满足条件就抛出CurrentLimitException异常。

2.判断策略的默认值。如果limit=-1或者second=-1直接跳过。

3.判断策略的默认值。如果limit=0或者second=0 则直接抛出CurrentLimitException异常。

4.从redis中根据当前值获取缓存值。然后判断当前redis存储的请求数值与限流策略中的limit的大小。

5.如果maxLimit为0，则缓存改请求。

6.当maxLimit < currentLimitStrategy.limit()时，则缓存改请求。

7.当超出限制，则抛出异常。

tips：

1.缓存时我们使用了redis的缓存时间当作被限流的请求的时间。

2.CurrentLimitStrategy.match()方法根据不同的限流策略实现。

5.注解标记

@CustomCurrentLimit(value = DbCustomCurrentLimitInjection.class, CustomReturnObject =
            DbCustomReturnInvocation.class)
public Msg getUser(String id, HttpServletResponse response, HttpServletRequest request) throws FileNotFoundException {
    ....
}
复制代码