框架级多种策略实现接口限流【阅读MallChat频控模块源代码】

概要

今天拉了Gitee上的MallChat项目的源代码，想学习一下代码编写。前不久刚好简单了解了使用自定义注解+Redis+aop实现接口限流的方法，恰好在MallChat源码中发现了类似的配置。因此记录一下个人分析读MallChat频控模块的代码的过程。仅供个人学习记录使用。如果有错误的地方欢迎大佬指正。

整体架构流程

首先是自定义注解类

/**
 * 频控注解
 */
@Repeatable(FrequencyControlContainer.class) // 可重复
@Retention(RetentionPolicy.RUNTIME)// 运行时生效
@Target(ElementType.METHOD)//作用在方法上
public @interface FrequencyControl {
    /**
     * 策略
     */
    String strategy() default FrequencyControlConstant.TOTAL_COUNT_WITH_IN_FIX_TIME;

    /**
     * 窗口大小，默认 5 个 period
     */
    int windowSize() default 5;

    /**
     * 窗口最小周期 1s (窗口大小是 5s， 1s一个小格子，共10个格子)
     */
    int period() default 1;


    /**
     * key的前缀，默认取方法全限定名，除非我们在不同方法上对同一个资源做频控，就自己指定
     *
     * @return key的前缀
     */
    String prefixKey() default "";

    /**
     * 频控对象，默认el表达指定具体的频控对象
     * 对于ip 和uid模式，需要是http入口的对象，保证RequestHolder里有值
     *
     * @return 对象
     */
    Target target() default Target.EL;

    /**
     * springEl 表达式，target=EL必填
     *
     * @return 表达式
     */
    String spEl() default "";

    /**
     * 频控时间范围，默认单位秒
     *
     * @return 时间范围
     */
    int time() default 10;

    /**
     * 频控时间单位，默认秒
     *
     * @return 单位
     */
    TimeUnit unit() default TimeUnit.SECONDS;

    /**
     * 单位时间内最大访问次数
     *
     * @return 次数
     */
    int count() default 1;

    long capacity() default 3; // 令牌桶容量

    double refillRate() default 0.5; // 每秒补充的令牌数

    enum Target {
        UID,
        IP,
        EL
    }
}

改注解作用于方法上，参数主要定义了当前方法使用的限流策略，如果不指定，默认使用FrequencyControlConstant.TOTAL_COUNT_WITH_IN_FIX_TIME策略。窗口大小、窗口周期、key的前缀、频控对象、频控时间、时间单位、单位时间内最大的访问次数、令牌桶容量(今天才了解到用于做限流的算法有令牌桶算法和漏桶算法)等。

然后创建AOP切面类FrequencyControlAspect，方法上需要使用
@Aspect 切面类
@Component 加入IOC容器
这两个注解
然后创建around方法，在方法上使用 @Around(“@annotation(com.abin.frequencycontrol.annotation.FrequencyControl)||@annotation(com.abin.frequencycontrol.annotation.FrequencyControlContainer)”)环绕增强打有@FrequencyControl或@FrequencyControlContainer自定义注解的方法。

/**
 * 频控实现
 */
