分布式限流不会用？一个注解简单搞定

背景

在一些高并发或者供外部访问的接口，在受服务器压力的背景下，我们需要根据自身服务器每分钟可提供的QPS对一些接口进行限流处理，来保障整个系统的稳定性，所以我们需要一个限流功能。最简单的限流方式就是使用Guava的 RateLimiter

public void testRateLimiter() {
 RateLimiter r = RateLimiter.create(10);
 while (true) {
 System.out.println("get 1 tokens: " + r.acquire() + "s");
 }

但是改方案不是一个分布式限流，现在都是分布式系统，多节点部署.我们希望基于IP或者自定义的key去分布式限流，比如一个用户在1分钟内只能访问接口100次。入股是这种方式限流，有三个接口，实际访问的次数就是300次

Redis分布式限流

Redis分布式限流自己实现一般是使用Lua脚本去实现，但是实际编写Lua脚本还是比较费劲，庆幸的是Redisson直接提供了基于Lua脚本实现的分布式限流类RRateLimiter

分布式限流sdk编写

为了使用简单方便，我们还是对Redisson进行简单封装，封装一个注解来使用分布式限流

定义注解

Limiter

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

 /**
  * 限流器的 key
  *
  * @return key
  */
 String key() default "";

 /**
  * 限制数量
  *
  * @return 许可数量
  */
 long rate() default 100;

 /**
  * 速率时间间隔
  *
  * @return 速率时间间隔
  */
 long rateInterval() default 1;

 /**
  * 时间单位
  *
  * @return 时间
  */
 RateIntervalUnit rateIntervalUnit() default RateIntervalUnit.MINUTES;

 RateType rateType() default RateType.OVERALL;

}

IP工具类

由于需要获取IP，所以我们写一个IP获取工具类

IpUtil

public class IpUtil {


 public static String getIpAddr(HttpServletRequest request) {
  String ipAddress = null;
  try {
   ipAddress = request.getHeader("x-forwarded-for");
   if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
    ipAddress = request.getHeader("Proxy-Client-IP");
   }
   if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
    ipAddress = request.getHeader("WL-Proxy-Client-IP");
   }
   if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
    ipAddress = request.getRemoteAddr();
    if (ipAddress.equals("127.0.0.1")) {
     // 根据网卡取本机配置的IP
     InetAddress inet = null;
     try {
      inet = InetAddress.getLocalHost();
     }
     catch (UnknownHostException e) {
      e.printStackTrace();
     }
     ipAddress = inet.getHostAddress();
    }
   }
   // 对于通过多个代理的情况，第一个IP为客户端真实IP，多个IP按照','分割
   if (ipAddress != null && ipAddress.length() > 15) { // "***.***.***.***".length()
    if (ipAddress.indexOf(",") > 0) {
     ipAddress = ipAddress.substring(0, ipAddress.indexOf(","));
    }
   }
  }
  catch (Exception e) {
   ipAddress = "";
  }
  return ipAddress;
 }

}

AOP切面

AnnotationAdvisor

public class AnnotationAdvisor extends AbstractPointcutAdvisor implements BeanFactoryAware {

    private final Advice advice;

    private final Pointcut pointcut;

    private final Class<? extends Annotation> annotation;

    public AnnotationAdvisor(@NonNull MethodInterceptor advice,
                                                 @NonNull Class<? extends Annotation> annotation) {
        this.advice = advice;
        this.annotation = annotation;
        this.pointcut = buildPointcut();
    }

    @Override
    public Pointcut getPointcut() {
        return this.pointcut;
    }

    @Override
    public Advice getAdvice() {
        return this.advice;
    }

    @Override
    public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
        if (this.advice instanceof BeanFactoryAware) {
            ((BeanFactoryAware) this.advice).setBeanFactory(beanFactory);
        }
    }

    private Pointcut buildPointcut() {
        Pointcut cpc = new AnnotationMatchingPointcut(annotation, true);
        Pointcut mpc = new AnnotationMethodPoint(annotation);
        return new ComposablePointcut(cpc).union(mpc);
    }

    /**
     * In order to be compatible with the spring lower than 5.0
     */
    private static class AnnotationMethodPoint implements Pointcut {

        private final Class<? extends Annotation> annotationType;

        public AnnotationMethodPoint(Class<? extends Annotation> annotationType) {
            Assert.notNull(annotationType, "Annotation type must not be null");
            this.annotationType = annotationType;
        }

        @Override
        public ClassFilter getClassFilter() {
            return ClassFilter.TRUE;
        }

        @Override
        public MethodMatcher getMethodMatcher() {
            return new AnnotationMethodMatcher(annotationType);
        }

        private static class AnnotationMethodMatcher extends StaticMethodMatcher {
            private final Class<? extends Annotation> annotationType;

            public AnnotationMethodMatcher(Class<? extends Annotation> annotationType) {
                this.annotationType = annotationType;
            }

            @Override
            public boolean matches(Method method, Class<?> targetClass) {
                if (matchesMethod(method)) {
                    return true;
                }
                // Proxy classes never have annotations on their redeclared methods.
                if (Proxy.isProxyClass(targetClass)) {
                    return false;
                }
                // The method may be on an interface, so let's check on the target class as well.
                Method specificMethod = AopUtils.getMostSpecificMethod(method, targetClass);
                return (specificMethod != method && matchesMethod(specificMethod));
            }

            private boolean matchesMethod(Method method) {
                return AnnotatedElementUtils.hasAnnotation(method, this.annotationType);
            }
        }
    }
}

