UHD_Media_Tool/packages/db/test/rate-limit.test.ts

import { fileURLToPath } from "node:url";
import { Kysely, PostgresDialect, sql } from "kysely";
import pg from "pg";
import { afterAll, beforeAll, describe, expect, it } from "vitest";
import { type DB, migrateToLatest, PgTokenBucket } from "../src/index.ts";

// 跨进程令牌桶(0041):单进程语义 + 三个 Deno 子进程并发对同一 key 的总速率不超限。
const url = process.env.DATABASE_URL ?? "postgres://umt:umt@localhost:55433/umt";
const skip = process.env.UMT_SKIP_DB_TESTS === "1";
const schema = `umt_rl_${process.pid}_${Date.now().toString(36)}`;

describe.skipIf(skip)("PgTokenBucket(真实 PG)", () => {
  let admin: Kysely<unknown>;
  let db: Kysely<DB>;

  beforeAll(async () => {
    admin = new Kysely({ dialect: new PostgresDialect({ pool: new pg.Pool({ connectionString: url, max: 1 }) }) });
    await sql.raw(`create schema ${schema}`).execute(admin);
    db = new Kysely<DB>({
      dialect: new PostgresDialect({
        pool: new pg.Pool({ connectionString: url, max: 4, options: `-c search_path=${schema},public` }),
      }),
    });
    await migrateToLatest(db, { migrationTableSchema: schema });
  });

  afterAll(async () => {
    await db?.destroy();
    await sql.raw(`drop schema if exists ${schema} cascade`).execute(admin);
    await admin?.destroy();
  });

  it("容量内立即发放,之后按速率补充;waitMs=0 拿不到即返回", async () => {
    const b = new PgTokenBucket(db);
    const spec = { capacity: 2, refillPerSec: 10 };
    expect(await b.tryTake("t:a", 1, spec)).toBe(0);
    expect(await b.tryTake("t:a", 1, spec)).toBe(0);
    const wait = await b.tryTake("t:a", 1, spec);
    expect(wait).toBeGreaterThan(0);
    expect(wait).toBeLessThanOrEqual(100);
    expect((await b.acquire("t:a", 1, { ...spec, waitMs: 0 })).ok).toBe(false);
    const r = await b.acquire("t:a", 1, { ...spec, waitMs: 2000 });
    expect(r.ok).toBe(true);
  });

  it("penalize 顺延:期间不发放令牌", async () => {
    const b = new PgTokenBucket(db);
    const spec = { capacity: 5, refillPerSec: 100 };
    await b.penalize("t:p", 400, spec);
    const wait = await b.tryTake("t:p", 1, spec);
    expect(wait).toBeGreaterThan(250);
    const t0 = Date.now();
    expect((await b.acquire("t:p", 1, { ...spec, waitMs: 3000 })).ok).toBe(true);
    expect(Date.now() - t0).toBeGreaterThanOrEqual(200);
  });

  it("取消等待", async () => {
    const b = new PgTokenBucket(db);
    const spec = { capacity: 1, refillPerSec: 0.1 };
    await b.tryTake("t:c", 1, spec);
    const ac = new AbortController();
    const p = b.acquire("t:c", 1, { ...spec, waitMs: 60_000, signal: ac.signal });
    setTimeout(() => ac.abort(new Error("cancelled")), 50);
    await expect(p).rejects.toThrow("cancelled");
  });

  it("三个子进程并发对同一 key:总速率不超过 容量 + 速率 × 时长", async () => {
    const script = fileURLToPath(new URL("./fixtures/bucket-worker.ts", import.meta.url));
    const capacity = 2;
    const rate = 20;
    const perProc = 12;
    const procs = [0, 1, 2].map(() =>
      new Deno.Command(Deno.execPath(), {
        args: ["run", "-A", script, url, schema, "t:procs", String(perProc), String(capacity), String(rate)],
        stdout: "piped",
        stderr: "piped",
      }).output(),
    );
    const outs = await Promise.all(procs);
    const stamps: number[] = [];
    for (const o of outs) {
      expect(o.code, new TextDecoder().decode(o.stderr)).toBe(0);
      for (const line of new TextDecoder().decode(o.stdout).split("\n")) if (line.trim()) stamps.push(Number(line));
    }
    stamps.sort((a, b) => a - b);
    expect(stamps).toHaveLength(perProc * 3);
    // 36 个令牌:首批 capacity 个立即发放,其余按 20/s → 至少 (36-2)/20 = 1.7s
    const span = stamps.at(-1)! - stamps[0]!;
    expect(span).toBeGreaterThanOrEqual(((perProc * 3 - capacity) / rate) * 1000 - 150);
    // 任意 500ms 窗口内发放数 ≤ 容量 + 速率 × 0.5s(+1 容忍子进程打印时间戳的误差)
    const win = 500;
    for (let i = 0; i < stamps.length; i++) {
      let j = i;
      while (j < stamps.length && stamps[j]! - stamps[i]! < win) j++;
      expect(j - i).toBeLessThanOrEqual(capacity + (rate * win) / 1000 + 1);
    }
  }, 60_000);
});