Node.js サーバーアプリケーションにおけるトークンバケットを用いたレートリミットの実現

はじめに

こんにちは、株式会社 ROXX で back check というサービスを開発しているぐっきー（@Area029S）です。

リクエストを受け付けるごとに重い処理が実行される機能の実装などにおいては一定の時間内に大量に処理を実行するとシステムのリソースを圧迫してしまう可能性があるため注意が必要です。
このような処理をリソースを圧迫させずに処理する方法の1つとして、レートリミット(時間あたりの実行数の制御)の導入があげられます。

今回はネットワークのトラフィックシェービングなどに利用されることの多いトークンバケットアルゴリズムを導入することで、サーバー側のアプリケーションでレートリミットを実現しようと思います。

レートリミットとは？

レートリミットとは主に時間あたりのネットワークトラフィックを制御するために用いられる戦略です。

一般的には下記のようなことに役立ちます。

• アプリケーションリソースの枯渇を防ぐ
  • 特定のエンドポイントへのリクエスト数を制限する
• 運用コストの超過を防ぐ
  • ChatGPT など利用ごとに課金される外部 API
    を利用したサービスのリクエスト数を制限することにより出費を管理する
• 悪意のあるユーザーから API を保護する
  • ブルートフォース攻撃を防ぐためエンドポイントへのリクエスト数を制限する

参照:
The Fundamentals of Rate Limiting: How it Works and Why You Need it

外部 API を利用したサービスのリクエスト数を制限する

今回のやりたいことであるリクエスト毎に重い処理を実行したい要件においては、上記の文脈でレートリミットを適用することでシステムリソースの負荷を分散させる用途に利用できそうです。

レートリミットを実装するための一般的なアルゴリズムには下記のようなものが存在します。

• token bucket
• leaky bucket
• fixed window counter
• sliding window log
• sliding window counter

参照：様々な rate limit アルゴリズム

今回は高負荷の掛かるリクエストが大量に来たときは処理量を制限しつつ、制限の範囲内においては短期間に多数のリクエストを処理（バースト）することの可能な方式を利用したいため token bucket がよさそうです。そこで token bucket の特徴について説明します。

token bucket（トークンバケットアルゴリズム）とは？

画像引用：様々な rate limit アルゴリズム より

概要

トークンバケットアルゴリズムは、リクエストのレートを一定の制限内に保ちつつ、短期的なトラフィックの増加（バースト）も許容するための手法です。このアルゴリズムは、トークンの生成速度とバケットのサイズを利用してリクエストの平均処理量に制限をかけます。柔軟かつ効率的なレートリミットの手段として、ネットワーク環境からアプリケーションまで幅広く採用されています。

主要用語

トークンバケットアルゴリズムを理解する上で重要な用語は以下の通りです。

1. トークン
   • トークンは、リクエストを処理するために消費される「通貨」のようなものです。トークンの有無がリクエストの実行可否を決定します。トークンは一定間隔でバケットに追加されます。
2. バケット
   • バケットはトークンを保持する仮想的な容器です。バケットの中のトークン数がレートリミットの基準となります。バケットが満杯の場合、新たに生成されるトークンは追加されず破棄されます。
3. トークンの生成速度
   • トークンの生成速度は、一定時間内にバケットに追加されるトークンの数を指します。トークンの生成速度が許可されるリクエストのレートを決定します。
4. トークンの消費
   • リクエストが行われるたび、バケットからトークンが消費されます。トークンが不足している場合、リクエストは待機状態になります。
5. バースト
   • バーストは、短期間に多数のリクエストを許可する能力です。バケットが満杯の場合、短時間で多数のリクエストを処理できますが、その後トークンが不足するとトークンの再生成を待たなければなりません。

仕組み

トークンバケットアルゴリズムでは、定められた速度でトークンがバケットに追加されます。バケットはこれらのトークンを保持し、そのサイズには上限があります。バケットが満杯の場合、新たに生成されるトークンは失われます。

