fix: org redis for lock flow

server/src/external/redis/customerRedisRouting.ts

@@ -64,9 +64,24 @@ export const assignCustomerRedisToCtx = ({
 	});
 
 	ctx.redisV2 = routingInfo.redis;
+
 	return routingInfo;
 };
 
+export const overrideCtxRedisV2 = ({
+	ctx,
+	redisV2,
+}: {
+	ctx: AutumnContext;
+	redisV2: Redis;
+}): AutumnContext => {
+	if (ctx.redisV2 === redisV2) return ctx;
+
+	const injectedCtx = Object.create(ctx) as AutumnContext;
+	injectedCtx.redisV2 = redisV2;
+	return injectedCtx;
+};
+
 export const getRedisUrlForCustomer = ({
 	org,
 	customerId,

server/src/external/redis/orgRedisUtils/orgRedisMigrationUtils.ts

@@ -0,0 +1,37 @@
+import type { Redis } from "ioredis";
+import type { AutumnContext } from "@/honoUtils/HonoEnv.js";
+import { getOrgRedis } from "../orgRedisPool.js";
+import { resolveRedisV2 } from "../resolveRedisV2.js";
+
+const dedupeRedisInstances = ({ candidates }: { candidates: Redis[] }) =>
+	candidates.filter(
+		(candidate, index) => candidates.indexOf(candidate) === index,
+	);
+
+export const getRedisV2LockReceiptCandidates = ({
+	ctx,
+}: {
+	ctx: AutumnContext;
+}): Redis[] => {
+	const candidates: Redis[] = [ctx.redisV2];
+
+	if (ctx.org.redis_config && ctx.org.redis_config.migrationPercent > 0) {
+		candidates.push(getOrgRedis({ org: ctx.org }), resolveRedisV2());
+	}
+
+	return dedupeRedisInstances({ candidates });
+};
+
+export const getRedisV2OrgCleanupCandidates = ({
+	ctx,
+}: {
+	ctx: AutumnContext;
+}): Redis[] => {
+	const candidates: Redis[] = [ctx.redisV2, resolveRedisV2()];
+
+	if (ctx.org.redis_config) {
+		candidates.push(getOrgRedis({ org: ctx.org }));
+	}
+
+	return dedupeRedisInstances({ candidates });
+};

server/src/internal/balances/finalizeLock/runFinalizeLock.ts

@@ -48,6 +48,7 @@ const runFinalizeLockInner = async ({ ctx, params }: RunFinalizeLockArgs) => {
 			receipt: fetchedReceipt.receipt,
 			lockReceiptKey: fetchedReceipt.lockReceiptKey,
 			claimed: fetchedReceipt.claimed,
+			lockRedisInstance: fetchedReceipt.redisInstance,
 		});
 	}
 

server/src/internal/balances/finalizeLock/runFinalizeLockV2.ts

@@ -5,6 +5,7 @@ import {
 	RecaseError,
 } from "@autumn/shared";
 import { Decimal } from "decimal.js";
+import type { Redis } from "ioredis";
 import type { AutumnContext } from "@/honoUtils/HonoEnv.js";
 import { cancelLockExpiry } from "@/internal/balances/utils/lock/cancelLockExpiry.js";
 import type { LockReceipt } from "@/internal/balances/utils/lock/fetchLockReceipt.js";
