【实践】SpringAOP + 自定义注解实现接口限流

【实践】SpringAOP + 自定义注解实现接口限流 

目录

【实践】SpringAOP + 自定义注解实现接口限流 

前置情景：

实践背景：

实践过程：

总结：

---

前置情景：

        1. 对 SpringAOP 相关知识有一定理解              

        2. 对 自定义注解 相关知识有一定理解     

        3. 对 Redis + lua脚本 相关知识有一定了解

实践背景：

        现某 ToC 接口存在 QPS 远高于设计预期，但内部逻辑较为复杂无法快速进行迭代优化情况

        现需对此接口改造，同时为了方便其他接口后续接入同样限流功能，采用AOP+自定义注解+Redis实现

实践过程：

        1. 准备redis

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import redis.clients.jedis.JedisPool;

import java.util.List;


@Slf4j
@Component
public class RedisComponent {

    protected JedisPool jedisPool;

    @SuppressWarnings("SpringJavaInjectionPointsAutowiringInspection")
    public RedisComponent(JedisPool jedisPool) {
        this.jedisPool = jedisPool;
    }

    public Long eval(final String script, final List<String> keys, final List<String> args) {
        Object result = RedisCaller.call(this.jedisPool, (jedis) -> jedis.eval(script, keys, args));
        return Long.parseLong(String.valueOf(result));
    }
}

注：RedisCaller为公司内部工具类，本质还是为了执行 jedis.eval()，使用其他工具类平替即可

        2. 准备自定义注解 + 枚举

import java.lang.annotation.*;


@Inherited
@Documented
@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.METHOD, ElementType.TYPE})
public @interface RateLimit {

    /**
     * 场景prefix
     */
    RateLimiterPrefixEnum prefix() default RateLimiterPrefixEnum.DEFAULT;

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

    /**
     * 给定的时间段
     * 单位秒
     */
    int period() default 15;

    /**
     * 最多的访问限制次数
     */
    int count() default 5;

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

import com.google.common.collect.Maps;

import java.util.Map;


public enum RateLimiterPrefixEnum {

    /**
     * 默认
     */
    DEFAULT(0, "default"),

    /**
     * TEST
     */
    TEST(1, "test");

    private final Integer code;

    private final String desc;

    private static Map<Integer, RateLimiterPrefixEnum> map;

    RateLimiterPrefixEnum(Integer code, String desc) {
        this.code = code;
        this.desc = desc;
    }

    static {
        map = Maps.newHashMap();
        for (RateLimiterPrefixEnum prefixEnum : RateLimiterPrefixEnum.values()) {
            map.put(prefixEnum.getCode(), prefixEnum);
        }
    }

    public Integer getCode() {
        return this.code;
    }

    public String getDesc() {
        return this.desc;
    }

    public static RateLimiterPrefixEnum parse(Integer code){
        return map.get(code);
    }
}

public enum LimitType {

    /**
     * 自定义key
     */
    CUSTOMER,

    /**
     * 根据请求者IP
     */
    IP;
}

        3. 准备切面

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
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.beans.factory.annotation.Value;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.List;
import java.util.Map;
import java.util.Objects;


@Slf4j
@Aspect
@Component
@Order(value = 1)
public class RateLimiterAspect {

    private static final String UNKNOWN = "unknown";

    private static final String PERIOD = "period";

    private static final String COUNT = "count";

    private static final String REDIS_SCRIPT;

    private final RedisComponent redisComponent;

    @Value("${app.limit.config.map}")
    private String limitConfigMapJson;

    static {
        // 构建lua脚本
        REDIS_SCRIPT = "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;";
    }

    public RateLimiterAspect(RedisComponent redisComponent) {
        this.redisComponent = redisComponent;
    }

    /**
     * 限流拦截器
     *
     * @param pjp       连接点
     * @return java.lang.Object
     */
    @Around("execution(public * a.b.c..*Controller.*(..)) && @annotation(自定义注解的全限定类名)")
    public Object limitInterceptor(ProceedingJoinPoint pjp) throws Throwable {
        MethodSignature signature = (MethodSignature) pjp.getSignature();
        Method method = signature.getMethod();
        RateLimit limit = method.getAnnotation(RateLimit.class);

        // 组装参数，优先取配置中对应场景的限流配置，没有就取接口上的兜底配置
        Map<String, Map<String, Integer>> configMap =
                new ObjectMapper().readValue(limitConfigMapJson, new TypeReference<Map<String, Map<String, Integer>>>(){});
        Map<String, Integer> periodCountMap = configMap.get(limit.prefix().getDesc() + "_" + limit.key());
        int limitPeriod = Objects.nonNull(periodCountMap) && Objects.nonNull(periodCountMap.get(PERIOD)) ?
                periodCountMap.get(PERIOD) : limit.period();
        int limitCount = Objects.nonNull(periodCountMap) && Objects.nonNull(periodCountMap.get(COUNT)) ?
                periodCountMap.get(COUNT) : limit.count();
        List<String> args = Lists.newArrayList(Integer.toString(limitCount), Integer.toString(limitPeriod));
        ImmutableList<String> keys = this.getLimitKey(limit);
        Long count = redisComponent.eval(REDIS_SCRIPT, keys, args);
        // 如果达到限流数则抛出异常
        if(count != null && count.intValue() <= limitCount) {
            return pjp.proceed();
        } else {
            throw new AppGatewayException(AppGatewayErrorCode.REQ_RATE_LIMIT_ERROR, "请求次数已达上限");
        }
    }

