UHD_Media_Tool/packages/adapters/src/sites/rate-limit.ts

/**
 * 站点级令牌桶:同一站点的列表/搜索/详情/下载(以及将来的浏览器操作)共享一个实例。
 * 容量 burst,每 minIntervalMs 补 1 个令牌;429/Retry-After 通过 penalize 整体顺延。
 */

export interface SiteRateLimiter {
  acquire(signal?: AbortSignal): Promise<void>;
  /** 不等待:有令牌返回 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<void>;
}

function defaultSleep(ms: number, signal?: AbortSignal): Promise<void> {
  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<void> = 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<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;
    },
  };
}