限流 RateLimit

介绍

限流框架是有的，比如Sentinel，但是不想集成到小项目，所以使用谷歌的guava实现令牌桶算法的限流，主要分两种，一种简单使用的针对所有人的接口限流，一种是针对不同IP实现接口QPS限流。

简单限流

对访问这个接口所有人总的QPS限制，且不设切面。
pom.xml

<dependency>
    <groupId>com.google.guava</groupId>
    <artifactId>guava</artifactId>
    <version>29.0-jre</version>
</dependency>

controller

//创建令牌桶，每秒0.5个令牌（每隔2秒允许请求一次）
private RateLimiter rl = RateLimiter.create(0.5);

@PostMapping("/insert")
@ResponseBody
public Result insert() {
    //获取令牌，如果没有则等待至超时，本代码超时时间为0，立刻返回错误信息
    boolean flag = rl.tryAcquire(0, TimeUnit.SECONDS);
    if (!flag) {
        return Result.fail("服务器忙");
    } else {
        xxx...
        return Result.ok("插入成功");
    }
}

针对不同IP实现

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

RateLimit.java 注解类

package com.rawchen.shorturl.limit;

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

/**
 * 限流注解
 *
 * @author RawChen
 * @date 2022-03-31 11:44
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {

    /**
     * 资源名称
     *
     * @return
     */
    String name() default "默认资源";

    /**
     * 每秒允许数（qps），默认每秒0.1次也就是10秒允许访问一次
     *
     * @return
     */
    double perSecond() default 0.1;

}

RateLimitAspect.java 切面类

package com.rawchen.shorturl.limit;

import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.util.concurrent.RateLimiter;
import com.rawchen.shorturl.entity.Result;
import org.aspectj.lang.ProceedingJoinPoint;
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 org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.concurrent.TimeUnit;

/**
 * 限流切面
 *
 * @author RawChen
 * @date 2022-03-31 11:46
 */
@Aspect
@Component
public class RateLimitAspect {

    private static final Logger logger = LoggerFactory.getLogger(RateLimitAspect.class);

    /**
     * 缓存
     * maximumSize 设置缓存个数
     * expireAfterWrite 写入后过期时间
     */
    private static LoadingCache<String, RateLimiter> limitCaches = CacheBuilder.newBuilder()
            .maximumSize(1000)
            .expireAfterWrite(1, TimeUnit.DAYS)
            .build(new CacheLoader<String, RateLimiter>() {
                @Override
                public RateLimiter load(String key) {
                    double perSecond = RateLimitUtil.getCacheKeyPerSecond(key);
                    return RateLimiter.create(perSecond);
                }
            });

    /**
     * 切点
     * 通过扫包切入 @Pointcut("execution(public * com.rawchen.shorturl.*.*(..))")
     * 带有指定注解切入 @Pointcut("@annotation(com.rawchen.shorturl.annotation.RateLimit)")
     */
    @Pointcut("@annotation(com.rawchen.shorturl.limit.RateLimit)")
    public void pointcut() {
    }

    @Around("pointcut()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        if (method.isAnnotationPresent(RateLimit.class)) {
            String cacheKey = RateLimitUtil.generateCacheKey(method, request);
            RateLimiter limiter = limitCaches.get(cacheKey);
            if (!limiter.tryAcquire()) {
                logger.info("限流{}方法，具体内容【{}】", point.getSignature().getName(), cacheKey);
                return Result.fail("你手速太快了");
            }
        }
        return point.proceed();
    }
}

RateLimitUtil.java 限流工具

package com.rawchen.shorturl.limit;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.net.InetAddress;
import java.net.UnknownHostException;

/**
 * 限流工具
 *
 * @author RawChen
 * @date 2022-03-31 11:48
 */
public class RateLimitUtil {
    /**
     * 获取唯一key根据注解类型，规则为： 资源名:ip:perSecond
     *
     * @param method
     * @param request
     * @return
     */
    public static String generateCacheKey(Method method, HttpServletRequest request) {
        //获取方法上的注解
        RateLimit rateLimit = method.getAnnotation(RateLimit.class);
        StringBuffer cacheKey = new StringBuffer(rateLimit.name() + ":");
        cacheKey.append(getIpAddr(request) + ":");
        cacheKey.append(rateLimit.perSecond());
        return cacheKey.toString();
    }

    /**
     * 获取缓存key的限制每秒访问次数
     * <p>
     * 规则 资源名:业务key:perSecond
     *
     * @param cacheKey
     * @return
     */
    public static double getCacheKeyPerSecond(String cacheKey) {
        String perSecond = cacheKey.split(":")[2];
        return Double.parseDouble(perSecond);
    }

    /**
     * 获取客户端IP地址
     *
     * @param request 请求
     * @return
     */
    public static String getIpAddr(HttpServletRequest request) {
        String ip = request.getHeader("x-forwarded-for");
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("WL-Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getRemoteAddr();
            if ("127.0.0.1".equals(ip)) {
                //根据网卡取本机配置的IP
                InetAddress inet = null;
                try {
                    inet = InetAddress.getLocalHost();
                } catch (UnknownHostException e) {
                    e.printStackTrace();
                }
                ip = inet.getHostAddress();
            }
        }
        // 对于通过多个代理的情况，第一个IP为客户端真实IP,多个IP按照','分割
        if (ip != null && ip.length() > 15) {
            if (ip.indexOf(",") > 0) {
                ip = ip.substring(0, ip.indexOf(","));
            }
        }
        if ("0:0:0:0:0:0:0:1".equals(ip)) {
            ip = "127.0.0.1";
        }
        return ip;
    }
}

Result.java

package com.rawchen.shorturl.entity;

/**
 * 封装结果
 *
 * @author RawChen
 * @date 2022-03-30 17:49
 */
public class Result {
    private Integer code;
    private String msg;
    private Object data;

    public Result() {
    }

    private Result(Integer code, String msg) {
        this.code = code;
        this.msg = msg;
        this.data = null;
    }

    private Result(Integer code, String msg, Object data) {
        this.code = code;
        this.msg = msg;
        this.data = data;
    }

    public static Result ok(String msg, Object data) {
        return new Result(200, msg, data);
    }

    public static Result ok(String msg) {
        return new Result(200, msg);
    }

    public static Result ok() {
        return new Result(200, "success");
    }

    public static Result ok(Object data) {
        return new Result(200, "success", data);
    }

    public static Result fail(String msg) {
        return new Result(400, msg);
    }

    public static Result fail() {
        return new Result(400, "fail");
    }

    public static Result error(String msg) {
        return new Result(500, msg);
    }

    public static Result error() {
        return new Result(500, "error");
    }

    public static Result create(Integer code, String msg, Object data) {
        return new Result(code, msg, data);
    }

    public static Result create(Integer code, String msg) {
        return new Result(code, msg);
    }

    public Integer getCode() {
        return code;
    }

    public void setCode(Integer code) {
        this.code = code;
    }

    public String getMsg() {
        return msg;
    }

    public void setMsg(String msg) {
        this.msg = msg;
    }

    public Object getData() {
        return data;
    }

    public void setData(Object data) {
        this.data = data;
    }

    @Override
    public String toString() {
        return "Result{" +
                "code=" + code +
                ", msg='" + msg + '\'' +
                ", data=" + data +
                '}';
    }
}