• 前言
• 实现
  • 依赖
  • 注解
  • 枚举
  • 配置项
  • AOP 切面
  • 容器装配

前言

本内容介绍使用 Redis 对接口进行限流访问。

主要使用 AOP 对接口进行切入，然后 Before 时获取相关信息，存入 Redis，并设置失效时间，这样在没有失效前，访问接口时会返回失败信息。

实现

依赖

<dependencies>
  <dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-redis</artifactId>
  </dependency>
  <dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-aop</artifactId>
  </dependency>
</dependencies>

注解

/**
 * 例： @RedisLimit(key = "redis-limit:test", permitsPerSecond = 4, expire = 1, msg = "请求太频繁，等一下啦")
 */
@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.METHOD})
@Documented
public @interface RateLimit {
    /**
     * 资源的 key，唯一
     * 作用：不同的接口，不同的流量控制
     */
    String key() default "";

    /**
     * 限流类型
     */
    RateLimitType rateLimitType() default RateLimitType.DEFAULT;

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

    /**
     * 过期时间也可以理解为单位时间，单位秒，默认 60
     */
    long expire() default 60;

    /**
     * 得不到令牌的提示语
     */
    String msg() default "系统繁忙，请稍后再试！";
}

枚举

public enum RateLimitType {
    /**
     * 默认策略全局限流
     */
    DEFAULT,

    /**
     * 根据用户 ID 限流
     */
    USER,

    /**
     * 根据请求者 IP 进行限流
     */
    IP,

    /**
     * 实例限流(集群多后端实例)
     */
    CLUSTER
}

配置项

@Getter
@Setter
@Component
@ConfigurationProperties(prefix = "redis-limit")
public class RateLimitProperties {
    
    /**
     * 访问次数
     */
    private Long limit;

    /**
     * 有效期,单位秒
     */
    private Long expire;

    /**
     * 限流 urls
     */
    private String[] urls;

}

AOP 切面

@Slf4j
@Aspect
@RequiredArgsConstructor
public class RateLimitAspect {

    private final StringRedisTemplate stringRedisTemplate;
    private final RateLimitProperties rateLimitProperties;

    private static final ExpressionParser EXPRESSION_PARSER = new SpelExpressionParser();
    private static final TemplateParserContext TEMPLATE_PARSER_CONTEXT = new TemplateParserContext();

    @Pointcut("execution(* com..controller..*.*(..)) || @annotation(cn.youngkbt.ratelimiter.annotations.RateLimit)")
    private void check() {
    }

    private DefaultRedisScript<Long> redisScript;

    @PostConstruct
    public void init() {
        redisScript = new DefaultRedisScript<>();
        redisScript.setResultType(Long.class);
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("rateLimiter.lua")));
    }

    @Before("check()")
    public void before(JoinPoint joinPoint) {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        // 拿到 RedisLimit 注解，如果存在则说明需要限流
        RateLimit rateLimit = method.getAnnotation(RateLimit.class);
        HttpServletRequest request = WebUtil.getRequest();
        if (Objects.isNull(request)) {
            return;
        }
        // StrUtil.containsAnyIgnoreCase 判断请求的 url 是否有配置文件限流的 urls
        if (Objects.nonNull(rateLimit) || StrUtil.containsAnyIgnoreCase(request.getRequestURI(), rateLimitProperties.getUrls())) {
            // 获取 redis 的 key
            String key;
            long limit;
            long expire;
            if (Objects.nonNull(rateLimit)) {
                key = getLimitKey(rateLimit.key(), rateLimit.rateLimitType(), joinPoint);
                limit = rateLimit.count();
                expire = rateLimit.expire();
            } else {
                // 根据 URI + 用户 ID 限流
                key = getLimitKey(request.getRequestURI(), RateLimitType.USER, joinPoint);
                limit = rateLimitProperties.getLimit();
                expire = rateLimitProperties.getExpire();
            }
            if (!StringUtil.hasText(key)) {
                throw new RedisLimitException("redis key cannot be null");
            }

            List<String> keys = new ArrayList<>();
            keys.add(key);
            Long count = stringRedisTemplate.execute(redisScript, keys, String.valueOf(limit), String.valueOf(expire));
            log.info("接口限流, 尝试访问次数为 {}，key：{}", count, key);

            if (Objects.nonNull(count) && count == 0) {
                log.debug("接口限流, 导致获取 key 失败，key 为 {}", key);
                throw new RedisLimitException(Objects.nonNull(rateLimit) ? rateLimit.msg() : "请求过于频繁！");
            }
        }
    }

    private String getLimitKey(String key, RateLimitType rateLimitType, JoinPoint joinPoint) {
        StringBuilder stringBuffer = new StringBuilder(RedisConstants.SERVER_REQUEST_LIMIT);

        if (StringUtil.hasText(key)) {
            key = getSpeElValue(key, joinPoint);
        }

        if (rateLimitType == RateLimitType.USER) {
            // 获取用户 ID
            stringBuffer.append(LoginHelper.getUsername()).append(":");
        } else if (rateLimitType == RateLimitType.IP) {
            // 获取请求 IP
            stringBuffer.append(ServletUtil.getClientIp()).append(":");
        } else if (rateLimitType == RateLimitType.CLUSTER) {
            // 获取客户端实例 ID
            // stringBuffer.append(getClientId()).append(":");
        }

        return stringBuffer.append(key).toString();
    }

    private String getSpeElValue(String key, JoinPoint joinPoint) {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method targetMethod = signature.getMethod();
        // 方法的参数
        Object[] args = joinPoint.getArgs();

        // 创建 MethodBasedEvaluationContext
        MethodBasedEvaluationContext context = new MethodBasedEvaluationContext("", targetMethod, args, new DefaultParameterNameDiscoverer());

        // 设置 ApplicationContext 到 Context 中
        context.setBeanResolver(new BeanFactoryResolver(SpringHelper.getBeanFactory()));

        Expression expression = null;
        // 如果 key 为 SpEl 表达式
        if (StringUtils.startsWith(key, TEMPLATE_PARSER_CONTEXT.getExpressionPrefix())
                && StringUtils.endsWith(key, TEMPLATE_PARSER_CONTEXT.getExpressionSuffix())) {
            expression = EXPRESSION_PARSER.parseExpression(key, TEMPLATE_PARSER_CONTEXT);
        } else if (StringUtils.startsWith(key, "#")) {
            expression = EXPRESSION_PARSER.parseExpression(key);
        }
        return Objects.nonNull(expression) ? expression.getValue(context, String.class) : key;
    }
}

容器装配

@AutoConfiguration(after = RedisConfiguration.class)
@EnableConfigurationProperties(RateLimitProperties.class)
public class RateLimitConfiguration {

    @Bean
    public RateLimitAspect rateLimiterAspect(StringRedisTemplate stringRedisTemplate, RateLimitProperties rateLimitProperties) {
        return new RateLimitAspect(stringRedisTemplate, rateLimitProperties);
    }
}

Spring Boot 3.x 需要在 resource 下建立 META-INF/spring 路径，然后创建 org.springframework.boot.autoconfigure.AutoConfiguration.imports 文件，内容为

cn.youngkbt.ratelimiter.config.RateLimitConfiguration

这样 Spring 会自动扫描该文件的容器装配类，将里面涉及的类注入到 Spring 容器。