    /**
     * 获取各场景限流key
     *
     * @param limit     方法注解配置
     * @return com.google.common.collect.ImmutableList<java.lang.String>
     */
    private ImmutableList<String> getLimitKey(RateLimit limit) {
        String key = null;
        LimitType limitType = limit.limitType();
        switch (limitType) {
            case IP:
                key = getIpAddress();
                break;
            case CUSTOMER:
                // 可以自定义限流
                key = limit.key();
                break;
            default:
                break;
        }
        return ImmutableList.of(StringUtils.join(limit.prefix(), key));
    }

    /**
     * 获取请求的ip
     *
     * @return 获取请求的ip
     */
    public String getIpAddress() {
        HttpServletRequest request = ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getRequest();
        String ip = null;
        //X-Forwarded-For：Squid 服务代理
        String ipAddresses = request.getHeader("X-Forwarded-For");

        if (ipAddresses == null || ipAddresses.length() == 0 || UNKNOWN.equalsIgnoreCase(ipAddresses)) {
            //Proxy-Client-IP：apache 服务代理
            ipAddresses = request.getHeader("Proxy-Client-IP");
        }

        if (ipAddresses == null || ipAddresses.length() == 0 || UNKNOWN.equalsIgnoreCase(ipAddresses)) {
            //WL-Proxy-Client-IP：weblogic 服务代理
            ipAddresses = request.getHeader("WL-Proxy-Client-IP");
        }

        if (ipAddresses == null || ipAddresses.length() == 0 || UNKNOWN.equalsIgnoreCase(ipAddresses)) {
            //HTTP_CLIENT_IP：有些代理服务器
            ipAddresses = request.getHeader("HTTP_CLIENT_IP");
        }

        if (ipAddresses == null || ipAddresses.length() == 0 || UNKNOWN.equalsIgnoreCase(ipAddresses)) {
            //X-Real-IP：nginx服务代理
            ipAddresses = request.getHeader("X-Real-IP");
        }

        //有些网络通过多层代理，那么获取到的ip就会有多个，一般都是通过逗号（,）分割开来，并且第一个ip为客户端的真实IP
        if (ipAddresses != null && ipAddresses.length() != 0) {
            ip = ipAddresses.split(",")[0];
        }

        //还是不能获取到，最后再通过request.getRemoteAddr();获取
        if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ipAddresses)) {
            ip = request.getRemoteAddr();
        }
        return ip;
    }
}

注：

    1. 配置文件中需添加 app.limit.config.map = {"test_generate":{"period":管控时间,"count":上限次数}}

    2. 若出现 "error Type referred to is not an annotation type: xxx" 报错，请检查环绕通知中路径是否正确

        4. 在controller接口上增加自定义注解

import io.swagger.annotations.ApiOperation;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestPart;
import org.springframework.web.bind.annotation.RestController;

import javax.annotation.Resource;


@Slf4j
@RestController
@RequestMapping("/test")
public class TestController {

    @Resource
    private TestBizService testBizService;

    /**
     * test（增加流量管控）
     *
     * @param req       请求json
     * @return TestDTO
     */
    @NoSign
    @PostMapping("/generate")
    @ApiOperation(value = "generate")
    @RateLimit(prefix = RateLimiterPrefixEnum.TEST, key = "generate", limitType = LimitType.CUSTOMER)
    public TestDTO generate(@RequestPart TestRequest req) {
        try {
            return testBizService.generate(req);
        } catch (Exception e) {
            log.error("generate test error", e);
            return new TestDTO();
        }
    }
}

注：若出现 "getInputStream() has already been called for this request" 说明请求中数据已被获取过了，请确认是否有其他拦截器等配置，并自行百度解决

        5. 使用postman验证效果

总结：

        需求其实很简单，但为了方便后续其他接口快速接入限流能力，还是使用了比较复杂但泛用的实现方式，正好修补一下自己的知识漏洞，也挺好的。

        愿这篇实践能对后来者有所帮助