基于Java编写一个限流工具类RateLimiter

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

public void limitTest(){  
DateTimeFormatter dtf = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSS");  
  
//创建令牌桶，一秒一个，容量为1  
RateLimiter rateLimiter = RateLimiter.create(1);  
//获取放令牌的速率  
System.out.println("放令牌的速率："+rateLimiter.getRate());  
for (int i=0;i<5;i++){  
//获取令牌，会阻塞，返回等待的时间  
double acquire = rateLimiter.acquire();  
System.out.println("第"+i+" 个令牌获取到的时间："+LocalDateTime.now().format(dtf)+"，等待时间："+acquire);  
}  
  
//是否会立即获取到令牌  
boolean tryAcquire = rateLimiter.tryAcquire();  
}

public @interface Limit {  
// 资源主键  
String key() default "";  
//最多访问次数,代表请求总数量  
double permitsPerSeconds();  
// 时间：即timeout时间内,只允许有permitsPerSeconds个请求总数量访问，超过的将被限制不能访问  
long timeout();  
//时间类型，默认秒
TimeUnit timeUnit() default TimeUnit.SECONDS;  
  
//提示信息  
String msg() default "系统繁忙,请稍后重试";  
}

@Aspect  
@Component  
public class LimitAop {  
  
private final Map<String, RateLimiter> limitMap = Maps.newConcurrentMap();  
  
private DateTimeFormatter dtf = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSS");  

@Around("@annotation(com.limit.Limit)")  
public Object around(ProceedingJoinPoint joinPoint) throws Throwable{  
    MethodSignature signature = (MethodSignature) joinPoint.getSignature();  
    Method method = signature.getMethod();  
    //拿limit的注解  
    Limit limit = method.getAnnotation(Limit.class);  
    if (limit != null) {  
        // key作用：不同的接口，不同的流量控制，相当于每个接口有一个对应的令牌桶  
        String key = limit.key();   
        RateLimiter rateLimiter;  
        //验证缓存是否有命中key  
        if (!limitMap.containsKey(key)) {  
            //创建令牌桶  
            rateLimiter = RateLimiter.create(limit.permitsPerSeconds());  
            limitMap.put(key, rateLimiter);  
            log.info("新建了令牌桶={},容量={}", key, limit.permitsPerSeconds());  
        }  
        rateLimiter = limitMap.get(key);  
        //拿一个令牌，拿不到会一直阻塞 
        double acquire = rateLimiter.acquire(1);  
        log.info("{},获取令牌时间{}", key,LocalDateTime.now().format(dtf));  
        
        /*
        //是否能立即拿到令牌，不能则桶里没有，还没到一秒钟，进行限流
        boolean acquire = rateLimiter.tryAcquire();  
        if (!acquire) {  
        log.info("令牌桶={},获取令牌失败", key);  
        throw new RuntimeException(limit.msg());  
        }*/  
        
        }  
    return joinPoint.proceed();  
    }  
}

@GetMapping("/limitTest")  
@Limit(key = "limitTest",permitsPerSeconds = 1,timeout = 1,msg = "触发接口限流，请重试")  
public void limitTest(){  
    //其它逻辑代码  
    //这里打印调用该接口执行的时间，以便观察限流  
    DateTimeFormatter dtf = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSS");  
    System.out.println(LocalDateTime.now().format(dtf));  
}