package com.fanfan.tuoren.common.commponent.RateLimit.TokenBucket;
import cn.hutool.core.date.DateUtil;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
/**
* ClassName:
* date: 2020/5/22
* 令牌桶
*
* 假如用户配置的平均发送速率为 rate/1000,则每隔 rate 毫秒一个令牌被加入到桶中;
* 假设桶最多可以存发 bucketSize 个令牌。如果令牌到达时令牌桶已经满了,那么这个令牌会被丢弃;
* 当一个消耗N个令牌的数据包到达时,就从令牌桶中删除n个令牌,并且数据包被发送到网络;
* 如果令牌桶中少于n个令牌,那么不会删除令牌,用户进入轮询状态,等待 waitMillsTime 后再次调用令牌桶,如果失败则认为这个数据包在流量限制之外;
* 算法允许最长 bucketSize 个令牌请求的突发,但从长期运行结果看,数据包的速率被限制成常量 rate/1000。对于在流量限制外的数据包可以以不同的方式处理:
*
* @author lifang
* @version 1.0
*/
@Data
@Slf4j
public class TokenBucketHandler {
//最大令牌数量
private int bucketSize;
//当前存储的令牌数
private AtomicLong storedPermits=new AtomicLong(0);
//上次填充令牌时间
private AtomicLong lastRefillTime=new AtomicLong(0);
//每个令牌产生时间
private long rate;
//此处以微秒为单位
private long millsWait;
private static volatile TokenBucketHandler obj;
private TokenBucketHandler(TokenBucketProperties tokenBucketProperties){
this.bucketSize=tokenBucketProperties.bucketSize;
this.rate=tokenBucketProperties.rate;
this.millsWait=TimeUnit.SECONDS.toMillis(tokenBucketProperties.getWaitTime());
this.storedPermits.set(bucketSize/2);
this.lastRefillTime.set(System.currentTimeMillis());
log.info("令牌桶参数设置成功,桶大小为--{},速率--{}ms/个,最大等待时间--{}ms,当前桶中存储令牌数量--{}",bucketSize,rate,millsWait,storedPermits.get());
}
public static TokenBucketHandler create(TokenBucketProperties tokenBucketProperties){
if(obj==null){
synchronized (TokenBucketHandler.class){
if(obj==null){
obj=new TokenBucketHandler(tokenBucketProperties);
}
}
}
return obj;
}
public boolean acquire(){
return tryAcquire(1,millsWait,System.currentTimeMillis());
}
public boolean acquire(long needPermits){
return tryAcquire(needPermits,millsWait,System.currentTimeMillis());
}
public boolean acquire(long needPermits, TimeUnit unit){
return tryAcquire(unit.toMillis(needPermits),0,System.currentTimeMillis());
}
/**
* 获取令牌
* @param needPermits 消耗令牌数量,1个bit位消耗1个令牌
* @return
*/
boolean tryAcquire(long needPermits,long millsToWait,long now){
if(!checkPermits(needPermits)) return false;
//填充令牌
refillPermits(bucketSize,rate,System.currentTimeMillis());
//消费令牌
if(consumerToken((int) needPermits)) return true;
//若允许等待
return waitToGetPermits(needPermits,now,millsToWait);
}
boolean checkPermits(long needPermits){
//最大令牌数量小于0,令牌产出速率小于0,参数初始化失败,则令牌桶调用失败
if(bucketSize<=0||rate<=0L||needPermits>bucketSize) {
return false;
}
//将初始容量改为桶大小的一半
if(storedPermits.get()==-1) storedPermits.set(bucketSize/2);
return true;
}
/**
* 向令牌桶中添加令牌
* @param bucketSize
* @param rate
* @param nowMills
*/
private void refillPermits(int bucketSize,long rate,long nowMills){
//上次填充时间
long lastTime = lastRefillTime.get();
//两次填充令牌桶的间隔时间
long interval = nowMills - lastTime;
//添加令牌数量
long newTokens = interval / rate;
if(newTokens>0){
long newFillTime=lastTime==0?nowMills:lastTime+newTokens*rate;
//CAS更新添加时间,若更新成功,则存放令牌
if(lastRefillTime.compareAndSet(lastTime,newFillTime)){
while (true){
//获取当前存储令牌数量
long l = storedPermits.get();
//设置调整后的令牌数
long adjust=Math.min(l+newTokens,bucketSize);
if(storedPermits.compareAndSet(l,adjust)) {
// log.info("更新后的令牌数量---{}",storedPermits.get());
return;
}
}
}
}
}
/**
* 消费令牌
* @param consumerTokens
* @return
*/
private boolean consumerToken(int consumerTokens) {
while (true){
long storePermits = storedPermits.get();
if(storePermits<consumerTokens) return false;
if(storedPermits.compareAndSet(storePermits,Math.max(0,storePermits-1))){
log.info("当前时间---{},消费令牌数量--{},剩余令牌数量--{}", DateUtil.now(),consumerTokens,storedPermits.get());
return true;
}
}
}
/**
* 等待时间内消费令牌
* @param now
* @param waitMillsTime
* @return
*/
private boolean waitToGetPermits(long needPermits,long now,long waitMillsTime){
if(waitMillsTime<=0) return false;
long targetTime = now + waitMillsTime;
for (;;){
//超过等待时间,按照waitTime时间计算令牌数
long timeMillis = System.currentTimeMillis();
if(timeMillis>=targetTime) {
return tryAcquire(needPermits,0,timeMillis);
}
}
}
}
令牌桶的实现
最新推荐文章于 2024-07-27 04:44:45 发布