使用 Bucket4j 限制 Spring API 請求速率

1. 概述

在本教程中，我們將重點介紹如何使用 Bucket4j 來對 Spring REST API 進行速率限制。

我們將探索 API 速率限制，瞭解 Bucket4j，並逐步處理在 Spring 應用程序中對 REST API 進行速率限制的幾種方法。

2. API 速率限制

速率限制是一種策略，用於限制 API 訪問。它限制客户端在一定時間範圍內可以進行的 API 調用次數。這有助於防禦 API 免受濫用，無論是有意還是惡意。

速率限制通常通過跟蹤 IP 地址，或者更具業務特定性的方式（例如 API 密鑰或訪問令牌）應用於 API。作為 API 開發人員，當客户端達到限制時，我們有幾個選項：

排隊請求，直到剩餘時間段結束後
立即允許請求，但為此請求收取額外費用
拒絕請求（HTTP 429 Too Many Requests）

3. Bucket4j Rate Limiting 庫

Bucket4j 是一個基於 Java 的輕量級、高性能的速率限制庫。它利用 Bucket4j 算法，提供靈活、可配置的速率限制功能，適用於各種應用場景。

核心功能：

基於桶的速率限制： Bucket4j 採用基於桶的算法，通過維護一個桶來跟蹤請求數量，從而實現速率限制。
靈活的配置： 可以根據需求自定義桶的大小、刷新時間、以及速率限制規則。
高併發支持： Bucket4j 能夠處理高併發請求，保證系統性能。
可擴展性： Bucket4j 易於集成到現有應用中，並可根據需要進行擴展。

主要特性：

支持多種速率限制模式： 包括固定窗口、滑動窗口等。
支持多種數據類型： 例如，整數、字符串等。
提供豐富的 API： 方便開發者進行速率限制配置和管理。
易於集成： Bucket4j 提供了多種集成方式，例如 Spring Boot 集成、自定義集成等。

示例代碼：

// 示例代碼，用於演示 Bucket4j 的使用
import org.bucket4j.Bucket4j;

// 創建一個 Bucket4j 實例
Bucket4j bucket = new Bucket4j(new StaticConfig(100, 10, 1000));

// 增加請求計數
bucket.add(1);

// 減少請求計數
bucket.remove(1);

3.1. Bucket4j 簡介

Bucket4j 是一個基於令牌桶算法的 Java 速率限制庫。 Bucket4j 是一種線程安全的庫，可用於獨立 JVM 應用程序或集羣環境。它還通過 JCache (JSR107) 規範支持內存或分佈式緩存。

3.2 令牌桶算法

讓我們直觀地從 API 速率限制的角度來看待該算法。

假設我們有一個桶，其容量定義為它可以容納的令牌數量。 當消費者想要訪問 API 端點時，必須從該桶中獲取一個令牌。如果桶中有可用令牌，我們則從桶中移除一個令牌並接受請求。反之，如果桶中沒有令牌，我們則拒絕請求。

隨着請求消耗令牌，我們也會以固定的速率對其進行補充，以確保桶永遠不會超過其容量。

假設我們有一個速率限制為每分鐘 100 個請求的 API。我們創建一個容量為 100 個令牌，補充速率為 100 個令牌/分鐘的桶。

如果我們在 1 分鐘內收到 70 個請求，則將消耗 70 個令牌，桶中剩餘 30 個令牌。在下個分鐘的開始時，桶將完全補充到其最大容量 100 個令牌，無論剩餘 30 個令牌如何。這確保了桶在每個時間窗口的開始時始終充滿最大容量。如果在分鐘結束之前所有 100 個令牌都被消耗掉，則桶將逐漸補充，直到有新的令牌可用時，請求將被拒絕。

4. 使用 Bucket4j 入門

4.1. Maven 配置

首先，我們添加 bucket4j 依賴到我們的 pom.xml 中：

<dependency>
    <groupId>com.bucket4j</groupId>
    <artifactId>bucket4j-core</artifactId>
    <version>8.1.0</version>
</dependency>

