使用令牌桶和漏桶实现请求限流逻辑

实现请求限流

令牌桶算法
- 原理
- 实现案例
- - 案例目的:
  - 实例demo
  - 运行结果:
漏桶算法
- 原理:
- 实现案例:
- - 案例目的:
  - 案例代码
  - 运行结果:

令牌桶算法和漏桶算法是两种常用的限流算法，用于控制系统对请求或数据的访问速率。下面分别详细解释这两种算法的原理.

令牌桶算法

原理

令牌桶算法的原理比较直观。在一个固定容量的桶中，以固定的速率持续放入令牌（token），每个令牌代表着允许通过的请求。当有请求到达时，需要获取一个令牌才能执行请求。如果没有令牌可用，请求将被限制。

实现案例

案例目的:

创建一个容量为10、速率为10个令牌/秒的令牌桶。然后循环进行20次请求，每次请求使用tryAcquire()方法尝试获取一个令牌。如果获取到令牌，则输出"Request x passed"；否则输出"Request x limited"。通过设置请求的时间间隔，可以观察到令牌桶限流的效果。

实例demo

package com.mianshi;

import java.util.concurrent.TimeUnit;

/**
 * 令牌桶案例
 */
public class TokenBucket {
    private final int capacity;              // 令牌桶容量
    private final int rate;                  // 令牌生成速率（令牌/秒）
    private int tokens;                      // 当前令牌数量
    private long lastRefillTime;             // 上次填充时间

    public TokenBucket(int capacity, int rate) {
        this.capacity = capacity;
        this.rate = rate;
        this.tokens = capacity;              // 初始时令牌桶是满的
        this.lastRefillTime = System.nanoTime();
    }

    public boolean tryAcquire() {
        refillTokens();                     // 填充令牌
        System.out.println("tokens === " + tokens);
        if (tokens > 0) {
            tokens--;                       // 消费一个令牌
            return true;                    // 返回true表示允许请求通过
        } else {
            return false;                   // 返回false表示请求限制
        }
    }

    private void refillTokens() {
        long currentTime = System.nanoTime();
        System.out.println("currentTime == " + currentTime);
        long elapsedTime = currentTime - lastRefillTime;
        System.out.println("elapsedTime == " + elapsedTime);
        int generatedTokens = (int) (elapsedTime * rate / TimeUnit.SECONDS.toNanos(1));
        System.out.println("generatedTokens == " + generatedTokens);

        if (generatedTokens > 0) {
            tokens = Math.min(capacity, tokens + generatedTokens); // 补充新生成的令牌
            lastRefillTime = currentTime;                         // 更新上次填充时间
        }
        System.out.println();
    }

    public static void main(String[] args) {
        TokenBucket tokenBucket = new TokenBucket(10, 10); // 创建一个容量为10、速率为10个令牌/秒的令牌桶

        for (int i = 1; i <= 20; i++) {
            if (tokenBucket.tryAcquire()) {
                System.out.println("Request " + i + " passed"); // 请求通过
            } else {
                System.out.println("Request " + i + " limited"); // 请求被限制
            }

            try {
                Thread.sleep(100); // 模拟请求间隔时间
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }
}

在上述代码中，TokenBucket类实现了令牌桶算法的限流器。构造函数中传入令牌桶的容量和生成速率（令牌/秒）。tryAcquire()方法用于尝试获取一个令牌，如果成功获取到令牌则返回true，表示允许请求通过；否则返回false，表示请求被限制。

在tryAcquire()方法中，首先调用refillTokens()方法来检查是否需要填充新的令牌。然后判断当前令牌数量是否大于0，如果是，则消费一个令牌并返回true；否则返回false，表示请求被限制。

refillTokens()方法根据时间间隔和生成速率计算应该生成的新令牌数量，并将其添加到当前令牌数量中。注意，这里使用纳秒级的时间戳来进行计算。如果生成的令牌数超过了令牌桶的容量，那么令牌桶的令牌数量将被限制在容量之内。

运行结果:

...


tokens === 0
Request 16 limited
generatedTokens == 0

tokens === 0
Request 17 limited
generatedTokens == 0

tokens === 0
Request 18 limited
generatedTokens == 0

tokens === 0
Request 19 limited
generatedTokens == 0

tokens === 0
Request 20 limited

Process finished with exit code 0

漏桶算法

原理:

漏桶算法的原理是将请求或数据按照恒定的速率从容器（桶）中处理或发送出去。当请求到达时，如果当前容器有足够的空闲容量，则可以处理该请求；如果没有足够的容量，则丢弃该请求或进行排队等待处理。

实现案例:

案例目的:

案例代码


import java.util.concurrent.TimeUnit;

/**
 * 漏桶原理
 */
public class LeakyBucket {
    private final int capacity;              // 漏桶容量
    private final int rate;                  // 处理速率（请求数/秒）
    private int waterLevel;                  // 当前水位
    private long lastLeakTime;               // 上次漏水时间

    public LeakyBucket(int capacity, int rate) {
        this.capacity = capacity;
        this.rate = rate;
        this.waterLevel = 0;                  // 初始时漏桶为空
        this.lastLeakTime = System.nanoTime();
    }

    public boolean tryConsume() {
        leakWater();                          // 漏水

        if (waterLevel < capacity) {
            waterLevel++;                     // 增加水位
            return true;                      // 返回true表示允许请求通过
        } else {
            return false;                     // 返回false表示请求被限制
        }
    }

    private void leakWater() {
        long currentTime = System.nanoTime();
        long elapsedTime = currentTime - lastLeakTime;
        int leakedRequests = (int) (elapsedTime * rate / TimeUnit.SECONDS.toNanos(1));

        if (leakedRequests > 0) {
            waterLevel = Math.max(0, waterLevel - leakedRequests); // 漏掉一定数量的请求
            lastLeakTime = currentTime;                           // 更新上次漏水时间
        }
    }

    public static void main(String[] args) {
        LeakyBucket leakyBucket = new LeakyBucket(10, 10); // 创建一个容量为10、速率为10个请求/秒的漏桶

        for (int i = 1; i <= 20; i++) {
            if (leakyBucket.tryConsume()) {
                System.out.println("Request " + i + " passed"); // 请求通过
            } else {
                System.out.println("Request " + i + " limited"); // 请求被限制
            }

            try {
                Thread.sleep(10); // 模拟请求间隔时间
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }
}