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

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;
}