本文共 3748 字，大约阅读时间需要 12 分钟。

令牌桶简单实现（Java）

简介

实现思路

参考了guava的RateLimiter之后，发现并没有“一个线程定时往桶里放token”，而是在请求进来时通过当前时间戳实时计算（下一次加令牌的时间为x，此时时间为m，平均n秒放一个令牌，则此时应当增加 (m - x) / n + 1 个令牌）

code

/** * 令牌桶 * 局限：最多1毫秒生成一个令牌 * * create time: 2020/7/15 9:42 */public class TokenBucket {
       private final double unitAddNum;    // 单位时间（1s）往桶中放令牌数量    private final long maxTokenNum;      // 桶中最大有多少令牌    private volatile long currentTokenCount = 0;  // 当前桶中有多少令牌    private volatile long nextRefreshTime = 0L;  // 下一次刷新桶中令牌数量的时间戳    private volatile long lastAcquireTime;       // 上一次从桶中获取令牌的时间戳（貌似用不到）    /**     *     * @param unitAddNum 1秒增加几个令牌     * @param maxToken 桶中最大令牌数     * @param isFullStart 一开始是否是满的     */    public TokenBucket(double unitAddNum, long maxToken, boolean isFullStart) {
           if (unitAddNum <= 0 || maxToken <= 0) {
               throw new RuntimeException("unitAddNum and maxToken can't less than 0");        }        if (unitAddNum > 1000) {
               throw new RuntimeException("unitAddNum max is 1000");        }        this.unitAddNum = unitAddNum;        this.maxTokenNum = maxToken;        if (isFullStart) {
               this.currentTokenCount = maxToken;        } else {
               this.currentTokenCount = 0;        }        this.nextRefreshTime = calculateNextRefreshTime(System.currentTimeMillis());    }    public boolean acquire(long needTokenNum) {
           if (needTokenNum > this.maxTokenNum) {
               return false;        }        synchronized (this) {
               long currentTimestamp = System.currentTimeMillis();            this.refreshCurrentTokenCount(currentTimestamp);            if (needTokenNum <= this.currentTokenCount) {
                   return this.doAquire(needTokenNum, currentTimestamp);            }            return false;        }    }    private boolean doAquire(long needTokenNum, long currentTimestamp) {
           this.currentTokenCount -= needTokenNum;        this.lastAcquireTime = currentTimestamp;        return true;    }    /**     * 刷新桶中令牌数量     * @param currentTimestamp     */    private void refreshCurrentTokenCount(long currentTimestamp) {
           if (this.nextRefreshTime > currentTimestamp) {
               return;        }        this.currentTokenCount = Math.min(this.maxTokenNum, this.currentTokenCount + calculateNeedAddTokenNum(currentTimestamp));        this.nextRefreshTime = calculateNextRefreshTime(currentTimestamp);    }    /**     * 计算当前需要添加多少令牌     * @param currentTimestamp     * @return     */    private long calculateNeedAddTokenNum(long currentTimestamp) {
           if (this.nextRefreshTime > currentTimestamp) {
               return 0;        }        long addOneMs = Math.round(1.0D / this.unitAddNum * 1000D); // 这么久才能加1个令牌        return (currentTimestamp - this.nextRefreshTime) / addOneMs + 1;    }    private long calculateNextRefreshTime(long currentTimestamp) {
           if (currentTimestamp < this.nextRefreshTime) {
               return this.nextRefreshTime;        }        long addOneMs = Math.round(1.0D / this.unitAddNum * 1000D); // 这么久才能加1个令牌        long result = 0;        if (this.nextRefreshTime <= 0) {
               result = currentTimestamp + addOneMs;        } else {
               result = this.nextRefreshTime + ((currentTimestamp - this.nextRefreshTime) / addOneMs + 1) * addOneMs;        }        return result;    }    public static void main(String[] args) throws InterruptedException {
           TokenBucket tokenBucket = new TokenBucket(1D, 1, true);        for (int i=0; i<10; i++) {
               System.out.println(tokenBucket.acquire(1));            Thread.sleep(500);        }    }}

main输出结果

truefalsetruefalsetruefalsetruefalsetruefalse

转载地址：https://blog.csdn.net/xxxxssss12/article/details/107359477 如侵犯您的版权，请留言回复原文章的地址，我们会给您删除此文章，给您带来不便请您谅解！