Merge pull request #95 from linzhihan/linzhihan

编程式限流
zongzibinbin · Jul 6, 2023 · 1bf6f9d · 1bf6f9d
2 parents 89d81fc + 58daaef
commit 1bf6f9d
Show file tree

Hide file tree

Showing 12 changed files with 506 additions and 75 deletions.
diff --git a/...t-common/src/main/java/com/abin/mallchat/common/common/aspect/FrequencyControlAspect.java b/...t-common/src/main/java/com/abin/mallchat/common/common/aspect/FrequencyControlAspect.java
@@ -2,9 +2,8 @@
 
 import cn.hutool.core.util.StrUtil;
 import com.abin.mallchat.common.common.annotation.FrequencyControl;
-import com.abin.mallchat.common.common.exception.BusinessException;
-import com.abin.mallchat.common.common.exception.CommonErrorEnum;
-import com.abin.mallchat.common.common.utils.RedisUtils;
+import com.abin.mallchat.common.common.domain.dto.FrequencyControlDTO;
+import com.abin.mallchat.common.common.service.frequencycontrol.FrequencyControlUtil;
 import com.abin.mallchat.common.common.utils.RequestHolder;
 import com.abin.mallchat.common.common.utils.SpElUtils;
 import lombok.extern.slf4j.Slf4j;
@@ -15,7 +14,12 @@
 import org.springframework.stereotype.Component;
 
 import java.lang.reflect.Method;
-import java.util.*;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import static com.abin.mallchat.common.common.service.frequencycontrol.FrequencyControlStrategyFactory.TOTAL_COUNT_WITH_IN_FIX_TIME_FREQUENCY_CONTROLLER;
 
 /**
  * Description: 频控实现
@@ -48,25 +52,25 @@ public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
             }
             keyMap.put(prefix + ":" + key, frequencyControl);
         }
-        //批量获取redis统计的值
-        ArrayList<String> keyList = new ArrayList<>(keyMap.keySet());
-        List<Integer> countList = RedisUtils.mget(keyList, Integer.class);
-        for (int i = 0; i < keyList.size(); i++) {
-            String key = keyList.get(i);
-            Integer count = countList.get(i);
-            FrequencyControl frequencyControl = keyMap.get(key);
-            if (Objects.nonNull(count) && count >= frequencyControl.count()) {//频率超过了
-                log.warn("frequencyControl limit key:{},count:{}", key, count);
-                throw new BusinessException(CommonErrorEnum.FREQUENCY_LIMIT);
-            }
-        }
-        try {
-            return joinPoint.proceed();
-        } finally {
-            //不管成功还是失败，都增加次数
-            keyMap.forEach((k, v) -> {
-                RedisUtils.inc(k, v.time(), v.unit());
-            });
-        }
+        // 将注解的参数转换为编程式调用需要的参数
+        List<FrequencyControlDTO> frequencyControlDTOS = keyMap.entrySet().stream().map(entrySet -> buildFrequencyControlDTO(entrySet.getKey(), entrySet.getValue())).collect(Collectors.toList());
+        // 调用编程式注解
+        return FrequencyControlUtil.executeWithFrequencyControlList(TOTAL_COUNT_WITH_IN_FIX_TIME_FREQUENCY_CONTROLLER, frequencyControlDTOS, joinPoint::proceed);
+    }
+
+    /**
+     * 将注解参数转换为编程式调用所需要的参数
+     *
+     * @param key              频率控制Key
+     * @param frequencyControl 注解
+     * @return 编程式调用所需要的参数-FrequencyControlDTO
+     */
+    private FrequencyControlDTO buildFrequencyControlDTO(String key, FrequencyControl frequencyControl) {
+        FrequencyControlDTO frequencyControlDTO = new FrequencyControlDTO();
+        frequencyControlDTO.setCount(frequencyControl.count());
+        frequencyControlDTO.setTime(frequencyControl.time());
+        frequencyControlDTO.setUnit(frequencyControl.unit());
+        frequencyControlDTO.setKey(key);
+        return frequencyControlDTO;
     }
 }
