Springboot-自定义注解

业务上我们使用注解实现限流和幂等性校验，下面已限流为例介绍

一、注解的定义

元注解

有一些注解可以修饰其他注解，这些注解就称为元注解（meta annotation）

@Inherited

使用@Inherited定义子类是否可继承父类定义的Annotation。@Inherited仅针对@Target(ElementType.TYPE)类型的annotation有效，并且仅针对class的继承，对interface的继承无效。

@Documented 注解

指明修饰的注解，可以被例如javadoc此类的工具文档化，只负责标记，没有成员取值。

@Target

指明了修饰的这个注解的使用范围，即被描述的注解可以用在哪里：

• 类或接口：ElementType.TYPE

• 字段：ElementType.FIELD

• 方法：ElementType.METHOD

• 构造方法：ElementType.CONSTRUCTOR

• 方法参数：ElementType.PARAMETER

@Retention

指明修饰的注解的生存周期，即会保留到哪个阶段：

• 仅编译期：RetentionPolicy.SOURCE

• 仅class文件：RetentionPolicy.CLASS

• 运行期：RetentionPolicy.RUNTIME

二、限流注解

Java语言使用@interface语法来定义注解（Annotation），它的格式如下：

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

/**
 * 限流注解
 * @author meng
 * @date 2021/11/17 9:38
 */
@Inherited
@Documented
@Target({ElementType.TYPE,ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
public @interface RequestRateLimiter {

    /**
     * 这里指吞吐率每秒多少许可数（通常是指QPS，每秒多少查询）
     * @return
     */
    double QPS() default 100D;

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

    /**
     * 获取令牌超时时间单位：默认为 毫秒
     * @return
     */
    TimeUnit timeunit() default TimeUnit.MILLISECONDS;

}

三、限流拦截器

import com.google.common.util.concurrent.RateLimiter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.annotation.AnnotatedElementUtils;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.stereotype.Component;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.lang.reflect.Method;
import java.util.Enumeration;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author meng
 * @date 2021/11/17 9:35
 */
@Component
public class RateLimiterInterceptor implements HandlerInterceptor {
    
    private static final Logger logger = LoggerFactory.getLogger(RateLimiterInterceptor.class);

    /**
     * 根据请求地址保存不同的令牌桶
     */
    private static final Map<String, RateLimiter> rateLimiterMap = new ConcurrentHashMap<String, RateLimiter>(16);

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
        if (handler instanceof HandlerMethod) {
            HandlerMethod handlerMethod = (HandlerMethod) handler;
            Method method = ((HandlerMethod) handler).getMethod();

            //没有标注注解，放行
            if (AnnotatedElementUtils.isAnnotated(method, RequestRateLimiter.class)) {
                //首先获取方法上的注解
                RequestRateLimiter requestRateLimiter = AnnotationUtils.findAnnotation(method, RequestRateLimiter.class);
                //方法上没有标注该注解，尝试获取类上的注解
                if (Objects.isNull(requestRateLimiter)) {
                    //获取类上的注解
                    requestRateLimiter = AnnotationUtils.findAnnotation(handlerMethod.getBean().getClass(), RequestRateLimiter.class);
                }
                // 获取请求 url
                String url = request.getRequestURI();
                // 定义令牌桶
                RateLimiter rateLimiter = null;
                if (!rateLimiterMap.containsKey(url)) {
                    // 为当前请求创建令牌桶
                    rateLimiter = RateLimiter.create(requestRateLimiter.QPS());
                    rateLimiterMap.put(url, rateLimiter);
                } else {
                    // 根据请求 url 获取令牌桶
                    rateLimiter = rateLimiterMap.get(url);
                }

                // 获取令牌
                boolean acquire = rateLimiter.tryAcquire(requestRateLimiter.acquireTokenTimeout(), requestRateLimiter.timeunit());
                logger.info("IP:{},url:{},acquire:{}", getIpAddr(request), url, acquire);
                if (!acquire) {
                    throw new RateLimiterValidatorException();
                }

            }
        }
        return true;
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) {
    }

    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;
    }

}

PS：

1. 使用Guava RateLimiter实现限流

2. RateLimiterValidatorException自定义异常，详见博客：Springboot - 统一异常处理