From ea66e470294cc1cf439d8ff70b4997e29b59c8ea Mon Sep 17 00:00:00 2001 From: Liam Phung Date: Fri, 16 Oct 2026 20:05:28 +0700 Subject: [PATCH] feat(services): start quota threshold notifier job at server boot Closes #1212 --- src/config/env.ts | 16 +++++ src/config/index.ts | 6 ++ src/index.ts | 28 ++++++++ src/services/quotaNotifier.test.ts | 94 +++++++++++++++++++++++++ src/services/quotaNotifier.ts | 107 +++++++++++++++++++++++++---- 5 files changed, 239 insertions(+), 12 deletions(-) diff --git a/src/config/env.ts b/src/config/env.ts index 95523758..58c7de27 100644 --- a/src/config/env.ts +++ b/src/config/env.ts @@ -461,6 +461,22 @@ export const envSchema = z .positive() .default(86400000), + // Quota threshold notifier job (issue #1212) + QUOTA_NOTIFIER_ENABLED: z + .enum(["true", "false"]) + .default("true") + .transform((v) => v === "true"), + QUOTA_NOTIFIER_INTERVAL_MS: z.coerce + .number() + .int() + .positive() + .default(60_000), + QUOTA_DEFAULT_MONTHLY_LIMIT: z.coerce + .number() + .int() + .nonnegative() + .default(0), + // ──────────────────────────────────────────────────────────────────────── // SLO burn-rate alerting (issue #706) // diff --git a/src/config/index.ts b/src/config/index.ts index 9094b5fd..0f8daad0 100644 --- a/src/config/index.ts +++ b/src/config/index.ts @@ -304,6 +304,12 @@ export const config = { intervalMs: env.MONTHLY_INVOICE_JOB_INTERVAL_MS, }, + quotaNotifier: { + enabled: env.QUOTA_NOTIFIER_ENABLED, + intervalMs: env.QUOTA_NOTIFIER_INTERVAL_MS, + defaultMonthlyLimit: env.QUOTA_DEFAULT_MONTHLY_LIMIT, + }, + sloAlert: { enabled: Boolean(env.SLO_ALERT_WEBHOOK_URL) && env.SLO_ROUTE_CONFIGS.length > 0, diff --git a/src/index.ts b/src/index.ts index b5b90c95..2bb64c99 100644 --- a/src/index.ts +++ b/src/index.ts @@ -58,6 +58,11 @@ import { import { createSloAlertJob } from "./workers/sloAlertJob.js"; import { createMonthlyInvoiceJob } from "./workers/monthlyInvoiceJob.js"; import { createSettlementReconWorker } from "./workers/settlementRecon.js"; +import { + createQuotaNotifierJob, + PgQuotaNotificationStore, + loadDeveloperQuotasFromDb, +} from "./services/quotaNotifier.js"; import { createDeveloperRouter } from './routes/developerRoutes.js'; import { createGatewayRouter } from './routes/gatewayRoutes.js'; import { createProxyRouter } from './routes/proxyRoutes.js'; @@ -239,6 +244,18 @@ if (isDirectExecution) { intervalMs: config.monthlyInvoiceJob.intervalMs, }); + const quotaNotificationStore = new PgQuotaNotificationStore(pool); + const quotaNotifierJob = config.quotaNotifier.enabled + ? createQuotaNotifierJob(usageEventsRepository, quotaNotificationStore, { + intervalMs: config.quotaNotifier.intervalMs, + getDeveloperQuotas: () => + loadDeveloperQuotasFromDb( + pool, + config.quotaNotifier.defaultMonthlyLimit, + ), + }) + : null; + const sloAlertJob = config.sloAlert.enabled ? createSloAlertJob({ webhookUrl: config.sloAlert.webhookUrl!, @@ -357,6 +374,15 @@ if (isDirectExecution) { beginShutdown: () => settlementReconJob.beginShutdown(), awaitIdle: () => settlementReconJob.awaitIdle(), }, + ...(quotaNotifierJob + ? [ + { + name: "quota-notifier", + beginShutdown: () => quotaNotifierJob.beginShutdown(), + awaitIdle: () => quotaNotifierJob.awaitIdle(), + }, + ] + : []), ]; // Mount the refresh-token router. The drain middleware is applied inside the @@ -385,6 +411,7 @@ if (isDirectExecution) { slowQueryAlerterJob?.stop(); anomalyDetectorJob?.stop(); monthlyInvoiceJob.stop(); + quotaNotifierJob?.stop(); sloAlertJob?.stop(); await closeDb(); await Promise.allSettled([ @@ -427,6 +454,7 @@ if (isDirectExecution) { slowQueryAlerterJob?.start(); anomalyDetectorJob?.start(); monthlyInvoiceJob.start(); + quotaNotifierJob?.start(); sloAlertJob?.start(); const server = app.listen(PORT, () => { diff --git a/src/services/quotaNotifier.test.ts b/src/services/quotaNotifier.test.ts index 7eebcd5d..9b8ebb51 100644 --- a/src/services/quotaNotifier.test.ts +++ b/src/services/quotaNotifier.test.ts @@ -650,4 +650,98 @@ describe('createQuotaNotifierJob', () => { expect(getDeveloperQuotas).toHaveBeenCalledTimes(1); }); + + it('beginShutdown() stops future ticks and awaitIdle() waits for in-flight tick', async () => { + let resolveQuotas!: (quotas: DeveloperQuota[]) => void; + const getDeveloperQuotas = jest.fn( + () => + new Promise((resolve) => { + resolveQuotas = resolve; + }), + ); + const store = new InMemoryQuotaNotificationStore(); + + const job = createQuotaNotifierJob(makeUsageRepo([]), store, { + intervalMs: 1000, + getDeveloperQuotas, + }); + + job.start(); + jest.advanceTimersByTime(1000); + await Promise.resolve(); + + expect(getDeveloperQuotas).toHaveBeenCalledTimes(1); + + job.beginShutdown(); + let idleDone = false; + const idlePromise = job.awaitIdle().then(() => { + idleDone = true; + }); + + await Promise.resolve(); + expect(idleDone).toBe(false); + + resolveQuotas([]); + await idlePromise; + expect(idleDone).toBe(true); + + // No further ticks after beginShutdown + jest.advanceTimersByTime(5000); + await Promise.resolve(); + expect(getDeveloperQuotas).toHaveBeenCalledTimes(1); + }); }); + +// --------------------------------------------------------------------------- +// PgQuotaNotificationStore & loadDeveloperQuotasFromDb (Issue #1212) +// --------------------------------------------------------------------------- + +describe('PgQuotaNotificationStore & loadDeveloperQuotasFromDb', () => { + it('queries and marks sent notifications in Postgres table', async () => { + const { PgQuotaNotificationStore } = await import('./quotaNotifier.js'); + const rows = new Set(); + const fakeDb = { + query: jest.fn(async (sql: string, params?: unknown[]) => { + const [dev, period, threshold] = params ?? []; + const key = `${dev}|${period}|${threshold}`; + if (sql.includes('SELECT EXISTS')) { + return { rows: [{ exists: rows.has(key) }] }; + } + if (sql.includes('INSERT INTO quota_notifications_sent')) { + rows.add(key); + return { rows: [] }; + } + return { rows: [] }; + }), + }; + + const store = new PgQuotaNotificationStore(fakeDb as never); + expect(await store.hasBeenSent('dev_pg', '2026-06', 80)).toBe(false); + await store.markSent('dev_pg', '2026-06', 80); + expect(await store.hasBeenSent('dev_pg', '2026-06', 80)).toBe(true); + }); + + it('loads developer quotas from developers.plan_overrides fallback', async () => { + const { loadDeveloperQuotasFromDb } = await import('./quotaNotifier.js'); + const fakeDb = { + query: jest.fn(async (sql: string) => { + if (sql.includes('developer_quotas')) { + throw new Error('relation "developer_quotas" does not exist'); + } + return { + rows: [ + { user_id: 'dev_a', plan_overrides: JSON.stringify({ monthly_call_limit: 5000 }) }, + { user_id: 'dev_b', plan_overrides: null }, + ], + }; + }), + }; + + const quotas = await loadDeveloperQuotasFromDb(fakeDb as never, 1000); + expect(quotas).toEqual([ + { developerId: 'dev_a', monthlyLimit: 5000 }, + { developerId: 'dev_b', monthlyLimit: 1000 }, + ]); + }); +}); + diff --git a/src/services/quotaNotifier.ts b/src/services/quotaNotifier.ts index fa47b658..62212561 100644 --- a/src/services/quotaNotifier.ts +++ b/src/services/quotaNotifier.ts @@ -54,6 +54,8 @@ export interface QuotaNotifierOptions { export interface QuotaNotifierJob { start(): void; stop(): void; + beginShutdown(): void; + awaitIdle(): Promise; } // --------------------------------------------------------------------------- @@ -224,28 +226,48 @@ export function createQuotaNotifierJob( let timer: NodeJS.Timeout | null = null; let running = false; + let shuttingDown = false; + let activeTick: Promise | null = null; const tick = async (): Promise => { - if (running) return; + if (running || shuttingDown) return; running = true; - try { - await runQuotaCheck(deps); - } catch (err) { - logger.error('[quotaNotifier] Unexpected error during quota check:', err); - } finally { - running = false; - } + const runPromise = (async () => { + try { + await runQuotaCheck(deps); + } catch (err) { + logger.error('[quotaNotifier] Unexpected error during quota check:', err); + } finally { + running = false; + activeTick = null; + } + })(); + activeTick = runPromise; + await runPromise; + }; + + const stopTimer = (): void => { + if (!timer) return; + clearInterval(timer); + timer = null; }; return { start() { - if (timer) return; + if (timer || shuttingDown) return; timer = setInterval(() => { void tick(); }, options.intervalMs); }, stop() { - if (!timer) return; - clearInterval(timer); - timer = null; + stopTimer(); + }, + beginShutdown() { + shuttingDown = true; + stopTimer(); + }, + async awaitIdle() { + if (activeTick) { + await activeTick; + } }, }; } @@ -306,3 +328,64 @@ export class PgQuotaNotificationStore implements QuotaNotificationStore { ); } } + +/** + * Loads active developer monthly quotas from PostgreSQL. + * Checks `developer_quotas` first, then falls back to `developers.plan_overrides`. + */ +export async function loadDeveloperQuotasFromDb( + db: QuotaNotificationDb, + defaultMonthlyLimit = 0, +): Promise { + try { + const res = await db.query<{ + developer_id?: string; + id?: string; + monthly_limit?: number | string; + }>('SELECT developer_id, monthly_limit FROM developer_quotas WHERE monthly_limit > 0'); + if (Array.isArray(res.rows) && res.rows.length > 0) { + return res.rows + .map((row) => ({ + developerId: String(row.developer_id ?? row.id ?? ''), + monthlyLimit: Number(row.monthly_limit ?? 0), + })) + .filter((q) => q.developerId.length > 0 && Number.isFinite(q.monthlyLimit) && q.monthlyLimit > 0); + } + } catch { + // Fallback to developers table with plan_overrides JSON + } + + try { + const devRes = await db.query<{ + user_id: string; + plan_overrides?: string | Record | null; + }>('SELECT user_id, plan_overrides FROM developers'); + + const quotas: DeveloperQuota[] = []; + for (const row of devRes.rows ?? []) { + if (!row.user_id) continue; + let limit = defaultMonthlyLimit; + if (row.plan_overrides) { + try { + const parsed = + typeof row.plan_overrides === 'string' + ? (JSON.parse(row.plan_overrides) as Record) + : row.plan_overrides; + const rawLimit = parsed.monthly_call_limit ?? parsed.monthlyCallLimit ?? parsed.monthlyLimit; + if (rawLimit !== undefined && Number.isFinite(Number(rawLimit))) { + limit = Number(rawLimit); + } + } catch { + // Ignore malformed plan_overrides JSON + } + } + if (limit > 0) { + quotas.push({ developerId: row.user_id, monthlyLimit: limit }); + } + } + return quotas; + } catch { + return []; + } +} +