feat(sites): cross-process site rate limit via PG token bucket

Migration 0041_rate_limit adds rate_limit_bucket. PgTokenBucket (@umt/db)
refills and takes tokens under SELECT ... FOR UPDATE on the DB clock and
polls with jitter within a wait budget; penalize() carries Retry-After.
Site adapters get createSharedSiteRateLimiter (key site:<id>), shared by
api / worker / worker:agent, with an in-process serial queue in front and
a local fallback when the bucket is unavailable.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

packages/db/migrations/0041_rate_limit.ts

@@ -0,0 +1,30 @@
+import type { Kysely } from "kysely";
+import { dropTables, execDdl } from "../src/migration-helpers.ts";
+
+// C2-A 跨进程站点限流:PG 令牌桶(api / worker / worker:agent 共用)。
+// - key:`site:<站点 id>` 等;tokens 为上次结算时的令牌数(可为负,表示被 Retry-After 顺延);
+// - capacity / refill_per_sec 由调用方每次扣令牌时带上(站点限流设置变化即时生效);
+// - blocked_until:站点要求暂停(429 Retry-After)时在此之前不发放令牌。
+// 扣令牌走 `select … for update`,补充按 DB 时钟(clock_timestamp)结算,多进程时钟不一致也不影响。
+
+// biome-ignore lint/suspicious/noExplicitAny: 迁移不依赖应用表类型
+export async function up(db: Kysely<any>): Promise<void> {
+  await execDdl(
+    db,
+    `
+create table rate_limit_bucket (
+  key text primary key,
+  tokens double precision not null,
+  capacity double precision not null check (capacity > 0),
+  refill_per_sec double precision not null check (refill_per_sec > 0),
+  blocked_until timestamptz,
+  updated_at timestamptz not null default now()
+);
+`,
+  );
+}
+
+// biome-ignore lint/suspicious/noExplicitAny: 同上
+export async function down(db: Kysely<any>): Promise<void> {
+  await execDdl(db, dropTables(["rate_limit_bucket"]));
+}

packages/db/src/index.ts

@@ -2,4 +2,11 @@
 export { type Kysely, sql, type Transaction } from "kysely";
 export { type CreateDbOptions, createDb } from "./client.ts";
 export { createMigrator, MIGRATIONS_DIR, type MigratorOptions, migrateToLatest, resetSchema } from "./migrator.ts";
+export {
+  type AcquireOptions,
+  type AcquireResult,
+  type BucketSpec,
+  PgTokenBucket,
+  pgTokenBucketFor,
+} from "./rate-limit.ts";
 export type * from "./types.ts";

packages/db/src/rate-limit.ts

