Zenn
🤠

Next.js Vercel KVでレートリミットとセッション管理の話

2023/06/08に公開
2件
Next.js
React
Redis
Vercel
KV
tech

Vercel KVを使用したレートリミットとセッション管理の実装方法の説明です。ソースコードはこちらにあります。

Vercel KVとは

Redisというインメモリのキー・バーリュー ストア(DB)です。裏ではUpstashのRedisを使用しています。主にショッピングカートのようなセッション管理やレートリミットの実装に使用します。もちろん、キャッシュとしても利用できます。

https://zenn.dev/tfutada/articles/10e9ef769144f9

レートリミットを設置する理由

例えば、OpenAI APIを使用したチャット機能をZennに追加するとします。その場合、30メッセージ/時間と制限を加えることで、クラウドの請求額にビクビクしないようにします。

インストール

App Router方式で説明します。upstashのレートリミットをインストールします。

インストール
yarn add npm i @upstash/ratelimit

コーディング

Upstashのレートリミット・ライブラリを使用した方法とスクラッチ実装の２通りを説明します。

middlewareにレートリミットを実装していきます。リクエストをインターセプトし、制限を超えた場合は429エラー(Too many requests)を返します。OKならそのままリクエストをディスパッチします。

/middleware.ts
import {NextResponse} from 'next/server';
import type {NextRequest} from 'next/server';
import {Ratelimit} from "@upstash/ratelimit";
import {Redis} from "@upstash/redis";

// プロジェクトにKVを接続するときにプレフィックス(KV_)を選べます。
const redis = new Redis({
    url: process.env.KV_REST_API_URL ?? 'expect redis url',
    token: process.env.KV_REST_API_TOKEN ?? 'expect redis token'
});

// 10秒間隔のウィンドウ。一つのウィンドウで最大１０リクエストまで許容。
const ratelimit = new Ratelimit({
    redis,
    limiter: Ratelimit.slidingWindow(10, "10 s")
});

export async function middleware(request: NextRequest) {
    const key = request.ip ?? "127.0.0.1"; // クライアントのIPでレートリミット
    // レートリミットの判定
    const resp = await ratelimit.limit(key);

    console.log(`success: ${resp.success}, remaining: ${resp.remaining}, reset: ${resp.reset}`)

    // リミットオーバーなら HTPステータス429で返す。
    return resp.success ? NextResponse.next()
        : new NextResponse("Too many requests", {status: 429});
}

yarn devで起動します。(ローカルからクラウドのVercel KVにアクセス可能です。)
vegetaを使用して負荷をかけてみます。1秒に2リクエスト(rate=2)で負荷をかけてみましょう。

Vegeta
echo 'GET http://localhost:3000/' | vegeta attack -rate=2 -duration=30s | tee results.bin | vegeta report

結果です。

5秒後(2リクエストx5秒=10リクエスト)にリミットに到達し、429エラーでリクエストはリジェクトされます。10秒後にウィンドウがスライドし、リミットが10リクエストにリセットされます。

Vegetaの出力結果です。

スループットが約1rpsになっています。成功率も約50%になっています。このことからレートリミットが機能していることが分かります。

コマンドラインから作成されたRedisキーを確認できます。

Redis CLI
redis-cli --tls -u redis://default:xxx@yyy.kv.vercel-storage.com:35306

自作版

スクラッチで実装してみました。RedisではZSCOREというsorted setが利用できます。リクエストのタイムスタンプでソートをすることでウィンドウ処理を実現できます。キーにはscoreというプロパティのようなものがあり、その値でソート処理します。そのscore値にタイムスタンプを使います。

自作版
import {kv} from '@vercel/kv';

export async function rateLimit(ip: string, maxRequests: number, windowInSecs: number): Promise<any> {
    const now = Date.now();
    const windowStart = now - windowInSecs * 1000;
    const key = `rate:${ip}`;

    try {
        // ウィンドウからハズれたキーを消します。
        await kv.zremrangebyscore(key, 0, windowStart);

        // ウィンドウ内のキーの数を数えます。
        const current = await kv.zcount(key, '-inf', '+inf');
	// 最大リクエスト数を超えたらリジェクトします。
        let resp = {success: current < maxRequests, remaining: maxRequests - current, reset: windowInSecs};

        if (resp.success) {
            // リクエストをredisに追加します。
            await kv.zadd(key, {score: now, member: `timestamp:${now}`});
            await kv.expire(key, windowInSecs);
        }

        return resp;
    } catch (err) {
        console.error('Redis error:', err);
        return {};
    }
}

同様にVegetaの実行結果です。

ご覧のように、ドンピシャでスループット1.00、成功率50%になりました。秒単位でスムーズにスライドしたためでしょう。

ちなみにリミッターを外すと綺麗にスケールします。

レイテンシーを見て分かるように、リミッター(Redis)を設置すると激おそになります。リミッターの部分のelapsed timeを測定してみました。

200msとかRedisではあり得ない遅さです。理由としては、まだVercel KVがベータ版であること、書き込み処理をしていることなど考えられます。また、VPNや時間帯でもレイテンシーが変わりました。ステーブル版がリリースされたらまた検証したいと思います。

Leacky Bucket

