fix(cache): cut per-request DB+Redis load in hot path

apps/gateway/src/lib/cached-queries.ts

@@ -21,7 +21,6 @@ import {
 	ne,
 	or,
 	cdb as db,
-	db as uncachedDb,
 	apiKey as apiKeyTable,
 	apiKeyIamRule as apiKeyIamRuleTable,
 	discount as discountTable,
@@ -270,26 +269,45 @@ export async function findProjectById(
 	});
 }
 
+// TTL for the "fresh" credit/balance refetch below. A zero-credit org or
+// zero-balance wallet otherwise refetches on EVERY request; under high
+// throughput that is one Postgres SELECT per request (and the DB pool, max 20,
+// saturates). A short TTL still reflects topups/debits within FRESH_TTL_SECONDS
+// while collapsing per-request DB load to at most one query per window per row.
+const FRESH_TTL_SECONDS = 2;
+
 /**
- * Find an organization by ID without cache (for fresh credit checks)
+ * Find an organization by ID with a short-TTL fresh read (for near-fresh credit
+ * checks when the org shows <= 0 credits). Uses a distinct Drizzle cache tag AND
+ * a distinct SWR mirror key (`org:fresh:${id}`) so it does not collide with the
+ * longer-lived `findOrganizationById` entry. The distinct mirror key matters:
+ * sharing it would let the (possibly stale-zero) regular read claim the mirror
+ * write-throttle slot and suppress this fresh value's mirror write, so a DB
+ * outage right after a topup could keep serving the stale-zero fallback.
  */
