Java令牌桶简单实现

令牌桶算法是一种反向的漏桶算法。在令牌桶算法中，桶中存放的不再是请求，而是
令牌。处理程序只有拿到令牌后，才能对请求进行处理。如果没有令牌，那么处理程序要
么丢弃请求，要么等待可用的令牌。为了限制流速，该算法在每个单位时间产生一定量的
令牌存入桶中。比如，要限定应用每秒只能处理1个请求，那么令牌桶就会每秒产生1个
令牌。通常，桶的容量是有限的，比如，当令牌没有被消耗掉时，只能累计有限单位时间
内的令牌数量，其基本原理如图。

令牌桶简单实现代码

import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

/**
 * 令牌桶简单实现
 *
 * @author zhangqi
 * @date 2022/5/9 15:47
 */
public class TokenLimiter {
    /**
     * 令牌
     */
    public static final String TOKEN = "lp";

    /**
     * 阻塞队列 用于存放令牌
     */
    private ArrayBlockingQueue<String> blockingQueue;
    /**
     * 令牌桶容量
     */
    private int limit;
    /**
     * 令牌的产生间隔时间 单位: 毫秒
     */
    private int period;
    /**
     * 令牌每次产生的个数
     */
    private int amount;


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

    /**
     * 生产令牌
     *
     * @param lock 锁对象 因为有可能多个地方调用 所以锁对象需要传过来
     */
    public void start(Object lock) {
        Executors.newScheduledThreadPool(1).scheduleAtFixedRate(() -> {
            synchronized (lock) {
                addToken();
                lock.notifyAll();
            }
        }, 500, this.period, TimeUnit.MILLISECONDS);
    }

    /**
     * 获取令牌
     *
     * @return 令牌
     */
    public boolean tryAcquire() {
        // 队首元素出队
        return blockingQueue.poll() != null;
    }

    /**
     * 创建时初始化令牌
     */
    private void init() {
        for (int i = 0; i < limit; i++) {
            blockingQueue.add(TOKEN);
        }
    }

    /**
     * 添加令牌
     */
    private void addToken() {
        for (int i = 0; i < this.amount; i++) {
            // 溢出返回false
            blockingQueue.offer(TOKEN);
        }
    }

}

测试代码:

/**
 * @author zhangqi
 * @date 2022/5/9 15:57
 */
public class TestTokenLimiter {
    final static Object LOCK = new Object();

    /**
     * 先生产2个令牌，减少4个令牌；再每500ms生产2个令牌，减少4个令牌
     */
    public static void main(String[] args) throws InterruptedException {
        int period = 500;
        TokenLimiter limiter = new TokenLimiter(2, period, 2);

        limiter.start(LOCK);

        // 让线程先产生2个令牌
        synchronized (LOCK) {
            LOCK.wait();
        }

        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 + ":没有令牌");
                    }
                    try {
                        Thread.sleep(period);
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                }
            }).start();
        }

    }

}

guava 的 RateLimiter

package com.example.demo.hmjuc.day13;

import com.google.common.util.concurrent.RateLimiter;

/**
 * @author zhangqi
 * @date 2022/5/13 16:16
 */
public class RateLimiterDemo {
    static RateLimiter limiter = RateLimiter.create(2);


    public static void main(String[] args) {
        for (int i = 0; i < 10; i++) {
            limiter.acquire();
            new Thread(new Task()).start();

            // if (!limiter.tryAcquire()) {
            //     continue;
            // }
            // new Thread(new Task()).start();
        }
    }
}


class Task implements Runnable {
    @Override
    public void run() {
        System.out.println(Thread.currentThread().getName() + " is running" + System.currentTimeMillis());
    }
}