Springboot+Guava实现单机令牌桶限流

令牌桶算法

系统会维护一个令牌（token）桶，以一个恒定的速度往桶里放入令牌（token），这时如果有请求进来想要被处理，则需要先从桶里获取一个令牌（token），当桶里没有令牌（token）可取时，则该请求将被拒绝服务。令牌桶算法通过控制桶的容量、发放令牌的速率，来达到对请求的限制。

==========================

Google开源工具包Guava提供了限流工具类RateLimiter，该类基于令牌桶算法实现流量限制，使用十分方便，而且十分高效

Pom中引入依赖

<properties>
  <spring-boot-start-version>2.6.9</spring-boot-start-version>
  <guava-version>31.1-jre</guava-version>
  </properties>
  <!--单机限流-->
        <dependency>
            <groupId>com.google.guava</groupId>
            <artifactId>guava</artifactId>
        </dependency>

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-aop</artifactId>
        </dependency>

编写注解，带上注解的接口实现限流控制

import java.lang.annotation.Documented;
import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
import java.util.concurrent.TimeUnit;

/**
 * @description: Guava限流单机
 * @author: GuoTong
 * @createTime: 2023-05-27 14:30
 * @since JDK 1.8 OR 11
 **/
@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.METHOD})
@Documented
public @interface GuavaAopLimit {

    /**
     * 资源的key,唯一：不同的接口，不同的流量控制
     */
    String key() default "";

    /**
     * 最多的访问限制次数
     */
    double permitsPerSecond();

    /**
     * 获取令牌最大等待时间
     */
    long timeout();

    /**
     * 获取令牌最大等待时间,单位(例:分钟/秒/毫秒) 默认:毫秒
     */
    TimeUnit timeunit() default TimeUnit.MILLISECONDS;

    /**
     * 得不到令牌的提示语
     */
    String msg() default "Guava单机限流令牌桶没了";
}

编写切面

import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import com.gton.annotation.GuavaAopLimit;
import com.gton.handler.RedisLimitException;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.Map;

/**
 * @description: ：使用AOP切面拦截限流注解
 * @author: GuoTong
 * @createTime: 2023-05-27 14:32
 * @since JDK 1.8 OR 11
 **/
@Slf4j
@Aspect
@Component
public class GuavaLimitAspect {
    /**
     * 不同的接口，不同的流量控制
     * map的key为 Limiter.key
     */
    private final Map<String, RateLimiter> limitMap = Maps.newConcurrentMap();

    @Around("@annotation(com.gton.annotation.GuavaAopLimit)")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        //拿limit的注解
        GuavaAopLimit limit = method.getAnnotation(GuavaAopLimit.class);
        if (limit != null) {
            //key作用：不同的接口，不同的流量控制
            String key = limit.key();
            RateLimiter rateLimiter = null;
            //验证缓存是否有命中key
            if (!limitMap.containsKey(key)) {
                // 创建令牌桶
                rateLimiter = RateLimiter.create(limit.permitsPerSecond());
                limitMap.put(key, rateLimiter);
                log.info("新建了令牌桶={}，容量={}", key, limit.permitsPerSecond());
            }
            rateLimiter = limitMap.get(key);
            // 拿令牌
            boolean acquire = rateLimiter.tryAcquire(limit.timeout(), limit.timeunit());
            // 拿不到命令，直接返回异常提示
            if (!acquire) {
                log.debug("令牌桶={}，获取令牌失败", key);
                throw new RedisLimitException(limit.msg());
            }
        }
        return joinPoint.proceed();
    }
}

定义统一响应结构

/**
 * @description: 通用返回对象
 * 贫血型模型
 * @author: GuoTong
 * @createTime: 2022-09-24 13:16
 * @since JDK 1.8 OR 11
 **/
@Data
@Builder
@NoArgsConstructor
@AllArgsConstructor
@Accessors(chain = true)
public class Resp<T> implements Serializable {
    //状态码
    protected String code;

