自定义aop思想实现RateLimiter令牌桶算法

1.自定义注解



package com.com.hangzhou.executor.ratelimiter.annotion;


import java.lang.annotation.*;

@Documented
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {
    double limitNum() default 10;
}

 

 

2.自定义aop实现类

package com.com.hangzhou.executor.ratelimiter;

import com.com.hangzhou.executor.ratelimiter.annotion.RateLimit;
import com.google.common.util.concurrent.RateLimiter;
import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;
import org.apache.commons.lang3.builder.ToStringBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
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.Charset;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 自定义注解限流切面 @RateLimit
 */
@Component
public class RateLimiterAop implements MethodInterceptor {

    /**
     * logger
     */
    private static final Logger LOGGER = LoggerFactory.getLogger(RateLimiterAop.class);

    /**
     * key -> methodName  value -> RateLimiter
     */
    private ConcurrentHashMap<String, RateLimiter> concurrentHashMap = new ConcurrentHashMap<>();

    @Autowired
    private HttpServletResponse response;

    @Override
    public Object invoke(MethodInvocation invocation) throws Throwable {
        //获取执行方法
        Method method = invocation.getMethod();
        //获取方法标注的注解@RateLimit
        RateLimit annotation = method.getAnnotation(RateLimit.class);
        //未标注的方法不需要拦截限流处理
        if (annotation == null) {
            return invocation.proceed();
        }
        //获取注解配置的限流值
        double limitNum = annotation.limitNum();

        //获取方法名
        String methodName = method.getName();

        if (!concurrentHashMap.containsKey(methodName)) {
            concurrentHashMap.put(methodName, RateLimiter.create(limitNum));
        }

        //获取限流RateLimiter令牌桶算法
        RateLimiter rateLimiter = concurrentHashMap.get(methodName);

        //尝试获取token  获取成功执行，打印日志，未获取到给于友好的提示
        if (rateLimiter.tryAcquire()) {
            Object result = invocation.proceed();
            LOGGER.info("执行的Controller" + method.getDeclaringClass().getName()
                    + "执行的方法名：" + methodName +
                    "执行的结果：" + ToStringBuilder.reflectionToString(result));
        } else {
            outWriteResultErrorMsg("网络繁忙，请稍后重试");
        }
        return invocation.proceed();

    }

    private void outWriteResultErrorMsg(String errorMsg) {
        ServletOutputStream outputStream = null;
        try {
            outputStream = response.getOutputStream();
            outputStream.write(errorMsg.getBytes(Charset.defaultCharset()));
        } catch (IOException e) {
            LOGGER.warn("IO异常" + e);
            throw new RuntimeException("自定义异常");
        } finally {
            if (outputStream != null) {
                try {
                    outputStream.flush();
                    outputStream.close();
                } catch (IOException e) {
                    LOGGER.warn("IO异常" + e);
                    throw new RuntimeException("自定义异常");
                }
            }
        }
    }
}

 

3.自定义controller

package com.com.hangzhou.executor.ratelimiter;

import com.com.hangzhou.executor.ratelimiter.annotion.RateLimit;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.ResponseBody;

@Controller
public class RateLimiterController {

    @RequestMapping("/test")
    @ResponseBody
    @RateLimit(limitNum = 20)
    public Object test() {
        return "success";
    }
}

 

4.自定义xml配置

<bean id="rateLimiterAop" class="com.com.hangzhou.executor.ratelimiter.RateLimiterAop"/>
<bean id="beanNameAutoProxyCreator" class="org.springframework.aop.framework.autoproxy.BeanNameAutoProxyCreator">
    <property name="interceptorNames">
        <list>
            <value>rateLimiterAop</value>
        </list>
    </property>
    <property name="beanNames">
        <list>
            <value>*Controller</value>
        </list>
    </property>
</bean>

 

spring aop自定义实现RateLimiter  @Pointcut  @Aroud  

package com.hangzhou.executor.aop.ratelimiter;

import com.com.hangzhou.executor.ratelimiter.annotion.RateLimit;
import com.google.common.util.concurrent.RateLimiter;
import org.apache.commons.lang3.builder.ToStringBuilder;
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.beans.factory.annotation.Autowired;
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.Charset;
import java.util.concurrent.ConcurrentHashMap;

@Aspect
@Component
@Scope
//配置@Scope默认为单例
public class RateLimiterAspect {

    private static final Logger LOGGER = LoggerFactory.getLogger(RateLimiterAspect.class);

    private ConcurrentHashMap<String, RateLimiter> concurrentHashMap = new ConcurrentHashMap<>();

    @Autowired
    private HttpServletResponse response;

    @Pointcut(value = "@annotation(com.com.hangzhou.executor.ratelimiter.annotion.RateLimit)")
    public void rateLimiterPointCut() {

    }

    @Around(value = "rateLimiterPointCut()")
    public Object executeRateLimiter(ProceedingJoinPoint joinPoint) {

        Object result = null;

        //获取执行的方法
        Signature signature = joinPoint.getSignature();

        //cast method type
        MethodSignature methodSignature = (MethodSignature) signature;

        //获取方法
        Method method = methodSignature.getMethod();

        //获取method中配置的注解
        RateLimit rateLimit = method.getAnnotation(RateLimit.class);

        //获取不到，说明不需要限流处理，直接放行
        if (rateLimit == null) {
            try {
                result = joinPoint.proceed();
                //打印日志
                LOGGER.info(ToStringBuilder.reflectionToString(result));
            } catch (Throwable throwable) {
                throwable.printStackTrace();
            }
        }
        //获取配置的限流数
        double limitNum = rateLimit.limitNum();

        //获取方法名 作为key
        String methodName = method.getName();

        if (!concurrentHashMap.containsKey(methodName)) {
            concurrentHashMap.put(methodName, RateLimiter.create(limitNum));
        }

        //获取限流RateLimiter令牌桶算法
        RateLimiter rateLimiter = concurrentHashMap.get(methodName);

        //尝试获取token  获取成功执行，打印日志，未获取到给于友好的提示
        if (rateLimiter.tryAcquire()) {
            try {
                result = joinPoint.proceed();
            } catch (Throwable throwable) {
                throwable.printStackTrace();
            }
            LOGGER.info("执行的Controller" + method.getDeclaringClass().getName()
                    + "执行的方法名：" + methodName +
                    "执行的结果：" + ToStringBuilder.reflectionToString(result));
        } else {
            outWriteResultErrorMsg("网络繁忙，请稍后重试");
        }

        return result;

    }

    private void outWriteResultErrorMsg(String msg) {
        ServletOutputStream outputStream = null;
        try {
            outputStream = response.getOutputStream();
            outputStream.write(msg.getBytes(Charset.defaultCharset()));
        } catch (IOException e) {
            LOGGER.warn("IO异常" + e);
            throw new RuntimeException("自定义异常");
        } finally {
            if (outputStream != null) {
                try {
                    outputStream.flush();
                    outputStream.close();
                } catch (IOException e) {
                    LOGGER.warn("IO异常" + e);
                    throw new RuntimeException("自定义异常");
                }
            }
        }
    }
}