@Slf4j
@Aspect
@Component
public class FrequencyControlAspect {
    @Around("@annotation(com.abin.frequencycontrol.annotation.FrequencyControl)||@annotation(com.abin.frequencycontrol.annotation.FrequencyControlContainer)")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        Method method = ((MethodSignature) joinPoint.getSignature()).getMethod();
        FrequencyControl[] annotationsByType = method.getAnnotationsByType(FrequencyControl.class);
        Map<String, FrequencyControl> keyMap = new HashMap<>();
        String strategy = FrequencyControlConstant.TOTAL_COUNT_WITH_IN_FIX_TIME;
        for (int i = 0; i < annotationsByType.length; i++) {
            // 获取频控注解
            FrequencyControl frequencyControl = annotationsByType[i];
            String prefix = StrUtil.isBlank(frequencyControl.prefixKey()) ? /* 默认方法限定名 + 注解排名（可能多个）*/method.toGenericString() + ":index:" + i : frequencyControl.prefixKey();
            String key = "";
            switch (frequencyControl.target()) {
                case EL:
                    key = SpElUtils.parseSpEl(method, joinPoint.getArgs(), frequencyControl.spEl());
                    break;
                case IP:
                    key = RequestHolder.get().getIp();
                    break;
                case UID:
                    key = RequestHolder.get().getUid().toString();
            }
            keyMap.put(prefix + ":" + key, frequencyControl);
            strategy = frequencyControl.strategy();
        }
        // 将注解的参数转换为编程式调用需要的参数
        if (FrequencyControlConstant.TOTAL_COUNT_WITH_IN_FIX_TIME.equals(strategy)) {
            // 调用编程式注解 固定窗口
            List<FrequencyControlDTO> frequencyControlDTOS = keyMap.entrySet().stream().map(entrySet -> buildFixedWindowDTO(entrySet.getKey(), entrySet.getValue())).collect(Collectors.toList());
            return FrequencyControlUtil.executeWithFrequencyControlList(strategy, frequencyControlDTOS, joinPoint::proceed);

        } else if (FrequencyControlConstant.TOKEN_BUCKET.equals(strategy)) {
            // 调用编程式注解 令牌桶
            List<TokenBucketDTO> frequencyControlDTOS = keyMap.entrySet().stream().map(entrySet -> buildTokenBucketDTO(entrySet.getKey(), entrySet.getValue())).collect(Collectors.toList());
            return FrequencyControlUtil.executeWithFrequencyControlList(strategy, frequencyControlDTOS, joinPoint::proceed);
        } else {
            // 调用编程式注解 滑动窗口
            List<SlidingWindowDTO> frequencyControlDTOS = keyMap.entrySet().stream().map(entrySet -> buildSlidingWindowFrequencyControlDTO(entrySet.getKey(), entrySet.getValue())).collect(Collectors.toList());
            return FrequencyControlUtil.executeWithFrequencyControlList(strategy, frequencyControlDTOS, joinPoint::proceed);
        }
    }

进入方法首先获取当前增强方法上@FrequencyControl注解的相关参数数据，因为@FrequencyControl是可以在一个方法上重复使用，所以需要循环获取。在循环内首先拼接前缀字符串prefix，然后通过注解内的Target类型获取对应的key。然后以prefix+“：”+key作为key，当前注解对象作为value存入keyMap。看代码好像最终该接口生效的限流策略取循环遍历中最后那个注解的strategy。然后在增强方法的最后三个if分别对应三种限流策略所做的逻辑处理。三种策略分别为“固定窗口”、“令牌桶”、“滑动窗口”。(我看代码里还有漏桶算法做限流的代码，但是在接口入口出好像没有使用该策略，不知道是后面用的还是没有启用)。

简单介绍一下两种限流算法
常用的限流算法有两种：漏桶算法和令牌桶算法。
漏桶算法思路很简单，水（请求）先进入到漏桶里，漏桶以一定的速度出水，当水流入速度过大会直接溢出，可以看出漏桶算法能强行限制数据的传输速率。我个人的理解是只要水桶内的水没满，任何速率的请求，无论是突增的密集请求还是匀速的请求都可以进入到水桶，但是出水的速率是恒定的，因此可以限制出口处。
令牌桶算法的原理是系统会以一个恒定的速度往桶里放入令牌，而如果请求需要被处理，则需要先从桶里获取一个令牌，当桶里没有令牌可取时，则拒绝服务。我个人的理解是令牌桶相较于漏桶，它可以限制入口处请求，因为补充令牌的速率是恒定的。