4.2. 術語

在探討如何使用 Bucket4j 之前，我們簡要地討論一下核心類及其在令牌桶算法正式模型中表示的不同元素。

<a href="https://javadoc.io/doc/com.github.vladimir-bukhtoyarov/bucket4j-core/4.10.0/io/github/bucket4j/Bucket.html">Bucket</a> 接口代表一個具有最大容量的令牌桶。它提供了諸如tryConsume 和 <a href="https://javadoc.io/doc/com.github.vladimir-bukhtoyarov/bucket4j-core/4.10.0/io/github/bucket4j/Bucket.html#tryConsumeAndReturnRemaining-long-">tryConsumeAndReturnRemaining</a> 方法，用於消耗令牌。這些方法將消耗結果作為true`，如果請求符合限制，並且令牌已消耗。

<a href="https://javadoc.io/doc/com.github.vladimir-bukhtoyarov/bucket4j-core/4.10.0/io/github/bucket4j/Bandwidth.html">Bandwidth</a> 類是桶的關鍵構建塊，因為它定義了桶的限制。我們使用Bandwidth 來配置桶的容量和令牌回充率。

`Refill 類用於定義令牌以固定的速率添加到桶中的方法。我們可以將速率配置為在給定時間段內添加的令牌數量，例如每秒 10 個桶或每 5 分鐘 200 個令牌，等等。

