支持方法动态参数的分布式限流实现

一、需求背景

核心诉求：

• 支持自定义统计窗口
  需要实现不定长统计窗口的限流，比如半小时一次。
• 支持分布式
  由于可能限定的流量为一次，需要分布式支持。
• 支持注解
  简化使用。
• 支持动态参数
  需要根据方法参数做细粒度的限流。

二、方案对比

方案	自定义统计窗口	分布式	注解	动态参数
Guava	简单	不支持	不支持	不支持
Sentinel	部分支持	较复杂	支持	部分支持
Redisson	简单	简单	不支持	不支持

对比写的比较浅显，仅体现当前需求，实际需要考虑的因素会很多，其它方案也有很多。

简单总结：

1. Guava 使用简单，缺少分布式和注解等支持。
2. Sentinel 功能强大，但是考虑到需求并非通用限流模式，导致整体实现起来较复杂。
3. Redisson 使用简单，支持分布式，缺少注解等支持。

最后选择 Redisson 方案，使用简单，补全缺少的功能的实现也不复杂。

三、Redisson 实现

• Spring Boot 2.x
• Jdk 1.8

1、添加依赖

注意依赖的版本，参考官网。

     <dependency>
         <groupId>org.redisson</groupId>
         <artifactId>redisson-spring-boot-starter</artifactId>
         <version>2.15.2</version>
     </dependency>

2、添加配置

支持多种配置方法，参考官网。

spring:
  redis:
    password: xxxxxxxx
    cluster:
      nodes:
        - xxx.xxx.xxx.xxx:6379
        - xxx.xxx.xxx.xxx:6379

3、添加自定义注解和AOP实现

自定义注解，key 值为限流全局唯一标识，支持 Spel 表达式。

@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {

    // 唯一标识
    String key() default "";

    // 限制流量
    int rate() default 1;

    // 统计窗口 (单位秒)
    int interval() default 1;

}

具体限流AOP实现如下，主要包括 key值解析，和限流逻辑。

限流逻辑：

• 逻辑比较简单，需要注意的是，如果程序限流配置变更，需要自行删除 Redis 中旧的配置，当然也是可以添加获取、对比和删除配置的逻辑，就是性能上有点影响。
• 可根据实际需要添加更多功能，比如分布式限流失败，降级为单机限流。

Key值解析：

• 将方法的所有参数封装到一个 Map 对象中。
• 在key值做进一步解析前，将其格式转变成存在于Map对象内，类似 user.id 转换为 ['user'].id。

@Aspect
@Component
@Slf4j
@RequiredArgsConstructor
public class RateLimitAspect {

    public static final String RATE_LIMITER_KEY_PREFIX = "rate-limiter:";
    private final RedissonClient redissonClient;


    @Around("@annotation(rateLimit)")
    private Object around(ProceedingJoinPoint joinPoint, RateLimit rateLimit) throws Throwable {
        String key = getKey(joinPoint, rateLimit);
        log.info("rate limiter key: {}", key);

        // Key为空，不做限流
        if (key == null) {
            return joinPoint.proceed();
        }

        RRateLimiter rateLimiter = redissonClient.getRateLimiter(key);
        if (!rateLimiter.isExists()) {
            // 如果配置变更，需要先删除Redis中旧的配置
            rateLimiter.trySetRate(RateType.OVERALL, rateLimit.rate(), rateLimit.interval(), RateIntervalUnit.SECONDS);
        }

        if (!rateLimiter.tryAcquire()) {
            throw new RRException("操作过于频繁，请稍后再试");
        }

        return joinPoint.proceed();
    }

    private String getKey(ProceedingJoinPoint joinPoint, RateLimit rateLimit) {
        String key = rateLimit.key();
        if (StringUtils.isEmpty(key)) {
            return null;
        }

        String parsedKey;
        try {
            parsedKey = ExprUtils.parse(key, getArgsMap(joinPoint));
        } catch (Throwable e) {
            log.error("parse rate limiter's key failed", e);
            return null;
        }
        return RATE_LIMITER_KEY_PREFIX + parsedKey;
    }

    private Map<String, Object> getArgsMap(ProceedingJoinPoint joinPoint) {
        String[] names = ((CodeSignature) joinPoint.getSignature()).getParameterNames();
        Object[] args = joinPoint.getArgs();

        Map<String, Object> argsMap = new HashMap<>(names.length);
        for (int i = 0; i < names.length; i++) {
            argsMap.put(names[i], args[i]);
        }
        return argsMap;
    }
}

public class ExprUtils {

    private static final ExpressionParser parser = new RateLimitExpressionParser();
    private static final SimpleEvaluationContext context = SimpleEvaluationContext.forReadOnlyDataBinding().build();


    public static String parse(String expressionString, Map<String, Object> rootObject) {
        return parser.parseExpression(expressionString, TEMPLATE_EXPRESSION).getValue(context, rootObject, String.class);
    }


    private static class RateLimitExpressionParser extends SpelExpressionParser {

        @Override
        protected SpelExpression doParseExpression(String expressionString, @Nullable ParserContext context) throws ParseException {
            return super.doParseExpression(wrapWithMap(expressionString), context);
        }

        private String wrapWithMap(String expressionString) {
            int dotIndex = expressionString.indexOf(".");
            if (dotIndex > 0) {
                return "['" + expressionString.substring(0, dotIndex) + "']" + expressionString.substring(dotIndex - 1);
            }
            return "['" + expressionString + "']";
        }
    }

}

4、配置使用

一个简单的查询用户信息的服务实现，限制每个用户每5分钟只能访问1次。

@Service
public class InspectionTriggerServiceImpl implements InspectionTriggerService {

    @Override
    @RateLimit(key = "test-demo:user-service:get-#{user.id}", rate = 1, interval = 5 * 60)
    public UserVO get(UserBO user) {
        // ...
    }

}

四、总结

该方案适合简单的场景，如果是复杂的场景，还是推荐使用 Sentinel 实现。