    //提示信息
    protected String msg;

    //返回的数据
    protected T data;


    /**
     * Description:常用返回数据抽离
     */
    public static <T> Resp<T> LoginOk() {
        return new Resp<T>().
                setCode(ContextCommonMsg.LOGIN_SUCCESS_CODE).
                setMsg(ContextCommonMsg.LOGIN_SUCCESS_MSG);
    }

    public static <T> Resp<T> LoginFail() {
        return new Resp<T>().
                setCode(ContextCommonMsg.LOGIN_FAIL_CODE).
                setMsg(ContextCommonMsg.LOGIN_FAIL_MSG);
    }

    public static <T> Resp<T> error(String errorMsg) {
        return new Resp<T>().
                setCode(ContextCommonMsg.FAIL_CODE).
                setMsg(errorMsg);
    }

    public static <T> Resp<T> error() {
        return new Resp<T>().
                setCode(ContextCommonMsg.FAIL_CODE).
                setMsg(ContextCommonMsg.FAIL_MSG);
    }

    public static <T> Resp<T> error(String errorMsg, String failCode) {
        return new Resp<T>().
                setCode(failCode).
                setMsg(errorMsg);
    }

    public static <T> Resp<T> error(String errorMsg, String failCode, T data) {
        return new Resp<T>().
                setCode(failCode).
                setData(data).
                setMsg(errorMsg);
    }

    public static <T> Resp<T> Ok(T data) {
        return new Resp<T>().
                setCode(ContextCommonMsg.SUCCESS_CODE).
                setMsg(ContextCommonMsg.SUCCESS_MSG).
                setData(data);
    }

    public static <T> Resp<T> Ok() {
        return new Resp<T>().
                setCode(ContextCommonMsg.SUCCESS_CODE).
                setMsg(ContextCommonMsg.SUCCESS_MSG);
    }

    public static <T> Resp<T> Ok(T data, String msg) {
        return new Resp<T>().
                setCode(ContextCommonMsg.SUCCESS_CODE).
                setMsg(msg).
                setData(data);
    }

    public static <T> Resp<T> Ok(T data, String msg, String successCode) {
        return new Resp<T>().
                setCode(successCode).
                setMsg(msg).
                setData(data);
    }

}

自定义异常：用于限流控制失败

/**
 * @description: 限流异常
 * @author: GuoTong
 * @createTime: 2023-05-27 13:46
 * @since JDK 1.8 OR 11
 **/
public class RedisLimitException extends RuntimeException {

    public RedisLimitException(String msg) {
        super(msg);
    }

}

springboot异常捕获

/**
 * @description: 全局异常处理：
 * @author: GuoTong
 * @createTime: 2022-11-27 11:17
 * @since JDK 1.8 OR 11
 **/
@RestControllerAdvice
@Slf4j
public class GlobalExceptionHandler {

    // 全局异常拦截
    @ExceptionHandler
    public Resp handlerException(Exception e) {
        e.printStackTrace();
        log.error("服务执行异常---->{}", e.getMessage());
        return Resp.error(e.getMessage());
    }
	
    @ExceptionHandler(value = RedisLimitException.class)
    public Resp handlerException(RedisLimitException e) {
        log.error("限流触发---->{}", e.getMessage());
        return Resp.error(e.getMessage());
    }

接口使用测试

/**
 * @description:
 * @author: GuoTong
 * @createTime: 2022-11-25 00:02
 * @since JDK 1.8 OR 11
 **/
@RestController
public class AA {

    @GetMapping("/guava")
    @GuavaAopLimit(key = "limit_guava", permitsPerSecond = 1, timeout = 500, timeunit = TimeUnit.MILLISECONDS, msg = "基于Guava+AOP实现的单机限流！")
    public Resp okGuava() {
        return Resp.Ok();
    }

}

测试

参考：https://zhuanlan.zhihu.com/p/621267375?utm_id=0