回到代码，首先第一种策略，如果接口注解上strategy字段值为TotalCountWithInFixTime，则使用固定窗口策略限流。进入固定窗口策略的if分支，首先将keyMap的value的list通过stream流调用buildFixedWindowDTO()方法封装成固定窗口所对应的FixedWindowDTO，然后调用FrequencyControlUtil.executeWithFrequencyControlList()

    /**
     * 多限流策略的编程式调用方法调用方法
     *
     * @param strategyName         策略名称
     * @param frequencyControlList 频控列表 包含每一个频率控制的定义以及顺序
     * @param supplier             函数式入参-代表每个频控方法执行的不同的业务逻辑
     * @return 业务方法执行的返回值
     * @throws Throwable 被限流或者限流策略定义错误
     */
    public static <T, K extends FrequencyControlDTO> T executeWithFrequencyControlList(String strategyName, List<K> frequencyControlList, AbstractFrequencyControlService.SupplierThrowWithoutParam<T> supplier) throws Throwable {
        boolean existsFrequencyControlHasNullKey = frequencyControlList.stream().anyMatch(frequencyControl -> ObjectUtils.isEmpty(frequencyControl.getKey()));
        AssertUtil.isFalse(existsFrequencyControlHasNullKey, "限流策略的Key字段不允许出现空值");
        AbstractFrequencyControlService<K> frequencyController = FrequencyControlStrategyFactory.getFrequencyControllerByName(strategyName);
        return frequencyController.executeWithFrequencyControlList(frequencyControlList, supplier);
    }

进入方法首先对list里面的数据进行判空，然后调用策略的工厂类FrequencyControlStrategyFactory的getFrequencyControllerByName方法，将策略值传入获取对应的策略类。

    /**
     * 根据名称获取策略类
     *
     * @param strategyName 策略名称
     * @return 对应的限流策略类
     */
    @SuppressWarnings("unchecked")
    public static <K extends FrequencyControlDTO> AbstractFrequencyControlService<K> getFrequencyControllerByName(String strategyName) {
        return (AbstractFrequencyControlService<K>) frequencyControlServiceStrategyMap.get(strategyName);
    }

可以看到直接从frequencyControlServiceStrategyMap获取对应key的value。那我想找一下什么时间将策略类放入这个map的。
可以看到这个方法将get到的value强转成了一个抽象类AbstractFrequencyControlService

/**
 * 抽象类频控服务 其他类如果要实现限流服务 直接注入使用通用限流类 后期会通过继承此类实现令牌桶等算法
 *
 * @param <K>
 */
@Slf4j
public abstract class AbstractFrequencyControlService<K extends FrequencyControlDTO> {

    @PostConstruct
    protected void registerMyselfToFactory() {
        FrequencyControlStrategyFactory.registerFrequencyController(getStrategyName(), this);
    }
 }

可以看到这个抽象类上的泛型K是继承FrequencyControlDTO类，其中registerMyselfToFactory方法在构造方法执行后执行将策略类注册到frequencyControlServiceStrategyMap(@PostConstruct在构造函数之后执行，init()方法之前执行。)。Ctrl+H看一下这个抽象类的树，发现三种策略对应的Controller类继承于它。因此，在这三个类构造方法执行之后以key为策略值，value为自己本身注册进FrequencyControlStrategyFactory的frequencyControlServiceStrategyMap。

