令牌桶 java_服务限流-令牌桶java实现

importjava.io.BufferedWriter;importjava.io.FileOutputStream;importjava.io.IOException;importjava.io.OutputStreamWriter;importjava.util.Random;importjava.util.concurrent.ArrayBlockingQueue;importjava.util.concurrent.Executors;importjava.util.concurrent.ScheduledExecutorService;importjava.util.concurrent.TimeUnit;importjava.util.concurrent.locks.ReentrantLock;/***

* Created by inter12 on 15-3-18.

*

*/

public classTokenBucket {//默认桶大小个数 即最大瞬间流量是64M

private static final int DEFAULT_BUCKET_SIZE = 1024 * 1024 * 64;//一个桶的单位是1字节

private int everyTokenSize = 1;//瞬间最大流量

private intmaxFlowRate;//平均流量

private intavgFlowRate;//队列来缓存桶数量：最大的流量峰值就是 = everyTokenSize*DEFAULT_BUCKET_SIZE 64M = 1 * 1024 * 1024//* 64

private ArrayBlockingQueue tokenQueue = new ArrayBlockingQueue(DEFAULT_BUCKET_SIZE);private ScheduledExecutorService scheduledExecutorService =Executors.newSingleThreadScheduledExecutor();private volatile boolean isStart = false;private ReentrantLock lock = new ReentrantLock(true);private static final byte A_CHAR = 'a';publicTokenBucket() {

}public TokenBucket(int maxFlowRate, intavgFlowRate) {this.maxFlowRate =maxFlowRate;this.avgFlowRate =avgFlowRate;

}public TokenBucket(int everyTokenSize, int maxFlowRate, intavgFlowRate) {this.everyTokenSize =everyTokenSize;this.maxFlowRate =maxFlowRate;this.avgFlowRate =avgFlowRate;

}public voidaddTokens(Integer tokenNum) {//若是桶已经满了，就不再加入新的令牌

for (int i = 0; i < tokenNum; i++) {

tokenQueue.offer(Byte.valueOf(A_CHAR));

}

}publicTokenBucket build() {

start();return this;

}/*** 获取足够的令牌个数

*

*@return

*/

public boolean getTokens(byte[] dataSize) {//Preconditions.checkNotNull(dataSize);//Preconditions.checkArgument(isStart,//"please invoke start method first !");

int needTokenNum = dataSize.length / everyTokenSize + 1;//传输内容大小对应的桶个数

final ReentrantLock lock = this.lock;

lock.lock();try{boolean result = needTokenNum <= tokenQueue.size(); //是否存在足够的桶数量

if (!result) {return false;

}int tokenCount = 0;for (int i = 0; i < needTokenNum; i++) {

Byte poll=tokenQueue.poll();if (poll != null) {

tokenCount++;

}

}return tokenCount ==needTokenNum;

}finally{

lock.unlock();

}

}public voidstart() {//初始化桶队列大小

if (maxFlowRate != 0) {

tokenQueue= new ArrayBlockingQueue(maxFlowRate);

}//初始化令牌生产者

TokenProducer tokenProducer = new TokenProducer(avgFlowRate, this);

scheduledExecutorService.scheduleAtFixedRate(tokenProducer,0, 1, TimeUnit.SECONDS);

isStart= true;

}public voidstop() {

isStart= false;

scheduledExecutorService.shutdown();

}public booleanisStarted() {returnisStart;

}class TokenProducer implementsRunnable {private intavgFlowRate;privateTokenBucket tokenBucket;public TokenProducer(intavgFlowRate, TokenBucket tokenBucket) {this.avgFlowRate =avgFlowRate;this.tokenBucket =tokenBucket;

}

@Overridepublic voidrun() {

tokenBucket.addTokens(avgFlowRate);

}

}public staticTokenBucket newBuilder() {return newTokenBucket();

}public TokenBucket everyTokenSize(inteveryTokenSize) {this.everyTokenSize =everyTokenSize;return this;

}public TokenBucket maxFlowRate(intmaxFlowRate) {this.maxFlowRate =maxFlowRate;return this;

}public TokenBucket avgFlowRate(intavgFlowRate) {this.avgFlowRate =avgFlowRate;return this;

}private String stringCopy(String data, intcopyNum) {

StringBuilder sbuilder= new StringBuilder(data.length() *copyNum);for (int i = 0; i < copyNum; i++) {

sbuilder.append(data);

}returnsbuilder.toString();

}public static void main(String[] args) throwsIOException, InterruptedException {

TokenBucket tokenBucket= TokenBucket.newBuilder().avgFlowRate(512).maxFlowRate(1024).build();

BufferedWriter bufferedWriter= new BufferedWriter(new OutputStreamWriter(new FileOutputStream("D:/ds_test")));

String data= "xxxx";//四个字节

for (int i = 1; i <= 1000; i++) {

Random random= newRandom();int i1 = random.nextInt(100);boolean tokens =tokenBucket.getTokens(tokenBucket.stringCopy(data, i1).getBytes());

TimeUnit.MILLISECONDS.sleep(100);if(tokens) {

bufferedWriter.write("token pass --- index:" +i1);

System.out.println("token pass --- index:" +i1);

}else{

bufferedWriter.write("token rejuect --- index" +i1);

System.out.println("token rejuect --- index" +i1);

}

bufferedWriter.newLine();

bufferedWriter.flush();

}

bufferedWriter.close();

}

}