令牌桶的实现

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);
            }
        }
    }
}