获取到AbstractFrequencyControlService frequencyController后，执行对应策略Controller的自己的executeWithFrequencyControlList()方法（其实其他两种策略的业务逻辑到目前为止都是一样，后面就不赘述了）。
主要限流逻辑都在每个对应的的executeWithFrequencyControlList方法里。executeWithFrequencyControlList方法每个类继承这个抽象类时没有重写，因此这个方法都走抽象父类里面的。

    /**
     * 多限流策略的编程式调用方法 无参的调用方法
     *
     * @param frequencyControlList 频控列表 包含每一个频率控制的定义以及顺序
     * @param supplier             函数式入参-代表每个频控方法执行的不同的业务逻辑
     * @return 业务方法执行的返回值
     * @throws Throwable 被限流或者限流策略定义错误
     */
    @SuppressWarnings("unchecked")
    public <T> T executeWithFrequencyControlList(List<K> frequencyControlList, SupplierThrowWithoutParam<T> supplier) throws Throwable {
        boolean existsFrequencyControlHasNullKey = frequencyControlList.stream().anyMatch(frequencyControl -> ObjectUtils.isEmpty(frequencyControl.getKey()));
        AssertUtil.isFalse(existsFrequencyControlHasNullKey, "限流策略的Key字段不允许出现空值");
        Map<String, K> frequencyControlDTOMap = frequencyControlList.stream().collect(Collectors.groupingBy(K::getKey, Collectors.collectingAndThen(Collectors.toList(), list -> list.get(0))));
        return executeWithFrequencyControlMap(frequencyControlDTOMap, supplier);
    }

进入方法首先也是先对DTO的key值做非空校验，然后对DTO的集合使用Collectors.groupingBy()处理成Map的形式，正常来说通过这个方法处理承德Map value是个List。Collectors.collectingAndThen()方法又将这个list只取第一个元素。因此这个Map最终变成value是一个DTO对象的形式。然后调用executeWithFrequencyControlMap方法。这个方法也是走的抽象父类的方法。

    /**
     * @param frequencyControlMap 定义的注解频控 Map中的Key-对应redis的单个频控的Key Map中的Value-对应redis的单个频控的Key限制的Value
     * @param supplier            函数式入参-代表每个频控方法执行的不同的业务逻辑
     * @return 业务方法执行的返回值
     * @throws Throwable
     */
    private <T> T executeWithFrequencyControlMap(Map<String, K> frequencyControlMap, SupplierThrowWithoutParam<T> supplier) throws Throwable {
        if (reachRateLimit(frequencyControlMap)) {
            throw new FrequencyControlException(CommonErrorEnum.FREQUENCY_LIMIT);
        }
        try {
            return supplier.get();
        } finally {
            //不管成功还是失败，都增加次数
            addFrequencyControlStatisticsCount(frequencyControlMap);
        }
    }

在这个方法内主要是判断是否达到限流阈值，达到了抛出异常，没达到测放行，执行被环绕增强的方法，其实就是最初传进来的joinPoint::proceed。判断是否限流的reachRateLimit()方法走各自策略controller自己的方法。下面就分别看一下三种策略控制类的限流判断逻辑。

首先是固定窗口控制类

/**
 * 抽象类频控服务 -使用redis实现 固定时间内不超过固定次数的限流类
 */
@Slf4j
@Service
public class TotalCountWithInFixTimeFrequencyController extends AbstractFrequencyControlService<FixedWindowDTO> {


    /**
     * 是否达到限流阈值 子类实现 每个子类都可以自定义自己的限流逻辑判断
     *
     * @param frequencyControlMap 定义的注解频控 Map中的Key-对应redis的单个频控的Key Map中的Value-对应redis的单个频控的Key限制的Value
     * @return true-方法被限流 false-方法没有被限流
     */
    @Override
    protected boolean reachRateLimit(Map<String, FixedWindowDTO> frequencyControlMap) {
        //批量获取redis统计的值
        List<String> frequencyKeys = new ArrayList<>(frequencyControlMap.keySet());
        List<Integer> countList = RedisUtils.mget(frequencyKeys, Integer.class);
        for (int i = 0; i < frequencyKeys.size(); i++) {
            String key = frequencyKeys.get(i);
            Integer count = countList.get(i);
            int frequencyControlCount = frequencyControlMap.get(key).getCount();
            if (Objects.nonNull(count) && count >= frequencyControlCount) {
                //频率超过了
                log.warn("frequencyControl limit key:{},count:{}", key, count);
                return true;
            }
        }
        return false;
    }