diff --git a/...-common/src/main/java/com/abin/mallchat/common/common/domain/dto/FrequencyControlDTO.java b/...-common/src/main/java/com/abin/mallchat/common/common/domain/dto/FrequencyControlDTO.java
@@ -0,0 +1,42 @@
+package com.abin.mallchat.common.common.domain.dto;
+
+import lombok.*;
+
+import java.util.concurrent.TimeUnit;
+
+@Data
+@ToString
+@Builder
+@NoArgsConstructor
+@AllArgsConstructor
+/** 限流策略定义
+ * @author linzhihan
+ * @date 2023/07/03
+ *
+ */
+public class FrequencyControlDTO {
+    /**
+     * 代表频控的Key 如果target为Key的话 这里要传值用于构建redis的Key target为Ip或者UID的话会从上下文取值 Key字段无需传值
+     */
+    private String key;
+    /**
+     * 频控时间范围，默认单位秒
+     *
+     * @return 时间范围
+     */
+    private Integer time;
+
+    /**
+     * 频控时间单位，默认秒
+     *
+     * @return 单位
+     */
+    private TimeUnit unit;
+
+    /**
+     * 单位时间内最大访问次数
+     *
+     * @return 次数
+     */
+    private Integer count;
+}
diff --git a/...on/src/main/java/com/abin/mallchat/common/common/exception/FrequencyControlException.java b/...on/src/main/java/com/abin/mallchat/common/common/exception/FrequencyControlException.java
@@ -0,0 +1,39 @@
+package com.abin.mallchat.common.common.exception;
+
+import lombok.Data;
+
+/**
+ * 自定义限流异常
+ *
+ * @author linzhihan
+ * @date 2023/07/034
+ */
+@Data
+public class FrequencyControlException extends RuntimeException {
+    private static final long serialVersionUID = 1L;
+
+    /**
+     *  错误码
+     */
+    protected Integer errorCode;
+
+    /**
+     *  错误信息
+     */
+    protected String errorMsg;
+
+    public FrequencyControlException() {
+        super();
+    }
+
+    public FrequencyControlException(String errorMsg) {
+        super(errorMsg);
+        this.errorMsg = errorMsg;
+    }
+
+    public FrequencyControlException(ErrorEnum error) {
+        super(error.getErrorMsg());
+        this.errorCode = error.getErrorCode();
+        this.errorMsg = error.getErrorMsg();
+    }
+}
diff --git a/...ommon/src/main/java/com/abin/mallchat/common/common/exception/GlobalExceptionHandler.java b/...ommon/src/main/java/com/abin/mallchat/common/common/exception/GlobalExceptionHandler.java
@@ -73,4 +73,12 @@ public ApiResult<Void> handleException(HttpRequestMethodNotSupportedException e)
         return ApiResult.fail(-1, String.format("不支持'%s'请求", e.getMethod()));
     }
 
+    /**
+     * 限流异常
+     */
+    @ExceptionHandler(value = FrequencyControlException.class)
+    public ApiResult frequencyControlExceptionHandler(FrequencyControlException e) {
+        log.info("frequencyControl exception！The reason is：{}", e.getMessage(), e);
+        return ApiResult.fail(e.getErrorCode(), e.getMessage());
+    }
 }
