漏桶算法和令牌桶算法

漏桶算法

• 这种算法很简单，我们把用户请求当作水流，水流到一个漏斗里面，无论有多少水流过来，漏斗的水流下去的速度是恒定的，如果流过来的水过多，多余的部分会溢出，这就是漏桶算法

package org.example.algorithm;

import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

public class LeakyBucket {
    private final long start = System.currentTimeMillis();

    private long capacity = 10000;// 桶的容量
    private long speed = 1;// 水流速
    private AtomicLong water = new AtomicLong(100);// 当前水量

    boolean add() {
//        添加令牌
        final long now = System.currentTimeMillis();
        water.set(Math.max(0, water.get() - speed * (now - start) * speed));
        long expected = water.get() + speed;
        return expected <= capacity && water.compareAndSet(expected, water.get() + speed);
    }

    public LeakyBucket setCapacity(int capacity) {
        this.capacity = capacity;
        return this;
    }

    public LeakyBucket setSpeed(int speed) {
        this.speed = speed;
        return this;
    }
}

令牌桶算法

• 这种算法就是有一个令牌桶，它会在桶的容量范围内源源不断的添加令牌，不同大小额数据包消耗的令牌数量不一样
• 通常情况下，令牌桶算法用来控制发送到网络上的数据的数目，并允许突发数据的传送

package org.example.algorithm;

import java.util.concurrent.atomic.AtomicLong;

public class TokenBucket {
    long capacity = 10000;// 桶容量
    long rate = 1;// 令牌添加速率
    AtomicLong tokens = new AtomicLong(100);// 桶内当前令牌数量
    long start = System.currentTimeMillis();

    boolean add(){
        // 添加令牌
        long now = System.currentTimeMillis();
        tokens.set(Math.min(capacity, tokens.get() + rate * (now - start)));
        if(tokens.get() >= 1){
            tokens.decrementAndGet();
            return true;
        }
        return false;
    }

    public TokenBucket setCapacity(int capacity) {
        this.capacity = capacity;
        return this;
    }

    public TokenBucket setRate(int rate) {
        this.rate = rate;
        return this;
    }
}