tryConsumeAndReturnRemaining 方法在 Bucket 中返回 <a href="https://javadoc.io/doc/com.github.vladimir-bukhtoyarov/bucket4j-core/4.10.0/io/github/bucket4j/ConsumptionProbe.html">ConsumptionProbe</a>。ConsumptionProbe` 包含消耗結果以及桶的狀態，例如剩餘令牌或請求的令牌再次可用在桶中的時間。

4.3. 基本用法

讓我們測試一些基本的速率限制模式。

對於每分鐘 10 個請求的速率限制，我們將創建一個容量為 10 個，刷新速率為每分鐘 10 個令牌的桶：

Refill refill = Refill.intervally(10, Duration.ofMinutes(1));
Bandwidth limit = Bandwidth.classic(10, refill);
Bucket bucket = Bucket.builder()
    .addLimit(limit)
    .build();

for (int i = 1; i <= 10; i++) {
    assertTrue(bucket.tryConsume(1));
}
assertFalse(bucket.tryConsume(1));

Refill.intervally 會在時間窗口的開始時填充桶，在本例中是每分鐘的 10 個 token。

接下來，讓我們看看 refill 的實際應用。

我們將 refill 速率設置為每 2 秒 1 個 token，並限制我們的請求以遵守速率限制：

Bandwidth limit = Bandwidth.classic(1, Refill.intervally(1, Duration.ofSeconds(2)));
Bucket bucket = Bucket.builder()
    .addLimit(limit)
    .build();
assertTrue(bucket.tryConsume(1));     // first request
Executors.newScheduledThreadPool(1)   // schedule another request for 2 seconds later
    .schedule(() -> assertTrue(bucket.tryConsume(1)), 2, TimeUnit.SECONDS);

假設我們有一個每分鐘10個請求的速率限制。同時，我們可能希望避免在第一 5 秒內耗盡所有令牌的突發情況。Bucket4j 允許我們在同一個桶上設置多個限制（帶寬）。讓我們添加另一個限制，允許在 20 秒的時間窗口內僅發送 5 個請求：

Bucket bucket = Bucket.builder()
    .addLimit(Bandwidth.classic(10, Refill.intervally(10, Duration.ofMinutes(1))))
    .addLimit(Bandwidth.classic(5, Refill.intervally(5, Duration.ofSeconds(20))))
    .build();

for (int i = 1; i <= 5; i++) {
    assertTrue(bucket.tryConsume(1));
}
assertFalse(bucket.tryConsume(1));

5. 使用 Bucket4j 限制 Spring API 的速率

讓我們使用 Bucket4j 在 Spring REST API 中應用速率限制。

5.1. 面積計算器 API

我們將實現一個簡單但非常流行的面積計算器 REST API。目前，它根據矩形的尺寸計算並返回其面積：

@RestController
class AreaCalculationController {

    @PostMapping(value = "/api/v1/area/rectangle")
    public ResponseEntity<AreaV1> rectangle(@RequestBody RectangleDimensionsV1 dimensions) {
        return ResponseEntity.ok(new AreaV1("rectangle", dimensions.getLength() * dimensions.getWidth()));
    }
}

讓我們確保我們的API已啓動並正常運行：

$ curl -X POST http://localhost:9001/api/v1/area/rectangle \
    -H "Content-Type: application/json" \
    -d '{ "length": 10, "width": 12 }'

{ "shape":"rectangle","area":120.0 }

5.2. 應用速率限制

現在我們將介紹一個簡單的速率限制，允許 API 每分鐘發送 20 個請求。換句話説，如果 API 在 1 分鐘的時間窗口內已經收到 20 個請求，它將拒絕該請求。

讓我們修改我們的 Controller 以創建 Bucket 並添加限制（帶寬）：

@RestController
class AreaCalculationController {

    private final Bucket bucket;

    public AreaCalculationController() {
        Bandwidth limit = Bandwidth.classic(20, Refill.greedy(20, Duration.ofMinutes(1)));
        this.bucket = Bucket.builder()
            .addLimit(limit)
            .build();
    }
    //..
}

在本 API 中，我們可以通過使用 tryConsume 方法，消耗一個令牌來檢查請求是否被允許。如果已達到限制，我們可以通過返回 HTTP 429 Too Many Requests 狀態碼來拒絕請求：

public ResponseEntity<AreaV1> rectangle(@RequestBody RectangleDimensionsV1 dimensions) {
    if (bucket.tryConsume(1)) {
        return ResponseEntity.ok(new AreaV1("rectangle", dimensions.getLength() * dimensions.getWidth()));
    }

    return ResponseEntity.status(HttpStatus.TOO_MANY_REQUESTS).build();
}

# 21st request within 1 minute
$ curl -v -X POST http://localhost:9001/api/v1/area/rectangle \
    -H "Content-Type: application/json" \
    -d '{ "length": 10, "width": 12 }'

< HTTP/1.1 429

5.3. API 客户端和計費方案

我們目前有一個簡單的速率限制，可能會導致 API 請求被限速。接下來，我們將介紹針對企業級場景的計費方案。

計費方案有助於我們實現 API 的盈利。假設我們為 API 客户端提供以下方案：

免費：每客户端每小時 20 個請求
基礎版：每客户端每小時 40 個請求
專業版：每客户端每小時 100 個請求

每個 API 客户端都擁有一個 唯一的 API 密鑰，必須在每個請求中發送。這有助於我們識別與 API 客户端關聯的計費方案。

讓我們定義每個計費方案的速率限制（帶寬）：

enum PricingPlan {
    FREE {
        Bandwidth getLimit() {
            return Bandwidth.classic(20, Refill.intervally(20, Duration.ofHours(1)));
        }
    },
    BASIC {
        Bandwidth getLimit() {
            return Bandwidth.classic(40, Refill.intervally(40, Duration.ofHours(1)));
        }
    },
    PROFESSIONAL {
        Bandwidth getLimit() {
            return Bandwidth.classic(100, Refill.intervally(100, Duration.ofHours(1)));
        }
    };
    //..
}

然後，我們添加一個方法來從提供的 API 密鑰中解析定價計劃：

enum PricingPlan {
    
    static PricingPlan resolvePlanFromApiKey(String apiKey) {
        if (apiKey == null || apiKey.isEmpty()) {
            return FREE;
        } else if (apiKey.startsWith("PX001-")) {
            return PROFESSIONAL;
        } else if (apiKey.startsWith("BX001-")) {
            return BASIC;
        }
        return FREE;
    }
    //..
}

接下來，我們需要為每個API密鑰存儲Bucket，並檢索用於限速的Bucket：

class PricingPlanService {

    private final Map<String, Bucket> cache = new ConcurrentHashMap<>();

    public Bucket resolveBucket(String apiKey) {
        return cache.computeIfAbsent(apiKey, this::newBucket);
    }

    private Bucket newBucket(String apiKey) {
        PricingPlan pricingPlan = PricingPlan.resolvePlanFromApiKey(apiKey);
        return Bucket.builder()
            .addLimit(pricingPlan.getLimit())
            .build();
    }
}

現在我們有一個基於內存的存儲，用於存儲每個API密鑰對應的桶。讓我們修改我們的Controller，使其使用PricingPlanService：

@RestController
class AreaCalculationController {

    private PricingPlanService pricingPlanService;

    public ResponseEntity<AreaV1> rectangle(@RequestHeader(value = "X-api-key") String apiKey,
        @RequestBody RectangleDimensionsV1 dimensions) {

        Bucket bucket = pricingPlanService.resolveBucket(apiKey);
        ConsumptionProbe probe = bucket.tryConsumeAndReturnRemaining(1);
        if (probe.isConsumed()) {
            return ResponseEntity.ok()
                .header("X-Rate-Limit-Remaining", Long.toString(probe.getRemainingTokens()))
                .body(new AreaV1("rectangle", dimensions.getLength() * dimensions.getWidth()));
        }
        
        long waitForRefill = probe.getNanosToWaitForRefill() / 1_000_000_000;
        return ResponseEntity.status(HttpStatus.TOO_MANY_REQUESTS)
            .header("X-Rate-Limit-Retry-After-Seconds", String.valueOf(waitForRefill))
            .build();
    }
}

讓我們逐步瞭解這些變更。API客户端通過X-api-key請求頭髮送API密鑰。我們使用PricingPlanService來獲取與該API密鑰關聯的存儲桶，並檢查請求是否允許通過消耗存儲桶中的令牌。

為了提升API客户端體驗，我們將使用以下附加響應頭來發送速率限制信息：

X-Rate-Limit-Remaining: 當前時間窗口中剩餘的令牌數量
X-Rate-Limit-Retry-After-Seconds: 存儲桶刷新剩餘時間，單位為秒

我們可以調用ConsumptionProbe中的getRemainingTokens和getNanosToWaitForRefill方法，分別獲取存儲桶中剩餘的令牌數量以及下一次刷新剩餘的時間。getNanosToWaitForRefill方法如果成功消耗令牌，則返回0。

## successful request
$ curl -v -X POST http://localhost:9001/api/v1/area/rectangle \
    -H "Content-Type: application/json" -H "X-api-key:FX001-99999" \
    -d '{ "length": 10, "width": 12 }'

< HTTP/1.1 200
< X-Rate-Limit-Remaining: 11
{"shape":"rectangle","area":120.0}

## rejected request
$ curl -v -X POST http://localhost:9001/api/v1/area/rectangle \
    -H "Content-Type: application/json" -H "X-api-key:FX001-99999" \
    -d '{ "length": 10, "width": 12 }'

< HTTP/1.1 429
< X-Rate-Limit-Retry-After-Seconds: 583

5.4. 使用 Spring MVC 攔截器

假設我們現在需要添加一個新的 API 端點，該端點根據給定的高度和底邊計算並返回三角形的面積：

@PostMapping(value = "/triangle")
public ResponseEntity<AreaV1> triangle(@RequestBody TriangleDimensionsV1 dimensions) {
    return ResponseEntity.ok(new AreaV1("triangle", 0.5d * dimensions.getHeight() * dimensions.getBase()));
}

實際上，我們還需要對新的端點進行限速處理。我們可以簡單地從我們之前的端點複製粘貼限速代碼。另一種方法是使用 Spring MVC 的 HandlerInterceptor 來解耦限速代碼與業務代碼。

讓我們創建一個 RateLimitInterceptor 並實現限速代碼在 preHandle 方法中：

public class RateLimitInterceptor implements HandlerInterceptor {

    private PricingPlanService pricingPlanService;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) 
      throws Exception {
        String apiKey = request.getHeader("X-api-key");
        if (apiKey == null || apiKey.isEmpty()) {
            response.sendError(HttpStatus.BAD_REQUEST.value(), "Missing Header: X-api-key");
            return false;
        }

        Bucket tokenBucket = pricingPlanService.resolveBucket(apiKey);
        ConsumptionProbe probe = tokenBucket.tryConsumeAndReturnRemaining(1);
        if (probe.isConsumed()) {
            response.addHeader("X-Rate-Limit-Remaining", String.valueOf(probe.getRemainingTokens()));
            return true;
        } else {
            long waitForRefill = probe.getNanosToWaitForRefill() / 1_000_000_000;
            response.addHeader("X-Rate-Limit-Retry-After-Seconds", String.valueOf(waitForRefill));
            response.sendError(HttpStatus.TOO_MANY_REQUESTS.value(),
              "You have exhausted your API Request Quota"); 
            return false;
        }
    }
}

最後，我們需要將攔截器添加到 InterceptorRegistry 中：

public class Bucket4jRateLimitApp implements WebMvcConfigurer {
    
    private RateLimitInterceptor interceptor;

    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(interceptor)
            .addPathPatterns("/api/v1/area/**");
    }
}

RateLimitInterceptor攔截所有對我們的區域計算API端點的請求。

讓我們嘗試一下我們的新端點：

## successful request
$ curl -v -X POST http://localhost:9001/api/v1/area/triangle \
    -H "Content-Type: application/json" -H "X-api-key:FX001-99999" \
    -d '{ "height": 15, "base": 8 }'

< HTTP/1.1 200
< X-Rate-Limit-Remaining: 9
{"shape":"triangle","area":60.0}

## rejected request
$ curl -v -X POST http://localhost:9001/api/v1/area/triangle \
    -H "Content-Type: application/json" -H "X-api-key:FX001-99999" \
    -d '{ "height": 15, "base": 8 }'

< HTTP/1.1 429
< X-Rate-Limit-Retry-After-Seconds: 299
{ "status": 429, "error": "Too Many Requests", "message": "You have exhausted your API Request Quota" }

看起來我們已經完成了。我們可以繼續添加端點，攔截器將為每個請求應用速率限制。

6. Bucket4j Spring Boot Starter

讓我們來看在 Spring 應用中使用 Bucket4j 的另一種方式。 Bucket4j Spring Boot Starter 提供對 Bucket4j 的自動配置，幫助我們通過 Spring Boot 應用屬性或配置實現 API 速率限制。

一旦我們將 Bucket4j starter 集成到我們的應用程序中，我們將擁有完全聲明式的 API 速率限制實現，無需任何應用程序代碼。

6.1 速率限制過濾器

在我們的示例中，我們使用請求頭 X-api-key 的值作為識別和應用速率限制的關鍵。

Bucket4j Spring Boot Starter 提供了一些預定義的配置，用於定義我們的速率限制關鍵：

一個簡單的速率限制過濾器，它是默認配置。
按 IP 地址過濾
基於表達式的過濾器

基於表達式的過濾器使用 Spring 表達式語言 (SpEL)。SpEL 允許訪問根對象，例如 HttpServletRequest，從而可以在 IP 地址 (getRemoteAddr())、請求頭 (getHeader(‘X-api-key’)) 等上構建過濾器表達式。

該庫還支持在過濾器表達式中使用自定義類，關於此內容的詳細信息請參閲文檔。

6.2. Maven 配置

讓我們首先將 <a href="https://mvnrepository.com/artifact/com.giffing.bucket4j.spring.boot.starter/bucket4j-spring-boot-starter">bucket4j-spring-boot-starter</a> 依賴添加到我們的 pom.xml 中：

<dependency>
    <groupId>com.giffing.bucket4j.spring.boot.starter</groupId>
    <artifactId>bucket4j-spring-boot-starter</artifactId>
    <version>0.8.1</version>
</dependency>

我們以前的實現中，使用了一種內存中的 Map 來存儲每個 API 密鑰（消費者）對應的 Bucket。在這裏，我們可以使用 Spring 的緩存抽象來配置一個內存存儲，例如 Caffeine 或 Guava。

讓我們添加緩存依賴項：

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-cache</artifactId>
    <version>3.3.2</version>
</dependency>
<dependency>
    <groupId>javax.cache</groupId>
    <artifactId>cache-api</artifactId>
</dependency>
<dependency>
    <groupId>com.github.ben-manes.caffeine</groupId>
    <artifactId>caffeine</artifactId>
    <version>2.8.2</version>
</dependency>
<dependency>
    <groupId>com.github.ben-manes.caffeine</groupId>
    <artifactId>jcache</artifactId>
    <version>2.8.2</version>
</dependency>

注意：我們還添加了 jcache 依賴項，以符合 Bucket4j 的緩存支持。

請務必通過在任何配置類中添加 @EnableCaching 註解來啓用緩存功能。

6.3. 應用配置

讓我們配置應用程序使用 Bucket4j starter 庫。首先，我們將配置 Caffeine 緩存，用於存儲 API 密鑰和 Bucket 在內存中：

spring:
  cache:
    cache-names:
    - rate-limit-buckets
    caffeine:
      spec: maximumSize=100000,expireAfterAccess=3600s

接下來，讓我們配置 Bucket4j：

bucket4j:
  enabled: true
  filters:
  - cache-name: rate-limit-buckets
    url: /api/v1/area.*
    strategy: first
    http-response-body: "{ \"status\": 429, \"error\": \"Too Many Requests\", \"message\": \"You have exhausted your API Request Quota\" }"
    rate-limits:
    - cache-key: "getHeader('X-api-key')"
      execute-condition: "getHeader('X-api-key').startsWith('PX001-')"
      bandwidths:
      - capacity: 100
        time: 1
        unit: hours
    - cache-key: "getHeader('X-api-key')"
      execute-condition: "getHeader('X-api-key').startsWith('BX001-')"
      bandwidths:
      - capacity: 40
        time: 1
        unit: hours
    - cache-key: "getHeader('X-api-key')"
      bandwidths:
      - capacity: 20
        time: 1
        unit: hours

那麼，我們剛才配置了什麼呢？

bucket4j.enabled=true – 啓用 Bucket4j 自動配置
bucket4j.filters.cache-name – 從緩存中獲取與 API 密鑰相關的 Bucket
bucket4j.filters.url – 指定用於應用速率限制的路徑表達式
bucket4j.filters.strategy=first – 在第一個匹配的速率限制配置中停止
bucket4j.filters.rate-limits.cache-key – 使用 Spring 表達式語言 (SpEL) 檢索鍵
bucket4j.filters.rate-limits.execute-condition – 使用 SpEL 決定是否執行速率限制
bucket4j.filters.rate-limits.bandwidths – 定義 Bucket4j 速率限制參數

我們替換了 PricingPlanService 和 RateLimitInterceptor，並使用一系列按順序評估的速率限制配置。

我們來試試看：

## successful request
$ curl -v -X POST http://localhost:9000/api/v1/area/triangle \
    -H "Content-Type: application/json" -H "X-api-key:FX001-99999" \
    -d '{ "height": 20, "base": 7 }'

< HTTP/1.1 200
< X-Rate-Limit-Remaining: 7
{"shape":"triangle","area":70.0}

## rejected request
$ curl -v -X POST http://localhost:9000/api/v1/area/triangle \
    -H "Content-Type: application/json" -H "X-api-key:FX001-99999" \
    -d '{ "height": 7, "base": 20 }'

< HTTP/1.1 429
< X-Rate-Limit-Retry-After-Seconds: 212
{ "status": 429, "error": "Too Many Requests", "message": "You have exhausted your API Request Quota" }

7. 結論

在本文中，我們演示了多種使用 Bucket4j 限制 Spring API 速率的方法。要了解更多信息，請務必查看官方文檔。

知識庫 / Spring / Spring Boot RSS 訂閱