diff --git a/server/src/external/redis/customerRedisRouting.ts b/server/src/external/redis/customerRedisRouting.ts index f4e69e174..f93a742fa 100644 --- a/server/src/external/redis/customerRedisRouting.ts +++ b/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, diff --git a/server/src/external/redis/orgRedisUtils/orgRedisMigrationUtils.ts b/server/src/external/redis/orgRedisUtils/orgRedisMigrationUtils.ts new file mode 100644 index 000000000..0ad421102 --- /dev/null +++ b/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 }); +}; diff --git a/server/src/internal/balances/finalizeLock/runFinalizeLock.ts b/server/src/internal/balances/finalizeLock/runFinalizeLock.ts index 292ff6748..067072f1b 100644 --- a/server/src/internal/balances/finalizeLock/runFinalizeLock.ts +++ b/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, }); } diff --git a/server/src/internal/balances/finalizeLock/runFinalizeLockV2.ts b/server/src/internal/balances/finalizeLock/runFinalizeLockV2.ts index 892db916f..bb9028d21 100644 --- a/server/src/internal/balances/finalizeLock/runFinalizeLockV2.ts +++ b/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; diff --git a/server/src/internal/balances/finalizeLock/runRedisFinalizeLockV2.ts b/server/src/internal/balances/finalizeLock/runRedisFinalizeLockV2.ts index 568c21465..1a4879d77 100644 --- a/server/src/internal/balances/finalizeLock/runRedisFinalizeLockV2.ts +++ b/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>; @@ -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()) { diff --git a/server/src/internal/balances/utils/lock/fetchLockReceipt.ts b/server/src/internal/balances/utils/lock/fetchLockReceipt.ts index 4d6c62359..af65bf2f6 100644 --- a/server/src/internal/balances/utils/lock/fetchLockReceipt.ts +++ b/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, 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, }; } diff --git a/server/src/internal/balances/utils/lockV2/buildFinalizeLockContextV2.ts b/server/src/internal/balances/utils/lockV2/buildFinalizeLockContextV2.ts index 9324257ce..3082640e5 100644 --- a/server/src/internal/balances/utils/lockV2/buildFinalizeLockContextV2.ts +++ b/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 => { + 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, diff --git a/server/src/internal/balances/utils/lockV2/fetchAndClaimLockReceiptV2.ts b/server/src/internal/balances/utils/lockV2/fetchAndClaimLockReceiptV2.ts index 041dbb511..67a402611 100644 --- a/server/src/internal/balances/utils/lockV2/fetchAndClaimLockReceiptV2.ts +++ b/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, }; }; diff --git a/server/src/internal/orgs/handlers/handleRedisConfig.ts b/server/src/internal/orgs/handlers/handleRedisConfig.ts index a663edd19..650a3652d 100644 --- a/server/src/internal/orgs/handlers/handleRedisConfig.ts +++ b/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( diff --git a/server/src/workers.ts b/server/src/workers.ts index bade4e121..ec0d9c4bd 100644 --- a/server/src/workers.ts +++ b/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 }); diff --git a/server/tests/integration/balances/lock/check-with-lock-expiry.test.ts b/server/tests/integration/balances/lock/check-with-lock-expiry.test.ts index a55950532..0600f88a6 100644 --- a/server/tests/integration/balances/lock/check-with-lock-expiry.test.ts +++ b/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 diff --git a/server/tests/integration/balances/utils/lockUtils/deleteLock.ts b/server/tests/integration/balances/utils/lockUtils/deleteLock.ts index d91b80ac6..a01576d1e 100644 --- a/server/tests/integration/balances/utils/lockUtils/deleteLock.ts +++ b/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), + ), ]); }; diff --git a/server/tests/integration/balances/utils/lockUtils/expectLockReceiptDeleted.ts b/server/tests/integration/balances/utils/lockUtils/expectLockReceiptDeleted.ts index 1d86dc66d..c77d51896 100644 --- a/server/tests/integration/balances/utils/lockUtils/expectLockReceiptDeleted.ts +++ b/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); + } }; diff --git a/server/tests/unit/balances/lock/fetchLockReceipt.test.ts b/server/tests/unit/balances/lock/fetchLockReceipt.test.ts new file mode 100644 index 000000000..1185e06d4 --- /dev/null +++ b/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(), + 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) => 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]); + }); +});