@@ -0,0 +1,127 @@
+import { type Kysely, sql } from "kysely";
+import type { DB } from "./types.ts";
+
+/**
+ * 跨进程令牌桶(rate_limit_bucket,迁移 0041)。api / worker / worker:agent 三个进程对同一 key 共用一个桶:
+ * 每次扣令牌在事务里 `select … for update` 锁住该行,按 DB 时钟补充后扣减;拿不到则算出还差多久,
+ * 调用方在 waitMs 预算内睡到那时(带抖动)再试。进程内的突发由调用方的本地串行队列挡住(见 adapters 的站点限流器)。
+ */
+
+export interface BucketSpec {
+  /** 桶容量(突发上限),≥ 1 */
+  capacity: number;
+  /** 每秒补充的令牌数,> 0 */
+  refillPerSec: number;
+}
+
+export interface AcquireOptions extends BucketSpec {
+  /** 最多等多久(毫秒);0 = 只试一次。缺省 60s */
+  waitMs?: number;
+  signal?: AbortSignal;
+}
+
+export interface AcquireResult {
+  ok: boolean;
+  /** 本次累计等待(毫秒) */
+  waitedMs: number;
+}
+
+// biome-ignore lint/suspicious/noExplicitAny: 测试会传入隔离 schema 的实例
+type AnyDb = Kysely<DB> | Kysely<any>;
+
+function sleep(ms: number, signal?: AbortSignal): Promise<void> {
+  return new Promise((resolve, reject) => {
+    if (signal?.aborted) return reject(signal.reason);
+    const t = setTimeout(() => {
+      signal?.removeEventListener("abort", onAbort);
+      resolve();
+    }, ms);
+    const onAbort = () => {
+      clearTimeout(t);
+      reject(signal?.reason);
+    };
+    signal?.addEventListener("abort", onAbort, { once: true });
+  });
+}
+
+export class PgTokenBucket {
+  constructor(
+    private readonly db: AnyDb,
+    private readonly opts: { jitterMs?: number; random?: () => number } = {},
+  ) {}
+
+  /**
+   * 试扣一次:成功返回 0,否则返回还需等待的毫秒数(被 penalize 顺延时为到期时间)。
+   * 行不存在时按满桶创建;容量 / 速率以本次传入为准(设置改了即时生效)。
+   */
+  async tryTake(key: string, cost: number, spec: BucketSpec): Promise<number> {
+    const capacity = Math.max(1, spec.capacity);
+    const rate = Math.max(1e-6, spec.refillPerSec);
+    const need = Math.min(cost, capacity);
+    return await this.db.transaction().execute(async (trx) => {
+      await sql`insert into rate_limit_bucket (key, tokens, capacity, refill_per_sec)
+        values (${key}, ${capacity}, ${capacity}, ${rate}) on conflict (key) do nothing`.execute(trx);
+      const row = (
+        await sql<{ tokens: number; elapsed: number; blocked_ms: number | null }>`
+          select tokens,
+                 greatest(0, extract(epoch from clock_timestamp() - updated_at))::float8 as elapsed,
+                 (extract(epoch from blocked_until - clock_timestamp()) * 1000)::float8 as blocked_ms
+            from rate_limit_bucket where key = ${key} for update`.execute(trx)
+      ).rows[0];
+      if (!row) return 1000;
+      let tokens = Math.min(capacity, Number(row.tokens) + Number(row.elapsed) * rate);
+      const blocked = row.blocked_ms !== null && Number(row.blocked_ms) > 0 ? Math.ceil(Number(row.blocked_ms)) : 0;
+      let wait = 0;
+      if (blocked > 0) wait = blocked;
+      else if (tokens >= need) tokens -= need;
+      else wait = Math.ceil(((need - tokens) / rate) * 1000);
+      await sql`update rate_limit_bucket set tokens = ${tokens}, capacity = ${capacity}, refill_per_sec = ${rate},
+          updated_at = clock_timestamp(),
+          blocked_until = case when blocked_until <= clock_timestamp() then null else blocked_until end
+        where key = ${key}`.execute(trx);
+      return wait;
+    });
+  }
+
+  /** 扣 cost 个令牌;waitMs 内拿不到返回 ok:false。signal 中止时抛出其 reason。 */
+  async acquire(key: string, cost: number, opts: AcquireOptions): Promise<AcquireResult> {
+    const budget = Math.max(0, opts.waitMs ?? 60_000);
+    const jitterMax = this.opts.jitterMs ?? 50;
+    const random = this.opts.random ?? Math.random;
+    const t0 = Date.now();
+    for (;;) {
+      opts.signal?.throwIfAborted();
+      const wait = await this.tryTake(key, cost, opts);
+      const spent = Date.now() - t0;
+      if (wait <= 0) return { ok: true, waitedMs: spent };
+      if (spent + wait > budget) return { ok: false, waitedMs: spent };
+      // 抖动:多个进程同时醒来时错开,避免一起撞锁再一起落空
+      const jitter = Math.floor(random() * Math.min(jitterMax, Math.max(1, wait * 0.2)));
+      await sleep(wait + jitter, opts.signal);
+    }
+  }
+
+  /** 站点要求暂停(Retry-After):ms 内不发放令牌,已有令牌清零 */
+  async penalize(key: string, ms: number, spec: BucketSpec): Promise<void> {
+    const capacity = Math.max(1, spec.capacity);
+    const rate = Math.max(1e-6, spec.refillPerSec);
+    const until = sql`clock_timestamp() + make_interval(secs => ${Math.max(0, ms) / 1000})`;
+    await sql`insert into rate_limit_bucket (key, tokens, capacity, refill_per_sec, blocked_until, updated_at)
+      values (${key}, 0, ${capacity}, ${rate}, ${until}, clock_timestamp())
+      on conflict (key) do update set tokens = least(rate_limit_bucket.tokens, 0),
+        blocked_until = greatest(coalesce(rate_limit_bucket.blocked_until, clock_timestamp()), ${until}),
+        updated_at = clock_timestamp()`.execute(this.db);
+  }
+}
+
+const buckets = new WeakMap<object, PgTokenBucket>();
+
+/** 同一个 Kysely 实例共用一个 PgTokenBucket */
+export function pgTokenBucketFor(db: AnyDb): PgTokenBucket {
+  let b = buckets.get(db);
+  if (!b) {
+    b = new PgTokenBucket(db);
+    buckets.set(db, b);
+  }
+  return b;
+}