leacky bucketアルゴリズムも実装してみました。その名の通り、バケツの底から水が漏れ出るイメージです。バケツから水が溢れる = リクエストをリジェクトするの意味になります。一定のスピードで水が減っていきます。スライディング・ウィンドウに比べるとスムーズなトラフィックになります。

leacky bucket
import {kv} from '@vercel/kv';

interface Bucket {
    level: number;
    lastUpdated: number;
}

class LeakyBucket {
    capacity: number; // バケツの容量
    leakRate: number; // 水が漏れるスピード

    constructor(capacity: number, leakRate: number) {
        this.capacity = capacity;
        this.leakRate = leakRate;
    }

    async increment(key: string) {
        const currentTime = Date.now() / 1000;

        // Update the bucket's last updated timestamp and its level
        let bucket = await kv.get<Bucket>(key);
        let level = 0;
        let lastUpdated = currentTime;
        if (bucket) {
            level = bucket.level; // 水かさ
            lastUpdated = bucket.lastUpdated; // 前回の計算のタイムスタンプ(秒)
        }

        // 前回計算した時からどれだけの量の水が漏れ出たか？
        const delta = this.leakRate * (currentTime - lastUpdated);
        level = Math.max(0, level - delta); // 抜けた分の水かさを減らす

        // If the resulting level is less than the capacity, increment the level
        let resp = {
            success: level < this.capacity, // 水かさ < バケツの容量
            remaining: this.capacity - level, // 満杯までどれくらい？
            reset: 0
        };

        if (resp.success) {
            await kv.set(key, {
                level: level + 1, // 水かさを増やす
                lastUpdated: currentTime
            });
        }

        return resp;
    }
}

export default LeakyBucket;

セッション管理

話題を変えまして、KVを使用したセッション管理の方法です。

以下の３つのプログラムから構成されます。

  1. app/api-cookie/route.ts APIです。クッキーがまだなければ仕込みます。
  2. app/components/mycookie.tsx クライアント・コンポーネント。補助的な部品
  3. app/mysession/page.tsx サーバ・コンポーネント。ショッピングカート本体

まずは、クッキーを仕込むロジックです。クッキー(MY_TOKEN)の存在チェックをし、なければセットします。このクッキーがショッピングカートを識別するセッションキーになります。

手抜きで、ここでカートにアイテムを追加しておきます。このセッションキーをRedisのキーにしてアイテムを追加します。

app/api-cookie/route.ts
import {cookies} from 'next/headers';
import {v4} from 'uuid';
import {kv} from "@vercel/kv";
import {Cart, MY_TOKEN} from "@/app/myconst";

export async function GET(request: Request) {
    const cookieStore = cookies();
    const token = cookieStore.get(MY_TOKEN); // クッキーを取得する

    const headers: { [key: string]: string } = {}
    
    // まだクッキーがなければ、セットします。
    if (!token) {
        const key = v4();
        headers['Set-Cookie'] = `${MY_TOKEN}=${key}; path=/; HttpOnly; SameSite=Strict`;

        // 手抜きでここでカートにアイテムを仕込みます。
        // 実際はボタンを設置して、サーバアクション等でカートに追加します。
        const vals: Cart[] = [{id: "1", quantity: 3}, {id: "2", quantity: 4}, {id: "3", quantity: 5}]
        await kv.set(key, vals);
    }

    return new Response('Hello, Next.js!', {
        status: 200,
        headers
    });
}

cookieStoreはイミュータブルなのでセットはできません。HTTPレスポンスでセットします。

次に、クライアント・コンポーネントです。このコンポーネントは先ほどのAPIを叩くためのコンポーネントです。

app/components/mycookie.tsx
'use client';

import {useEffect} from "react";
import {useRouter} from 'next/navigation';


export default function SessionCookie() {
    const router = useRouter()

    useEffect(
        () => {
            fetch('/api-cookie')
                .then((res) => res.text())
                .then((data) => {
                        console.log('data', data)
                        router.refresh()
                    }
                )
        }, []
    )

    return (
        <div>
            setting cookie...
        </div>
    )
}

routerでリフレッシュします。

最後に、サーバ・コンポーネントです。ここにショッピングカートを実装します。
クッキーからセッションキーを取得し、Redisからアイテム一覧を取得します。
もしキーがない場合は、先ほどのクライアント・コンポーネントをレンダリングし、クッキー作成のAPIを叩きます。

app/mysession/page.tsx
import dynamic from 'next/dynamic';
import {cookies} from 'next/headers';
import {kv} from "@vercel/kv";
import {Cart, MY_TOKEN} from "@/app/myconst";

const SessionCookie = dynamic(() => import('@/app/components/mycookie'), {ssr: false});


export default async function MySession() {
    const cookieStore = cookies();
    const token = cookieStore.get(MY_TOKEN);

    let items: Cart[] = [];

    // Redisからカートのアイテム一覧を取得する。
    if (token?.value) {
        items = await kv.get<Cart[]>(token.value) ?? [];
    }

    return (
        <div>
            {!token && <SessionCookie/>}
            <div>
                {items.map((item) => (
                    <div key={item.id}>
                        アイテムID {item.id} : {item.quantity}個
                    </div>
                ))}
            </div>
        </div>
    )
}

Discussion

TFTF

読んでいただきありがとうございます。サーバレスはスケールしてしまうので、従量課金だとヒヤヒヤです。個人で使うにはやはりインスタンスを自分で立ち上げた方が安心感はあるかなと感じます。