国际惯例原理图
代码实现
package Thread; import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong; /**
* @ProjectName: cutter-point
* @Package: Thread
* @ClassName: RateLimiter
* @Author: xiaof
* @Description: 令牌桶，限流
* @Date: 2019/6/21 11:41
* @Version: 1.0
*/
public class RateLimiter { //限流消费的令牌
private final AtomicInteger consumedTokens = new AtomicInteger(); private final AtomicLong lastRefrushTokenTime = new AtomicLong(0); //限流类型，是秒，还是分
private final long rateType; public RateLimiter(TimeUnit averageRateUnit) {
switch (averageRateUnit) {
case SECONDS:
rateType = 1000;
break;
case MINUTES:
rateType = 60 * 1000;
break;
default:
throw new IllegalArgumentException("TimeUnit of " + averageRateUnit + " is not supported");
}
} //请求令牌，判断是否可以获取到新的令牌
public boolean acquire(int bucketSize, long avgRate) {
return acquire(bucketSize, avgRate, System.currentTimeMillis());
} public boolean acquire(int bucketSize, long avgRate, long curentTimeMillis) { if(bucketSize <= 0 || avgRate <= 0) {
return true;//如果这2个参数，任意一个为0 ，我们就认为没有上限
} //刷新令牌桶
refillToken(bucketSize, avgRate, curentTimeMillis);
//开始消费令牌
return consumToken(bucketSize);
} private void refillToken(int bucketSize, long avgRate, long currentTimeMillis) {
//获取上次最后以后更新令牌时间
long freshTime = lastRefrushTokenTime.get();
//获取当前间隔时间
long timeDelta = currentTimeMillis - freshTime; //计算这次需要填充的token数
long newToken = timeDelta * avgRate /rateType;
if(newToken > 0) {
//新的更新时间
long newFillTime = freshTime == 0 ? currentTimeMillis : freshTime + timeDelta;
//用cas操作，以保证只有一个线程注入新令牌
if(lastRefrushTokenTime.compareAndSet(freshTime, newFillTime)) {
//死循环，直到设置成功新的令牌
while(true) {
//1.获取当前消费的令牌
int currentConsumToken = consumedTokens.get();
//2.获取消费的令牌容量，跟桶极限大小比较，取小的那个
int realConsumTokens = Math.min(currentConsumToken, bucketSize);
//3.计算填充之后剩余的被消费的容量,计算新增容量，用来填充被消费的令牌数
//剩余的消费容量,但是不能比0还小，这个要取值
int newConsumSize = (int) Math.max(0, realConsumTokens - newToken);
//然后设置进去
if(consumedTokens.compareAndSet(currentConsumToken, newConsumSize)) {
return;
}
}
}
}
} //消费令牌
private boolean consumToken(int bucketSize) {
while (true) {
int currentLevel = consumedTokens.get();
//如果超出负载
if (currentLevel >= bucketSize) {
return false;
}
//每次消费一个
if (consumedTokens.compareAndSet(currentLevel, currentLevel + 1)) {
return true;
}
}
} public void reset() {
consumedTokens.set(0);
lastRefrushTokenTime.set(0);
} }
到这里可能有的人不清楚怎么用，来我们测试一波
我们假设有100个用户同时请求，然后令牌恢复速率调成10，然后速率单位改为秒，也就是1秒恢复10个令牌
这样同时100个请求过来，马上令牌就会被用完，那么就会被限流，比如我们拦截器这个时候可以返回404，或者503
public static void main(String args[]) {
RateLimiter rateLimiter = new RateLimiter(TimeUnit.SECONDS);
final int bucketSize = 10;
//回复令牌生产速率
final long avgRate = 10;
//判断是否流量达到上限
ExecutorService pool = Executors.newCachedThreadPool();
for(int i = 0; i < 100; ++i) {
pool.submit(new Runnable() {
@Override
public void run() {
while(true) {
try {
Thread.sleep(((int) (Math.random() * 10)) * 1000) ;
if(!rateLimiter.acquire(bucketSize, avgRate)) {
System.err.println(Thread.currentThread().getName() + "已经限流成功----response.setStatus(404)");
} else {
System.out.println(Thread.currentThread().getName() + "正常执行");
}
} catch (InterruptedException e) {
e.printStackTrace();
}
}
}
});
}
while(true) {
}
}
效果展示
