RateLimit接口限流

引入依赖

        <!--guava-->
        <dependency>
            <groupId>com.google.guava</groupId>
            <artifactId>guava</artifactId>
            <version>30.1-jre</version>
        </dependency>
        <!--aop-->
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-aop</artifactId>
        </dependency>

创建自定义注解

import java.lang.annotation.*;
import java.util.concurrent.TimeUnit;

@Documented
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {
    /**
     * 每秒创建令牌个数,默认为10
     * @return
     */
    double permitsPerSecond() default 10D;

    /**
     * 获取令牌超时时间
     * @return
     */
    long timeout() default 0;

    /**
     * 超时时间单位
     * @return
     */
    TimeUnit timeUnit() default TimeUnit.SECONDS;


}

创建aop

import com.google.common.util.concurrent.RateLimiter;
import com.kaying.luck.annotation.ratelimit.RateLimit;
import com.kaying.luck.controller.draw.UserDrawController;
import com.kaying.luck.exception.ServiceException;
import com.kaying.luck.response.enums.ApiResponseCodeEnum;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

@Aspect
@Component
public class RateLimiterAOP {
    private static final Logger logger = LoggerFactory.getLogger(UserDrawController.class);
    /**
     * 不同的方法存放不同的令牌桶
     */
    private final Map<String, RateLimiter> map = new ConcurrentHashMap<>();

    /**
     * 定义切入点,自定义RateLimit
     */
    @Pointcut("@annotation(com.kaying.luck.annotation.ratelimit.RateLimit)")
    public void pointCut() {}

    @Around(value = "pointCut()")
    public Object  around(ProceedingJoinPoint joinPoint) throws Throwable {
        logger.info("进入限流控制");
        Object obj = null;
        Signature signature = joinPoint.getSignature();
        String methodName = signature.getName();
        Method method = ((MethodSignature) signature).getMethod();
        //获取注解对象
        RateLimit rateLimit = method.getAnnotation(RateLimit.class);
        if (rateLimit != null) {
            double permitsPerSecond = rateLimit.permitsPerSecond();
            long timeout = rateLimit.timeout();
            TimeUnit timeUnit = rateLimit.timeUnit();
            RateLimiter rateLimiter = null;
            if (!map.containsKey(methodName)) {
                // 创建令牌桶
                rateLimiter = RateLimiter.create(permitsPerSecond);
                map.put(methodName, rateLimiter);
            }
            rateLimiter = map.get(methodName);
            if (rateLimiter.tryAcquire(timeout, timeUnit)) {
                logger.info("成功访问!");
                obj = joinPoint.proceed();
                return obj;
            } else {
                logger.info("当前访问人数众多,请稍后再来");
                throw new ServiceException(ApiResponseCodeEnum.FAIL,"当前访问人数众多,请稍后再来");
            }
        } else {
            logger.info("aop未知错误");
            throw new ServiceException(ApiResponseCodeEnum.FAIL,"当前访问人数众多,请稍后再来");
        }
    }

}

在方法上添加自定义注解

    @ApiOperation("RateLimit测试")
    @GetMapping("rate/limit/test/{name}")
    @RateLimit(permitsPerSecond = 1D, timeout = 0, timeUnit = TimeUnit.SECONDS)
    public String sendMessage(@PathVariable("name") String name) {
        System.out.println("name = " + name);
        return "通过";
    }

结果

猜你喜欢

转载自blog.csdn.net/m0_58709145/article/details/129334065