packages/db/src/types.ts

@@ -1156,6 +1156,15 @@ export interface ProcessHeartbeatSampleTable {
   sse_connections: number | null;
 }
 
+export interface RateLimitBucketTable {
+  key: string;
+  tokens: number;
+  capacity: number;
+  refill_per_sec: number;
+  blocked_until: Timestamptz | null;
+  updated_at: GeneratedTimestamptz;
+}
+
 export interface ReconcileRunTable {
   id: Generated<string>;
   kind: ReconcileRunKind;
@@ -1653,6 +1662,7 @@ export interface DB {
   preference_tag: PreferenceTagTable;
   process_heartbeat: ProcessHeartbeatTable;
   process_heartbeat_sample: ProcessHeartbeatSampleTable;
+  rate_limit_bucket: RateLimitBucketTable;
   reconcile_run: ReconcileRunTable;
   resolution_attempt: ResolutionAttemptTable;
   rule_atom: RuleAtomTable;
@@ -1837,6 +1847,9 @@ export type ProcessHeartbeatUpdate = Updateable<ProcessHeartbeatTable>;
 export type ProcessHeartbeatSampleRow = Selectable<ProcessHeartbeatSampleTable>;
 export type NewProcessHeartbeatSample = Insertable<ProcessHeartbeatSampleTable>;
 export type ProcessHeartbeatSampleUpdate = Updateable<ProcessHeartbeatSampleTable>;
+export type RateLimitBucketRow = Selectable<RateLimitBucketTable>;
+export type NewRateLimitBucket = Insertable<RateLimitBucketTable>;
+export type RateLimitBucketUpdate = Updateable<RateLimitBucketTable>;
 export type ReconcileRunRow = Selectable<ReconcileRunTable>;
 export type NewReconcileRun = Insertable<ReconcileRunTable>;
 export type ReconcileRunUpdate = Updateable<ReconcileRunTable>;

packages/db/test/fixtures/bucket-worker.ts

@@ -0,0 +1,21 @@
+// 子进程:对同一 key 连续扣 N 个令牌,每次成功输出一行时间戳(毫秒)。由 rate-limit.test.ts 以 Deno 子进程启动。
+import { Kysely, PostgresDialect } from "kysely";
+import pg from "pg";
+import { PgTokenBucket } from "../../src/rate-limit.ts";
+
+const [url, schema, key, n, capacity, rate] = Deno.args as [string, string, string, string, string, string];
+const db = new Kysely<unknown>({
+  dialect: new PostgresDialect({
+    pool: new pg.Pool({ connectionString: url, max: 1, options: `-c search_path=${schema},public` }),
+  }),
+});
+const bucket = new PgTokenBucket(db);
+try {
+  for (let i = 0; i < Number(n); i++) {
+    const r = await bucket.acquire(key, 1, { capacity: Number(capacity), refillPerSec: Number(rate), waitMs: 30_000 });
+    if (!r.ok) throw new Error("等待超时");
+    console.log(Date.now());
+  }
+} finally {
+  await db.destroy();
+}

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

@@ -0,0 +1,98 @@
+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);
+});

packages/db/test/schema.test.ts

@@ -115,6 +115,8 @@ const EXPECTED_TABLES = [
   "admin_recovery_code",
   "admin_passkey",
   "webauthn_challenge",
+  // C2-A 跨进程站点限流(0041)
+  "rate_limit_bucket",
 ].sort();
 
 /** 不走 entityTable() 的表:追加型/引擎自有表,没有 uuid id + created_at/updated_at 三件套 */
@@ -139,6 +141,8 @@ const NON_ENTITY_TABLES = [
   // 一次性凭据:只记创建 / 用掉的时间
   "admin_recovery_code",
   "webauthn_challenge",
+  // 限流桶:key 主键,只有 updated_at
+  "rate_limit_bucket",
 ];
 
 describe.skipIf(skip)("schema migrations (real PG)", () => {