LimiterAnnotationInterceptor

核心实现类

@RequiredArgsConstructor
@Slf4j
public class LimiterAnnotationInterceptor implements MethodInterceptor {


 private final RedissonClient redisson;

 private static final Map<RateIntervalUnit, String> INSTANCE = Map.ofEntries(
   entry(RateIntervalUnit.SECONDS, "秒"),
   entry(RateIntervalUnit.MINUTES, "分钟"),
   entry(RateIntervalUnit.HOURS, "小时"),
   entry(RateIntervalUnit.DAYS, "天"));


 @Nullable
 @Override
 public Object invoke(@NotNull MethodInvocation invocation) throws Throwable {

  Method method = invocation.getMethod();
  Limiter limiter = method.getAnnotation(Limiter.class);
  long limitNum = limiter.rate();
  long limitTimeInterval = limiter.rateInterval();

  ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
  HttpServletRequest request = attributes.getRequest();
  String ip = IpUtil.getIpAddr(request);

  String key = DataUtils.isEmpty(limiter.key()) ? "limit:" + ip + "-" + request.getRequestURI() : limiter.key();

  RateIntervalUnit rateIntervalUnit = limiter.rateIntervalUnit();
  RRateLimiter rateLimiter = redisson.getRateLimiter(key);
  if (rateLimiter.isExists()) {
   RateLimiterConfig config = rateLimiter.getConfig();
   if (!Objects.equals(limiter.rate(), config.getRate())
     || !Objects.equals(limiter.rateIntervalUnit()
     .toMillis(limiter.rateInterval()), config.getRateInterval())
     || !Objects.equals(limiter.rateType(), config.getRateType())) {
    rateLimiter.delete();
    rateLimiter.trySetRate(limiter.rateType(), limiter.rate(), limiter.rateInterval(), limiter.rateIntervalUnit());
   }
  }
  else {
   rateLimiter.trySetRate(RateType.OVERALL, limiter.rate(), limiter.rateInterval(), limiter.rateIntervalUnit());
  }

  boolean allow = rateLimiter.tryAcquire();
  if (!allow) {
   String url = request.getRequestURL().toString();
   String unit = getInstance().get(rateIntervalUnit);
   String tooManyRequestMsg = String.format("用户IP[%s]访问地址[%s]时间间隔[%s %s]超过了限定的次数[%s]", ip, url, limitTimeInterval, unit, limitNum);
   log.info(tooManyRequestMsg);
   throw new BizException("访问速度过于频繁,请稍后再试");
  }
  return invocation.proceed();
 }

 public static Map<RateIntervalUnit, String> getInstance() {
  return INSTANCE;
 }

}

自动装载AOP Bean

AutoConfiguration

@Slf4j
@Configuration(proxyBeanMethods = false)
public class AutoConfiguration {
    @Bean
    public Advisor limiterAdvisor(RedissonClient redissonClient) {
        LimiterAnnotationInterceptor advisor = new LimiterAnnotationInterceptor(redissonClient);
        return new AnnotationAdvisor(advisor, Limiter.class);
    }
}

定义一个开启功能的注解

EnableLimiter

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.TYPE)
@Documented
@Import(AutoConfiguration.class)
public @interface EnableLimiter {
}

使用

@SpringBootApplication
@EnableLimiter
public class Application {

    public static void main(String[] args) {
        TimeZone.setDefault(TimeZone.getTimeZone("Asia/Shanghai"));
        SpringApplication.run(Application.class, args);
    }
}

配置一个RedissonClient

RedissonClient

@Configuration
public class RedissonConfig {

    @Value("${redis.host}")
    private String redisLoginHost;
    @Value("${redis.port}")
    private Integer redisLoginPort;
    @Value("${redis.password}")
    private String redisLoginPassword;


    @Bean
    public RedissonClient redissonClient() {
        return createRedis(redisLoginHost, redisLoginPort, redisLoginPassword);
    }

    private RedissonClient createRedis(String redisHost, Integer redisPort, String redisPassword) {
        Config config = new Config();
        SingleServerConfig singleServerConfig = config.useSingleServer();
        singleServerConfig.setAddress("redis://" + redisHost + ":" + redisPort + "");
        if (DataUtils.isNotEmpty(redisPassword)) {
            singleServerConfig.setPassword(redisPassword);
        }
        config.setCodec(new JsonJacksonCodec());
        return Redisson.create(config);
    }

}

controller使用注解

    @GetMapping("/testLimiter")
    @Limiter(rate = 2, rateInterval = 10, rateIntervalUnit = RateIntervalUnit.SECONDS)
    public ActionEnum testLimiter(String name) {
        log.info("testLimiter {}", name);
        return ActionEnum.SUCCESS;
    }