AOP+策略模式实现接口限流器

1 引言 

        在日常项目开发中，一些接口可能需要对用户进行一些访问的速率限制。比如：保存某类表单时，如果用户恶意请求，狂点保存，前端没做校验的话，可能会存在同一份数据被保存多次的问题。这个时候需要对接口进行一定的请求限制，比如，对单个用户进行限制等操作。

本篇笔记实现了通过注解+AOP+策略模式的方法来进行接口速率限制的过程。                                                                                  （主要记录对同一用户进行限制和对方法参数进行限制）

2 实现过程

2.1 思路分析

        要对接口进行速率限制，则需要存储接口的访问记录，来对短时间内连续访问的请求做限制。对于接口访问记录的记录选择使用Redis进行存储。使用SetNx来对请求进行获取锁，如果获取到锁就放行。

        实现策略模式，首先定义一个接口，接口定义速率限制的方法，返回一个成功或者失败。定义一个策略的TAG。

最后使用方法注解加AOP的形式就可以实现一个接口过滤器了。

2.2 代码实现

首先定义一个接口，代码如下：

public interface IRateLimiterStrategy {
    RateLimiterType type();

    boolean rateLimiterCheck(ProceedingJoinPoint joinPoint, RateLimiter rateLimiter);
}

        其中rateLimiterCheck返回一个boolean类型，方法用于鉴别请求是否通过。type方法返回一个限制类型的枚举，用于不同策略的选择。

定义一个限制类型枚举类，代码如下：

public enum RateLimiterType {
    /**
     * 通过参数限制
     */
    ARGS,
    /**
     * 通过用户id限制
     */
    USER_ID
}

定义一个用于标记限流方法注解，代码如下：

/**
 * 方法速率限制器
 */
@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimiter {
    RateLimiterType type(); //限制策略类型

    int time() default 10; //限制时间 单位/s

    String result() default "请稍后重试"; //限制后返回提示语

    String[] args() default {};
}

        其中type属性是策略选择Tag，time是接口速率限制的时间，result是接口限流后，响应的值，args是使用参数进行限流的参数名称。

接下来定义一个对同一用户进行限制的策略类，代码中，每个请求的用户ID被存在ThreadLocal中，所以只需对方法的签名和用户id进行组合便可以得到Redis的key。

代码如下：

@Component
@Slf4j
public class UserIdRateLimiter implements IRateLimiterStrategy {

  @Autowired
  private StringRedisTemplate stringRedisTemplate;

  @Override
  public boolean rateLimiterCheck(ProceedingJoinPoint joinPoint, RateLimiter rateLimiter) {
    int time = rateLimiter.time();

    //获取用户ID
    String userId = IdentityContext.getUserId();

    //如果用户ID为空，抛出未登录异常
    if (StringUtils.isEmpty(userId)) {
      throw new BizException(ErrorEnum.NO_LOGIN_ERROR);
    }

    // 获取方法的全局唯一签名，作为限流器的键
    String signatureStr = joinPoint.getSignature().toLongString() + ":" + userId;

    //拼接redisKey
    String redisKey = String.format(RedisKeyPrefix.RATE_LIMITER, signatureStr);

    //判断redis中是否存在该key
    Boolean isHas = stringRedisTemplate.hasKey(redisKey);

    //如果存在，说明请求过于频繁，返回false
    if (Boolean.TRUE.equals(isHas)) {
      log.info("用户 {} 请求{}过于频繁,接口频率限制:{}s", redisKey, signatureStr, time);
      return false;
    }

    //如果不存在，将该key存入redis，并设置过期时间
    stringRedisTemplate.opsForValue().set(redisKey, LocalDateTime.now().toString(), time, TimeUnit.SECONDS);

    //返回true
    return true;
  }

  @Override
  public RateLimiterType type() {
    return RateLimiterType.USER_ID;
  }
}

接下来定义一个对参数进行过滤的类，首先获取注解中需要限制的参数名称，再获取方法对应的参数值，组合成Redis的key。代码如下：

@Slf4j
@Component
public class ArgsRateLimiter implements IRateLimiterStrategy {

    @Autowired
    private StringRedisTemplate stringRedisTemplate;

    @Override
    public RateLimiterType type() {
        return RateLimiterType.ARGS;
    }

