令牌桶与漏桶

1 令牌桶原理

令牌桶（Token Bucket）是一种流量控制算法。

• 系统以固定速率往桶里放入令牌，每个令牌代表发送一个请求或处理一个数据包的许可。

• 桶有容量上限，如果满了，新生成的令牌会被丢弃。

• 当请求到来时，需要消耗一个令牌才能被处理；如果没有令牌，请求要么等待，要么被丢弃。

• 特点：允许短时突发流量，但长期平均速率受到令牌产生速率的限制。

2 漏桶原理

漏桶（Leaky Bucket）是一种流量整形算法。

• 把请求想象成往桶里倒水，桶的底部有个小孔，水以固定速率流出。

• 输入流可以是突发的，但输出流是严格平滑的。

• 如果桶满了还继续倒水，多余的水就会溢出（对应请求被丢弃）。

• 特点：保证稳定的输出速率，不允许突发流量超过孔的速率。

3 令牌桶和漏桶的用途

• 令牌桶：常用于限流，特别是在需要允许突发的场景，例如：

  • 接口 QPS 限制（用户短时间内可发起多个请求）。

  • API 网关流量控制。

• 漏桶：常用于流量整形，例如：

  • 网络 QoS，保证数据包以稳定速率输出。

  • 消息队列消费者限速，避免下游过载。

4 令牌桶和漏桶的对比

令牌桶是系统以固定速率产生令牌，请求到来时需要消耗令牌才能被处理，因此在令牌积累的情况下可以一次性处理突发请求；漏桶是请求先进入桶中，然后系统按照固定速率取出并处理，请求的消费速率始终恒定，多余的请求会被丢弃，所以令牌桶允许突发流量而漏桶只能输出平滑流量。下面列举了他们的一些区别

特性	令牌桶 (Token Bucket)	漏桶 (Leaky Bucket)
控制方式	以速率生成令牌，允许突发流量	以固定速率漏水，严格平滑
是否允许突发	允许短时突发	不允许突发
输出速率	平均受控，可能有瞬时波动	固定速率，严格平滑
常见场景	限流（接口、API 网关）	流量整形（网络传输、QoS）

5 怎么选择令牌桶和漏桶

• 如果场景需要支持突发请求（如秒杀、短时峰值），选择 令牌桶。

• 如果场景要求输出稳定（如视频流、网络 QoS），选择 漏桶。

• 实际工程中：

  • 令牌桶更常用，因为多数业务允许突发。

  • 漏桶更多用于网络协议栈、硬件 QoS。

6 实现

一般来说，令牌桶和漏桶都会有专门的任务来负责发放令牌或者在桶中按照固定速率获取请求，但是这里为了方便采用了懒发放和懒漏的方式

6.1 令牌桶

package com.oraen.box.common.structure;

import java.util.concurrent.atomic.AtomicLong;

public class TokenBucket {
    //容量
    private final long capacity;
    //每次填充的令牌数量
    private final long refillTokens;
    //填充令牌的时间间隔
    private final long refillInterval;
    //当前令牌数量
    private final AtomicLong tokens;
    //上次填充令牌的时间戳
    private final AtomicLong lastRefillTimestamp;

    public TokenBucket(long capacity, long refillTokens, long refillInterval) {
        this.capacity = capacity;
        this.refillTokens = refillTokens;
        this.refillInterval = refillInterval;
        this.tokens = new AtomicLong(capacity);
        this.lastRefillTimestamp = new AtomicLong(System.currentTimeMillis());
    }

    private void refill() {
        long now = System.currentTimeMillis();
        long lastRefill = lastRefillTimestamp.get();
        long elapsed = now - lastRefill;
        if (elapsed > refillInterval) {
            long refillCount = elapsed / refillInterval;
            long adjustNow = lastRefill + refillCount * refillInterval;

            if (lastRefillTimestamp.compareAndSet(lastRefill, adjustNow)) {
                long add = refillCount * refillTokens;
                long currentTokens = tokens.get();
                long newTokens = Math.min(capacity, currentTokens + add);
                while (!tokens.compareAndSet(currentTokens, newTokens)) {
                    currentTokens = tokens.get();
                    newTokens = Math.min(capacity, currentTokens + add);
                }
            }
        }
    }

    public long tryTake(long num) {
        if(num <= 0) {
            return 0;
        }

        refill();
        long current = tokens.get();
        long takeNum = Math.min(num, current);
        while(! tokens.compareAndSet(current, current - takeNum)) {
            current = tokens.get();
            takeNum = Math.min(num, current);
        }

        return takeNum;
    }

    public boolean takeAll(long num) {
        if(num <= 0) {
            return true;
        }

        refill();
        long current = tokens.get();
        if(current < num) {
            return false;
        }

        while(! tokens.compareAndSet(current, current - num)) {
            current = tokens.get();
            if(current < num) {
                return false;
            }
        }

        return true;
    }
}

6.2 漏桶

package com.oraen.box.common.structure;

import java.util.concurrent.atomic.AtomicLong;

public class LeakyBucket {

    private final long capacity;          // 桶容量
    private final long leakInterval;          // 漏水间隔
    private final long leakCount;          // 每次漏水量
    private final AtomicLong count;       // 当前桶里的水量
    private volatile AtomicLong lastLeakTime;   // 上次漏水时间

    public LeakyBucket(long capacity, long leakInterval, long leakCount) {
        this.capacity = capacity;
        this.leakInterval = leakInterval;
        this.leakCount = leakCount;
        this.count = new AtomicLong(0);
        this.lastLeakTime = new AtomicLong(System.currentTimeMillis());
    }

    public boolean addAll(long num) {
        if (num <= 0) {
            return true;
        }

        long current = count.get();
        long targetNum = current + num;
        if (targetNum > capacity) {
            // 桶满，无法添加
            return false;
        }

        while (! count.compareAndSet(current, targetNum)) {
            current = count.get();
            targetNum = current + num;
            if (targetNum > capacity) {
                // 桶满，无法添加
                return false;
            }
        }

        // 成功添加
        return true;
    }

    public long tryAdd(long num) {
        if (num <= 0) {
            return 0;
        }

        long current = count.get();
        long targetNum = Math.min(current + num, capacity);

        while (! count.compareAndSet(current, targetNum)) {
            current = count.get();
            targetNum = Math.min(current + num, capacity);
        }

        // 成功添加
        return targetNum - current;
    }
}