解决：请求限流AOP+RateLimiter+注解

╰+哭是因爲堅強的太久メ 2023-10-09 04:14 12阅读 0赞

1、注解

import java.lang.annotation.*;
    
    /**
     * 限流 一段时间内，允许请求的最大数量
     *
     * @author 
     */
    @Inherited
    @Documented
    @Target({ElementType.METHOD})
    @Retention(RetentionPolicy.RUNTIME)
    public @interface RateLimit {
    
        /**
         * 请求数 1秒最多处理请求数   默认-1，不限制
         */
        long value() default -1;
    }

2、切面

import com.alibaba.fastjson.JSON;
    import com.google.common.util.concurrent.RateLimiter;
    import com.myfutech.common.util.Responses;
    import com.myfutech.common.util.enums.ResponseCode;
    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.springframework.beans.factory.annotation.Autowired;
    import org.springframework.context.annotation.Scope;
    import org.springframework.http.MediaType;
    import org.springframework.stereotype.Component;
    
    import javax.servlet.ServletOutputStream;
    import javax.servlet.http.HttpServletResponse;
    import java.io.IOException;
    import java.lang.reflect.Method;
    import java.nio.charset.StandardCharsets;
    import java.util.HashMap;
    import java.util.Map;
    
    @Component
    @Scope
    @Aspect
    public class RateLimitAop {
     
        @Autowired
        private HttpServletResponse response;
     
        private final Map<Method, RateLimiter> rateLimiterMap = new HashMap<>();
     
        @Pointcut("@annotation(com.unifgroup.system.callback.limit.RateLimit)")
        public void serviceLimit() {
     
        }
     
        @Around("serviceLimit()")
        public Object around(ProceedingJoinPoint joinPoint) throws Throwable{
            Signature signature = joinPoint.getSignature();
            MethodSignature methodSignature = (MethodSignature)signature;
            Method targetMethod = methodSignature.getMethod();
            RateLimit annotation = targetMethod.getDeclaredAnnotation(RateLimit.class);
            long requestNum = annotation.value();
    
            if (requestNum <= 0) {
                return joinPoint.proceed();
            }
    
            RateLimiter rateLimiter = rateLimiterMap.get(targetMethod);
            if (rateLimiter == null) {
                synchronized (rateLimiterMap) {
                    rateLimiter = rateLimiterMap.get(targetMethod);
                    if (rateLimiter == null) {
                        rateLimiter = RateLimiter.create(requestNum);
                        rateLimiterMap.put(targetMethod, rateLimiter);
                    }
                }
            }
    
            boolean acquire = rateLimiter.tryAcquire();
    
            if (acquire){
                return joinPoint.proceed();
            }
            String result = JSON.toJSONString(Responses.error(ResponseCode.WARN_CODE, "服务器压力过大，请稍后再试"));
            output(response, result);
            return null;
        }
        
        private void output(HttpServletResponse response, String msg) throws IOException {
            response.setContentType(MediaType.APPLICATION_JSON_UTF8_VALUE);
            ServletOutputStream os = null;
            try {
                os = response.getOutputStream();
                os.write(msg.getBytes(StandardCharsets.UTF_8));
            } finally {
                if (os != null) {
                    os.flush();
                    os.close();
                }
            }
        }
    }

3、controller

import com.myfutech.common.util.http.HttpClientUtils;
    import com.unifgroup.system.callback.limit.RateLimit;
    import io.swagger.annotations.ApiOperation;
    import lombok.extern.slf4j.Slf4j;
    import org.springframework.http.MediaType;
    import org.springframework.stereotype.Controller;
    import org.springframework.web.bind.annotation.PostMapping;
    import org.springframework.web.bind.annotation.RequestMapping;
    import org.springframework.web.bind.annotation.ResponseBody;
    
    import java.util.HashMap;
    import java.util.Map;
    import java.util.Random;
    
    /**
     * 测试请求
     */
    @Slf4j
    @Controller
    @RequestMapping("/limit")
    public class LimitCtrl {
    
    
        @RateLimit(1)
        @ApiOperation(value = "call")
        @PostMapping(value = "/call", produces = MediaType.APPLICATION_XML_VALUE)
        @ResponseBody
        public String call() {
            log.info("正常接收请求...");
            return "success";
        }
    }