diff --git a/...abin/mallchat/common/common/service/frequencycontrol/AbstractFrequencyControlService.java b/...abin/mallchat/common/common/service/frequencycontrol/AbstractFrequencyControlService.java
@@ -0,0 +1,113 @@
+package com.abin.mallchat.common.common.service.frequencycontrol;
+
+import com.abin.mallchat.common.common.domain.dto.FrequencyControlDTO;
+import com.abin.mallchat.common.common.exception.CommonErrorEnum;
+import com.abin.mallchat.common.common.exception.FrequencyControlException;
+import com.abin.mallchat.common.common.utils.AssertUtil;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.lang3.ObjectUtils;
+
+import javax.annotation.PostConstruct;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+/**
+ * 抽象类频控服务 其他类如果要实现限流服务 直接注入使用通用限流类
+ * 后期会通过继承此类实现令牌桶等算法
+ *
+ * @author linzhihan
+ * @date 2023/07/03
+ * @see TotalCountWithInFixTimeFrequencyController 通用限流类
+ */
+@Slf4j
+public abstract class AbstractFrequencyControlService<K extends FrequencyControlDTO> {
+
+    @PostConstruct
+    protected void registerMyselfToFactory() {
+        FrequencyControlStrategyFactory.registerFrequencyController(getStrategyName(), this);
+    }
+
+    /**
+     * @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);
+        }
+    }
+
+
+    /**
+     * 多限流策略的编程式调用方法 无参的调用方法
+     *
+     * @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, FrequencyControlDTO> frequencyControlDTOMap = frequencyControlList.stream().collect(Collectors.groupingBy(FrequencyControlDTO::getKey, Collectors.collectingAndThen(Collectors.toList(), list -> list.get(0))));
+        return executeWithFrequencyControlMap((Map<String, K>) frequencyControlDTOMap, supplier);
+    }
+
+    /**
+     * 单限流策略的调用方法-编程式调用
+     *
+     * @param frequencyControl 单个频控对象
+     * @param supplier         服务提供着
+     * @return 业务方法执行结果
+     * @throws Throwable
+     */
+    public <T> T executeWithFrequencyControl(K frequencyControl, SupplierThrowWithoutParam<T> supplier) throws Throwable {
+        return executeWithFrequencyControlList(Collections.singletonList(frequencyControl), supplier);
+    }
+
+
+    @FunctionalInterface
+    public interface SupplierThrowWithoutParam<T> {
+
+        /**
+         * Gets a result.
+         *
+         * @return a result
+         */
+        T get() throws Throwable;
+    }
+
+    /**
+     * 是否达到限流阈值 子类实现 每个子类都可以自定义自己的限流逻辑判断
+     *
+     * @param frequencyControlMap 定义的注解频控 Map中的Key-对应redis的单个频控的Key Map中的Value-对应redis的单个频控的Key限制的Value
+     * @return true-方法被限流 false-方法没有被限流
+     */
+    protected abstract boolean reachRateLimit(Map<String, K> frequencyControlMap);
+
+    /**
+     * 增加限流统计次数 子类实现 每个子类都可以自定义自己的限流统计信息增加的逻辑
+     *
+     * @param frequencyControlMap 定义的注解频控 Map中的Key-对应redis的单个频控的Key Map中的Value-对应redis的单个频控的Key限制的Value
+     */
+    protected abstract void addFrequencyControlStatisticsCount(Map<String, K> frequencyControlMap);
+
+    /**
+     * 获取策略名称
+     *
+     * @return 策略名称
+     */
+    protected abstract String getStrategyName();
+
+}
diff --git a/...abin/mallchat/common/common/service/frequencycontrol/FrequencyControlStrategyFactory.java b/...abin/mallchat/common/common/service/frequencycontrol/FrequencyControlStrategyFactory.java
@@ -0,0 +1,51 @@
+package com.abin.mallchat.common.common.service.frequencycontrol;
+
+import com.abin.mallchat.common.common.domain.dto.FrequencyControlDTO;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * 限流策略工厂
+ *
+ * @author linzhihan
+ * @date 2023/07/03
+ */
+public class FrequencyControlStrategyFactory {
+    /**
+     * 指定时间内总次数限流
+     */
+    public static final String TOTAL_COUNT_WITH_IN_FIX_TIME_FREQUENCY_CONTROLLER = "TotalCountWithInFixTime";
+    /**
+     * 限流策略集合
+     */
+    static Map<String, AbstractFrequencyControlService<?>> frequencyControlServiceStrategyMap = new ConcurrentHashMap<>(8);
+
+    /**
+     * 将策略类放入工厂
+     *
+     * @param strategyName                    策略名称
+     * @param abstractFrequencyControlService 策略类
+     */
+    public static <K extends FrequencyControlDTO> void registerFrequencyController(String strategyName, AbstractFrequencyControlService<K> abstractFrequencyControlService) {
+        frequencyControlServiceStrategyMap.put(strategyName, abstractFrequencyControlService);
+    }
+
+    /**
+     * 根据名称获取策略类
+     *
+     * @param strategyName 策略名称
+     * @return 对应的限流策略类
+     */
+    @SuppressWarnings("unchecked")
+    public static <K extends FrequencyControlDTO> AbstractFrequencyControlService<K> getFrequencyControllerByName(String strategyName) {
+        return (AbstractFrequencyControlService<K>) frequencyControlServiceStrategyMap.get(strategyName);
+    }
+
+    /**
+     * 构造器私有
+     */
+    private FrequencyControlStrategyFactory() {
+
+    }
+}
diff --git a/...n/java/com/abin/mallchat/common/common/service/frequencycontrol/FrequencyControlUtil.java b/...n/java/com/abin/mallchat/common/common/service/frequencycontrol/FrequencyControlUtil.java
@@ -0,0 +1,55 @@
+package com.abin.mallchat.common.common.service.frequencycontrol;
+
+import com.abin.mallchat.common.common.domain.dto.FrequencyControlDTO;
+import com.abin.mallchat.common.common.utils.AssertUtil;
+import org.apache.commons.lang3.ObjectUtils;
+
+import java.util.List;
+
+/**
+ * 限流工具类 提供编程式的限流调用方法
+ *
+ * @author linzhihan
+ * @date 2023/07/03
+ */
+public class FrequencyControlUtil {
+
+    /**
+     * 单限流策略的调用方法-编程式调用
+     *
+     * @param strategyName     策略名称
+     * @param frequencyControl 单个频控对象
+     * @param supplier         服务提供着
+     * @return 业务方法执行结果
+     * @throws Throwable
+     */
+    public static <T, K extends FrequencyControlDTO> T executeWithFrequencyControl(String strategyName, K frequencyControl, AbstractFrequencyControlService.SupplierThrowWithoutParam<T> supplier) throws Throwable {
+        AbstractFrequencyControlService<K> frequencyController = FrequencyControlStrategyFactory.getFrequencyControllerByName(strategyName);
+        return frequencyController.executeWithFrequencyControl(frequencyControl, supplier);
+    }
+
+
+    /**
+     * 多限流策略的编程式调用方法调用方法
+     *
+     * @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);
+    }
+
+    /**
+     * 构造器私有
+     */
+    private FrequencyControlUtil() {
+
+    }
+
+}