Java 简单实现令牌桶

一、原理

     令牌桶可用作流量控制，令牌桶控制流量的原理：单位时间内只发放固定数量的令牌到令牌桶中，规定服务在处理请求之前，必须先从令牌桶中拿出一个令牌，如果令牌桶中没有令牌，则拒绝请求。这样就保证单位时间内能处理的请求不超过发放令牌的数量，起到流量控制的作用。

二、动手前分析

     我理解的步骤拆分：
（1）得有算法支持单位时间内生成固定数量令牌
（2）获取令牌的抽象化，若获取不到令牌，丢弃请求或是等待可用令牌（等下一波单位时间到来，再试图获取令牌了）
（3）数据结构选型
① 并发工具包里的阻塞队列 ArrayBlockingQueue，先进先出（这个倒是不重要，我感觉栈也行，先进先出还是先进后出对令牌桶没啥影响），线程安全，默认不保证线程公平（这个蛮重要❗ 即 等待的线程若因获取锁而竞争，无论是先等待还是后等待，都有可能获取到锁，而非先等待的先获取，后等待的后获取）
②

---

     我理解的类与属性与方法：
（1）类：令牌桶类、客户类；
（2）属性：桶容量（没啥实在用处）、令牌们、时间单位；状态？——是否获取令牌成功；
（3）方法：生成令牌（称之为 初始化init~ 可拆分成先 new 出来一堆令牌，完了再往桶里加）、获取令牌、发令牌（给予令牌or拒绝请求or下一波重试）；

---

三、代码实现

令牌桶类：

public class TokenLimiter {

    /**
     * 令牌
     */
    private static final String TOKEN = "token";

    /**
     * 阻塞队列，线程安全，非公平容量有限
     */
    private ArrayBlockingQueue<String> arrayBlockingQueue;

    /**
     * 令牌桶容量
     */
    private int limit;

    /**
     * 令牌每次生成的数量
     */
    private int amount;

    /**
     * 令牌桶间隔时间——单位时间
     */
    private int period;

    public TokenLimiter(int limit, int period, int amount) {
        this.limit = limit;
        this.amount = amount;
        this.period = period;
        arrayBlockingQueue = new ArrayBlockingQueue<>(limit);
        init();
    }

    /**
     * 初始化生成令牌，按桶容量生成
     */
    private void init() {
        for (int i = 0; i < limit; i++) {
            arrayBlockingQueue.offer(TOKEN);
        }
    }
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

生成令牌的方法：

	/**
     * 生成令牌，周期性线程池，启动后延迟500毫秒，每500毫秒执行一次
     * （可以理解为启动后每隔500毫秒执行一次）
     * @param lock
     */
  public void start(Object lock) {
        Executors.newScheduledThreadPool(1).scheduleAtFixedRate(() -> {
            synchronized (lock) {
                // 往令牌桶添加令牌
                addToken();
                lock.notifyAll();
            }
        }, 500, this.period, TimeUnit.MILLISECONDS);
    }
1
2
3
4
5
6
7
8
9
10
11
12
13
14

往令牌桶添加令牌方法：

    /**
     * 添加令牌
     */
    private void addToken() {
        for (int i = 0; i < this.amount; i++) {
            // 溢出返回false
            arrayBlockingQueue.offer(TOKEN);
        }
    }
1
2
3
4
5
6
7
8
9

获取令牌方法：

    public boolean tryAcquire() {
        // 队首元素出队
        return arrayBlockingQueue.poll() != null;
    }
1
2
3
4

main方法：

public class TestTokenLimiter {
    final static Object LOCK = new Object();
   
    public static void main(String[] args) throws InterruptedException {
        int period = 500;
        // 创建令牌桶
        TokenLimiter limiter = new TokenLimiter(2, period, 2);
        // 生产令牌
        limiter.start(LOCK);

        // 这里主要是保证让线程池初始先产生2个令牌，
        // 其实init方法里已经加了2个令牌了，而队列长度限定为2，
        // 所以初始时addToken()并不会成功。
        // 如果这里不使用wait和notify配合，可能客户端获取令牌会先执行，打破节奏了
        synchronized (LOCK) {
         // 该线程被阻塞挂起，直到其他线程调用了该共享对象的notify()或者notifyAll()方法，才返回
         LOCK.wait();          
        }
   
        // 启动4个线程，模拟4个客户端各自每隔500ms获取令牌
        for (int i = 0; i < 4; i++) {
            new Thread(() -> {
                while (true) {
                    String name = Thread.currentThread().getName();
                    if (limiter.tryAcquire()) {
                        System.out.println(name + ":拿到令牌");
                    } else {
                        System.out.println(name + ":没有令牌");
                    }

                  // 如果不sleep  线程没法切到生产的去执行
                  try {
                        Thread.sleep(period);
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                }
            }).start();
        }
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

     整个的流程就是令牌桶里先初始化了2个令牌，然后开始500ms一个周期的往复——客户端取令牌，4个客户端中2个成功2个失败；令牌桶生成2个令牌…

参考：@https://blog.csdn.net/weixin_43939924/article/details/124668935，
特别鸣谢！~