Merge pull request #10977 from dataease/pr@dev-v2@feat_traffic

feat: 增加api限流机制

de-xpack

@@ -1 +1 @@
-Subproject commit d5537b40bc950fe703b3b726762b34b3f1f57474
+Subproject commit 839f7d618a87a336fa4cbddb3dbe5d589b24e54b

sdk/common/src/main/java/io/dataease/traffic/DeTraffic.java

@@ -0,0 +1,12 @@
+package io.dataease.traffic;
+
+import java.lang.annotation.*;
+
+@Target(ElementType.METHOD)
+@Retention(RetentionPolicy.RUNTIME)
+@Documented
+public @interface DeTraffic {
+    int value() default 2;
+
+    String api();
+}

sdk/common/src/main/java/io/dataease/traffic/DeTrafficAop.java

@@ -0,0 +1,63 @@
+package io.dataease.traffic;
+
+import io.dataease.exception.DEException;
+import io.dataease.traffic.dao.entity.CoreApiTraffic;
+import io.dataease.traffic.dao.mapper.CoreApiTrafficMapper;
+import io.dataease.utils.IDUtils;
+import io.dataease.utils.LogUtil;
+import jakarta.annotation.Resource;
+import org.apache.commons.lang3.ObjectUtils;
+import org.aspectj.lang.ProceedingJoinPoint;
+import org.aspectj.lang.annotation.Around;
+import org.aspectj.lang.annotation.Aspect;
+import org.aspectj.lang.reflect.MethodSignature;
+import org.springframework.stereotype.Component;
+
+import java.lang.reflect.Method;
+
+@Aspect
+@Component
+public class DeTrafficAop {
+
+    @Resource
+    private CoreApiTrafficMapper coreApiTrafficMapper;
+
+    final private static String errorMsg = "当前API【%s】设定并发阈值为【%s】，现已经达到限流阈值，请稍后再试！";
+
+    @Around(value = "@annotation(io.dataease.traffic.DeTraffic)")
+    public Object trafficAround(ProceedingJoinPoint point) throws Throwable {
+        MethodSignature ms = (MethodSignature) point.getSignature();
+        Method method = ms.getMethod();
+        DeTraffic traffic = method.getAnnotation(DeTraffic.class);
+        int value = traffic.value();
+        String api = traffic.api();
+        Object result = null;
+        try {
+            Integer count = coreApiTrafficMapper.apiCount(api);
+            if (count == 0) {
+                CoreApiTraffic apiTraffic = new CoreApiTraffic();
+                apiTraffic.setId(IDUtils.snowID());
+                apiTraffic.setAlive(1);
+                apiTraffic.setThreshold(value);
+                apiTraffic.setApi(api);
+                coreApiTrafficMapper.insert(apiTraffic);
+                result = point.proceed();
+                return result;
+            }
+            int alive = coreApiTrafficMapper.getAlive(api);
+            if (alive < value) {
+                coreApiTrafficMapper.upgrade(api);
+                result = point.proceed();
+                return result;
+            }
+        } catch (Exception e) {
+            LogUtil.error(e.getMessage(), e);
+        } finally {
+            if (ObjectUtils.isNotEmpty(result)) {
+                coreApiTrafficMapper.releaseAlive(api);
+            }
+        }
+        DEException.throwException(errorMsg);
+        return null;
+    }
+}

sdk/common/src/main/java/io/dataease/traffic/dao/entity/CoreApiTraffic.java

@@ -0,0 +1,22 @@
+package io.dataease.traffic.dao.entity;
+
+import com.baomidou.mybatisplus.annotation.TableName;
+import lombok.Data;
+
+import java.io.Serial;
+import java.io.Serializable;
+
+@TableName("core_api_traffic")
+@Data
+public class CoreApiTraffic implements Serializable {
+    @Serial
+    private static final long serialVersionUID = -9130425144350145905L;
+
+    private Long id;
+
+    private String api;
+
+    private Integer threshold;
+
+    private Integer alive;
+}

sdk/common/src/main/java/io/dataease/traffic/dao/mapper/CoreApiTrafficMapper.java

@@ -0,0 +1,31 @@
+package io.dataease.traffic.dao.mapper;
+
+import io.dataease.traffic.dao.entity.CoreApiTraffic;
+import org.apache.ibatis.annotations.*;
+
+@Mapper
+public interface CoreApiTrafficMapper {
+
+    @Select("select `alive` from `core_api_traffic` where `api` = #{api}")
+    int getAlive(@Param("api") String api);
+
+    @Update("update `core_api_traffic` set alive = alive + 1 where `api` = #{api}")
+    void upgrade(@Param("api") String api);
+
+    @Insert("insert into core_api_traffic values(#{id}, #{api}, #{threshold}, 0)")
+    void insert(CoreApiTraffic traffic);
+
+    @Select("select count(*) from core_api_traffic where api = #{api}")
+    Integer apiCount(@Param("api") String api);
+
+    @Update("""
+        update `core_api_traffic` set alive = 
+        CASE WHEN alive > 0 THEN alive - 1
+        ELSE alive END 
+        where `api` = #{api}
+    """)
+    void releaseAlive(@Param("api") String api);
+
+    @Delete("delete from core_api_traffic")
+    void cleanTraffic();
+}

sdk/common/src/main/java/io/dataease/traffic/starter/DeTrafficStarter.java

@@ -0,0 +1,23 @@
+package io.dataease.traffic.starter;
+
+import io.dataease.traffic.dao.mapper.CoreApiTrafficMapper;
+import io.dataease.utils.LogUtil;
+import jakarta.annotation.Resource;
+import org.springframework.boot.ApplicationArguments;
+import org.springframework.boot.ApplicationRunner;
+import org.springframework.stereotype.Component;
+
+@Component
+public class DeTrafficStarter implements ApplicationRunner {
+
+    @Resource
+    private CoreApiTrafficMapper coreApiTrafficMapper;
+    @Override
+    public void run(ApplicationArguments args) throws Exception {
+        try {
+            coreApiTrafficMapper.cleanTraffic();
+        } catch (Exception e) {
+            LogUtil.error(e.getMessage(), new Throwable(e));
+        }
+    }
+}