    /**
     * 增加限流统计次数 子类实现 每个子类都可以自定义自己的限流统计信息增加的逻辑
     *
     * @param frequencyControlMap 定义的注解频控 Map中的Key-对应redis的单个频控的Key Map中的Value-对应redis的单个频控的Key限制的Value
     */
    @Override
    protected void addFrequencyControlStatisticsCount(Map<String, FixedWindowDTO> frequencyControlMap) {
        frequencyControlMap.forEach((k, v) -> RedisUtils.inc(k, v.getTime(), v.getUnit()));
    }

    @Override
    protected String getStrategyName() {
        return FrequencyControlConstant.TOTAL_COUNT_WITH_IN_FIX_TIME;
    }
}

首先获取传入进来的Map的key值集合，然后使用Redis批量获取这些key对应的value值并转成Integer。然后遍历这些Key，获取对应注解上的count值，如果redis里这个值存的大于注解上的count，意味着限流，返回true，否则不限流返回false。不管成功还是失败，都增加次数调用 addFrequencyControlStatisticsCount()，使对应key的value自增1。自增这块是用Lua脚本实现，保证原子性。如果有key就自增，如果没有key就设置，过期时间为注解上的time参数。
假设第一个请求过来，如果redis里没有对应key，则先放行，然后redis存入这个key，value为1，过期时间假定10s。10s内如果还有相同方法的请求，则先判断当前value值有没有超过设置的最大的请求数，如果没有，则放行然后再自增加一，如果大于等于则限流。个人理解是从第一个请求到后面的固定一段时间内，期间固定接口访问次数。

第二种是滑动窗口控制类，比固定窗口更灵活

/**
 * 抽象类频控服务 -使用redis实现 滑动窗口是一种更加灵活的频率控制策略，它在一个滑动的时间窗口内限制操作的发生次数
 */
@Slf4j
@Service
public class SlidingWindowFrequencyController extends AbstractFrequencyControlService<SlidingWindowDTO> {
    @Override
    protected boolean reachRateLimit(Map<String, SlidingWindowDTO> frequencyControlMap) {
        // 批量获取redis统计的值
        List<String> frequencyKeys = new ArrayList<>(frequencyControlMap.keySet());
        for (int i = 0; i < frequencyKeys.size(); i++) {
            String key = frequencyKeys.get(i);
            SlidingWindowDTO controlDTO = frequencyControlMap.get(key);
            // 获取窗口时间内计数
            Long count = RedisUtils.ZSetGet(key);
            int frequencyControlCount = controlDTO.getCount();
            if (Objects.nonNull(count) && count >= frequencyControlCount) {
                //频率超过了
                log.warn("frequencyControl limit key:{},count:{}", key, count);
                return true;
            }
        }
        return false;
    }

    @Override
    protected void addFrequencyControlStatisticsCount(Map<String, SlidingWindowDTO> frequencyControlMap) {
        List<String> frequencyKeys = new ArrayList<>(frequencyControlMap.keySet());
        for (int i = 0; i < frequencyKeys.size(); i++) {
            String key = frequencyKeys.get(i);
            SlidingWindowDTO controlDTO = frequencyControlMap.get(key);
            // 窗口最小周期转秒
            long period = controlDTO.getUnit().toMillis(controlDTO.getPeriod());
            long current = System.currentTimeMillis();
            // 窗口大小 单位 秒
            long length = period * controlDTO.getWindowSize();
            long start = current - length;
//            long expireTime = length + period;
            RedisUtils.ZSetAddAndExpire(key, start, length, current);
        }
    }

    @Override
    protected String getStrategyName() {
        return FrequencyControlConstant.SLIDING_WINDOW;

    }
}

滑动窗口是使用Redis的ZSet实现的。首先看判断是否限流的方法。首先还是循环遍历key的list，然后通过Redis.zCard()方法获取当前key的元素集合长度。判断如果集合长度大于等于最大count则限流，否则通过。再看一下增加次数的方法。进入方法还是循环遍历Key，然后将注解上的period字段值根据unit值转换成毫秒(我看注释写的是转秒好像是有点不对，刚开始看注释被误导了)，然后给key以当前时间添加元素，删除当前时间减去窗口大小之前的所有数据。最后给当前key设置窗口大小的过期时间。