    @Override
    public boolean rateLimiterCheck(ProceedingJoinPoint joinPoint, RateLimiter rateLimiter) {
        // 获取速率限制器的参数
        String[] args = rateLimiter.args();

        // 获取速率限制器的时长
        int time = rateLimiter.time();

        // 如果参数为空
        if (args.length == 0){
            return true;
        }

        // 获取方法签名
        Signature signature = joinPoint.getSignature();

        // 获取方法参数拼接的字符串
        StringBuilder argsSignatureStr = getSignatureStr(joinPoint, (MethodSignature) signature, args);

        // 如果字符串为空
        if (StringUtils.isEmpty(argsSignatureStr.toString())){
            return true;
        }

        // 在字符串最前面插入方法签名
        argsSignatureStr.insert(0,signature);

        // 获取最终的rediskey
        String redisKey = String.format(RedisKeyPrefix.RATE_LIMITER, argsSignatureStr);


        //判断redis中是否存在该key
        Boolean isHas = stringRedisTemplate.hasKey(redisKey);

        //如果存在，说明请求过于频繁，返回false
        if (Boolean.TRUE.equals(isHas)) {
            log.info("用户 {} 请求{}过于频繁,接口频率限制:{}s", redisKey, signature, time);
            return false;
        }

        //如果不存在，将该key存入redis，并设置过期时间
        stringRedisTemplate.opsForValue().set(redisKey, LocalDateTime.now().toString(), time, TimeUnit.SECONDS);

        //返回true
        return true;
    }

    @NotNull
    private static StringBuilder getSignatureStr(ProceedingJoinPoint joinPoint, MethodSignature methodSignature, String[] args) {

        String[] parameterNames = methodSignature.getParameterNames(); // 获取参数名称

        Object[] parameterValues = joinPoint.getArgs(); // 获取参数值

        List<String> argList = Arrays.asList(args);

        return getArgsSignatureStr(parameterNames, argList, parameterValues);
    }

    @NotNull
    private static StringBuilder getArgsSignatureStr(String[] parameterNames, List<String> argList, Object[] parameterValues) {
        StringBuilder argsSignatureStr = new StringBuilder();

        for (int i = 0; i < parameterNames.length; i++) {

            String parameterName = parameterNames[i];

            if (!argList.contains(parameterName)){
                continue;
            }

            Object parameterValue = parameterValues[i];

            argsSignatureStr.append(":");
            argsSignatureStr.append(parameterName);
            argsSignatureStr.append(":");
            argsSignatureStr.append(parameterValue);
        }
        return argsSignatureStr;
    }
}

定义一个用于选择策略的类，该类用于通过Tag来选择对应的策略，代码如下

@Component
public class RateLimiterHandler {
    @Autowired
    private List<IRateLimiterStrategy> rateLimiterStrategyList;

    public IRateLimiterStrategy get(RateLimiterType type) {
        return rateLimiterStrategyList.stream()
                .filter(item -> item.type() != null)
                .filter(item -> item.type().equals(type))
                .findFirst().orElse(null);
    }
}

最后定义一个处理被注解标注方法的AOP就可以了，代码如下：

@Aspect
@Slf4j
@Component
public class RateLimiterAspect {

    @Autowired
    private RateLimiterHandler rateLimiterHandler;

    @Around("@annotation(rateLimiter)")
    public Object rateLimit(ProceedingJoinPoint joinPoint, RateLimiter rateLimiter) throws Throwable {

        IRateLimiterStrategy limiterStrategy = rateLimiterHandler.get(rateLimiter.type());

        boolean isPass = limiterStrategy.rateLimiterCheck(joinPoint, rateLimiter);

        if (!isPass){
            throw new BizException(rateLimiter.result());
        }
        // 继续执行原方法
        return joinPoint.proceed();
    }
}

2.3 使用方法
接下来只需要将注解添加在需要限流的方法上，即可进行限流。

对用户id进行限流：

@RateLimiter(type = RateLimiterType.USER_ID, time = 2, result = "请稍后重试哦")
@PostMapping("/testRate")
public BaseResponse<String> testRate() {
      return BaseResponse.success("成功");
 }

对参数进行限流：

@RateLimiter(type = RateLimiterType.ARGS, args = {"dialogueRecordId"}, time = 2, result = "请稍后重试哦")
@ClientResponseHandler
@GetMapping("/wtb/generateVoice")
public BaseResponse<String> wtbGenerateVoice(@RequestParam Long dialogueRecordId) {
    return BaseResponse.success("成功");
}