chore: further redis cleanup

server/src/external/redis/initRedis.ts

@@ -28,6 +28,7 @@ import {
 	UPSERT_INVOICE_IN_CUSTOMER_SCRIPT,
 } from "../../_luaScriptsV2/luaScriptsV2.js";
 import { instrumentRedis } from "../../utils/otel/instrumentRedis.js";
+import { withTimeout } from "../../utils/withTimeout.js";
 
 // if (!process.env.CACHE_URL) {
 // 	throw new Error("CACHE_URL (redis) is not set");
@@ -82,31 +83,6 @@ let redisTickInFlight = false;
 
 let redisAvailabilityState: RedisAvailabilityState = "degraded";
 
-const withTimeout = async <T>({
-	timeoutMs,
-	fn,
-}: {
-	timeoutMs: number;
-	fn: () => Promise<T>;
-}): Promise<T> => {
-	let timeoutId: ReturnType<typeof setTimeout> | undefined;
-
-	try {
-		return await Promise.race([
-			fn(),
-			new Promise<never>((_, reject) => {
-				timeoutId = setTimeout(
-					() => reject(new Error(`timed out after ${timeoutMs}ms`)),
-					timeoutMs,
-				);
-				timeoutId.unref?.();
-			}),
-		]);
-	} finally {
-		clearTimeout(timeoutId);
-	}
-};
-
 const attachRedisErrorHandler = ({
 	redisInstance,
 	label,

server/src/honoUtils/handleReadyCheck.ts

@@ -3,6 +3,7 @@ import type { Context } from "hono";
 import { clientCritical } from "@/db/initDrizzle.js";
 import { getPgHealthState } from "@/db/pgHealthMonitor.js";
 import { getRedisAvailability } from "@/external/redis/initRedis.js";
+import { withTimeout } from "@/utils/withTimeout.js";
 import type { HonoEnv } from "./HonoEnv.js";
 
 const POSTGRES_TIMEOUT_MS = 1_000;
@@ -10,18 +11,13 @@ const READY_CHECK_TOKEN = process.env.READY_CHECK_TOKEN?.trim();
 
 const checkPostgresReady = async () => {
 	const query = clientCritical`SELECT 1`;
-	let timeoutId: ReturnType<typeof setTimeout> | undefined;
 
 	try {
-		await Promise.race([
-			query,
-			new Promise<never>((_, reject) => {
-				timeoutId = setTimeout(() => {
-					void query.cancel();
-					reject(new Error(`timed out after ${POSTGRES_TIMEOUT_MS}ms`));
-				}, POSTGRES_TIMEOUT_MS);
-			}),
-		]);
+		await withTimeout({
+			timeoutMs: POSTGRES_TIMEOUT_MS,
+			fn: () => query,
+			onTimeout: () => query.cancel(),
+		});
 
 		return {
 			ok: true,
@@ -33,8 +29,6 @@ const checkPostgresReady = async () => {
 			error: error instanceof Error ? error.message : "unknown postgres error",
 			...getPgHealthState(),
 		};
-	} finally {
-		clearTimeout(timeoutId);
 	}
 };
 

server/src/internal/api/check/checkUtils/getRetryableCheckFallbackResponse.ts

@@ -0,0 +1,39 @@
+import {
+	type CheckParams,
+	CheckResponseV3Schema,
+	type ParsedCheckParams,
+} from "@autumn/shared";
+import type { AutumnContext } from "@/honoUtils/HonoEnv.js";
+import { transformCheckResponse } from "./transformCheckResponse.js";
+
+export const getRetryableCheckFallbackResponse = ({
+	ctx,
+	body,
+	requiredBalance,
+}: {
+	ctx: AutumnContext;
+	body: ParsedCheckParams | (CheckParams & { feature_id: string });
+	requiredBalance: number;
+}) => {
+	const fallbackResponse = CheckResponseV3Schema.parse({
+		allowed: true,
+		customer_id: body.customer_id || "",
+		entity_id: body.entity_id,
+		required_balance: requiredBalance,
+		balance: null,
+		flag: null,
+	});
+
+	const featureToUse = ctx.features.find(
+		(feature) => feature.id === body.feature_id,
+	);
+
+	return featureToUse
+		? transformCheckResponse({
+				ctx,
+				response: fallbackResponse,
+				featureToUse,
+				noCusEnts: false,
+			})
+		: fallbackResponse;
+};

server/src/internal/api/check/checkUtils/transformCheckResponse.ts

@@ -0,0 +1,29 @@
+import {
+	AffectedResource,
+	applyResponseVersionChanges,
+	type CheckResponseV3,
+	type Feature,
+} from "@autumn/shared";
+import type { AutumnContext } from "@/honoUtils/HonoEnv.js";
+
+export const transformCheckResponse = ({
+	ctx,
+	response,
+	featureToUse,
+	noCusEnts,
+}: {
+	ctx: AutumnContext;
+	response: CheckResponseV3;
+	featureToUse: Feature;
+	noCusEnts: boolean;
+}) =>
+	applyResponseVersionChanges<CheckResponseV3>({
+		input: response,
+		targetVersion: ctx.apiVersion,
+		resource: AffectedResource.Check,
+		legacyData: {
+			noCusEnts,
+			featureToUse,
+		},
+		ctx,
+	});

server/src/internal/api/check/handleCheck.ts

@@ -1,19 +1,19 @@
 import {
 	AffectedResource,
 	ApiVersion,
-	applyResponseVersionChanges,
 	type CheckParams,
 	CheckParamsSchema,
 	CheckQuerySchema,
 	type CheckResponseV3,
-	CheckResponseV3Schema,
 	type ParsedCheckParams,
 } from "@autumn/shared";
 import { isRetryableDbError } from "@/db/dbUtils.js";
 import { createRoute } from "@/honoMiddlewares/routeHandler.js";
 import { parseCheckParamsForLock } from "@/internal/balances/utils/lock/parseCheckParamsForLock.js";
 import { getCheckData } from "./checkUtils/getCheckData.js";
+import { getRetryableCheckFallbackResponse } from "./checkUtils/getRetryableCheckFallbackResponse.js";
 import { getV2CheckResponse } from "./checkUtils/getV2CheckResponse.js";
+import { transformCheckResponse } from "./checkUtils/transformCheckResponse.js";
 import { getCheckPreview } from "./getCheckPreview.js";
 import { handleProductCheck } from "./handlers/handleProductCheck.js";
 import { runCheckWithTrack } from "./runCheckWithTrack.js";
@@ -68,33 +68,13 @@ export const handleCheck = createRoute({
 				throw error;
 			}
 
-			const fallbackResponse = CheckResponseV3Schema.parse({
-				allowed: true,
-				customer_id: customer_id || "",
-				entity_id: entity_id,
-				required_balance: requiredBalance,
-				balance: null,
-				flag: null,
-			});
-
-			const featureToUse = ctx.features.find(
-				(feature) => feature.id === body.feature_id,
+			return c.json(
+				getRetryableCheckFallbackResponse({
+					ctx,
+					body,
+					requiredBalance,
+				}),
 			);
-
-			const transformedFallbackResponse = featureToUse
-				? applyResponseVersionChanges<CheckResponseV3>({
-						input: fallbackResponse,
-						targetVersion: ctx.apiVersion,
-						resource: AffectedResource.Check,
-						legacyData: {
-							noCusEnts: false,
-							featureToUse,
-						},
-						ctx,
-					})
-				: fallbackResponse;
-
-			return c.json(transformedFallbackResponse);
 		}
 
 		let response: CheckResponseV3;
@@ -122,17 +102,12 @@ export const handleCheck = createRoute({
 				})
 			: undefined;
 
-		// Version changes will transform V3 -> V2 -> V1 -> V0 based on target API version
-		const transformedResponse = applyResponseVersionChanges<CheckResponseV3>({
-			input: response,
-			targetVersion: ctx.apiVersion,
-			resource: AffectedResource.Check,
-			legacyData: {
-				noCusEnts:
-					checkData.apiBalance === undefined && checkData.apiFlag === undefined,
-				featureToUse: checkData.featureToUse,
-			},
+		const transformedResponse = transformCheckResponse({
 			ctx,
+			response,
+			featureToUse: checkData.featureToUse,
+			noCusEnts:
+				checkData.apiBalance === undefined && checkData.apiFlag === undefined,
 		});
 
 		return c.json({

server/src/queue/initWorkers.ts

@@ -13,6 +13,7 @@ import { type DrizzleCli, initDrizzle } from "@/db/initDrizzle.js";
 import { logger } from "@/external/logtail/logtailUtils.js";
 import { verifyCacheConsistency } from "@/internal/billing/v2/workflows/verifyCacheConsistency/verifyCacheConsistency.js";
 import { generateId } from "@/utils/genUtils.js";
+import { withTimeout } from "@/utils/withTimeout.js";
 import { hatchet } from "../external/hatchet/initHatchet.js";
 import { getSqsClient, QUEUE_URL, recreateSqsClient } from "./initSqs.js";
 import { JobName } from "./JobName.js";
@@ -41,27 +42,6 @@ const ZERO_MESSAGE_ALERT_THRESHOLD = 20; // ~20 min of 0 messages
 
 // ============ Helper Functions ============
 
-const withTimeout = <T>(promise: Promise<T>, timeoutMs: number): Promise<T> =>
-	new Promise((resolve, reject) => {
-		const timeout = setTimeout(() => {
-			reject(new Error(`Processing timed out after ${timeoutMs}ms`));
-		}, timeoutMs);
-
-		if (timeout.unref) {
-			timeout.unref();
-		}
-
-		promise
-			.then((result) => {
-				clearTimeout(timeout);
-				resolve(result);
-			})
-			.catch((error) => {
-				clearTimeout(timeout);
-				reject(error);
-			});
-	});
-
 const logPrefix = ({ queueUrl }: { queueUrl: string }) =>
 	`[SQS Worker ${process.pid}][${queueUrl.split("/").pop()}]`;
 
@@ -209,7 +189,11 @@ const startPollingLoop = async ({
 		if (isMigration) {
 			await processMessage({ message, db });
 		} else {
-			await withTimeout(processMessage({ message, db }), MESSAGE_TIMEOUT_MS);
+			await withTimeout({
+				timeoutMs: MESSAGE_TIMEOUT_MS,
+				timeoutMessage: `Processing timed out after ${MESSAGE_TIMEOUT_MS}ms`,
+				fn: () => processMessage({ message, db }),
+			});
 		}
 
 		messagesProcessed++;

server/src/utils/withTimeout.ts

@@ -0,0 +1,28 @@
+export const withTimeout = async <T>({
+	timeoutMs,
+	fn,
+	timeoutMessage,
+	onTimeout,
+}: {
+	timeoutMs: number;
+	fn: () => Promise<T>;
+	timeoutMessage?: string;
+	onTimeout?: () => void | Promise<void>;
+}): Promise<T> => {
+	let timeoutId: ReturnType<typeof setTimeout> | undefined;
+
+	try {
+		return await Promise.race([
+			fn(),
+			new Promise<never>((_, reject) => {
+				timeoutId = setTimeout(() => {
+					void Promise.resolve(onTimeout?.());
+					reject(new Error(timeoutMessage || `timed out after ${timeoutMs}ms`));
+				}, timeoutMs);
+				timeoutId.unref?.();
+			}),
+		]);
+	} finally {
+		clearTimeout(timeoutId);
+	}
+};

server/tests/unit/utils/with-timeout.test.ts

@@ -0,0 +1,29 @@
+import { describe, expect, test } from "bun:test";
+import { withTimeout } from "@/utils/withTimeout.js";
+
+describe("withTimeout", () => {
+	test("returns the wrapped result before the timeout", async () => {
+		const result = await withTimeout({
+			timeoutMs: 50,
+			fn: async () => "ok",
+		});
+
+		expect(result).toBe("ok");
+	});
+
+	test("rejects and runs onTimeout when the timeout elapses", async () => {
+		let timedOut = false;
+
+		await expect(
+			withTimeout({
+				timeoutMs: 10,
+				fn: () => new Promise<string>(() => {}),
+				onTimeout: () => {
+					timedOut = true;
+				},
+			}),
+		).rejects.toThrow("timed out after 10ms");
+
+		expect(timedOut).toBe(true);
+	});
+});