全网最简洁Java实现多线程安全的令牌桶限流算法

问题描述

在许多应用中，我们需要限制某些操作的频率，例如，限制API调用的速率，防止系统被过度使用，这种需求就需要一个限流算法来满足

令牌桶算法是一种常用的限流算法

它的基本思想是：系统以恒定的速率向桶中添加令牌，当桶满时，新添加的令牌会被丢弃。当一个请求到达时，系统会尝试从桶中取出一个令牌，如果桶中没有令牌，则请求被拒绝，这种算法常用于网络流量整形和流量控制。

limiter

实现原理

我们需要实现一个令牌桶限流算法，主要包括以下几个步骤：

初始化桶中的令牌数和桶的容量
每隔一段时间向桶中添加一个令牌
当有请求到来时，尝试从桶中取出一个令牌，如果桶中没有令牌，则拒绝请求
根据请求的时间，更新桶中的令牌数
处理多线程并发请求，保证线程安全

Java实现

我们要求这个令牌桶限流算法是多线程安全的，因此我们需要使用线程安全的数据结构来存储桶中的令牌数，这里我们使用AtomicInteger来表示桶中的令牌数。

获取请求的时间可以使用System.currentTimeMillis()，这个时间戳可以精确到毫秒级别。

所以我们的代码如下：


    private final AtomicInteger permits;
    private final AtomicLong lastRequestTime;
    private final int maxBurst;
    private final int permitsPerSecond;

    public boolean allowN(int n) {
        long currentTime = System.currentTimeMillis();
        long timePassed = currentTime - lastRequestTime.getAndSet(currentTime);
        int permitsToAdd = (int) (timePassed / 1000.0 * permitsPerSecond);
        permits.addAndGet(Math.min(permitsToAdd, maxBurst - permits.get()));

        return permits.getAndAdd(-n) >= n;
    }

这个代码就是基本的流程，但是缺少多线程的处理，当多个线程同时请求时，可能会出现问题，主要是因为：

lastRequestTime和permits的更新虽然是原子操作，但是lastRequestTime, permits 这两个原子的操作过程并不是原子的

如果两个线程同时访问allowN这个函数，那么就会出现并发的冲突，剩余的令牌计算并不准确，所以我们用最简单的synchronized来解决这个问题

    public synchronized boolean allowN(int n) {
        ///...
        return permits.getAndAdd(-n) >= n;
    }

测试

    RateLimiter rateLimiter = new RateLimiter(10, 10);
    long now = System.currentTimeMillis();
    for (int i = 0; i < 20; i++) {
        if (rateLimiter.allow()) {
            System.out.println("Request " + i + " allowed");
            if (System.currentTimeMillis() - now > 1000) {
                System.out.println("Request " + i + " delayed by " + (System.currentTimeMillis() - now) + "ms");
            }
        } else {
            System.out.println("Request " + i + " not allowed");
            TimeUnit.MILLISECONDS.sleep(1000);
            now = System.currentTimeMillis();
        }
    }

在这个测试中，我们创建了一个速率限制器，每秒可以处理10个请求，最大突发处理量也是10个请求。尝试进行20次请求，每次请求前，都会调用 allow() 方法检查是否允许进行请求，如果允许则输出请求允许，否则输出请求不允许，并且等待1秒后再次尝试请求。

这里有完整的代码: ratelimiter

总结

后端经常有些最基础的算法，比如限流、缓存、排序等，这些算法虽然简单，但是在实际应用中却非常重要，因为它们直接影响到系统的性能和稳定性。

后端的工程师对这些基础算法的原理和应用一定要非常的熟悉，因为不光你懂算法，你还需要考虑到并发这种场景，如果你没处理好并发访问，就会留下漏洞。

最后，希望这篇文章对你有所帮助，如果有任何问题，欢迎留言讨论。可以加入我们的实战项目群，一起学习，一起进步。

所有的后端面试常见的问题，我们每天都会在我们的编程群里面讨论和Code review, 欢迎大家加入我们的编程群，一起学习和进步。

编程交流群