-export async function findOrganizationByIdUncached(
+export async function findOrganizationByIdFresh(
 	id: string,
 ): Promise<Organization | undefined> {
-	return await swrWrap(`org:${id}`, [organizationTableName], async () => {
-		const results = await uncachedDb
+	return await swrWrap(`org:fresh:${id}`, [organizationTableName], async () => {
+		const results = await db
 			.select()
 			.from(organizationTable)
 			.where(eq(organizationTable.id, id))
-			.limit(1);
+			.limit(1)
+			.$withCache({
+				tag: `org-fresh:${id}`,
+				autoInvalidate: false,
+				config: { ex: FRESH_TTL_SECONDS },
+			});
 		return results[0];
 	});
 }
 
 /**
  * Find an organization by ID (cacheable)
- * When the organization has 0 credits, refetch without cache to ensure
- * no delay in reflecting topups and usage updates.
+ * When the organization has 0 credits, refetch via a short-TTL fresh read so
+ * topups and usage updates are reflected within FRESH_TTL_SECONDS without
+ * hitting Postgres on every request.
  */
 export async function findOrganizationById(
 	id: string,
@@ -303,8 +321,8 @@ export async function findOrganizationById(
 		return results[0];
 	});
 
-	// If org has 0 or negative credits, refetch without cache
-	// to ensure topups are reflected immediately
+	// If org has 0 or negative credits, refetch via the short-TTL fresh read
+	// so topups are reflected promptly without a per-request Postgres hit
 	if (org) {
 		const regularCredits = parseFloat(org.credits || "0");
 		const devPlanCreditsUsed = parseFloat(org.devPlanCreditsUsed || "0");
@@ -319,33 +337,42 @@ export async function findOrganizationById(
 			regularCredits + devPlanCreditsRemaining + chatPlanCreditsRemaining;
 
 		if (totalCredits <= 0) {
-			return await findOrganizationByIdUncached(id);
+			return await findOrganizationByIdFresh(id);
 		}
 	}
 
 	return org;
 }
 
 /**
- * Find an end-user wallet by ID without cache (for fresh balance checks)
+ * Find an end-user wallet by ID with a short-TTL fresh read (for near-fresh
+ * balance checks when the wallet shows <= 0 balance). Uses a distinct Drizzle
+ * cache tag AND a distinct SWR mirror key (`wallet:fresh:${id}`) so it does not
+ * collide with the longer-lived `findWalletById` entry — see
+ * `findOrganizationByIdFresh` for why the distinct mirror key matters.
  */
-export async function findWalletByIdUncached(
+export async function findWalletByIdFresh(
 	id: string,
 ): Promise<Wallet | undefined> {
-	return await swrWrap(`wallet:${id}`, [walletTableName], async () => {
-		const results = await uncachedDb
+	return await swrWrap(`wallet:fresh:${id}`, [walletTableName], async () => {
+		const results = await db
 			.select()
 			.from(walletTable)
 			.where(eq(walletTable.id, id))
-			.limit(1);
+			.limit(1)
+			.$withCache({
+				tag: `wallet-fresh:${id}`,
+				autoInvalidate: false,
+				config: { ex: FRESH_TTL_SECONDS },
+			});
 		return results[0];
 	});
 }
 
 /**
  * Find an end-user wallet by ID (cacheable). Mirrors findOrganizationById: when
- * the wallet balance is 0 or negative, refetch without cache so top-ups and
- * usage debits are reflected immediately.
+ * the wallet balance is 0 or negative, refetch via a short-TTL fresh read so
+ * top-ups and usage debits are reflected promptly without a per-request DB hit.
  */
 export async function findWalletById(id: string): Promise<Wallet | undefined> {
 	const w = await swrWrap(`wallet:${id}`, [walletTableName], async () => {
@@ -358,7 +385,7 @@ export async function findWalletById(id: string): Promise<Wallet | undefined> {
 	});
 
 	if (w && parseFloat(w.balance || "0") <= 0) {
-		return await findWalletByIdUncached(id);
+		return await findWalletByIdFresh(id);
 	}
 
 	return w;

packages/cache/src/swr.spec.ts

@@ -1,9 +1,10 @@
-import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 
 import { redisClient } from "./redis.js";
 import {
 	SWR_PREFIX,
 	SWR_TABLE_INDEX_PREFIX,
+	SWR_THROTTLE_PREFIX,
 	getSwrStaleTtlSeconds,
 	invalidateSwrByTables,
 	swrWrap,
@@ -107,6 +108,72 @@ describe("swrWrap", () => {
 		);
 	});
 
+	it("throttles mirror rewrites for the same key within the window", async () => {
+		await swrWrap("test:key:throttle", ["table_a"], () =>
+			Promise.resolve({ v: 1 }),
+		);
+		expect(
+			await redisClient.get(`${SWR_PREFIX}test:key:throttle`),
+		).not.toBeNull();
+
+		// Drop the mirror, then call again within the throttle window. The fresh
+		// fetcher value is still returned, but the mirror is NOT rewritten.
+		await redisClient.del(`${SWR_PREFIX}test:key:throttle`);
+		const value = await swrWrap("test:key:throttle", ["table_a"], () =>
+			Promise.resolve({ v: 2 }),
+		);
+		expect(value).toEqual({ v: 2 });
+		expect(await redisClient.get(`${SWR_PREFIX}test:key:throttle`)).toBeNull();
+	});
+
+	it("releases the throttle slot so the next call retries when the mirror write fails", async () => {
+		const realPipeline = redisClient.pipeline.bind(redisClient);
+		const pipelineSpy = vi
+			.spyOn(redisClient, "pipeline")
+			.mockImplementationOnce(() => {
+				const pipeline = realPipeline();
+				vi.spyOn(pipeline, "exec").mockRejectedValueOnce(
+					new Error("redis write failed"),
+				);
+				return pipeline;
+			});
+
+		await swrWrap("test:key:relfail", ["table_a"], () =>
+			Promise.resolve({ v: 1 }),
+		);
+
+		// Write failed, so neither the mirror nor a lingering throttle marker
+		// should remain — the slot must be released for a retry.
+		expect(await redisClient.get(`${SWR_PREFIX}test:key:relfail`)).toBeNull();
+		expect(
+			await redisClient.get(`${SWR_THROTTLE_PREFIX}test:key:relfail`),
+		).toBeNull();
+
+		pipelineSpy.mockRestore();
+
+		// The next call (still within the window) now succeeds in writing the
+		// mirror instead of being suppressed by a stuck throttle marker.
+		await swrWrap("test:key:relfail", ["table_a"], () =>
+			Promise.resolve({ v: 2 }),
+		);
+		expect(
+			await redisClient.get(`${SWR_PREFIX}test:key:relfail`),
+		).not.toBeNull();
+	});
+
+	it("repopulates mirror on next fetch after invalidation despite throttle", async () => {
+		await swrWrap("test:key:reinv", ["table_a"], () =>
+			Promise.resolve({ v: 1 }),
+		);
+		await invalidateSwrByTables(["table_a"]);
+		expect(await redisClient.get(`${SWR_PREFIX}test:key:reinv`)).toBeNull();
+
+		await swrWrap("test:key:reinv", ["table_a"], () =>
+			Promise.resolve({ v: 2 }),
+		);
+		expect(await redisClient.get(`${SWR_PREFIX}test:key:reinv`)).not.toBeNull();
+	});
+
 	it("honors SWR_STALE_TTL_SECONDS env var", async () => {
 		process.env.SWR_STALE_TTL_SECONDS = "120";
 		expect(getSwrStaleTtlSeconds()).toBe(120);

packages/cache/src/swr.ts

@@ -4,9 +4,55 @@ import { redisClient } from "./redis.js";
 
 export const SWR_PREFIX = "swr:";
 export const SWR_TABLE_INDEX_PREFIX = "swr:tables:";
+export const SWR_THROTTLE_PREFIX = "swr:throttle:";
 export const SWR_DEFAULT_TTL_SECONDS = 14400;
 export const SWR_BATCH_SIZE = 500;
 
+// The SWR mirror is only a fallback served when the underlying fetcher (a
+// Postgres query) throws. The happy path never reads it, so it does NOT need to
+// be rewritten on every successful request — doing so adds a Redis pipeline
+// (SET + SADD/EXPIRE per table) to every hot-path query and, at hundreds of
+// req/s with ~15-30 cached queries each, becomes a dominant source of Redis
+// load. We instead refresh a given key's mirror at most once per throttle
+// window, collapsing the per-request pipeline to a single conditional SET.
+//
+// The throttle marker lives in Redis (not in process memory) on purpose: the
+// thing it gates — the mirror — also lives in Redis and can disappear out from
+// under us (eviction, TTL, FLUSHDB, failover). An in-memory marker would
+// desync from that and suppress re-priming for a full window even though the
+// mirror is gone, silently removing the disaster fallback. A Redis-side marker
+// is cleared by the same events that drop the mirror, so the next request
+// re-primes immediately.
+export const SWR_MIRROR_WRITE_THROTTLE_SECONDS = 30;
+
+function throttleKey(key: string): string {
+	return SWR_THROTTLE_PREFIX + key;
+}
+
+// Returns true if this caller won the throttle slot and should (re)write the
+// mirror. Uses SET NX so exactly one caller per window per key wins. Fails open
+// (returns true) on Redis error so an unavailable Redis never suppresses a
+// write the caller would otherwise make.
+async function claimMirrorWrite(key: string): Promise<boolean> {
+	try {
+		const result = await redisClient.set(
+			throttleKey(key),
+			"1",
+			"EX",
+			SWR_MIRROR_WRITE_THROTTLE_SECONDS,
+			"NX",
+		);
+		return result === "OK";
+	} catch (error) {
+		logger.error(
+			"Error claiming SWR mirror write throttle",
+			error instanceof Error ? error : new Error(String(error)),
+			{ key },
+		);
+		return true;
+	}
+}
+
 const SWR_NONE_SENTINEL = "__swrNone" as const;
 
 interface SwrNoneSentinel {
@@ -45,7 +91,7 @@ async function writeMirror<T>(
 	key: string,
 	tables: string[],
 	value: T,
-): Promise<void> {
+): Promise<boolean> {
 	try {
 		const ttl = getSwrStaleTtlSeconds();
 		const cacheKey = swrKey(key);
@@ -60,12 +106,26 @@ async function writeMirror<T>(
 			pipeline.expire(indexKey, ttl + 60);
 		}
 		await pipeline.exec();
+		return true;
 	} catch (error) {
 		logger.error(
 			"Error writing SWR mirror",
 			error instanceof Error ? error : new Error(String(error)),
 			{ key },
 		);
+		return false;
+	}
+}
+
+// Release the throttle slot so the next request retries the write. Used when a
+// claimed write fails, so a transient Redis write error does not suppress the
+// mirror (and weaken the stale fallback) for the whole throttle window.
+async function releaseMirrorThrottle(key: string): Promise<void> {
+	try {
+		await redisClient.del(throttleKey(key));
+	} catch {
+		// Best-effort: the throttle key's own TTL bounds how long it can wrongly
+		// suppress writes, so a failed release is non-fatal.
 	}
 }
 
@@ -110,7 +170,12 @@ export async function swrWrap<T>(
 		throw error;
 	}
 
-	await writeMirror(key, tables, value);
+	if (await claimMirrorWrite(key)) {
+		const wrote = await writeMirror(key, tables, value);
+		if (!wrote) {
+			await releaseMirrorThrottle(key);
+		}
+	}
 	return value;
 }
 
@@ -134,9 +199,16 @@ export async function invalidateSwrByTables(tables: string[]): Promise<void> {
 			return;
 		}
 
+		// Also drop the write-throttle markers for the invalidated keys so the
+		// next fetch repopulates the mirror immediately instead of waiting out
+		// the throttle window.
 		const keysArray = Array.from(allKeysToDelete);
-		for (let i = 0; i < keysArray.length; i += SWR_BATCH_SIZE) {
-			const batch = keysArray.slice(i, i + SWR_BATCH_SIZE);
+		const throttleKeys = keysArray.map((cacheKey) =>
+			throttleKey(cacheKey.slice(SWR_PREFIX.length)),
+		);
+		const allUnlinkKeys = [...keysArray, ...throttleKeys];
+		for (let i = 0; i < allUnlinkKeys.length; i += SWR_BATCH_SIZE) {
+			const batch = allUnlinkKeys.slice(i, i + SWR_BATCH_SIZE);
 			if (batch.length > 0) {
 				await redisClient.unlink(...batch);
 			}