リクエストが行われると、バケットからトークンが消費されます。処理するためには、必要な数のトークンがバケット内に存在する必要があります。十分なトークンがない場合、リクエストは待機状態に入ります。

このアルゴリズムによるレートリミットは、トークンの生成速度とバケットのサイズによって決定されます。これにより、長期的には一定のレートが保たれる一方で、バケットが満杯の状態では短期間のバーストが可能になります。

もっとわかりやすく説明してほしい！という方には下記のブログをおすすめします。

トークンバケットアルゴリズムではない、MP アルゴリズムと呼べ

Node.jsにおける実装例

Node.js での実装例を簡単にまとめます。

ライブラリの活用

トークンバケットアルゴリズムを自前で実装するとそこそこの工数はかかりそうだということで、ライブラリがないか探してみます。
npm trends をみてみるとレートリミットを実装するものがいくつか見つかり、 limiter というライブラリが最も利用されてることがわかりました。
(比較的レートリミット系ライブラリ自体あまり出回っていないようです)

参照: npm trends - bottleneck vs leaky-bucket vs limiter vs tokenbucket

そこで今回は limiter を利用します。

実装例

コードの概要としては、リクエストを受け付けるごとに重い処理が非同期に実行される機能を想定しています。
そのため express でサーバーをたて、 bull というライブラリを用いて local 環境で動作するキューワーカーを立ち上げます。

キューに追加されたジョブを処理する箇所でレートリミットを設定することで、クライアントからのリクエストは受け取りつつ、ミッションクリティカルな処理の実行はトークンバケットを利用したレートリミット制限下で実行できる実装となっています。

https://github.com/SotaYamaguchi/express-latelimiter

const express = require("express");
const Queue = require("bull");
const { TokenBucket } = require("limiter");

// トークンバケットによるレートリミッターを設定
const bucket = new TokenBucket({
  bucketSize: 2, // トークンの最大保持可能数
  tokensPerInterval: 1, // 1分間に生成されるトークン数
  interval: "minute", // トークン生成の間隔
});

// キューの作成
const myQueue = new Queue("myQueue");

const addJobToQueue = (data) => {
  myQueue.add(data);
}

const someProcess = () => {
  const currentTime = new Date().toLocaleString();
  console.log(`Processing at: ${currentTime}`);
};

// キューでのジョブの処理
myQueue.process(async function(job, done) {
  console.log("Processing job");
  await bucket.removeTokens(1);
  someProcess();
  done();
});

const app = express();
app.use(express.json());

app.post("/job", (req, res) => {
  addJobToQueue(req.body);
  res.send("Job added to queue");
});

const PORT = 3000;
app.listen(PORT, () => {
  console.log(`Server is running on port ${PORT}`);
});

それでは具体的に下記の設定でトークンバケットを動かしてみます。

バケットサイズ: 2
トークンの生成速度: 1 個/分
処理実行時のトークンの消費数: 1 個

下記の振る舞いが確認できれば機能していることがわかります。

1. 3分間待機してバケットにトークンが溜まるのを待つ。バケットでトークンを保持できる最大数（バケットサイズ）を2に設定しているため、3つめに生成されたトークンは破棄される。
2. curlにて連続で3回リクエストを送信する。2個は即座に処理され、1分後にバケットに空きがある状態でトークンが生成されるため、生成されたトークンを使い残りの処理が実行される。

実行結果:

想定した動作をすることが確認できました。

おわりに

今回は受け付けたリクエストに対して非同期に処理を行う Node.js サーバーアプリケーションにおいてトークンバケットを用いたレートリミットを実装してみました。
レートリミットのアルゴリズムはとっつきづらいイメージを持ちがちですが、1度理解できると便利なので読んでいただいた方の学習の促進に少しでも繋がると幸いです。
また、本記事内でリンクを貼らせていただいている関連の記事もわかりやすいのでぜひご覧ください。