diff --git a/apps/server/src/modules/sources-identity/runtime.ts b/apps/server/src/modules/sources-identity/runtime.ts index b4062a73..a62e0021 100644 --- a/apps/server/src/modules/sources-identity/runtime.ts +++ b/apps/server/src/modules/sources-identity/runtime.ts @@ -1,5 +1,6 @@ import { createRssFetcher, + createSharedSiteRateLimiter, createSiteAdapter, createSiteRateLimiter, createTmdbClient, @@ -12,10 +13,10 @@ import { type TmdbClient, } from "@umt/adapters"; import { DEFAULT_SETTINGS, type SitesSettings } from "@umt/contracts"; -import type { SiteRow } from "@umt/db"; +import { pgTokenBucketFor, type SiteRow } from "@umt/db"; import type { ServerContext } from "../types.ts"; -// 模块运行时依赖:站点适配器与限流器(每站一个实例)、RSS 抓取器、TMDB 客户端、时钟。 +// 模块运行时依赖:站点适配器与限流器(每站一个实例,令牌在 PG 跨进程共享)、RSS 抓取器、TMDB 客户端、时钟。 // 测试通过 createSourcesIdentityModule(overrides) 注入 fetch / TMDB / 时钟。 export interface RuntimeOverrides { @@ -28,6 +29,8 @@ export interface RuntimeOverrides { rss?: Parameters[0]; /** 调度 tick 间隔(毫秒),默认 30s */ schedulerIntervalMs?: number; + /** false = 站点限流只用进程内令牌桶(缺省走 PG 跨进程令牌桶 rate_limit_bucket) */ + sharedRateLimit?: boolean; } interface CachedAdapter { @@ -158,7 +161,21 @@ export class SourcesRuntime { ]); const hit = this.adapters.get(site.id); if (hit && hit.sig === sig) return hit.adapter; - const limiter = createSiteRateLimiter(limiterOptions(def, site.rate_limit, defaults.defaultRateLimit)); + const lo = limiterOptions(def, site.rate_limit, defaults.defaultRateLimit); + // 跨进程:api / worker / worker:agent 对同一站点共用 PG 令牌桶(key = site:);进程内先串行排队 + const limiter = + this.overrides.sharedRateLimit === false + ? createSiteRateLimiter(lo) + : createSharedSiteRateLimiter({ + ...lo, + key: `site:${site.id}`, + bucket: pgTokenBucketFor(this.ctx.db), + onError: (err) => + this.ctx.logger.warn( + { siteId: site.id, err: { message: (err as Error).message } }, + "站点共享限流不可用,本次退回进程内限流", + ), + }); const cfg = (site.config ?? {}) as { apiBaseUrl?: string }; const baseUrl = def.umt.kind === "mteam" ? cfg.apiBaseUrl : sameUrl(site.base_url, def.domain) ? undefined : site.base_url; diff --git a/apps/server/test/sources-identity/api.test.ts b/apps/server/test/sources-identity/api.test.ts index 8e4693d8..7b099684 100644 --- a/apps/server/test/sources-identity/api.test.ts +++ b/apps/server/test/sources-identity/api.test.ts @@ -131,6 +131,19 @@ describe.skipIf(!up)("来源管理与站点工作台 API", () => { expect(observed.length).toBe(r.newCount); const site = await json(await h.req(`/sites/${siteId}`)); expect(site.body.data.lastSnapshotAt).not.toBeNull(); + // 限流间隔 > 0 时令牌在 PG 跨进程共享(rate_limit_bucket,key = site:);0 = 不限流只走本地 + await h.db + .updateTable("site") + .set({ rate_limit: JSON.stringify({ minIntervalMs: 5, concurrency: 5 }) }) + .where("id", "=", siteId) + .execute(); + expect((await snapshotSite(h.rt, { siteId, kind: "manual" })).status).toBe("ok"); + const bucket = await h.db + .selectFrom("rate_limit_bucket") + .selectAll() + .where("key", "=", `site:${siteId}`) + .executeTakeFirst(); + expect(bucket).toMatchObject({ capacity: 5, refill_per_sec: 200 }); }); it("工作台列表:快照优先,标注来源与更新时间;筛选与分页", async () => { diff --git a/packages/adapters/src/sites/index.ts b/packages/adapters/src/sites/index.ts index ff32bdcd..f6ff55a7 100644 --- a/packages/adapters/src/sites/index.ts +++ b/packages/adapters/src/sites/index.ts @@ -38,7 +38,14 @@ export { mteamSearchParams, parseMteamResult, } from "./mteam/spider.ts"; -export { createSiteRateLimiter, type RateLimiterOptions, type SiteRateLimiter } from "./rate-limit.ts"; +export { + createSharedSiteRateLimiter, + createSiteRateLimiter, + type RateLimiterOptions, + type SharedBucketPort, + type SharedRateLimiterOptions, + type SiteRateLimiter, +} from "./rate-limit.ts"; export { createSiteAdapter, getSiteAdapter, diff --git a/packages/adapters/src/sites/rate-limit.test.ts b/packages/adapters/src/sites/rate-limit.test.ts index a92be361..fb5d98b7 100644 --- a/packages/adapters/src/sites/rate-limit.test.ts +++ b/packages/adapters/src/sites/rate-limit.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { createSiteRateLimiter } from "./rate-limit.ts"; +import { createSharedSiteRateLimiter, createSiteRateLimiter, type SharedBucketPort } from "./rate-limit.ts"; import { parseSizeMp as parseSize } from "./spider/spider.ts"; import { parseCount, parseSiteTime } from "./util.ts"; @@ -34,6 +34,52 @@ describe("站点令牌桶", () => { }); }); +describe("跨进程站点限流器", () => { + it("同进程串行去共享桶扣令牌,key 与速率按站点参数;Retry-After 同时顺延共享桶", async () => { + const calls: { key: string; capacity: number; refillPerSec: number }[] = []; + let inflight = 0; + let maxInflight = 0; + const penalties: number[] = []; + const bucket: SharedBucketPort = { + async acquire(key, _cost, o) { + inflight++; + maxInflight = Math.max(maxInflight, inflight); + await new Promise((r) => setTimeout(r, 5)); + inflight--; + calls.push({ key, capacity: o.capacity, refillPerSec: o.refillPerSec }); + return { ok: true, waitedMs: 5 }; + }, + async penalize(_k, ms) { + penalties.push(ms); + }, + }; + const l = createSharedSiteRateLimiter({ key: "site:x", bucket, minIntervalMs: 500, burst: 2 }); + await Promise.all([l.acquire(), l.acquire(), l.acquire()]); + expect(maxInflight).toBe(1); + expect(calls).toEqual(Array(3).fill({ key: "site:x", capacity: 2, refillPerSec: 2 })); + expect(l.stats.granted).toBe(3); + l.penalize(3000); + expect(penalties).toEqual([3000]); + }); + + it("共享桶不可用时退回进程内令牌桶", async () => { + const errors: unknown[] = []; + const bucket: SharedBucketPort = { + acquire: () => Promise.reject(new Error("db down")), + penalize: () => Promise.reject(new Error("db down")), + }; + const l = createSharedSiteRateLimiter({ + key: "site:y", + bucket, + minIntervalMs: 10, + onError: (e) => errors.push(e), + }); + await l.acquire(); + expect(l.stats.granted).toBe(1); + expect(errors).toHaveLength(1); + }); +}); + describe("解析辅助", () => { it("parseSize/parseCount/parseSiteTime", () => { expect(parseSize("78.42 GB")).toBe(84202833838); diff --git a/packages/adapters/src/sites/rate-limit.ts b/packages/adapters/src/sites/rate-limit.ts index 2ae5f81b..0cb2b4c4 100644 --- a/packages/adapters/src/sites/rate-limit.ts +++ b/packages/adapters/src/sites/rate-limit.ts @@ -93,3 +93,81 @@ export function createSiteRateLimiter(opts: RateLimiterOptions): SiteRateLimiter }, }; } + +/** 跨进程令牌桶端口(@umt/db 的 PgTokenBucket 结构上满足;adapters 不依赖 db 包) */ +export interface SharedBucketPort { + acquire( + key: string, + cost: number, + opts: { capacity: number; refillPerSec: number; waitMs?: number; signal?: AbortSignal }, + ): Promise<{ ok: boolean; waitedMs: number }>; + penalize(key: string, ms: number, spec: { capacity: number; refillPerSec: number }): Promise; +} + +export interface SharedRateLimiterOptions extends RateLimiterOptions { + /** 桶 key,站点用 `site:` */ + key: string; + bucket: SharedBucketPort; + /** 单轮等待预算,缺省 60s;超出后重新排队(直到 signal 中止) */ + waitMs?: number; + /** 共享桶出错(DB 不可达等)时回调;本次退回进程内令牌桶,不让站点请求失败 */ + onError?: (err: unknown) => void; +} + +/** + * 跨进程站点限流器:进程内先串行排队(同进程并发不去 PG 抢锁),再到共享桶扣令牌; + * 429 Retry-After 同时顺延本地与共享桶。minIntervalMs = 0(不限流)时只走本地。 + */ +export function createSharedSiteRateLimiter(opts: SharedRateLimiterOptions): SiteRateLimiter { + const local = createSiteRateLimiter(opts); + const interval = Math.max(0, opts.minIntervalMs); + const spec = { capacity: Math.max(1, opts.burst ?? 1), refillPerSec: interval > 0 ? 1000 / interval : 0 }; + const stats = { granted: 0, waitedMs: 0 }; + let queue: Promise = Promise.resolve(); + const onError = opts.onError ?? (() => {}); + + const viaLocal = async (signal?: AbortSignal) => { + const before = local.stats.waitedMs; + await local.acquire(signal); + stats.waitedMs += local.stats.waitedMs - before; + }; + + return { + stats, + tryAcquire() { + if (!local.tryAcquire()) return false; + stats.granted++; + return true; + }, + penalize(ms) { + local.penalize(ms); + if (interval > 0) opts.bucket.penalize(opts.key, ms, spec).catch(onError); + }, + acquire(signal) { + const run = queue.then(async () => { + if (interval === 0) { + await viaLocal(signal); + } else { + try { + for (;;) { + const r = await opts.bucket.acquire(opts.key, 1, { + ...spec, + waitMs: opts.waitMs ?? 60_000, + ...(signal ? { signal } : {}), + }); + stats.waitedMs += r.waitedMs; + if (r.ok) break; + } + } catch (e) { + if (signal?.aborted) throw e; + onError(e); + await viaLocal(signal); + } + } + stats.granted++; + }); + queue = run.catch(() => {}); + return run; + }, + }; +} diff --git a/packages/db/migrations/0041_rate_limit.ts b/packages/db/migrations/0041_rate_limit.ts new file mode 100644 index 00000000..854def31 --- /dev/null +++ b/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): Promise { + 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): Promise { + await execDdl(db, dropTables(["rate_limit_bucket"])); +} diff --git a/packages/db/src/index.ts b/packages/db/src/index.ts index fb8141c8..05b34b56 100644 --- a/packages/db/src/index.ts +++ b/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"; diff --git a/packages/db/src/rate-limit.ts b/packages/db/src/rate-limit.ts new file mode 100644 index 00000000..e49ef0e2 --- /dev/null +++ b/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 | Kysely; + +function sleep(ms: number, signal?: AbortSignal): Promise { + 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 { + 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 { + 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 { + 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(); + +/** 同一个 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; +} diff --git a/packages/db/src/types.ts b/packages/db/src/types.ts index 985e168f..f76d35e3 100644 --- a/packages/db/src/types.ts +++ 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; 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; export type ProcessHeartbeatSampleRow = Selectable; export type NewProcessHeartbeatSample = Insertable; export type ProcessHeartbeatSampleUpdate = Updateable; +export type RateLimitBucketRow = Selectable; +export type NewRateLimitBucket = Insertable; +export type RateLimitBucketUpdate = Updateable; export type ReconcileRunRow = Selectable; export type NewReconcileRun = Insertable; export type ReconcileRunUpdate = Updateable; diff --git a/packages/db/test/fixtures/bucket-worker.ts b/packages/db/test/fixtures/bucket-worker.ts new file mode 100644 index 00000000..06168fed --- /dev/null +++ b/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({ + 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(); +} diff --git a/packages/db/test/rate-limit.test.ts b/packages/db/test/rate-limit.test.ts new file mode 100644 index 00000000..4156fd8f --- /dev/null +++ b/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; + let db: Kysely; + + 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({ + 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); +}); diff --git a/packages/db/test/schema.test.ts b/packages/db/test/schema.test.ts index 1962e803..6c733c36 100644 --- a/packages/db/test/schema.test.ts +++ b/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)", () => {