podcastdistributiona/lib/ratelimit/index.ts

import { Pool } from "pg";
import {
  RateLimiterMemory,
  RateLimiterPostgres,
  type RateLimiterAbstract,
} from "rate-limiter-flexible";

/**
 * Backend selection (default = in-memory):
 *
 *   - DEFAULT: RateLimiterMemory — per-process counters. Fine for a single
 *     instance (the Plesk VPS). No DB connection is ever opened in this mode.
 *   - OPT-IN: set RATE_LIMIT_BACKEND="postgres" *and* provide DATABASE_URL to
 *     share limits across multiple instances via RateLimiterPostgres. The
 *     limiter table ("rate_limits") is created automatically by
 *     rate-limiter-flexible on first use, so no migration is required.
 *
 * The pg Pool is created lazily on first consume — importing this module never
 * connects to the database, so memory mode stays connection-free.
 */
const usePostgres =
  process.env.RATE_LIMIT_BACKEND === "postgres" && !!process.env.DATABASE_URL;

const limiters = new Map<string, RateLimiterAbstract>();

// Lazily-created shared pg Pool. Constructing it does NOT open a connection
// (pg connects on first query), and it is only built when a Postgres-backed
// limiter is first needed — so memory mode never creates a pool.
let pgPool: Pool | null = null;
function getPool(): Pool {
  if (!pgPool) {
    pgPool = new Pool({ connectionString: process.env.DATABASE_URL });
  }
  return pgPool;
}

function getLimiter(name: string, points: number, durationSec: number): RateLimiterAbstract {
  let limiter = limiters.get(name);
  if (!limiter) {
    if (usePostgres) {
      // The constructor opens the (lazy) pool and ensures the backing table
      // exists; ready before the first consume() thanks to the insurance limiter.
      limiter = new RateLimiterPostgres({
        storeClient: getPool(),
        tableName: "rate_limits",
        keyPrefix: name,
        points,
        duration: durationSec,
        // If Postgres is briefly unreachable, degrade to per-process counting
        // instead of failing the request outright.
        insuranceLimiter: new RateLimiterMemory({ points, duration: durationSec }),
      });
    } else {
      limiter = new RateLimiterMemory({ points, duration: durationSec });
    }
    limiters.set(name, limiter);
  }
  return limiter;
}

export interface RateLimitResult {
  ok: boolean;
  retryAfterSec?: number;
}

/** Consume one unit for `key` under a named bucket; returns ok=false when throttled. */
export async function rateLimit(
  bucket: string,
  key: string,
  opts: { points: number; durationSec: number }
): Promise<RateLimitResult> {
  const limiter = getLimiter(bucket, opts.points, opts.durationSec);
  try {
    await limiter.consume(key);
    return { ok: true };
  } catch (res) {
    const retryAfterSec = Math.ceil((res as { msBeforeNext?: number }).msBeforeNext ?? 1000) / 1000;
    return { ok: false, retryAfterSec: Math.ceil(retryAfterSec) };
  }
}

// Common presets.
export const LIMITS = {
  generation: { points: 10, durationSec: 60 }, // 10 generations / min / user
  repurpose: { points: 15, durationSec: 60 },
  api: { points: 60, durationSec: 60 }, // 60 API calls / min / key (writes)
  read: { points: 120, durationSec: 60 }, // 120 read/list calls / min / key
  stream: { points: 30, durationSec: 60 }, // SSE (re)connects / min / user
  publicMedia: { points: 120, durationSec: 60 }, // anon audio/cover (Range) reqs / min / IP
} as const;