@@ -24,12 +25,14 @@ export const runFinalizeLockV2 = async ({
 	receipt,
 	lockReceiptKey,
 	claimed,
+	lockRedisInstance,
 }: {
 	ctx: AutumnContext;
 	params: FinalizeLockParamsV0;
 	receipt: LockReceipt;
 	lockReceiptKey: string;
 	claimed: boolean;
+	lockRedisInstance: Redis;
 }) => {
 	if (!claimed) {
 		throw new RecaseError({
@@ -45,6 +48,7 @@ export const runFinalizeLockV2 = async ({
 		params,
 		receipt,
 		lockReceiptKey,
+		redisInstance: lockRedisInstance,
 	});
 	const { redisInstance, finalValue, lockValue } = finalizeLockContext;
 

server/src/internal/balances/finalizeLock/runRedisFinalizeLockV2.ts

@@ -15,7 +15,7 @@ export const runRedisFinalizeLockV2 = async ({
 	ctx: AutumnContext;
 	finalizeLockContext: FinalizeLockContextV2;
 }) => {
-	const { receipt, fullSubject, deduction, deductionOptions } =
+	const { receipt, fullSubject, deduction, deductionOptions, redisInstance } =
 		finalizeLockContext;
 
 	let redisResult: Awaited<ReturnType<typeof executeRedisDeductionV2>>;
@@ -27,6 +27,7 @@ export const runRedisFinalizeLockV2 = async ({
 			entityId: receipt.entity_id ?? undefined,
 			deductions: [deduction],
 			deductionOptions,
+			redisInstance,
 		});
 	} catch (error) {
 		if (error instanceof RedisDeductionError && error.shouldFallback()) {

server/src/internal/balances/utils/lock/fetchLockReceipt.ts

@@ -1,5 +1,6 @@
 import { ErrCode, RecaseError } from "@autumn/shared";
 import { redis } from "@/external/redis/initRedis.js";
+import { getRedisV2LockReceiptCandidates } from "@/external/redis/orgRedisUtils/orgRedisMigrationUtils.js";
 import type { AutumnContext } from "@/honoUtils/HonoEnv";
 import { fetchAndClaimLockReceiptV2 } from "@/internal/balances/utils/lockV2/fetchAndClaimLockReceiptV2.js";
 import type { MutationLogItem } from "@/internal/balances/utils/types/mutationLogItem.js";
@@ -39,6 +40,28 @@ const normalizeLockReceiptItems = ({
 	});
 };
 
+const fetchAndClaimLockReceiptV2FromCandidates = async ({
+	ctx,
+	lockId,
+}: {
+	ctx: AutumnContext;
+	lockId: string;
+}) => {
+	const candidates = getRedisV2LockReceiptCandidates({ ctx });
+
+	for (const redisInstance of candidates) {
+		const result = await fetchAndClaimLockReceiptV2({
+			ctx,
+			lockId,
+			redisInstance,
+		});
+
+		if (result.found) return result;
+	}
+
+	return { found: false as const };
+};
+
 export const fetchLockReceipt = async ({
 	ctx,
 	lockId,
@@ -46,7 +69,6 @@ export const fetchLockReceipt = async ({
 	ctx: AutumnContext;
 	lockId: string;
 }) => {
-	const { redisV2 } = ctx;
 	const hashedKey = Bun.hash(lockId).toString();
 	const lockReceiptKey = buildLockReceiptKey({
 		orgId: ctx.org.id,
@@ -56,6 +78,7 @@ export const fetchLockReceipt = async ({
 
 	// V2 half doubles as a fetch+claim (pipelined GET + SET NX on a marker key)
 	// so the dispatcher can route to runFinalizeLockV2 without a follow-up claim RT.
+	// During org Redis migrations, V2 checks both shared and dedicated Redis.
 	// V1 half stays a plain JSON.GET — V1 finalize still claims via Lua afterwards.
 	const [rawReceiptV1, v2Result] = await Promise.all([
 		tryRedisRead(
@@ -63,10 +86,9 @@ export const fetchLockReceipt = async ({
 				redis.call("JSON.GET", lockReceiptKey, "$") as Promise<string | null>,
 			redis,
 		),
-		fetchAndClaimLockReceiptV2({
+		fetchAndClaimLockReceiptV2FromCandidates({
 			ctx,
 			lockId,
-			redisInstance: redisV2,
 		}),
 	]);
 
@@ -76,6 +98,7 @@ export const fetchLockReceipt = async ({
 			lockReceiptKey: v2Result.lockReceiptKey,
 			source: "redis_v2" as const,
 			claimed: v2Result.claimed,
+			redisInstance: v2Result.redisInstance,
 		};
 	}
 

server/src/internal/balances/utils/lockV2/buildFinalizeLockContextV2.ts

@@ -1,6 +1,7 @@
 import type { Feature, FullSubject } from "@autumn/shared";
 import { type FinalizeLockParamsV0, findFeatureById } from "@autumn/shared";
 import type { Redis } from "ioredis";
+import { overrideCtxRedisV2 } from "@/external/redis/customerRedisRouting.js";
 import type { AutumnContext } from "@/honoUtils/HonoEnv.js";
 import type { LockReceipt } from "@/internal/balances/utils/lock/fetchLockReceipt.js";
 import {
@@ -30,14 +31,18 @@ export const buildFinalizeLockContextV2 = async ({
 	params,
 	receipt,
 	lockReceiptKey,
+	redisInstance,
 }: {
 	ctx: AutumnContext;
 	params: FinalizeLockParamsV0;
 	receipt: LockReceipt;
 	lockReceiptKey: string;
+	redisInstance: Redis;
 }): Promise<FinalizeLockContextV2> => {
+	const ctxWithRedisV2 = overrideCtxRedisV2({ ctx, redisV2: redisInstance });
+
 	const fullSubject = await getOrSetCachedFullSubject({
-		ctx,
+		ctx: ctxWithRedisV2,
 		customerId: receipt.customer_id,
 		entityId: receipt.entity_id ?? undefined,
 		source: "runFinalizeLockV2",
@@ -61,7 +66,7 @@ export const buildFinalizeLockContextV2 = async ({
 	return {
 		receipt,
 		lockReceiptKey,
-		redisInstance: ctx.redisV2,
+		redisInstance,
 		fullSubject,
 		feature,
 		lockValue,

server/src/internal/balances/utils/lockV2/fetchAndClaimLockReceiptV2.ts

@@ -21,6 +21,7 @@ type FetchAndClaimResult =
 			claimed: boolean;
 			receipt: LockReceipt;
 			lockReceiptKey: string;
+			redisInstance: Redis;
 	  };
 
 const normalizeLockReceiptItems = ({
@@ -118,5 +119,6 @@ export const fetchAndClaimLockReceiptV2 = async ({
 		claimed: claimResult === "OK",
 		receipt,
 		lockReceiptKey,
+		redisInstance,
 	};
 };

server/src/internal/orgs/handlers/handleRedisConfig.ts

@@ -4,6 +4,7 @@ import { getOrgRedis, removeOrgRedis } from "@/external/redis/orgRedisPool.js";
 import { createRoute } from "@/honoMiddlewares/routeHandler.js";
 import { encryptData } from "@/utils/encryptUtils.js";
 import { OrgService } from "../OrgService.js";
+import { clearOrgCache } from "../orgUtils/clearOrgCache.js";
 
 const REDIS_PROTOCOLS = new Set(["redis:", "rediss:"]);
 
@@ -70,6 +71,7 @@ export const handleUpsertRedisConfig = createRoute({
 
 		if (updatedOrg) {
 			getOrgRedis({ org: updatedOrg });
+			await clearOrgCache({ db, orgId: org.id, env: ctx.env, logger });
 			logger.info(
 				`[handleUpsertRedisConfig] org=${org.id}: redis_config created, url=${redisUrl.host}, actor=${ctx.user?.email ?? ctx.userId ?? "unknown"}`,
 			);
@@ -109,6 +111,7 @@ export const handleUpdateRedisMigration = createRoute({
 				},
 			},
 		});
+		await clearOrgCache({ db, orgId: org.id, env: ctx.env, logger });
 
 		logger.info(
 			`[handleUpdateRedisMigration] org=${org.id}: ${org.redis_config.migrationPercent}% -> ${migrationPercent}%`,
@@ -139,6 +142,7 @@ export const handleDeleteRedisConfig = createRoute({
 			orgId: org.id,
 			updates: { redis_config: null },
 		});
+		await clearOrgCache({ db, orgId: org.id, env: ctx.env, logger });
 		removeOrgRedis({ orgId: org.id });
 
 		logger.info(

server/src/workers.ts

@@ -100,7 +100,35 @@ if (cluster.isPrimary) {
 	startMemoryMonitor("worker", 60_000);
 	await startAllEdgeConfigPolling({ logger });
 
-	process.once("exit", stopAllEdgeConfigPolling);
+	const { db } = await import("./db/initDrizzle.js");
+	const { primeRedisMonitor } = await import(
+		"./external/redis/initUtils/redisAvailability.js"
+	);
+	const {
+		primeRedisV2Monitor,
+		startRedisV2Monitor,
+		stopRedisV2Monitor,
+	} = await import("./external/redis/initUtils/redisV2Availability.js");
+	const { startRedisMonitor, stopRedisMonitor } = await import(
+		"./external/redis/initRedis.js"
+	);
+	const { preWarmOrgRedisConnections } = await import(
+		"./external/redis/orgRedisPool.js"
+	);
+
+	await Promise.all([primeRedisMonitor(), primeRedisV2Monitor()]);
+	startRedisMonitor();
+	startRedisV2Monitor();
+
+	void preWarmOrgRedisConnections({ db }).catch((error) => {
+		logger.warn("[OrgRedis] Warmup failed", { error });
+	});
+
+	process.once("exit", () => {
+		stopAllEdgeConfigPolling();
+		stopRedisMonitor();
+		stopRedisV2Monitor();
+	});
 
 	const { initWorkers } = await import("./queue/initWorkers.js");
 	await initWorkers({ startupStartedAt, queueImplementation });

server/tests/integration/balances/lock/check-with-lock-expiry.test.ts

@@ -201,10 +201,13 @@ test.concurrent(`${chalk.yellowBright("check-lock-expiry-4: no expires_at sets T
 		lock: { enabled: true, lock_id: customerId },
 	});
 
-	const { lockReceiptKey, source } = await fetchLockReceipt({ ctx, lockId: customerId });
-	const redisInstance = source === "redis_v2" ? ctx.redisV2 : redis;
+	const fetchedReceipt = await fetchLockReceipt({ ctx, lockId: customerId });
+	const redisInstance =
+		fetchedReceipt.source === "redis_v2" ? fetchedReceipt.redisInstance : redis;
 
-	const expireAt = await redisInstance.expiretime(lockReceiptKey);
+	const expireAt = await redisInstance.expiretime(
+		fetchedReceipt.lockReceiptKey,
+	);
 	const expectedTtl = beforeCheck + 24 * 60 * 60;
 
 	// TTL should be within 5s of now + 1 day
@@ -238,10 +241,13 @@ test.concurrent(`${chalk.yellowBright("check-lock-expiry-5: expires_at set, TTL
 		lock: { enabled: true, lock_id: customerId, expires_at: expiresAt },
 	});
 
-	const { lockReceiptKey, source } = await fetchLockReceipt({ ctx, lockId: customerId });
-	const redisInstance = source === "redis_v2" ? ctx.redisV2 : redis;
+	const fetchedReceipt = await fetchLockReceipt({ ctx, lockId: customerId });
+	const redisInstance =
+		fetchedReceipt.source === "redis_v2" ? fetchedReceipt.redisInstance : redis;
 
-	const expireAt = await redisInstance.expiretime(lockReceiptKey);
+	const expireAt = await redisInstance.expiretime(
+		fetchedReceipt.lockReceiptKey,
+	);
 	const expectedTtl = Math.ceil(expiresAt / 1000) + 60 * 60;
 
 	// TTL should be within 5s of expires_at + 1 hour

server/tests/integration/balances/utils/lockUtils/deleteLock.ts

@@ -1,5 +1,6 @@
 import type { TestContext } from "@tests/utils/testInitUtils/createTestContext.js";
 import { redis } from "@/external/redis/initRedis.js";
+import { getRedisV2OrgCleanupCandidates } from "@/external/redis/orgRedisUtils/orgRedisMigrationUtils.js";
 import { buildLockReceiptKey } from "@/internal/balances/utils/lock/buildLockReceiptKey.js";
 import { buildClaimMarkerKey } from "@/internal/balances/utils/lockV2/buildClaimMarkerKey.js";
 
@@ -20,6 +21,8 @@ export const deleteLock = async ({
 
 	await Promise.all([
 		redis.del(redisReceiptKey),
-		ctx.redisV2.del(redisReceiptKey, claimMarkerKey),
+		...getRedisV2OrgCleanupCandidates({ ctx }).map((redisInstance) =>
+			redisInstance.del(redisReceiptKey, claimMarkerKey),
+		),
 	]);
 };

server/tests/integration/balances/utils/lockUtils/expectLockReceiptDeleted.ts

@@ -1,6 +1,7 @@
 import { expect } from "bun:test";
 import type { TestContext } from "@tests/utils/testInitUtils/createTestContext.js";
 import { redis } from "@/external/redis/initRedis.js";
+import { getRedisV2OrgCleanupCandidates } from "@/external/redis/orgRedisUtils/orgRedisMigrationUtils.js";
 import { buildLockReceiptKey } from "@/internal/balances/utils/lock/buildLockReceiptKey.js";
 
 /** Asserts that the lock receipt for the given ID no longer exists in Redis. */
@@ -20,4 +21,13 @@ export const expectLockReceiptDeleted = async ({
 
 	const receipt = await redis.call("JSON.GET", redisReceiptKey, "$");
 	expect(receipt).toBeNull();
+
+	const v2ReceiptCounts = await Promise.all(
+		getRedisV2OrgCleanupCandidates({ ctx }).map((redisInstance) =>
+			redisInstance.exists(redisReceiptKey),
+		),
+	);
+	for (const receiptCount of v2ReceiptCounts) {
+		expect(receiptCount).toBe(0);
+	}
 };

server/tests/unit/balances/lock/fetchLockReceipt.test.ts

@@ -0,0 +1,114 @@
+import { beforeEach, describe, expect, mock, test } from "bun:test";
+import type { Redis } from "ioredis";
+
+const sharedRedis = { name: "shared" } as unknown as Redis;
+const dedicatedRedis = { name: "dedicated" } as unknown as Redis;
+
+const mockState = {
+	jsonGetResult: null as string | null,
+	v2ResultsByRedis: new Map<object, { found: boolean }>(),
+	v2Calls: [] as object[],
+};
+
+mock.module("@/external/redis/initRedis.js", () => ({
+	redis: {
+		call: async () => mockState.jsonGetResult,
+	},
+}));
+
+mock.module("@/external/redis/orgRedisPool.js", () => ({
+	getOrgRedis: () => dedicatedRedis,
+}));
+
+mock.module("@/external/redis/resolveRedisV2.js", () => ({
+	resolveRedisV2: () => sharedRedis,
+}));
+
+mock.module("@/utils/cacheUtils/cacheUtils.js", () => ({
+	tryRedisRead: async (operation: () => Promise<unknown>) => operation(),
+}));
+
+mock.module(
+	"@/internal/balances/utils/lockV2/fetchAndClaimLockReceiptV2.js",
+	() => ({
+		fetchAndClaimLockReceiptV2: async ({
+			redisInstance,
+		}: {
+			redisInstance: object;
+		}) => {
+			mockState.v2Calls.push(redisInstance);
+			const result = mockState.v2ResultsByRedis.get(redisInstance);
+
+			if (!result?.found) return { found: false };
+
+			return {
+				found: true,
+				claimed: true,
+				lockReceiptKey: "lock:receipt",
+				redisInstance,
+				receipt: {
+					customer_id: "customer_123",
+					feature_id: "messages",
+					items: [],
+				},
+			};
+		},
+	}),
+);
+
+import { fetchLockReceipt } from "@/internal/balances/utils/lock/fetchLockReceipt.js";
+
+const makeContext = ({ migrationPercent }: { migrationPercent: number }) =>
+	({
+		org: {
+			id: "org_123",
+			redis_config: {
+				connectionString: "encrypted",
+				url: "dragonfly.internal:6379",
+				migrationPercent,
+				previousMigrationPercent: 0,
+				migrationChangedAt: 1,
+			},
+		},
+		env: "sandbox",
+		redisV2: sharedRedis,
+	}) as never;
+
+beforeEach(() => {
+	mockState.jsonGetResult = null;
+	mockState.v2ResultsByRedis.clear();
+	mockState.v2Calls = [];
+});
+
+describe("fetchLockReceipt", () => {
+	test("checks dedicated Redis during an org Redis migration and returns the source instance", async () => {
+		mockState.v2ResultsByRedis.set(sharedRedis, { found: false });
+		mockState.v2ResultsByRedis.set(dedicatedRedis, { found: true });
+
+		const result = await fetchLockReceipt({
+			ctx: makeContext({ migrationPercent: 50 }),
+			lockId: "lock_123",
+		});
+
+		expect(result.source).toBe("redis_v2");
+		if (result.source !== "redis_v2")
+			throw new Error("Expected Redis V2 receipt");
+		expect(result.redisInstance).toBe(dedicatedRedis);
+		expect(mockState.v2Calls).toEqual([sharedRedis, dedicatedRedis]);
+	});
+
+	test("does not check dedicated Redis when the org migration is complete", async () => {
+		mockState.v2ResultsByRedis.set(sharedRedis, { found: true });
+
+		const result = await fetchLockReceipt({
+			ctx: makeContext({ migrationPercent: 100 }),
+			lockId: "lock_123",
+		});
+
+		expect(result.source).toBe("redis_v2");
+		if (result.source !== "redis_v2")
+			throw new Error("Expected Redis V2 receipt");
+		expect(result.redisInstance).toBe(sharedRedis);
+		expect(mockState.v2Calls).toEqual([sharedRedis]);
+	});
+});