/** * 站点级令牌桶:同一站点的列表/搜索/详情/下载(以及将来的浏览器操作)共享一个实例。 * 容量 burst,每 minIntervalMs 补 1 个令牌;429/Retry-After 通过 penalize 整体顺延。 */ export interface SiteRateLimiter { acquire(signal?: AbortSignal): Promise; /** 不等待:有令牌返回 true 并消耗 */ tryAcquire(): boolean; /** 站点要求暂停(Retry-After 等),在 ms 内不发放令牌 */ penalize(ms: number): void; readonly stats: { granted: number; waitedMs: number }; } export interface RateLimiterOptions { minIntervalMs: number; burst?: number; now?: () => number; sleep?: (ms: number, signal?: AbortSignal) => Promise; } function defaultSleep(ms: number, signal?: AbortSignal): Promise { return new Promise((resolve, reject) => { if (signal?.aborted) { reject(signal.reason); return; } const t = setTimeout(() => { signal?.removeEventListener("abort", onAbort); resolve(); }, ms); const onAbort = () => { clearTimeout(t); reject(signal?.reason); }; signal?.addEventListener("abort", onAbort, { once: true }); }); } export function createSiteRateLimiter(opts: RateLimiterOptions): SiteRateLimiter { const interval = Math.max(0, opts.minIntervalMs); const burst = Math.max(1, opts.burst ?? 1); const now = opts.now ?? Date.now; const sleep = opts.sleep ?? defaultSleep; let tokens = burst; let last = now(); let blockedUntil = 0; let queue: Promise = Promise.resolve(); const stats = { granted: 0, waitedMs: 0 }; const refill = () => { const t = now(); if (interval === 0) tokens = burst; else tokens = Math.min(burst, tokens + (t - last) / interval); last = t; }; const waitMs = (): number => { refill(); const t = now(); if (t < blockedUntil) return blockedUntil - t; if (tokens >= 1) return 0; return Math.ceil((1 - tokens) * interval); }; return { stats, tryAcquire() { if (waitMs() > 0) return false; tokens -= 1; stats.granted++; return true; }, penalize(ms) { blockedUntil = Math.max(blockedUntil, now() + ms); tokens = Math.min(tokens, 0); }, acquire(signal) { // 串行排队,保证先到先得、不会并发抢同一个令牌 const run = queue.then(async () => { for (;;) { signal?.throwIfAborted(); const w = waitMs(); if (w <= 0) break; stats.waitedMs += w; await sleep(w, signal); } tokens -= 1; stats.granted++; }); queue = run.catch(() => {}); return run; }, }; } /** 跨进程令牌桶端口(@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; }, }; }