From f15bbde76adbba2a4385a264759f8e7326a31b09 Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Fri, 24 Apr 2026 17:12:24 +0100 Subject: [PATCH] chore: cleanup --- .../balances/track/utils/queueTrack.ts | 2 +- .../track/v3/handleRedisTrackErrorV3.ts | 7 +- server/src/queue/initWorkers.ts | 12 ++-- .../queue/start-polling-loop.test.ts | 36 +++++++++- .../track-v3/handleRedisTrackErrorV3.test.ts | 9 ++- .../track/handle-track-queue-fallback.test.ts | 71 ++++++++++--------- .../unit/balances/track/queueTrack.test.ts | 38 ++++++---- 7 files changed, 120 insertions(+), 55 deletions(-) diff --git a/server/src/internal/balances/track/utils/queueTrack.ts b/server/src/internal/balances/track/utils/queueTrack.ts index 8d25a588f..fbc77186a 100644 --- a/server/src/internal/balances/track/utils/queueTrack.ts +++ b/server/src/internal/balances/track/utils/queueTrack.ts @@ -24,7 +24,7 @@ export const queueTrack = async ({ jobName: JobName.Track, queueUrl, messageGroupId: `${ctx.org.id}:${ctx.env}:${body.customer_id}:${body.entity_id ?? "none"}`, - messageDeduplicationId: body.idempotency_key || ctx.id, + messageDeduplicationId: ctx.id, payload: { orgId: ctx.org.id, env: ctx.env, diff --git a/server/src/internal/balances/track/v3/handleRedisTrackErrorV3.ts b/server/src/internal/balances/track/v3/handleRedisTrackErrorV3.ts index c91bac399..5df489ac0 100644 --- a/server/src/internal/balances/track/v3/handleRedisTrackErrorV3.ts +++ b/server/src/internal/balances/track/v3/handleRedisTrackErrorV3.ts @@ -7,6 +7,7 @@ import { type TrackResponseV3, } from "@autumn/shared"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { RedisUnavailableError } from "@/external/redis/utils/errors.js"; import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; import { RedisDeductionError, @@ -57,7 +58,11 @@ export const handleRedisTrackErrorV3 = async ({ } if (error.isRedisUnavailable()) { - throw error; + throw new RedisUnavailableError({ + source: "runTrackV3", + reason: "other", + cause: error, + }); } if (error.shouldFallback()) { diff --git a/server/src/queue/initWorkers.ts b/server/src/queue/initWorkers.ts index 75e896175..556c7be72 100644 --- a/server/src/queue/initWorkers.ts +++ b/server/src/queue/initWorkers.ts @@ -26,6 +26,7 @@ import { processMessage, type SqsJob } from "./processMessage.js"; // ============ Shared State ============ let isRunning = true; const abortControllers = new Set(); +export const getAbortControllerCountForTesting = () => abortControllers.size; // Process recycling — exit after processing this many messages to prevent memory leaks const MAX_MESSAGES_BEFORE_RECYCLE = 50_000; @@ -78,6 +79,11 @@ export const startPollingLoop = async ({ const prefix = logPrefix({ queueUrl }); let abortController = new AbortController(); abortControllers.add(abortController); + const replaceAbortController = () => { + abortControllers.delete(abortController); + abortController = new AbortController(); + abortControllers.add(abortController); + }; const alertZeroMessages = () => { const minutes = consecutiveZeroMessageIntervals; @@ -259,8 +265,7 @@ export const startPollingLoop = async ({ `${prefix} ${consecutiveEmptyPolls} consecutive empty polls - recreating SQS client`, ); consecutiveEmptyPolls = 0; - abortController = new AbortController(); - abortControllers.add(abortController); + replaceAbortController(); return recreateSqsClientFn(); } @@ -285,8 +290,7 @@ export const startPollingLoop = async ({ if (consecutiveEmptyPolls >= EMPTY_POLL_THRESHOLD) { console.warn(`${prefix} Repeated errors - recreating SQS client`); consecutiveEmptyPolls = 0; - abortController = new AbortController(); - abortControllers.add(abortController); + replaceAbortController(); await new Promise((resolve) => setTimeout(resolve, 5000)); return recreateSqsClientFn(); } diff --git a/server/tests/integration/queue/start-polling-loop.test.ts b/server/tests/integration/queue/start-polling-loop.test.ts index 5ea4dbb9f..570c86600 100644 --- a/server/tests/integration/queue/start-polling-loop.test.ts +++ b/server/tests/integration/queue/start-polling-loop.test.ts @@ -1,5 +1,8 @@ import { afterEach, beforeEach, describe, expect, test } from "bun:test"; -import { startPollingLoop } from "@/queue/initWorkers.js"; +import { + getAbortControllerCountForTesting, + startPollingLoop, +} from "@/queue/initWorkers.js"; const originalSetTimeout = globalThis.setTimeout; @@ -86,4 +89,35 @@ describe("startPollingLoop", () => { expect(shouldPollCalls).toBe(1); expect(sendCalls).toBe(1); }); + + test("does not leak abort controllers when the SQS client is recreated", async () => { + let sendCalls = 0; + let recreateCalls = 0; + const makeClient = (abortAfterRecreate: boolean) => + ({ + send: async () => { + sendCalls++; + if (abortAfterRecreate) { + throw makeAbortError(); + } + return { Messages: [] }; + }, + }) as never; + + await startPollingLoop({ + db: {} as never, + queueUrl: "https://sqs.eu-west-1.amazonaws.com/123/track.fifo", + isFifo: true, + getSqsClientFn: () => makeClient(false), + recreateSqsClientFn: () => { + recreateCalls++; + return makeClient(true); + }, + shouldPoll: () => true, + }); + + expect(recreateCalls).toBe(1); + expect(sendCalls).toBeGreaterThan(9); + expect(getAbortControllerCountForTesting()).toBe(0); + }); }); diff --git a/server/tests/unit/balances/track-v3/handleRedisTrackErrorV3.test.ts b/server/tests/unit/balances/track-v3/handleRedisTrackErrorV3.test.ts index 0602f1bf1..9ca2f8558 100644 --- a/server/tests/unit/balances/track-v3/handleRedisTrackErrorV3.test.ts +++ b/server/tests/unit/balances/track-v3/handleRedisTrackErrorV3.test.ts @@ -1,5 +1,6 @@ import { beforeEach, describe, expect, mock, test } from "bun:test"; import { ApiVersion, ApiVersionClass, AppEnv } from "@autumn/shared"; +import { RedisUnavailableError } from "@/external/redis/utils/errors.js"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { RedisDeductionError, @@ -36,7 +37,7 @@ describe("handleRedisTrackErrorV3", () => { mockState.postgresCalls = []; }); - test("rethrows when Redis is unavailable", async () => { + test("normalizes Redis unavailable to RedisUnavailableError", async () => { const error = new RedisDeductionError({ message: "Redis not ready for deduction", code: RedisDeductionErrorCode.RedisUnavailable, @@ -54,7 +55,11 @@ describe("handleRedisTrackErrorV3", () => { fullSubject: {} as never, featureDeductions: [], }), - ).rejects.toBe(error); + ).rejects.toMatchObject({ + name: "RedisUnavailableError", + source: "runTrackV3", + reason: "other", + } satisfies Partial); expect(mockState.postgresCalls).toHaveLength(0); }); diff --git a/server/tests/unit/balances/track/handle-track-queue-fallback.test.ts b/server/tests/unit/balances/track/handle-track-queue-fallback.test.ts index 16621a77d..360fa676c 100644 --- a/server/tests/unit/balances/track/handle-track-queue-fallback.test.ts +++ b/server/tests/unit/balances/track/handle-track-queue-fallback.test.ts @@ -1,23 +1,20 @@ -import { beforeEach, describe, expect, mock, test } from "bun:test"; +import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test"; import { ApiVersion, ApiVersionClass, AppEnv, - type TrackResponseV3, } from "@autumn/shared"; import { RedisUnavailableError } from "@/external/redis/utils/errors.js"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { getSqsClient } from "@/queue/initSqs.js"; const mockState = { - queueCalls: [] as Record[], + queueCommands: [] as Record[], + queueError: null as Error | null, + originalSend: null as ReturnType["send"] | null, runTrackV2Calls: [] as Record[], runTrackV3Calls: [] as Record[], v3Error: null as unknown, - queuedResponse: { - customer_id: "cus_123", - value: 2, - balance: null, - } as TrackResponseV3 | null, }; mock.module("@/internal/balances/track/runTrackV2.js", () => ({ @@ -35,13 +32,6 @@ mock.module("@/internal/balances/track/v3/runTrackV3.js", () => ({ }, })); -mock.module("@/internal/balances/track/utils/queueTrack.js", () => ({ - queueTrack: async (args: Record) => { - mockState.queueCalls.push(args); - return mockState.queuedResponse; - }, -})); - mock.module("@/external/redis/initUtils/redisV2Availability.js", () => ({ shouldUseRedisV2: () => true, })); @@ -74,15 +64,24 @@ const body = { describe("track queue fallback", () => { beforeEach(() => { - mockState.queueCalls = []; + mockState.queueCommands = []; + mockState.queueError = null; mockState.runTrackV2Calls = []; mockState.runTrackV3Calls = []; mockState.v3Error = null; - mockState.queuedResponse = { - customer_id: "cus_123", - value: 2, - balance: null, - }; + process.env.TRACK_SQS_QUEUE_URL = + "https://sqs.eu-west-1.amazonaws.com/123456789012/track-dev.fifo"; + + const sqsClient = getSqsClient(); + mockState.originalSend = sqsClient.send.bind(sqsClient); + sqsClient.send = (async (command: { input: Record }) => { + if (mockState.queueError) { + throw mockState.queueError; + } + + mockState.queueCommands.push(command.input); + return {}; + }) as typeof sqsClient.send; }); test("queues track when rollout path hits a retryable Redis failure", async () => { @@ -99,16 +98,17 @@ describe("track queue fallback", () => { expect(mockState.runTrackV3Calls).toHaveLength(1); expect(mockState.runTrackV2Calls).toHaveLength(0); - expect(mockState.queueCalls).toHaveLength(1); - expect(mockState.queueCalls[0]).toMatchObject({ - body: { - customer_id: "cus_123", - feature_id: "messages", - }, + expect(mockState.queueCommands).toHaveLength(1); + expect(mockState.queueCommands[0]).toMatchObject({ + QueueUrl: + "https://sqs.eu-west-1.amazonaws.com/123456789012/track-dev.fifo", + }); + expect(response).toEqual({ + customer_id: "cus_123", + entity_id: undefined, + value: 2, + balance: null, }); - - if (!mockState.queuedResponse) throw new Error("expected queued response"); - expect(response).toEqual(mockState.queuedResponse); }); test("throws retryable Redis failure when queue fallback is unavailable", async () => { @@ -117,7 +117,7 @@ describe("track queue fallback", () => { reason: "timeout", }); mockState.v3Error = error; - mockState.queuedResponse = null; + mockState.queueError = new Error("sqs unavailable"); await expect( runTrackWithRollout({ @@ -127,6 +127,13 @@ describe("track queue fallback", () => { }), ).rejects.toBe(error); - expect(mockState.queueCalls).toHaveLength(1); + expect(mockState.queueCommands).toHaveLength(0); + }); + + afterEach(() => { + const sqsClient = getSqsClient(); + if (mockState.originalSend) { + sqsClient.send = mockState.originalSend; + } }); }); diff --git a/server/tests/unit/balances/track/queueTrack.test.ts b/server/tests/unit/balances/track/queueTrack.test.ts index ea604fd47..62ca8d912 100644 --- a/server/tests/unit/balances/track/queueTrack.test.ts +++ b/server/tests/unit/balances/track/queueTrack.test.ts @@ -1,17 +1,14 @@ import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test"; import { ApiVersion, ApiVersionClass, AppEnv } from "@autumn/shared"; +import type { SQSClient } from "@aws-sdk/client-sqs"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { getSqsClient } from "@/queue/initSqs.js"; const mockState = { - queueCalls: [] as Record[], + queueCommands: [] as Record[], + originalSend: null as null | SQSClient["send"], }; -mock.module("@/queue/queueUtils.js", () => ({ - addTaskToQueue: async (args: Record) => { - mockState.queueCalls.push(args); - }, -})); - mock.module("@/internal/balances/track/utils/getQueuedTrackResponse.js", () => ({ getQueuedTrackResponse: () => ({ customer_id: "cus_123", @@ -26,9 +23,15 @@ describe("queueTrack", () => { const originalTrackQueueUrl = process.env.TRACK_SQS_QUEUE_URL; beforeEach(() => { - mockState.queueCalls = []; + mockState.queueCommands = []; process.env.TRACK_SQS_QUEUE_URL = "https://sqs.eu-west-1.amazonaws.com/123456789012/track-dev.fifo"; + const sqsClient = getSqsClient(); + mockState.originalSend = sqsClient.send.bind(sqsClient); + sqsClient.send = (async (command: { input: Record }) => { + mockState.queueCommands.push(command.input); + return {}; + }) as typeof sqsClient.send; }); test("queues track with request identity and entity-scoped grouping", async () => { @@ -52,13 +55,16 @@ describe("queueTrack", () => { }, }); - expect(mockState.queueCalls).toHaveLength(1); - expect(mockState.queueCalls[0]).toMatchObject({ - queueUrl: + expect(mockState.queueCommands).toHaveLength(1); + expect(mockState.queueCommands[0]).toMatchObject({ + QueueUrl: "https://sqs.eu-west-1.amazonaws.com/123456789012/track-dev.fifo", - messageGroupId: "org_123:sandbox:cus_123:ent_123", - messageDeduplicationId: "req_123", - payload: { + MessageGroupId: "org_123:sandbox:cus_123:ent_123", + MessageDeduplicationId: "req_123", + }); + expect(JSON.parse(mockState.queueCommands[0]?.MessageBody as string)).toMatchObject({ + name: "track", + data: { orgId: "org_123", env: AppEnv.Sandbox, customerId: "cus_123", @@ -70,6 +76,10 @@ describe("queueTrack", () => { }); afterEach(() => { + if (mockState.originalSend) { + const sqsClient = getSqsClient(); + sqsClient.send = mockState.originalSend as typeof sqsClient.send; + } process.env.TRACK_SQS_QUEUE_URL = originalTrackQueueUrl; }); });