最后一种就是令牌桶策略

/**
 * 抽象类频控服务 -使用redis实现 维护一个令牌桶来限制操作的发生次数
 */
@Slf4j
@Service
public class TokenBucketFrequencyController extends AbstractFrequencyControlService<TokenBucketDTO> {

    @Autowired
    private TokenBucketManager tokenBucketManager;

    @Override
    protected boolean reachRateLimit(Map<String, TokenBucketDTO> frequencyControlMap) {
        // 批量获取redis统计的值
        List<String> frequencyKeys = new ArrayList<>(frequencyControlMap.keySet());
        for (int i = 0; i < frequencyKeys.size(); i++) {
            String key = frequencyKeys.get(i);
            // 获取 1 个令牌
            return tokenBucketManager.tryAcquire(key, 1);
        }
        return false;
    }

    @Override
    protected void addFrequencyControlStatisticsCount(Map<String, TokenBucketDTO> frequencyControlMap) {
        List<String> frequencyKeys = new ArrayList<>(frequencyControlMap.keySet());
        for (int i = 0; i < frequencyKeys.size(); i++) {
            String key = frequencyKeys.get(i);
            TokenBucketDTO tokenBucketDTO = frequencyControlMap.get(key);
            tokenBucketManager.createTokenBucket(key, tokenBucketDTO.getCapacity(), tokenBucketDTO.getRefillRate());
            // 扣减 1 个令牌
            tokenBucketManager.deductionToken(key, 1);
        }
    }

    @Override
    protected String getStrategyName() {
        return FrequencyControlConstant.TOKEN_BUCKET;
    }
}

令牌桶策略通过调用TokenBucketManager的tryAcquire方法尝试获取令牌，tryAcquire中调用TokenBucketDTO的tryAcquire方法。TokenBucketManager类中有一个final的ConcurrentHashMap。程序启动后TokenBucketManager被放到IOC容器，每当有接口第一次被调用，TokenBucketManager把对应的key和DTO存到这个ConcurrentHashMap中。

    public boolean tryAcquire(int permits) {
        lock.lock();
        try {
            refillTokens();
            if (tokens < permits) {
                return true;
            }
            return false;
        } finally {
            lock.unlock();
        }
    }

可以看到进入这个方法首先加锁，这是为了保证在多线程的时候一次只有一个线程获取令牌。获取锁之后调用 refillTokens()方法。

    /**
     * 补充令牌
     */
    private void refillTokens() {
        long currentTime = System.nanoTime();
        // 转换为秒
        double elapsedTime = (currentTime - lastRefillTime) / 1e9;
        double tokensToAdd = elapsedTime * refillRate;
        log.info("tokensToAdd is {}", tokensToAdd);
        // 令牌总数不能超过令牌桶容量
        tokens = Math.min(capacity, tokens + tokensToAdd);
        log.info("current tokens is {}", tokens);
        lastRefillTime = currentTime;
    }

进入refillTokens方法首先获取当前时间，然后将当前这个接口最后一次调用的时间和当前时间做减法算出来的是以毫秒为单位的，再转换成秒做单位的。然后计算这段时间内能补充几个令牌。然后比较令牌桶最大容量和当前令牌数量加补充的令牌，取二者最小值，最后把当前时间赋值给lastRefillTime。refillTokens方法执行完后去比较当前需要的令牌数和令牌桶内剩余的令牌数，如果小于则放行，大于等于则限流。最后执行addFrequencyControlStatisticsCount方法扣减令牌。
总结令牌桶策略，相当于是每当有接口请求进来之后每次在获取令牌之前才去计算一次令牌数量。

小结

看到哪写到哪了，写的比较乱。