当前位置：首页 > news >正文

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

news 来源：原创 2024/9/30 1:42:13

这里是weihubeats,觉得文章不错可以关注公众号小奏技术，文章首发。拒绝营销号，拒绝标题党

背景

在一些高并发或者供外部访问的接口，在受服务器压力的背景下，我们需要根据自身服务器每分钟可提供的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;
    }