From e75965d922886ceaf18c5e34cea2f59a4d89677b Mon Sep 17 00:00:00 2001 From: John Yeo Date: Wed, 22 Apr 2026 12:09:33 +0100 Subject: [PATCH] fix: optimized queries / dedupe --- .vscode/settings.json | 2 +- bun.lock | 2 +- server/experiments/experimentEnv.ts | 4 +- server/experiments/explainEntityAggregate.ts | 22 +- server/experiments/explainGetFullSubject.ts | 61 ++++ .../fullSubjectDeduction/contextUtilsV2.lua | 4 +- .../deductFromSubjectBalances.lua | 1 + .../lock/lockReceiptV2.lua | 118 +++++++ server/src/db/initDrizzle.ts | 1 + server/src/db/redactDatabaseUrl.ts | 7 +- .../balances/autoTopUp/triggerAutoTopUp.ts | 4 - .../balances/check/runCheckWithRollout.ts | 2 +- .../deductionV2/executeRedisDeductionV2.ts | 4 + .../RefreshEntityAggregateBatchingManager.ts | 196 ++++++++++++ .../utils/refreshEntityAggregate/index.ts | 11 + .../queueRefreshEntityAggregate.ts | 40 +++ .../refreshEntityAggregateCache.ts | 53 +++- .../utils/sync/SyncBatchingManagerV3.ts | 1 + .../balances/utils/sync/syncItemV4.ts | 20 +- .../getEntityAggregateForSync.ts | 3 + .../getEntityAggregateFragments.ts | 299 ++++++------------ .../getEntityOptionsAggregateFragments.ts | 5 +- .../getFullSubject/getFullSubjectQuery.ts | 2 + server/src/queue/JobName.ts | 1 + server/src/queue/processMessage.ts | 18 +- server/src/queue/queueUtils.ts | 7 + server/src/workers.ts | 1 + .../balances/track/basic/track-basic.test.ts | 1 - .../refresh-entity-aggregate-dedup.test.ts | 220 +++++++++++++ ...fresh-entity-aggregate-track-dedup.test.ts | 127 ++++++++ 30 files changed, 994 insertions(+), 243 deletions(-) create mode 100644 server/experiments/explainGetFullSubject.ts create mode 100644 server/src/_luaScriptsV2/fullSubjectDeduction/lock/lockReceiptV2.lua create mode 100644 server/src/internal/balances/utils/refreshEntityAggregate/RefreshEntityAggregateBatchingManager.ts create mode 100644 server/src/internal/balances/utils/refreshEntityAggregate/index.ts create mode 100644 server/src/internal/balances/utils/refreshEntityAggregate/queueRefreshEntityAggregate.ts rename server/src/internal/balances/utils/{sync => refreshEntityAggregate}/refreshEntityAggregateCache.ts (56%) create mode 100644 server/tests/integration/others/refresh-entity-aggregate/refresh-entity-aggregate-dedup.test.ts create mode 100644 server/tests/integration/others/refresh-entity-aggregate/refresh-entity-aggregate-track-dedup.test.ts diff --git a/.vscode/settings.json b/.vscode/settings.json index d927606a7..cf6a231c4 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -33,7 +33,7 @@ "**/.cursor": true, // "**/.github": true, "**/.superset": true, - ".claude": true, + // ".claude": true, ".codex": true, ".agents": true, ".cursor": true, diff --git a/bun.lock b/bun.lock index 1d209f708..f4f82b95c 100644 --- a/bun.lock +++ b/bun.lock @@ -206,7 +206,7 @@ }, "packages/autumn-js": { "name": "autumn-js", - "version": "1.1.7", + "version": "1.2.10", "dependencies": { "query-string": "^9.2.2", "rou3": "^0.6.1", diff --git a/server/experiments/experimentEnv.ts b/server/experiments/experimentEnv.ts index 95add6752..c53bd00ec 100644 --- a/server/experiments/experimentEnv.ts +++ b/server/experiments/experimentEnv.ts @@ -17,8 +17,6 @@ export const prodTestOrgId = requireEnv({ key: "PROD_TEST_ORG_ID" }); export const prodTestCustomerId = requireEnv({ key: "PROD_TEST_CUSTOMER_ID", }); -export const prodTestEntityId = requireEnv({ - key: "PROD_TEST_ENTITY_ID", -}); +export const prodTestEntityId = process.env.PROD_TEST_ENTITY_ID || undefined; export const { initDrizzle } = await import("../src/db/initDrizzle"); diff --git a/server/experiments/explainEntityAggregate.ts b/server/experiments/explainEntityAggregate.ts index 8dc0caa28..eb70cde9a 100644 --- a/server/experiments/explainEntityAggregate.ts +++ b/server/experiments/explainEntityAggregate.ts @@ -1,10 +1,11 @@ import { AppEnv } from "@autumn/shared"; import { sql } from "drizzle-orm"; import { - initDrizzle, - prodTestCustomerId, - prodTestOrgId, + initDrizzle, + prodTestCustomerId, + prodTestOrgId, } from "./experimentEnv"; +import { OrgService } from "@/internal/orgs/OrgService"; const { getEntityAggregateForSync } = await import( "../src/internal/customers/repos/getFullSubject/getEntityAggregateForSync" @@ -19,6 +20,15 @@ const main = async () => { const { db } = initDrizzle(); + const org = await OrgService.getWithFeatures({ + db, + orgId, + env, + }); + + const features = org?.features.filter((feature) => feature.id.toLowerCase().includes("credit")); + + console.log("--- Running entity aggregate query ---"); const start = performance.now(); const result = await getEntityAggregateForSync({ @@ -26,6 +36,7 @@ const main = async () => { orgId, env, customerId, + internalFeatureIds: features?.map((feature) => feature.internal_id), }); const elapsed = performance.now() - start; console.log(`Rows returned: ${result.length}`); @@ -39,7 +50,10 @@ const main = async () => { ); const statusFilter = sql`AND cp.status = ANY(ARRAY['active', 'past_due', 'scheduled'])`; - const entityFragments = getEntityAggregateFragments({ statusFilter }); + const entityFragments = getEntityAggregateFragments({ + statusFilter, + internalFeatureIds: features?.map((feature) => feature.internal_id), + }); const query = sql` WITH subject_customer_records AS ( diff --git a/server/experiments/explainGetFullSubject.ts b/server/experiments/explainGetFullSubject.ts new file mode 100644 index 000000000..b4bcdaa2e --- /dev/null +++ b/server/experiments/explainGetFullSubject.ts @@ -0,0 +1,61 @@ +import { AppEnv } from "@autumn/shared"; +import { sql } from "drizzle-orm"; +import { + initDrizzle, + prodTestCustomerId, + prodTestEntityId, + prodTestOrgId, +} from "./experimentEnv"; + +const { getFullSubjectQuery } = await import( + "../src/internal/customers/repos/getFullSubject/getFullSubjectQuery" +); +const { RELEVANT_STATUSES } = await import( + "../src/internal/customers/cusProducts/CusProductService" +); + +// Run with: +// bun run experiments/explainGetFullSubject.ts +// Or scoped to an entity: +// PROD_TEST_ENTITY_ID=... bun run experiments/explainGetFullSubject.ts + +const main = async () => { + const orgId = prodTestOrgId; + const env = AppEnv.Live; + const customerId = prodTestCustomerId; + const entityId = prodTestEntityId; + + const { db } = initDrizzle(); + + const query = getFullSubjectQuery({ + orgId, + env, + customerId, + entityId, + inStatuses: RELEVANT_STATUSES, + }); + + console.log( + `--- Running full subject query (customer=${customerId}${entityId ? `, entity=${entityId}` : ""}) ---`, + ); + const start = performance.now(); + const result = await db.execute(query); + const elapsed = performance.now() - start; + console.log(`Rows returned: ${result.length}`); + console.log(`Wall-clock time: ${elapsed.toFixed(2)}ms`); + console.log("Result:", JSON.stringify(result, null, 2)); + console.log(); + + console.log("--- EXPLAIN (ANALYZE, BUFFERS) ---\n"); + const explainQuery = sql`EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) ${query}`; + const explainResult = await db.execute(explainQuery); + + for (const row of explainResult) { + const line = (row as Record)["QUERY PLAN"]; + console.log(line); + } + + process.exit(0); +}; + +await main(); diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/contextUtilsV2.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/contextUtilsV2.lua index bad7de72b..8e30a8a22 100644 --- a/server/src/_luaScriptsV2/fullSubjectDeduction/contextUtilsV2.lua +++ b/server/src/_luaScriptsV2/fullSubjectDeduction/contextUtilsV2.lua @@ -39,10 +39,12 @@ local function init_context(params) missing_customer_entitlement_ids = read_result.missing_customer_entitlement_ids or {}, logs = logs, - logger = { + logger = params.debug and { log = function(fmt, ...) table.insert(logs, string.format(fmt, ...)) end, + } or { + log = function() end, }, } diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromSubjectBalances.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromSubjectBalances.lua index 428500e5f..ed7dfa1cf 100644 --- a/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromSubjectBalances.lua +++ b/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromSubjectBalances.lua @@ -124,6 +124,7 @@ local context = init_context({ customer_id = customer_id, customer_entitlement_deductions = customer_entitlement_deductions, balance_keys_by_feature_id = params.balance_keys_by_feature_id, + debug = params.debug, }) if #(context.missing_customer_entitlement_ids or {}) > 0 then diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/lock/lockReceiptV2.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/lock/lockReceiptV2.lua new file mode 100644 index 000000000..ca9a2ee6b --- /dev/null +++ b/server/src/_luaScriptsV2/fullSubjectDeduction/lock/lockReceiptV2.lua @@ -0,0 +1,118 @@ +-- ============================================================================ +-- LOCK RECEIPT HELPERS (V2 — plain string storage) +-- V2 variant of deduction/lock/lockReceipt.lua. Stores the lock receipt as a +-- plain string via core Redis GET/SET instead of RedisJSON JSON.GET/JSON.SET. +-- +-- Why: fewer Redis calls per lock save (combined SET NX EXAT), no RedisJSON +-- module parse/serialize overhead. V2 full-subject cache has no legacy +-- RedisJSON lock receipts so plain GET is safe. +-- +-- Function names intentionally match deduction/lock/lockReceipt.lua so scripts +-- that import these helpers (deductFromSubjectBalances.lua, unwindLockV2.lua, +-- claimLockReceiptV2.lua) don't need call-site changes — only the bundle swaps +-- which helper file is included. +-- +-- Depends on: +-- - is_nil (luaUtils.lua) +-- ============================================================================ + +-- ============================================================================ +-- HELPER: Load lock receipt from Redis +-- Returns the decoded receipt table, or nil if the key does not exist / is +-- unparseable. +-- ============================================================================ +local function load_lock_receipt(lock_receipt_key) + local raw = redis.call('GET', lock_receipt_key) + if is_nil(raw) or raw == false then + return nil + end + + local ok, decoded = pcall(cjson.decode, raw) + if not ok or type(decoded) ~= 'table' then + return nil + end + + return decoded +end + +-- ============================================================================ +-- HELPER: Store lock receipt in Redis (full overwrite, no TTL change) +-- Use `save_lock_receipt_from_updates` when you want to set a TTL atomically. +-- Use this when updating an existing receipt whose TTL should be preserved — +-- callers should pair with `SET ... KEEPTTL` instead of this helper if they +-- want to keep the existing TTL; plain SET here clears it. +-- ============================================================================ +local function store_lock_receipt(lock_receipt_key, receipt) + redis.call('SET', lock_receipt_key, cjson.encode(receipt)) + return receipt +end + +-- ============================================================================ +-- HELPER: Overwrite an existing lock receipt and preserve its TTL. +-- Used by the claim path (pending -> processing) where we mutate one field and +-- write the full receipt back without resetting the expiry. +-- ============================================================================ +local function store_lock_receipt_keep_ttl(lock_receipt_key, receipt) + redis.call('SET', lock_receipt_key, cjson.encode(receipt), 'KEEPTTL') + return receipt +end + +-- ============================================================================ +-- HELPER: Save a lock receipt from deduction update objects. +-- +-- Combines the previous `JSON.SET + EXPIREAT` into a single `SET ... EXAT` so +-- the whole operation is one Redis round trip. When `ttl_at` is nil the TTL is +-- omitted (the key will persist until explicitly deleted or overwritten). +-- +-- params: +-- lock_receipt_key: string +-- receipt: table (base receipt metadata to persist) +-- mutation_logs: table | nil +-- ttl_at: number | nil (Unix seconds for EXAT) +-- ============================================================================ +local function save_lock_receipt_from_updates(params) + local receipt = params.receipt or {} + local mutation_logs = params.mutation_logs or {} + receipt.items = #mutation_logs > 0 and mutation_logs or cjson.decode('[]') + + local encoded = cjson.encode(receipt) + + if not is_nil(params.ttl_at) then + redis.call('SET', params.lock_receipt_key, encoded, 'EXAT', params.ttl_at) + else + redis.call('SET', params.lock_receipt_key, encoded) + end + + return receipt +end + +-- ============================================================================ +-- HELPER: Atomically create a lock receipt if no receipt exists at that key. +-- Returns true on success, false if a receipt is already present. +-- +-- This is the Lua analogue of the TS `SET NX EXAT` path in saveLockReceipt.ts. +-- Use when you want a single-RT "create if absent" (replaces EXISTS + SET + +-- EXPIREAT). +-- +-- params: +-- lock_receipt_key: string +-- receipt: table +-- mutation_logs: table | nil +-- ttl_at: number | nil (Unix seconds for EXAT) +-- ============================================================================ +local function create_lock_receipt_if_absent(params) + local receipt = params.receipt or {} + local mutation_logs = params.mutation_logs or {} + receipt.items = #mutation_logs > 0 and mutation_logs or cjson.decode('[]') + + local encoded = cjson.encode(receipt) + + local result + if not is_nil(params.ttl_at) then + result = redis.call('SET', params.lock_receipt_key, encoded, 'NX', 'EXAT', params.ttl_at) + else + result = redis.call('SET', params.lock_receipt_key, encoded, 'NX') + end + + return result == 'OK' +end diff --git a/server/src/db/initDrizzle.ts b/server/src/db/initDrizzle.ts index 136efb769..743f77189 100644 --- a/server/src/db/initDrizzle.ts +++ b/server/src/db/initDrizzle.ts @@ -55,6 +55,7 @@ export const initDrizzle = ({ const envDbUrl = replica ? process.env.DATABASE_REPLICA_URL : process.env.DATABASE_URL; + const dbUrl = databaseUrl || envDbUrl || ""; const client = new pg.Pool({ diff --git a/server/src/db/redactDatabaseUrl.ts b/server/src/db/redactDatabaseUrl.ts index 635d81099..c5095939e 100644 --- a/server/src/db/redactDatabaseUrl.ts +++ b/server/src/db/redactDatabaseUrl.ts @@ -3,17 +3,12 @@ import { createHash } from "node:crypto"; const hash = (value: string) => createHash("sha256").update(value).digest("hex").slice(0, 12); -const mask = (value = "") => - value.length > 6 - ? `${value.slice(0, 3)}***${value.slice(-3)}` - : `${value[0] ?? ""}***${value.slice(-1)}`; - const auth = (url: URL) => { const username = decodeURIComponent(url.username); const password = decodeURIComponent(url.password); if (!username && !password) return ""; - return `${mask(username)}${password ? `:${mask(password)}` : ""}@`; + return `${username}${password ? ":***" : ""}@`; }; const formatUrl = (value: string) => { diff --git a/server/src/internal/balances/autoTopUp/triggerAutoTopUp.ts b/server/src/internal/balances/autoTopUp/triggerAutoTopUp.ts index 487bb6721..f9c2fb3a3 100644 --- a/server/src/internal/balances/autoTopUp/triggerAutoTopUp.ts +++ b/server/src/internal/balances/autoTopUp/triggerAutoTopUp.ts @@ -28,10 +28,6 @@ export const triggerAutoTopUp = async ({ featureId: relevantFeature.id, }); - console.log( - `resolved, feature ${relevantFeature.id}, balance below threshold: ${resolved?.balanceBelowThreshold}, customerEntitlement balance: ${resolved?.customerEntitlement.balance}`, - ); - if (!resolved?.balanceBelowThreshold) continue; // Enqueue the auto top-up job diff --git a/server/src/internal/balances/check/runCheckWithRollout.ts b/server/src/internal/balances/check/runCheckWithRollout.ts index d6702cbc7..08346bb7d 100644 --- a/server/src/internal/balances/check/runCheckWithRollout.ts +++ b/server/src/internal/balances/check/runCheckWithRollout.ts @@ -29,7 +29,7 @@ export const runCheckWithRollout = async ({ body, requiredBalance, }), - catch: (error) => error, + catch: (error: unknown) => error, }); if (Result.isOk(result)) return result.value; diff --git a/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts b/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts index 7b94a1e2c..e1b3a769f 100644 --- a/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts +++ b/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts @@ -160,6 +160,7 @@ export const executeRedisDeductionV2 = async ({ } : null, unwind_value: unwindValue ?? null, + debug: process.env.NODE_ENV !== "production", }; const targetRedis = redisInstance ?? ctx.redisV2; @@ -262,6 +263,9 @@ export const executeRedisDeductionV2 = async ({ }); } } catch (error) { + // if (error.message?.includes("declined")) { + // return; + // } if (error instanceof Error && !error?.message?.includes("declined")) { ctx.logger.error( `[executeRedisDeductionV2] Attempting rollback due to error: ${error}`, diff --git a/server/src/internal/balances/utils/refreshEntityAggregate/RefreshEntityAggregateBatchingManager.ts b/server/src/internal/balances/utils/refreshEntityAggregate/RefreshEntityAggregateBatchingManager.ts new file mode 100644 index 000000000..6fa4f597a --- /dev/null +++ b/server/src/internal/balances/utils/refreshEntityAggregate/RefreshEntityAggregateBatchingManager.ts @@ -0,0 +1,196 @@ +import type { AppEnv } from "@autumn/shared"; +import { logger } from "@/external/logtail/logtailUtils.js"; +import { currentRegion } from "@/external/redis/initRedis.js"; +import { JobName } from "@/queue/JobName.js"; +import { addTaskToQueue } from "@/queue/queueUtils.js"; +import { + buildRefreshEntityAggregateDedupId, + REFRESH_ENTITY_AGGREGATE_DEDUP_BUCKET_MS, + REFRESH_ENTITY_AGGREGATE_SETTLE_BUFFER_MS, +} from "./queueRefreshEntityAggregate.js"; + +export type QueueRefreshEntityAggregatePayload = { + jobName: string; + payload: { + customerId: string; + orgId: string; + env: AppEnv; + region: string; + internalFeatureIds: string[]; + }; + messageGroupId: string; + messageDeduplicationId: string; +}; + +interface RefreshEntry { + customerId: string; + orgId: string; + env: AppEnv; + internalFeatureIds: Set; + timer: NodeJS.Timeout | null; +} + +/** + * Coalesces `RefreshEntityAggregate` enqueues per (org, env, customer). + * + * All `schedule()` calls that arrive inside the same 5s bucket share one + * timer aimed at the bucket's trailing edge (+ settle buffer). When the timer + * fires, the manager emits a single SQS message with the merged + * `internalFeatureIds`. Because we're in the same worker process that + * consumes the customer's sync-v4 stream (FIFO MessageGroupId already + * serializes that), this in-memory dedup is effectively per-customer global. + */ +export class RefreshEntityAggregateBatchingManager { + private entries: Map = new Map(); + + private readonly bucketMs: number; + private readonly settleBufferMs: number; + + private readonly _addTaskToQueue: ( + args: QueueRefreshEntityAggregatePayload, + ) => Promise; + private readonly _now: () => number; + + constructor({ + addTaskToQueueFn, + bucketMs, + settleBufferMs, + now, + }: { + addTaskToQueueFn?: ( + args: QueueRefreshEntityAggregatePayload, + ) => Promise; + bucketMs?: number; + settleBufferMs?: number; + now?: () => number; + } = {}) { + this._addTaskToQueue = + addTaskToQueueFn ?? + (addTaskToQueue as unknown as ( + args: QueueRefreshEntityAggregatePayload, + ) => Promise); + this.bucketMs = bucketMs ?? REFRESH_ENTITY_AGGREGATE_DEDUP_BUCKET_MS; + this.settleBufferMs = + settleBufferMs ?? REFRESH_ENTITY_AGGREGATE_SETTLE_BUFFER_MS; + this._now = now ?? Date.now; + } + + schedule({ + orgId, + env, + customerId, + internalFeatureIds, + }: { + orgId: string; + env: AppEnv; + customerId: string; + internalFeatureIds: string[]; + }): void { + const key = this.buildKey({ orgId, env, customerId }); + const existing = this.entries.get(key); + if (existing) { + for (const id of internalFeatureIds) { + existing.internalFeatureIds.add(id); + } + return; + } + + const entry: RefreshEntry = { + customerId, + orgId, + env, + internalFeatureIds: new Set(internalFeatureIds), + timer: null, + }; + this.entries.set(key, entry); + this.scheduleTimer({ key, entry }); + } + + async flush(): Promise { + const keys = Array.from(this.entries.keys()); + await Promise.all(keys.map((key) => this.fire({ key }))); + } + + getStats(): { totalPending: number } { + return { totalPending: this.entries.size }; + } + + private buildKey({ + orgId, + env, + customerId, + }: { + orgId: string; + env: AppEnv; + customerId: string; + }): string { + return `${orgId}:${env}:${customerId}`; + } + + private scheduleTimer({ + key, + entry, + }: { + key: string; + entry: RefreshEntry; + }): void { + const nowMs = this._now(); + const bucketEndMs = + (Math.floor(nowMs / this.bucketMs) + 1) * this.bucketMs; + const delayMs = bucketEndMs - nowMs + this.settleBufferMs; + + entry.timer = setTimeout(() => { + this.fire({ key }); + }, delayMs); + + if (entry.timer.unref) entry.timer.unref(); + } + + private async fire({ key }: { key: string }): Promise { + const entry = this.entries.get(key); + if (!entry) return; + + if (entry.timer) { + clearTimeout(entry.timer); + entry.timer = null; + } + this.entries.delete(key); + + if (entry.internalFeatureIds.size === 0) return; + + const internalFeatureIds = Array.from(entry.internalFeatureIds).sort(); + const messageDeduplicationId = buildRefreshEntityAggregateDedupId({ + orgId: entry.orgId, + env: entry.env, + customerId: entry.customerId, + nowMs: this._now(), + bucketMs: this.bucketMs, + }); + + try { + await this._addTaskToQueue({ + jobName: JobName.RefreshEntityAggregate, + payload: { + customerId: entry.customerId, + orgId: entry.orgId, + env: entry.env, + region: currentRegion, + internalFeatureIds, + }, + messageGroupId: `refresh-agg:${entry.orgId}:${entry.env}:${entry.customerId}`, + messageDeduplicationId, + }); + + logger.info( + `[RefreshEntityAggregate] Queued refresh for ${entry.customerId}, ${internalFeatureIds.length} features`, + ); + } catch (error) { + logger.error( + `[RefreshEntityAggregate] Failed to queue refresh for ${entry.customerId}: ${error}`, + ); + } + } +} + +export const globalRefreshEntityAggregateBatchingManager = + new RefreshEntityAggregateBatchingManager(); diff --git a/server/src/internal/balances/utils/refreshEntityAggregate/index.ts b/server/src/internal/balances/utils/refreshEntityAggregate/index.ts new file mode 100644 index 000000000..a82dfa5c8 --- /dev/null +++ b/server/src/internal/balances/utils/refreshEntityAggregate/index.ts @@ -0,0 +1,11 @@ +export { + buildRefreshEntityAggregateDedupId, + REFRESH_ENTITY_AGGREGATE_DEDUP_BUCKET_MS, + REFRESH_ENTITY_AGGREGATE_SETTLE_BUFFER_MS, +} from "./queueRefreshEntityAggregate.js"; +export { refreshEntityAggregateCache } from "./refreshEntityAggregateCache.js"; +export { + globalRefreshEntityAggregateBatchingManager, + RefreshEntityAggregateBatchingManager, + type QueueRefreshEntityAggregatePayload, +} from "./RefreshEntityAggregateBatchingManager.js"; diff --git a/server/src/internal/balances/utils/refreshEntityAggregate/queueRefreshEntityAggregate.ts b/server/src/internal/balances/utils/refreshEntityAggregate/queueRefreshEntityAggregate.ts new file mode 100644 index 000000000..d08a84080 --- /dev/null +++ b/server/src/internal/balances/utils/refreshEntityAggregate/queueRefreshEntityAggregate.ts @@ -0,0 +1,40 @@ +import type { AppEnv } from "@autumn/shared"; +import { JobName } from "@/queue/JobName.js"; + +export const REFRESH_ENTITY_AGGREGATE_DEDUP_BUCKET_MS = + process.env.NODE_ENV === "development" ? 1000 : 5000; + +/** + * Buffer added after the bucket boundary so the trailing enqueue fires *after* + * the final sync batch of the bucket has settled into Postgres. + * `SyncBatchingManagerV3` uses a 1s tumbling window, so 1.5s is enough. + */ +export const REFRESH_ENTITY_AGGREGATE_SETTLE_BUFFER_MS = 1500; + +/** + * Deterministic dedup id per (org, env, customer) within a time bucket. + * Rapid schedules inside the same bucket collapse to a single SQS message. + */ +export const buildRefreshEntityAggregateDedupId = ({ + orgId, + env, + customerId, + nowMs, + bucketMs = REFRESH_ENTITY_AGGREGATE_DEDUP_BUCKET_MS, +}: { + orgId: string; + env: AppEnv; + customerId: string; + nowMs: number; + bucketMs?: number; +}): string => { + const bucket = Math.floor(nowMs / bucketMs); + const key = JSON.stringify({ + jobName: JobName.RefreshEntityAggregate, + orgId, + env, + customerId, + bucket, + }); + return Bun.hash(key).toString(); +}; diff --git a/server/src/internal/balances/utils/sync/refreshEntityAggregateCache.ts b/server/src/internal/balances/utils/refreshEntityAggregate/refreshEntityAggregateCache.ts similarity index 56% rename from server/src/internal/balances/utils/sync/refreshEntityAggregateCache.ts rename to server/src/internal/balances/utils/refreshEntityAggregate/refreshEntityAggregateCache.ts index 55c54ce07..b940c09bf 100644 --- a/server/src/internal/balances/utils/sync/refreshEntityAggregateCache.ts +++ b/server/src/internal/balances/utils/refreshEntityAggregate/refreshEntityAggregateCache.ts @@ -1,9 +1,8 @@ -import type { AppEnv } from "@autumn/shared"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { buildSharedFullSubjectBalanceKey } from "@/internal/customers/cache/fullSubject/builders/buildSharedFullSubjectBalanceKey.js"; import { AGGREGATED_BALANCE_FIELD } from "@/internal/customers/cache/fullSubject/config/fullSubjectCacheConfig.js"; -import { tryRedisWrite } from "@/utils/cacheUtils/cacheUtils.js"; import { getEntityAggregateForSync } from "@/internal/customers/repos/getFullSubject/getEntityAggregateForSync.js"; +import { tryRedisWrite } from "@/utils/cacheUtils/cacheUtils.js"; /** * After DB sync, recompute entity aggregation from the now-authoritative DB @@ -12,33 +11,65 @@ import { getEntityAggregateForSync } from "@/internal/customers/repos/getFullSub export const refreshEntityAggregateCache = async ({ ctx, customerId, - orgId, - env, - featureIds, + internalFeatureIds, }: { ctx: AutumnContext; customerId: string; - orgId: string; - env: AppEnv; - featureIds: string[]; + internalFeatureIds: string[]; }): Promise => { try { + const orgId = ctx.org.id; + const env = ctx.env; + const internalIdSet = new Set(internalFeatureIds); + const featureIds = ctx.features + .filter((feature) => internalIdSet.has(feature.internal_id)) + .map((feature) => feature.id); + + if (featureIds.length === 0) return; + + const { redisV2 } = ctx; + + // Only refresh features whose balance hash already has `_aggregated`. + // If none of the hashes have it cached, there is nothing to refresh — + // avoid the expensive CTE entirely. + const existsPipeline = redisV2.pipeline(); + for (const featureId of featureIds) { + const balanceKey = buildSharedFullSubjectBalanceKey({ + orgId, + env, + customerId, + featureId, + }); + existsPipeline.hexists(balanceKey, AGGREGATED_BALANCE_FIELD); + } + const existsResults = (await existsPipeline.exec()) ?? []; + const featuresWithAggregated = new Set(); + existsResults.forEach(([, exists], idx) => { + if (exists === 1) featuresWithAggregated.add(featureIds[idx]); + }); + + if (featuresWithAggregated.size === 0) { + ctx.logger.info( + `[SYNC V4] (${customerId}) No _aggregated fields cached — skipping refresh`, + ); + return; + } + const aggregated = await getEntityAggregateForSync({ db: ctx.db, orgId, env, customerId, + internalFeatureIds, }); if (aggregated.length === 0) return; - const affectedFeatureIds = new Set(featureIds); - const { redisV2 } = ctx; const pipeline = redisV2.pipeline(); let writeCount = 0; for (const entry of aggregated) { - if (!affectedFeatureIds.has(entry.feature_id)) continue; + if (!featuresWithAggregated.has(entry.feature_id)) continue; const balanceKey = buildSharedFullSubjectBalanceKey({ orgId, diff --git a/server/src/internal/balances/utils/sync/SyncBatchingManagerV3.ts b/server/src/internal/balances/utils/sync/SyncBatchingManagerV3.ts index 7c55d7790..02ed944b0 100644 --- a/server/src/internal/balances/utils/sync/SyncBatchingManagerV3.ts +++ b/server/src/internal/balances/utils/sync/SyncBatchingManagerV3.ts @@ -293,6 +293,7 @@ export class SyncBatchingManagerV3 { entityId: context.entityId, modifiedCusEntIdsByFeatureId: context.modifiedCusEntIdsByFeatureId, }, + messageGroupId: `sync-v4:${context.orgId}:${context.env}:${context.customerId}`, messageDeduplicationId, }); diff --git a/server/src/internal/balances/utils/sync/syncItemV4.ts b/server/src/internal/balances/utils/sync/syncItemV4.ts index 6987b026c..09da61f45 100644 --- a/server/src/internal/balances/utils/sync/syncItemV4.ts +++ b/server/src/internal/balances/utils/sync/syncItemV4.ts @@ -7,9 +7,9 @@ import { } from "@autumn/shared"; import { sql } from "drizzle-orm"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; -import { deleteCachedFullCustomer } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/deleteCachedFullCustomer.js"; import { getCachedFeatureBalance } from "@/internal/customers/cache/fullSubject/balances/getCachedFeatureBalances.js"; -import { refreshEntityAggregateCache } from "./refreshEntityAggregateCache.js"; +import { deleteCachedFullCustomer } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/deleteCachedFullCustomer.js"; +import { globalRefreshEntityAggregateBatchingManager } from "../refreshEntityAggregate/index.js"; const SYNC_CONFLICT_CODES = { ResetAtMismatch: "RESET_AT_MISMATCH", @@ -135,8 +135,6 @@ export const syncItemV4 = async ({ const { customerId, entityId, - orgId, - env, rolloverIds, modifiedCusEntIdsByFeatureId, } = payload; @@ -244,12 +242,16 @@ export const syncItemV4 = async ({ (subjectBalance) => subjectBalance.isEntityLevel, ); if (hasEntityLevel) { - await refreshEntityAggregateCache({ - ctx, + const featureIds = Object.keys(modifiedCusEntIdsByFeatureId); + const internalFeatureIds = ctx.features + .filter((feature) => featureIds.includes(feature.id)) + .map((feature) => feature.internal_id); + + globalRefreshEntityAggregateBatchingManager.schedule({ + orgId: ctx.org.id, + env: ctx.env, customerId, - orgId, - env, - featureIds: Object.keys(modifiedCusEntIdsByFeatureId), + internalFeatureIds, }); } }; diff --git a/server/src/internal/customers/repos/getFullSubject/getEntityAggregateForSync.ts b/server/src/internal/customers/repos/getFullSubject/getEntityAggregateForSync.ts index 9ccfc18d4..406d01645 100644 --- a/server/src/internal/customers/repos/getFullSubject/getEntityAggregateForSync.ts +++ b/server/src/internal/customers/repos/getFullSubject/getEntityAggregateForSync.ts @@ -20,12 +20,14 @@ export const getEntityAggregateForSync = async ({ orgId, env, customerId, + internalFeatureIds, inStatuses = RELEVANT_STATUSES, }: { db: DrizzleCli; orgId: string; env: AppEnv; customerId: string; + internalFeatureIds?: string[]; inStatuses?: CusProductStatus[]; }): Promise => { const statusFilter = @@ -38,6 +40,7 @@ export const getEntityAggregateForSync = async ({ const entityFragments = getEntityAggregateFragments({ statusFilter, + internalFeatureIds, }); const query = sql` diff --git a/server/src/internal/customers/repos/getFullSubject/getEntityAggregateFragments.ts b/server/src/internal/customers/repos/getFullSubject/getEntityAggregateFragments.ts index 5410346e1..cae7defa7 100644 --- a/server/src/internal/customers/repos/getFullSubject/getEntityAggregateFragments.ts +++ b/server/src/internal/customers/repos/getFullSubject/getEntityAggregateFragments.ts @@ -2,20 +2,27 @@ import { type SQL, sql } from "drizzle-orm"; import { getEntityOptionsAggregateFragments } from "./getEntityOptionsAggregateFragments.js"; /** - * Per-entity rollover rows sourced from the `rollovers` table, mirroring the - * four branches in `entity_balance_rows` (product-attached / loose, per-entity - * jsonb / top-level). Top-level rollovers are attributed to the owning entity - * via `cp.internal_entity_id` (product-attached) or `ce.internal_entity_id` - * (loose), matching main-balance behaviour. Also exposes per-entity and - * per-feature rollups used by the outer aggregate CTEs. + * Rollover CTEs driven from the shared `entity_product_cus_ents` and + * `entity_loose_cus_ents` base CTEs — avoids re-scanning `customer_entitlements` + * and `customer_products` per branch. */ -const buildEntityRolloverCtes = ({ - statusFilter, -}: { - statusFilter: SQL; -}) => sql` +const buildEntityRolloverCtes = () => sql` entity_rollover_rows AS ( - -- Product-attached cusEnt, per-entity rollover (rollovers.entities jsonb) + -- Top-level rollover: one row per (rollover × cus_ent). + -- entity_key = cp.internal_entity_id for product-attached, ce.internal_entity_id for loose. + SELECT + ce.internal_feature_id, + ce.internal_customer_id, + COALESCE(ce.cp_entity_key, ce.internal_entity_id) AS entity_key, + r.balance::numeric AS rollover_balance, + COALESCE(r.usage, 0)::numeric AS rollover_usage + FROM rollovers r + JOIN entity_level_cus_ents ce ON r.cus_ent_id = ce.id + WHERE (r.expires_at IS NULL OR r.expires_at > EXTRACT(EPOCH FROM now()) * 1000) + + UNION ALL + + -- Per-entity rollover from jsonb_each(r.entities). SELECT ce.internal_feature_id, ce.internal_customer_id, @@ -23,67 +30,10 @@ const buildEntityRolloverCtes = ({ COALESCE((kv.entity_value->>'balance')::numeric, 0) AS rollover_balance, COALESCE((kv.entity_value->>'usage')::numeric, 0) AS rollover_usage FROM rollovers r - JOIN customer_entitlements ce ON r.cus_ent_id = ce.id - JOIN customer_products cp ON ce.customer_product_id = cp.id + JOIN entity_level_cus_ents ce ON r.cus_ent_id = ce.id CROSS JOIN LATERAL jsonb_each(r.entities) AS kv(entity_key, entity_value) - WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records) - AND cp.internal_entity_id IS NOT NULL - AND jsonb_typeof(r.entities) = 'object' + WHERE jsonb_typeof(r.entities) = 'object' AND (r.expires_at IS NULL OR r.expires_at > EXTRACT(EPOCH FROM now()) * 1000) - ${statusFilter} - - UNION ALL - - -- Product-attached cusEnt, top-level rollover (attributed to cp.internal_entity_id) - SELECT - ce.internal_feature_id, - ce.internal_customer_id, - cp.internal_entity_id AS entity_key, - r.balance::numeric AS rollover_balance, - COALESCE(r.usage, 0)::numeric AS rollover_usage - FROM rollovers r - JOIN customer_entitlements ce ON r.cus_ent_id = ce.id - JOIN customer_products cp ON ce.customer_product_id = cp.id - WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records) - AND cp.internal_entity_id IS NOT NULL - AND (r.expires_at IS NULL OR r.expires_at > EXTRACT(EPOCH FROM now()) * 1000) - ${statusFilter} - - UNION ALL - - -- Loose cusEnt (no customer_product), per-entity rollover - SELECT - ce.internal_feature_id, - ce.internal_customer_id, - kv.entity_key AS entity_key, - COALESCE((kv.entity_value->>'balance')::numeric, 0) AS rollover_balance, - COALESCE((kv.entity_value->>'usage')::numeric, 0) AS rollover_usage - FROM rollovers r - JOIN customer_entitlements ce ON r.cus_ent_id = ce.id - CROSS JOIN LATERAL jsonb_each(r.entities) AS kv(entity_key, entity_value) - WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records) - AND ce.customer_product_id IS NULL - AND ce.internal_entity_id IS NOT NULL - AND jsonb_typeof(r.entities) = 'object' - AND (r.expires_at IS NULL OR r.expires_at > EXTRACT(EPOCH FROM now()) * 1000) - AND (ce.expires_at IS NULL OR ce.expires_at > EXTRACT(EPOCH FROM now()) * 1000) - - UNION ALL - - -- Loose cusEnt, top-level rollover (attributed to ce.internal_entity_id) - SELECT - ce.internal_feature_id, - ce.internal_customer_id, - ce.internal_entity_id AS entity_key, - r.balance::numeric AS rollover_balance, - COALESCE(r.usage, 0)::numeric AS rollover_usage - FROM rollovers r - JOIN customer_entitlements ce ON r.cus_ent_id = ce.id - WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records) - AND ce.customer_product_id IS NULL - AND ce.internal_entity_id IS NOT NULL - AND (r.expires_at IS NULL OR r.expires_at > EXTRACT(EPOCH FROM now()) * 1000) - AND (ce.expires_at IS NULL OR ce.expires_at > EXTRACT(EPOCH FROM now()) * 1000) ), entity_rollover_keys AS ( @@ -112,10 +62,20 @@ const buildEntityRolloverCtes = ({ export const getEntityAggregateFragments = ({ entityId, statusFilter, + internalFeatureIds, }: { entityId?: string; statusFilter: SQL; + internalFeatureIds?: string[]; }) => { + const featureFilter = + internalFeatureIds && internalFeatureIds.length > 0 + ? sql`AND ce.internal_feature_id = ANY(ARRAY[${sql.join( + internalFeatureIds.map((internalFeatureId) => sql`${internalFeatureId}`), + sql`, `, + )}])` + : sql``; + if (entityId) { return { ctes: sql``, @@ -131,39 +91,52 @@ export const getEntityAggregateFragments = ({ const ctes = sql`, - entity_distinct_product_ids AS ( - SELECT DISTINCT cp.internal_product_id, cp.internal_customer_id - FROM customer_products cp - JOIN subject_customer_records scr - ON cp.internal_customer_id = scr.internal_id - WHERE cp.internal_entity_id IS NOT NULL - ${statusFilter} - ), - - entity_distinct_cus_products AS ( - SELECT sub.* - FROM entity_distinct_product_ids edpi - JOIN LATERAL ( - SELECT cp.* - FROM customer_products cp - WHERE cp.internal_customer_id = edpi.internal_customer_id - AND cp.internal_product_id = edpi.internal_product_id - AND cp.internal_entity_id IS NOT NULL - ${statusFilter} - ORDER BY cp.created_at DESC - LIMIT 1 - ) sub ON true - ), - - entity_cus_products_for_options AS ( + -- (A) Subject's entity-level customer_products, filtered once. + entity_cus_products AS ( SELECT cp.* FROM customer_products cp - JOIN subject_customer_records scr - ON cp.internal_customer_id = scr.internal_id - WHERE cp.internal_entity_id IS NOT NULL + WHERE cp.internal_customer_id IN (SELECT internal_id FROM subject_customer_records) + AND cp.internal_entity_id IS NOT NULL ${statusFilter} ), + -- (B) All entity-level cus_ents for the subject — product-attached (with + -- cp_entity_key set) and loose (cp_entity_key = NULL). Filtered once. + entity_level_cus_ents AS ( + SELECT ce.*, cp.internal_entity_id AS cp_entity_key + FROM entity_cus_products cp + JOIN customer_entitlements ce ON ce.customer_product_id = cp.id + WHERE 1 = 1 + ${featureFilter} + + UNION ALL + + SELECT ce.*, NULL::text AS cp_entity_key + FROM customer_entitlements ce + WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records) + AND ce.customer_product_id IS NULL + AND ce.internal_entity_id IS NOT NULL + AND (ce.expires_at IS NULL OR ce.expires_at > EXTRACT(EPOCH FROM now()) * 1000) + AND ce.balance != 0 + ${featureFilter} + ), + + -- Most-recent customer_product per (customer, product) — replaces the old + -- DISTINCT + LATERAL dance in entity_distinct_cus_products. + entity_distinct_cus_products AS ( + SELECT * + FROM ( + SELECT + cp.*, + ROW_NUMBER() OVER ( + PARTITION BY cp.internal_customer_id, cp.internal_product_id + ORDER BY cp.created_at DESC + ) AS rn + FROM entity_cus_products cp + ) ranked + WHERE ranked.rn = 1 + ), + entity_cus_prices AS ( SELECT cpr.* FROM customer_prices cpr @@ -173,82 +146,8 @@ export const getEntityAggregateFragments = ({ ${entityOptionsAggregateFragments.ctes}, entity_balance_rows AS ( - SELECT - COALESCE(ce.external_id, ce.id) AS api_id, - ce.internal_feature_id, - ce.internal_customer_id, - ce.feature_id, - COALESCE(ent.allowance, 0)::numeric AS allowance, - ce.balance::numeric AS balance, - ce.adjustment::numeric AS adjustment, - COALESCE(ce.additional_balance, 0)::numeric AS additional_balance, - ce.unlimited, - ce.usage_allowed, - cp.internal_entity_id AS entity_key, - ce.balance::numeric AS entity_balance, - COALESCE(ce.adjustment, 0)::numeric AS entity_adjustment, - COALESCE(ce.additional_balance, 0)::numeric AS entity_additional_balance - FROM customer_entitlements ce - JOIN customer_products cp ON ce.customer_product_id = cp.id - JOIN entitlements ent ON ce.entitlement_id = ent.id - WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records) - AND cp.internal_entity_id IS NOT NULL - ${statusFilter} - - UNION ALL - - SELECT - COALESCE(ce.external_id, ce.id) AS api_id, - ce.internal_feature_id, - ce.internal_customer_id, - ce.feature_id, - COALESCE(ent.allowance, 0)::numeric AS allowance, - 0::numeric AS balance, - 0::numeric AS adjustment, - 0::numeric AS additional_balance, - ce.unlimited, - ce.usage_allowed, - kv.entity_key AS entity_key, - (kv.entity_value->>'balance')::numeric AS entity_balance, - COALESCE((kv.entity_value->>'adjustment')::numeric, 0) AS entity_adjustment, - COALESCE((kv.entity_value->>'additional_balance')::numeric, 0) AS entity_additional_balance - FROM customer_entitlements ce - JOIN customer_products cp ON ce.customer_product_id = cp.id - JOIN entitlements ent ON ce.entitlement_id = ent.id - CROSS JOIN LATERAL jsonb_each(ce.entities) AS kv(entity_key, entity_value) - WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records) - AND cp.internal_entity_id IS NOT NULL - AND jsonb_typeof(ce.entities) = 'object' - ${statusFilter} - - UNION ALL - - SELECT - COALESCE(ce.external_id, ce.id) AS api_id, - ce.internal_feature_id, - ce.internal_customer_id, - ce.feature_id, - COALESCE(ent.allowance, 0)::numeric AS allowance, - 0::numeric AS balance, - 0::numeric AS adjustment, - 0::numeric AS additional_balance, - ce.unlimited, - ce.usage_allowed, - kv.entity_key AS entity_key, - (kv.entity_value->>'balance')::numeric AS entity_balance, - COALESCE((kv.entity_value->>'adjustment')::numeric, 0) AS entity_adjustment, - COALESCE((kv.entity_value->>'additional_balance')::numeric, 0) AS entity_additional_balance - FROM customer_entitlements ce - JOIN entitlements ent ON ce.entitlement_id = ent.id - CROSS JOIN LATERAL jsonb_each(ce.entities) AS kv(entity_key, entity_value) - WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records) - AND ce.customer_product_id IS NULL - AND ce.internal_entity_id IS NOT NULL - AND jsonb_typeof(ce.entities) = 'object' - AND (ce.expires_at IS NULL OR ce.expires_at > EXTRACT(EPOCH FROM now()) * 1000) - - UNION ALL - + -- Top-level: one row per cus_ent (product-attached or loose). + -- entity_key = cp.internal_entity_id for product-attached, ce.internal_entity_id for loose. SELECT COALESCE(ce.external_id, ce.id) AS api_id, ce.internal_feature_id, @@ -260,19 +159,39 @@ export const getEntityAggregateFragments = ({ COALESCE(ce.additional_balance, 0)::numeric AS additional_balance, ce.unlimited, ce.usage_allowed, - ce.internal_entity_id AS entity_key, + COALESCE(ce.cp_entity_key, ce.internal_entity_id) AS entity_key, ce.balance::numeric AS entity_balance, COALESCE(ce.adjustment, 0)::numeric AS entity_adjustment, COALESCE(ce.additional_balance, 0)::numeric AS entity_additional_balance - FROM customer_entitlements ce + FROM entity_level_cus_ents ce JOIN entitlements ent ON ce.entitlement_id = ent.id - WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records) - AND ce.customer_product_id IS NULL - AND ce.internal_entity_id IS NOT NULL - AND (ce.expires_at IS NULL OR ce.expires_at > EXTRACT(EPOCH FROM now()) * 1000) + + UNION ALL + + -- Per-entity: N rows per cus_ent from jsonb_each(ce.entities). + -- balance/adj/additional = 0 to avoid double-counting at the aggregate level. + SELECT + COALESCE(ce.external_id, ce.id) AS api_id, + ce.internal_feature_id, + ce.internal_customer_id, + ce.feature_id, + COALESCE(ent.allowance, 0)::numeric AS allowance, + 0::numeric AS balance, + 0::numeric AS adjustment, + 0::numeric AS additional_balance, + ce.unlimited, + ce.usage_allowed, + kv.entity_key AS entity_key, + (kv.entity_value->>'balance')::numeric AS entity_balance, + COALESCE((kv.entity_value->>'adjustment')::numeric, 0) AS entity_adjustment, + COALESCE((kv.entity_value->>'additional_balance')::numeric, 0) AS entity_additional_balance + FROM entity_level_cus_ents ce + JOIN entitlements ent ON ce.entitlement_id = ent.id + CROSS JOIN LATERAL jsonb_each(ce.entities) AS kv(entity_key, entity_value) + WHERE jsonb_typeof(ce.entities) = 'object' ), - ${buildEntityRolloverCtes({ statusFilter })}, + ${buildEntityRolloverCtes()}, entity_balance_keys AS ( SELECT @@ -389,21 +308,7 @@ export const getEntityAggregateFragments = ({ SELECT DISTINCT ce.internal_customer_id, ce.entitlement_id - FROM customer_entitlements ce - JOIN customer_products cp ON ce.customer_product_id = cp.id - WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records) - AND cp.internal_entity_id IS NOT NULL - ${statusFilter} - - UNION - SELECT DISTINCT - ce.internal_customer_id, - ce.entitlement_id - FROM customer_entitlements ce - WHERE ce.internal_customer_id IN (SELECT internal_id FROM subject_customer_records) - AND ce.customer_product_id IS NULL - AND ce.internal_entity_id IS NOT NULL - AND (ce.expires_at IS NULL OR ce.expires_at > EXTRACT(EPOCH FROM now()) * 1000) + FROM entity_level_cus_ents ce `; const priceRefsUnion = sql` diff --git a/server/src/internal/customers/repos/getFullSubject/getEntityOptionsAggregateFragments.ts b/server/src/internal/customers/repos/getFullSubject/getEntityOptionsAggregateFragments.ts index bb91420a7..50efbb16a 100644 --- a/server/src/internal/customers/repos/getFullSubject/getEntityOptionsAggregateFragments.ts +++ b/server/src/internal/customers/repos/getFullSubject/getEntityOptionsAggregateFragments.ts @@ -19,7 +19,7 @@ export const getEntityOptionsAggregateFragments = () => { NULLIF(option_row.option_value->>'internal_feature_id', '') AS option_internal_feature_id, NULLIF(option_row.option_value->>'feature_id', '') AS option_feature_id, COALESCE((option_row.option_value->>'quantity')::numeric, 0) AS option_quantity - FROM entity_cus_products_for_options ecp + FROM entity_cus_products ecp CROSS JOIN LATERAL unnest( COALESCE(ecp.options, ARRAY[]::jsonb[]) ) AS option_row(option_value) @@ -35,7 +35,7 @@ export const getEntityOptionsAggregateFragments = () => { * COALESCE((prepaid_price.config->>'billing_units')::numeric, 1) AS prepaid_grant FROM entity_option_rows eor - JOIN customer_entitlements ce + JOIN entity_level_cus_ents ce ON ce.customer_product_id = eor.customer_product_id JOIN entitlements ent ON ent.id = ce.entitlement_id @@ -83,4 +83,3 @@ export const getEntityOptionsAggregateFragments = () => { ctes, }; }; - diff --git a/server/src/internal/customers/repos/getFullSubject/getFullSubjectQuery.ts b/server/src/internal/customers/repos/getFullSubject/getFullSubjectQuery.ts index 573280bc2..a11e6d83c 100644 --- a/server/src/internal/customers/repos/getFullSubject/getFullSubjectQuery.ts +++ b/server/src/internal/customers/repos/getFullSubject/getFullSubjectQuery.ts @@ -232,7 +232,9 @@ export const getFullSubjectQuery = ({ ON ce.internal_customer_id = scr.internal_id WHERE ce.customer_product_id IS NULL AND (ce.expires_at IS NULL OR ce.expires_at > EXTRACT(EPOCH FROM now()) * 1000) + AND ce.balance != 0 ${extraCustomerEntitlementEntityFilter} + LIMIT 20 ), all_cus_ent_ids AS ( diff --git a/server/src/queue/JobName.ts b/server/src/queue/JobName.ts index d8df93c52..bcf11058b 100644 --- a/server/src/queue/JobName.ts +++ b/server/src/queue/JobName.ts @@ -19,6 +19,7 @@ export enum JobName { SyncBalanceBatchV2 = "sync-balance-batch-v2", SyncBalanceBatchV3 = "sync-balance-batch-v3", SyncBalanceBatchV4 = "sync-balance-batch-v4", + RefreshEntityAggregate = "refresh-entity-aggregate", InsertEventBatch = "insert-event-batch", Track = "track", diff --git a/server/src/queue/processMessage.ts b/server/src/queue/processMessage.ts index 76ed7111f..54079df59 100644 --- a/server/src/queue/processMessage.ts +++ b/server/src/queue/processMessage.ts @@ -10,6 +10,7 @@ import { runActionHandlerTask } from "@/internal/analytics/runActionHandlerTask. import { autoTopup } from "@/internal/balances/autoTopUp/autoTopup.js"; import { runInsertEventBatch } from "@/internal/balances/events/runInsertEventBatch.js"; import { expireLock } from "@/internal/balances/finalizeLock/expireLock.js"; +import { refreshEntityAggregateCache } from "@/internal/balances/utils/refreshEntityAggregate/index.js"; import { syncItemV3 } from "@/internal/balances/utils/sync/syncItemV3.js"; import { syncItemV4 } from "@/internal/balances/utils/sync/syncItemV4.js"; import { grantCheckoutReward } from "@/internal/billing/v2/workflows/grantCheckoutReward/grantCheckoutReward.js"; @@ -174,6 +175,20 @@ export const processMessage = async ({ return; } + if (job.name === JobName.RefreshEntityAggregate) { + if (!ctx) { + workerLogger.error("No context found for refresh entity aggregate job"); + return; + } + + await refreshEntityAggregateCache({ + ctx, + customerId: job.data.customerId, + internalFeatureIds: job.data.internalFeatureIds, + }); + return; + } + if (job.name === JobName.InsertEventBatch) { await runInsertEventBatch({ db, @@ -278,7 +293,8 @@ export const processMessage = async ({ // won't fix on retry. DB errors (connection, timeout) will. if ( (job.name === JobName.SyncBalanceBatchV3 || - job.name === JobName.SyncBalanceBatchV4) && + job.name === JobName.SyncBalanceBatchV4 || + job.name === JobName.RefreshEntityAggregate) && isRetryableDbError({ error }) ) { Sentry.captureException(error); diff --git a/server/src/queue/queueUtils.ts b/server/src/queue/queueUtils.ts index 8012a3841..34bf6c92d 100644 --- a/server/src/queue/queueUtils.ts +++ b/server/src/queue/queueUtils.ts @@ -56,6 +56,13 @@ export interface Payloads { entityId?: string; modifiedCusEntIdsByFeatureId: Record; }; + [JobName.RefreshEntityAggregate]: { + customerId: string; + orgId: string; + env: AppEnv; + region?: string; + internalFeatureIds: string[]; + }; [JobName.InsertEventBatch]: { events: EventInsert[]; }; diff --git a/server/src/workers.ts b/server/src/workers.ts index 49b8bd031..bac4ac87f 100644 --- a/server/src/workers.ts +++ b/server/src/workers.ts @@ -26,6 +26,7 @@ if (cluster.isPrimary) { // await initHatchetWorker(); console.log(`Starting ${NUM_PROCESSES} worker processes`); + console.log(`SQS URL: ${process.env.SQS_QUEUE_URL}`); // Fork workers for (let i = 0; i < NUM_PROCESSES; i++) { diff --git a/server/tests/integration/balances/track/basic/track-basic.test.ts b/server/tests/integration/balances/track/basic/track-basic.test.ts index 22e1fd89c..97d8e64d6 100644 --- a/server/tests/integration/balances/track/basic/track-basic.test.ts +++ b/server/tests/integration/balances/track/basic/track-basic.test.ts @@ -1,7 +1,6 @@ import { expect, test } from "bun:test"; import type { ApiCustomerV3, TrackResponseV2 } from "@autumn/shared"; -import { track } from "@tests/_groups/domains/balances/track"; import { TestFeature } from "@tests/setup/v2Features.js"; import { items } from "@tests/utils/fixtures/items.js"; import { products } from "@tests/utils/fixtures/products.js"; diff --git a/server/tests/integration/others/refresh-entity-aggregate/refresh-entity-aggregate-dedup.test.ts b/server/tests/integration/others/refresh-entity-aggregate/refresh-entity-aggregate-dedup.test.ts new file mode 100644 index 000000000..3c9fd0c97 --- /dev/null +++ b/server/tests/integration/others/refresh-entity-aggregate/refresh-entity-aggregate-dedup.test.ts @@ -0,0 +1,220 @@ +import { describe, expect, test } from "bun:test"; +import { AppEnv } from "@autumn/shared"; +import chalk from "chalk"; +import { + buildRefreshEntityAggregateDedupId, + type QueueRefreshEntityAggregatePayload, + REFRESH_ENTITY_AGGREGATE_DEDUP_BUCKET_MS, + REFRESH_ENTITY_AGGREGATE_SETTLE_BUFFER_MS, + RefreshEntityAggregateBatchingManager, +} from "@/internal/balances/utils/refreshEntityAggregate/index.js"; +import { JobName } from "@/queue/JobName.js"; + +const wait = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); + +const createMockQueue = () => { + const calls: QueueRefreshEntityAggregatePayload[] = []; + const fn = async (args: QueueRefreshEntityAggregatePayload) => { + calls.push(structuredClone(args)); + }; + return { fn, calls }; +}; + +const baseArgs = { + orgId: "org-1", + env: AppEnv.Sandbox, + customerId: "cust-1", + internalFeatureIds: ["if-seats"], +}; + +describe("buildRefreshEntityAggregateDedupId", () => { + test( + `${chalk.yellowBright("dedup-id-1: same bucket → same id")}`, + () => { + const t0 = 1_700_000_000_000; + const a = buildRefreshEntityAggregateDedupId({ + orgId: "org-1", + env: AppEnv.Sandbox, + customerId: "cust-1", + nowMs: t0, + }); + const b = buildRefreshEntityAggregateDedupId({ + orgId: "org-1", + env: AppEnv.Sandbox, + customerId: "cust-1", + nowMs: t0 + 4999, + }); + expect(a).toBe(b); + }, + ); + + test( + `${chalk.yellowBright("dedup-id-2: across bucket boundary → different ids")}`, + () => { + const t0 = 1_700_000_000_000; + const a = buildRefreshEntityAggregateDedupId({ + orgId: "org-1", + env: AppEnv.Sandbox, + customerId: "cust-1", + nowMs: t0, + }); + const b = buildRefreshEntityAggregateDedupId({ + orgId: "org-1", + env: AppEnv.Sandbox, + customerId: "cust-1", + nowMs: t0 + REFRESH_ENTITY_AGGREGATE_DEDUP_BUCKET_MS, + }); + expect(a).not.toBe(b); + }, + ); + + test( + `${chalk.yellowBright("dedup-id-3: different orgs/envs/customers → different ids")}`, + () => { + const nowMs = 1_700_000_000_000; + const ids = new Set([ + buildRefreshEntityAggregateDedupId({ + orgId: "org-1", + env: AppEnv.Sandbox, + customerId: "cust-1", + nowMs, + }), + buildRefreshEntityAggregateDedupId({ + orgId: "org-2", + env: AppEnv.Sandbox, + customerId: "cust-1", + nowMs, + }), + buildRefreshEntityAggregateDedupId({ + orgId: "org-1", + env: AppEnv.Live, + customerId: "cust-1", + nowMs, + }), + buildRefreshEntityAggregateDedupId({ + orgId: "org-1", + env: AppEnv.Sandbox, + customerId: "cust-2", + nowMs, + }), + ]); + expect(ids.size).toBe(4); + }, + ); +}); + +describe("RefreshEntityAggregateBatchingManager", () => { + test( + `${chalk.yellowBright("batch-1: rapid schedule() calls in one bucket → 1 enqueue with merged features")}`, + async () => { + const { fn, calls } = createMockQueue(); + const manager = new RefreshEntityAggregateBatchingManager({ + addTaskToQueueFn: fn, + bucketMs: 100, + settleBufferMs: 20, + }); + + for (let i = 0; i < 20; i++) { + manager.schedule({ + ...baseArgs, + internalFeatureIds: [`if-${i % 3}`], + }); + } + + expect(calls.length).toBe(0); + + await wait(200); + + expect(calls.length).toBe(1); + expect(calls[0].jobName).toBe(JobName.RefreshEntityAggregate); + expect(calls[0].payload.customerId).toBe(baseArgs.customerId); + expect(calls[0].payload.internalFeatureIds.sort()).toEqual([ + "if-0", + "if-1", + "if-2", + ]); + expect(calls[0].messageGroupId).toBe( + `refresh-agg:${baseArgs.orgId}:${baseArgs.env}:${baseArgs.customerId}`, + ); + }, + { timeout: 5_000 }, + ); + + test( + `${chalk.yellowBright("batch-2: schedules across a bucket boundary → 2 enqueues")}`, + async () => { + const { fn, calls } = createMockQueue(); + const manager = new RefreshEntityAggregateBatchingManager({ + addTaskToQueueFn: fn, + bucketMs: 100, + settleBufferMs: 20, + }); + + manager.schedule(baseArgs); + await wait(180); + manager.schedule(baseArgs); + await wait(180); + + expect(calls.length).toBe(2); + expect(calls[0].messageDeduplicationId).not.toBe( + calls[1].messageDeduplicationId, + ); + }, + { timeout: 5_000 }, + ); + + test( + `${chalk.yellowBright("batch-3: different customers fire independently")}`, + async () => { + const { fn, calls } = createMockQueue(); + const manager = new RefreshEntityAggregateBatchingManager({ + addTaskToQueueFn: fn, + bucketMs: 100, + settleBufferMs: 20, + }); + + manager.schedule({ ...baseArgs, customerId: "cust-A" }); + manager.schedule({ ...baseArgs, customerId: "cust-B" }); + + expect(manager.getStats().totalPending).toBe(2); + + await wait(200); + + expect(calls.length).toBe(2); + const customerIds = calls.map((c) => c.payload.customerId).sort(); + expect(customerIds).toEqual(["cust-A", "cust-B"]); + }, + { timeout: 5_000 }, + ); + + test( + `${chalk.yellowBright("batch-4: flush() drains all pending immediately")}`, + async () => { + const { fn, calls } = createMockQueue(); + const manager = new RefreshEntityAggregateBatchingManager({ + addTaskToQueueFn: fn, + bucketMs: 10_000, // Will never fire naturally + settleBufferMs: 1000, + }); + + manager.schedule({ ...baseArgs, customerId: "cust-A" }); + manager.schedule({ ...baseArgs, customerId: "cust-B" }); + + expect(calls.length).toBe(0); + + await manager.flush(); + + expect(calls.length).toBe(2); + expect(manager.getStats().totalPending).toBe(0); + }, + { timeout: 5_000 }, + ); + + test( + `${chalk.yellowBright("batch-5: settle buffer honored")}`, + () => { + expect(REFRESH_ENTITY_AGGREGATE_DEDUP_BUCKET_MS).toBe(5000); + expect(REFRESH_ENTITY_AGGREGATE_SETTLE_BUFFER_MS).toBe(1500); + }, + ); +}); diff --git a/server/tests/integration/others/refresh-entity-aggregate/refresh-entity-aggregate-track-dedup.test.ts b/server/tests/integration/others/refresh-entity-aggregate/refresh-entity-aggregate-track-dedup.test.ts new file mode 100644 index 000000000..f7fa06706 --- /dev/null +++ b/server/tests/integration/others/refresh-entity-aggregate/refresh-entity-aggregate-track-dedup.test.ts @@ -0,0 +1,127 @@ +/** + * Verifies that 20 concurrent tracks spread across ~3 seconds result in + * exactly ONE `RefreshEntityAggregate` enqueue to SQS at the end of the + * bucket, thanks to the producer-side `RefreshEntityAggregateBatchingManager`. + */ + +import { expect, mock, test } from "bun:test"; +import chalk from "chalk"; + +type QueueCall = { + jobName: string; + payload: Record; + messageGroupId?: string; + messageDeduplicationId?: string; +}; + +const queueCalls: QueueCall[] = []; + +mock.module("@/queue/queueUtils.js", () => ({ + addTaskToQueue: async (args: QueueCall) => { + queueCalls.push(structuredClone(args)); + }, +})); + +const { JobName } = await import("@/queue/JobName.js"); +const { TestFeature } = await import("@tests/setup/v2Features.js"); +const { items } = await import("@tests/utils/fixtures/items.js"); +const { products } = await import("@tests/utils/fixtures/products.js"); +const { timeout } = await import("@tests/utils/genUtils.js"); +const { initScenario, s } = await import( + "@tests/utils/testInitUtils/initScenario.js" +); +const { globalRefreshEntityAggregateBatchingManager } = await import( + "@/internal/balances/utils/refreshEntityAggregate/index.js" +); +const { globalSyncBatchingManagerV3 } = await import( + "@/internal/balances/utils/sync/SyncBatchingManagerV3.js" +); +const { syncItemV4 } = await import( + "@/internal/balances/utils/sync/syncItemV4.js" +); + +test( + `${chalk.yellowBright( + "refresh-dedup-track: 20 concurrent tracks across 3s → exactly 1 RefreshEntityAggregate enqueue", + )}`, + async () => { + const customerId = "refresh-agg-dedup-cus"; + + const perEntityMessages = items.monthlyMessages({ + includedUsage: 500, + entityFeatureId: TestFeature.Users, + }); + const prod = products.base({ + id: "refresh-agg-dedup", + items: [perEntityMessages], + }); + + const { autumnV2_1, ctx, entities } = await initScenario({ + customerId, + setup: [ + s.customer({ testClock: false }), + s.products({ list: [prod] }), + s.entities({ count: 1, featureId: TestFeature.Users }), + ], + actions: [s.attach({ productId: prod.id, entityIndex: 0 })], + }); + + // Discard enqueues produced during setup. + queueCalls.length = 0; + + const trackPromises: Promise[] = []; + for (let i = 0; i < 20; i++) { + trackPromises.push( + (async () => { + await timeout(i * 150); // 20 * 150ms ≈ 3s spread + return autumnV2_1.track({ + customer_id: customerId, + entity_id: entities[0].id, + feature_id: TestFeature.Messages, + value: 1, + }); + })(), + ); + } + + await Promise.all(trackPromises); + + // Drain pending sync batches → SyncBalanceBatchV4 enqueues. + await globalSyncBatchingManagerV3.flush(); + + // Simulate the worker processing each sync-v4 job in this process + // (there's no live worker inside `bun test`). Each call schedules a + // refresh with the batching manager. + const syncJobs = queueCalls.filter( + (call) => call.jobName === JobName.SyncBalanceBatchV4, + ); + expect(syncJobs.length).toBeGreaterThanOrEqual(1); + + for (const job of syncJobs) { + await syncItemV4({ + ctx: ctx as never, + payload: job.payload as never, + }); + } + + // Now drain the refresh batching manager — this is what the trailing + // timer would otherwise do at bucket end + settle buffer. Flushing + // directly keeps the test fast and deterministic. + await globalRefreshEntityAggregateBatchingManager.flush(); + + const refreshCalls = queueCalls.filter( + (call) => call.jobName === JobName.RefreshEntityAggregate, + ); + + expect(refreshCalls.length).toBe(1); + expect(refreshCalls[0].payload).toMatchObject({ + customerId, + orgId: ctx.org.id, + env: ctx.env, + }); + expect(refreshCalls[0].messageGroupId).toBe( + `refresh-agg:${ctx.org.id}:${ctx.env}:${customerId}`, + ); + }, + 60_000, +);