1. 首页
  2. 系列专栏
  3. 限流算法教程 - 讲解

限流算法之令牌桶算法

1.简介

令牌桶算法的原理是系统以一个恒定的速度往桶里放入令牌，请求需要被处理时，需要先从桶里获取一个令牌。当桶里没有令牌可取时，则拒绝服务。令牌桶算法可以用以下几个概念来描述：

  • 令牌将按照固定的速率被放入令牌桶中。
  • 桶中最多存放一定数量的令牌，当桶满时，新添加的令牌被丢弃。
  • 当一个数据包到达，将从桶中删除对应数量的令牌，然后数据包被发送到网络上。
  • 如果桶中的令牌不足，则不会删除令牌，且该数据包将被限流。

image-20241025145945261

2.代码实现

package top.houry.limit;

import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

public class TokenBucketRateLimiter {
    private final long capacity; // 桶的容量
    private final long tokens;   // 初始令牌数
    private final long refillRate; // 令牌填充速率（每秒添加的令牌数）
    private final long refillInterval; // 令牌填充间隔（毫秒）
    private long availableTokens; // 当前可用的令牌数
    private long lastRefillTimestamp; // 上次填充令牌的时间戳

    private final Lock lock = new ReentrantLock();

    public TokenBucketRateLimiter(long capacity, long tokens, long refillRate) {
        this.capacity = capacity;
        this.tokens = tokens;
        this.refillRate = refillRate;
        this.refillInterval = 1000 / refillRate; // 计算填充间隔
        this.availableTokens = tokens;
        this.lastRefillTimestamp = System.currentTimeMillis();
    }

    public boolean tryAcquire() {
        lock.lock();
        try {
            refillTokens();
            if (availableTokens >= 1) {
                availableTokens--;
                return true;
            } else {
                return false;
            }
        } finally {
            lock.unlock();
        }
    }

    private void refillTokens() {
        long now = System.currentTimeMillis();
        if (now - lastRefillTimestamp > refillInterval) {
            long tokensToAdd = (now - lastRefillTimestamp) / refillInterval * refillRate;
            availableTokens = Math.min(capacity, availableTokens + tokensToAdd);
            lastRefillTimestamp = now;
        }
    }

    public static void main(String[] args) throws InterruptedException {
        TokenBucketRateLimiter rateLimiter = new TokenBucketRateLimiter(100, 100, 100);

        for (int i = 0; i < 150; i++) {
            boolean result = rateLimiter.tryAcquire();
            System.out.println("请求 " + i + ": " + (result ? "通过" : "拒绝"));
//            if (i % 10 == 0) { // 模拟请求间隔
//                Thread.sleep(50); // 每10个请求暂停50毫秒
//            }
        }
    }
}

3.运行结果

image-20241025150023337