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]>
This commit is contained in:
1 parent
57346d438c
commit
61f585a707
12 files changed
+466
-5
No files matched your search
@@ -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<typeof createRssFetcher>[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:<id>);进程内先串行排队
|
||||
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;
|
||||
|
||||
@@ -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:<id>);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 () => {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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<void>;
|
||||
}
|
||||
|
||||
export interface SharedRateLimiterOptions extends RateLimiterOptions {
|
||||
/** 桶 key,站点用 `site:<id>` */
|
||||
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<void> = 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;
|
||||
},
|
||||
};
|
||||
}
|
||||
@@ -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"]));
|
||||
}
|
||||
@@ -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";
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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>;
|
||||
|
||||
+21
@@ -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();
|
||||
}
|
||||
@@ -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);
|
||||
});
|
||||
@@ -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)", () => {
|
||||
|
||||
Reference in new issue
Block a user