From 292d7e16d851b56e3ff3344111f2fa77d74969a3 Mon Sep 17 00:00:00 2001 From: John Yeo Date: Mon, 23 Feb 2026 15:44:59 +0000 Subject: [PATCH] feat: lazy resets implementation and tests --- server/src/_luaScriptsV2/luaScriptsV2.ts | 18 ++ .../resetCustomerEntitlements.lua | 112 ++++++++ server/src/db/initializeDatabaseFunctions.ts | 1 + server/src/external/redis/initRedis.ts | 10 + .../src/internal/balances/utils/sql/client.ts | 64 +++++ .../balances/utils/sql/resetCusEnts.sql | 133 +++++++++ .../balances/utils/sync/syncItemV3.ts | 5 +- .../verifyCacheConsistency.ts | 3 +- .../src/internal/customers/CusBatchService.ts | 52 ++-- server/src/internal/customers/CusService.ts | 11 +- .../applyResetResults.ts | 76 +++++ .../batchResetCustomerEntitlements.ts | 62 ++-- .../executeResetCache.ts | 34 +++ .../resetCustomerEntitlements/processReset.ts | 48 +++- .../resetCustomerEntitlements.ts | 142 +++++----- .../runBatchResetCusEntsTask.ts | 59 ---- .../triggerBatchResetCustomerEntitlements.ts | 40 +++ .../cusEnts/cusRollovers/RolloverService.ts | 33 ++- .../cusProducts/cusEnts/groupByUtils.ts | 23 +- .../getCachedFullCustomer.ts | 20 +- .../getOrCreateCachedFullCustomer.ts | 36 +-- .../getOrSetCachedFullCustomer.ts | 38 +-- .../handleGetCustomerEvents.ts | 4 +- .../handleGetFullCustomers.ts | 8 +- server/src/queue/processMessage.ts | 4 +- server/src/queue/workflows.ts | 6 +- .../reset-customer-entitlements.test.ts | 21 -- .../get-customer-reset-concurrent.test.ts | 251 +++++++++++++++++ .../get-customer-reset.test.ts | 265 ++++++++++++++++++ .../list-customers-reset.test.ts | 131 +++++++++ .../utils/cusProductUtils/resetTestUtils.ts | 110 ++++++++ .../tests/utils/testInitUtils/initScenario.ts | 2 +- 32 files changed, 1489 insertions(+), 333 deletions(-) create mode 100644 server/src/_luaScriptsV2/resetCustomerEntitlements/resetCustomerEntitlements.lua create mode 100644 server/src/internal/balances/utils/sql/client.ts create mode 100644 server/src/internal/balances/utils/sql/resetCusEnts.sql create mode 100644 server/src/internal/customers/actions/resetCustomerEntitlements/applyResetResults.ts create mode 100644 server/src/internal/customers/actions/resetCustomerEntitlements/executeResetCache.ts delete mode 100644 server/src/internal/customers/actions/resetCustomerEntitlements/runBatchResetCusEntsTask.ts create mode 100644 server/src/internal/customers/actions/resetCustomerEntitlements/triggerBatchResetCustomerEntitlements.ts delete mode 100644 server/tests/integration/crud/customers/reset-customer-entitlements.test.ts create mode 100644 server/tests/integration/crud/customers/reset-customer-entitlements/get-customer-reset-concurrent.test.ts create mode 100644 server/tests/integration/crud/customers/reset-customer-entitlements/get-customer-reset.test.ts create mode 100644 server/tests/integration/crud/customers/reset-customer-entitlements/list-customers-reset.test.ts create mode 100644 server/tests/utils/cusProductUtils/resetTestUtils.ts diff --git a/server/src/_luaScriptsV2/luaScriptsV2.ts b/server/src/_luaScriptsV2/luaScriptsV2.ts index b04919b7c..bca1abd42 100644 --- a/server/src/_luaScriptsV2/luaScriptsV2.ts +++ b/server/src/_luaScriptsV2/luaScriptsV2.ts @@ -8,6 +8,7 @@ const __dirname = dirname(__filename); // Path to script folders const DEDUCT_DIR = join(__dirname, "deductFromCustomerEntitlements"); const DELETE_CACHE_DIR = join(__dirname, "deleteFullCustomerCache"); +const RESET_DIR = join(__dirname, "resetCustomerEntitlements"); // ============================================================================ // HELPER MODULES @@ -93,3 +94,20 @@ export const BATCH_DELETE_FULL_CUSTOMER_CACHE_SCRIPT = readFileSync( join(DELETE_CACHE_DIR, "batchDeleteFullCustomerCache.lua"), "utf-8", ); + +// ============================================================================ +// RESET CUSTOMER ENTITLEMENTS SCRIPT +// ============================================================================ + +const resetMainScript = readFileSync( + join(RESET_DIR, "resetCustomerEntitlements.lua"), + "utf-8", +); + +/** + * Lua script for atomically resetting cusEnt fields in the cached FullCustomer. + * Reuses luaUtils helpers for find_entitlement navigation. + * Skips if cache doesn't exist or cusEnt already reset (optimistic guard). + */ +export const RESET_CUSTOMER_ENTITLEMENTS_SCRIPT = `${LUA_UTILS} +${resetMainScript}`; diff --git a/server/src/_luaScriptsV2/resetCustomerEntitlements/resetCustomerEntitlements.lua b/server/src/_luaScriptsV2/resetCustomerEntitlements/resetCustomerEntitlements.lua new file mode 100644 index 000000000..c14465aa5 --- /dev/null +++ b/server/src/_luaScriptsV2/resetCustomerEntitlements/resetCustomerEntitlements.lua @@ -0,0 +1,112 @@ +--[[ + Lua Script: Reset Customer Entitlements in Redis Cache + + Atomically updates cached cusEnt fields after a DB reset. + Skips if cache doesn't exist or if the cusEnt's next_reset_at already + equals the new value (same optimistic guard as the SQL function). + + Helper functions prepended via string interpolation from: + - luaUtils.lua (find_entitlement, safe_number, is_nil) + + KEYS[1] = FullCustomer cache key + + ARGV[1] = JSON params: + { + resets: [{ + cus_ent_id: string, + balance: number | null, + additional_balance: number | null, + adjustment: number, + entities: object | null, + next_reset_at: number, + rollover_insert: { id, cus_ent_id, balance, usage, expires_at, entities } | null + }] + } + + Returns JSON: + { "applied": { "": true }, "skipped": ["id1"] } +]] + +local cache_key = KEYS[1] +local params = cjson.decode(ARGV[1]) +local resets = params.resets or {} + +-- Early return if no resets +if #resets == 0 then + return cjson.encode({ applied = {}, skipped = {} }) +end + +-- Check if cache exists +local key_exists = redis.call('EXISTS', cache_key) +if key_exists == 0 then + return cjson.encode({ applied = {}, skipped = {}, cache_miss = true }) +end + +-- Read the full customer structure for entitlement path lookups +local full_customer_json = redis.call('JSON.GET', cache_key, '.') +if not full_customer_json then + return cjson.encode({ applied = {}, skipped = {}, cache_miss = true }) +end + +local full_customer = cjson.decode(full_customer_json) + +local applied = {} +local skipped = {} + +for _, reset in ipairs(resets) do + local ent_id = reset.cus_ent_id + local new_next_reset_at = reset.next_reset_at + + -- Find the cusEnt in the FullCustomer structure + local cus_ent, cus_product, ce_idx, cp_idx = find_entitlement(full_customer, ent_id) + + if not cus_ent then + table.insert(skipped, ent_id) + else + -- Build the JSON path to this cusEnt + local base_path + local is_loose = (cp_idx == nil) + + if is_loose then + base_path = '$.extra_customer_entitlements[' .. (ce_idx - 1) .. ']' + else + base_path = '$.customer_products[' .. (cp_idx - 1) .. '].customer_entitlements[' .. (ce_idx - 1) .. ']' + end + + -- Optimistic guard: skip if next_reset_at already equals the new value + local current_reset_at = safe_number(cus_ent.next_reset_at) + if current_reset_at == new_next_reset_at then + table.insert(skipped, ent_id) + else + -- Apply reset fields using JSON.SET for absolute values + if not is_nil(reset.balance) then + redis.call('JSON.SET', cache_key, base_path .. '.balance', tostring(reset.balance)) + end + + if not is_nil(reset.additional_balance) then + redis.call('JSON.SET', cache_key, base_path .. '.additional_balance', tostring(reset.additional_balance)) + end + + redis.call('JSON.SET', cache_key, base_path .. '.adjustment', tostring(reset.adjustment)) + redis.call('JSON.SET', cache_key, base_path .. '.next_reset_at', tostring(new_next_reset_at)) + + -- Set entities if provided (entity-scoped entitlement) + if not is_nil(reset.entities) then + redis.call('JSON.SET', cache_key, base_path .. '.entities', cjson.encode(reset.entities)) + end + + -- Increment cache_version + redis.call('JSON.NUMINCRBY', cache_key, base_path .. '.cache_version', 1) + + -- Append rollover if provided + if not is_nil(reset.rollover_insert) then + local rollover_json = cjson.encode(reset.rollover_insert) + redis.call('JSON.ARRAPPEND', cache_key, base_path .. '.rollovers', rollover_json) + end + + applied[ent_id] = true + end + end +end + +return cjson.encode({ applied = applied, skipped = skipped }) diff --git a/server/src/db/initializeDatabaseFunctions.ts b/server/src/db/initializeDatabaseFunctions.ts index 811397c14..d6033c655 100644 --- a/server/src/db/initializeDatabaseFunctions.ts +++ b/server/src/db/initializeDatabaseFunctions.ts @@ -28,6 +28,7 @@ export const initializeDatabaseFunctions = async () => { "performDeduction.sql", "syncBalances.sql", "syncBalancesV2.sql", + "resetCusEnts.sql", ]; for (const file of sqlFiles) { diff --git a/server/src/external/redis/initRedis.ts b/server/src/external/redis/initRedis.ts index c689ca9d5..e6fa3fb54 100644 --- a/server/src/external/redis/initRedis.ts +++ b/server/src/external/redis/initRedis.ts @@ -17,6 +17,7 @@ import { BATCH_DELETE_FULL_CUSTOMER_CACHE_SCRIPT, DEDUCT_FROM_CUSTOMER_ENTITLEMENTS_SCRIPT, DELETE_FULL_CUSTOMER_CACHE_SCRIPT, + RESET_CUSTOMER_ENTITLEMENTS_SCRIPT, SET_FULL_CUSTOMER_CACHE_SCRIPT, } from "../../_luaScriptsV2/luaScriptsV2.js"; @@ -179,6 +180,11 @@ const configureRedisInstance = (redisInstance: Redis): Redis => { lua: SET_FULL_CUSTOMER_CACHE_SCRIPT, }); + redisInstance.defineCommand("resetCustomerEntitlements", { + numberOfKeys: 1, + lua: RESET_CUSTOMER_ENTITLEMENTS_SCRIPT, + }); + redisInstance.on("error", (error) => { console.error(`[Redis] Connection error:`, error.message); }); @@ -353,6 +359,10 @@ declare module "ioredis" { serializedData: string, overwrite: string, ): Promise<"STALE_WRITE" | "CACHE_EXISTS" | "OK">; + resetCustomerEntitlements( + cacheKey: string, + paramsJson: string, + ): Promise; } } diff --git a/server/src/internal/balances/utils/sql/client.ts b/server/src/internal/balances/utils/sql/client.ts new file mode 100644 index 000000000..e16986ffc --- /dev/null +++ b/server/src/internal/balances/utils/sql/client.ts @@ -0,0 +1,64 @@ +import type { EntityBalance, Rollover } from "@autumn/shared"; +import { sql } from "drizzle-orm"; +import type { DrizzleCli } from "@/db/initDrizzle.js"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; + +export type ResetCusEntParam = { + cus_ent_id: string; + balance: number | null; + additional_balance: number | null; + adjustment: number; + entities: Record | null; + next_reset_at: number; + rollover_insert: Pick< + Rollover, + "id" | "cus_ent_id" | "balance" | "usage" | "expires_at" | "entities" + > | null; +}; + +export type AppliedCusEntReset = { + balance: number; + additional_balance: number; + adjustment: number; + entities: Record | null; + next_reset_at: number; + cache_version: number; + rollover: Pick< + Rollover, + "id" | "cus_ent_id" | "balance" | "usage" | "expires_at" | "entities" + > | null; +}; + +type ResetCusEntsResult = { + applied: Record; + skipped: string[]; +}; + +/** Calls the `reset_customer_entitlements` PL/pgSQL function atomically. */ +export const resetCusEnts = async ({ + ctx, + resets, +}: { + ctx: AutumnContext; + resets: ResetCusEntParam[]; +}): Promise => { + const { db } = ctx; + if (resets.length === 0) { + return { applied: {}, skipped: [] }; + } + + const result = await db.execute( + sql`SELECT * FROM reset_customer_entitlements(${JSON.stringify({ + resets, + })}::jsonb)`, + ); + + const raw = result[0]?.reset_customer_entitlements as + | ResetCusEntsResult + | undefined; + + return { + applied: raw?.applied ?? {}, + skipped: raw?.skipped ?? [], + }; +}; diff --git a/server/src/internal/balances/utils/sql/resetCusEnts.sql b/server/src/internal/balances/utils/sql/resetCusEnts.sql new file mode 100644 index 000000000..93024a681 --- /dev/null +++ b/server/src/internal/balances/utils/sql/resetCusEnts.sql @@ -0,0 +1,133 @@ +-- Atomically reset customer entitlements that have passed their next_reset_at. +-- Uses per-row locking + optimistic check: only resets a cusEnt if its +-- next_reset_at does NOT already equal the new value (prevents double-resets). +-- +-- Params (JSONB): +-- resets: array of objects with: +-- - cus_ent_id: text +-- - balance: numeric (null if entity-scoped) +-- - additional_balance: numeric (null if entity-scoped) +-- - adjustment: numeric +-- - entities: jsonb (null if non-entity) +-- - next_reset_at: bigint (new next_reset_at value) +-- - rollover_insert: jsonb object or null, with fields: +-- id, cus_ent_id, balance, usage, expires_at, entities +-- +-- Returns JSONB: +-- { +-- "applied": { +-- "": { +-- "balance": number, +-- "additional_balance": number, +-- "adjustment": number, +-- "entities": jsonb, +-- "next_reset_at": number, +-- "cache_version": number, +-- "rollover": jsonb or null +-- } +-- }, +-- "skipped": ["id1", "id2"] +-- } +-- +DROP FUNCTION IF EXISTS reset_customer_entitlements(jsonb); + +CREATE FUNCTION reset_customer_entitlements(params jsonb) +RETURNS jsonb +LANGUAGE plpgsql +AS $$ +DECLARE + resets_param jsonb := params->'resets'; + + reset_obj jsonb; + ent_id text; + new_balance numeric; + new_additional_balance numeric; + new_adjustment numeric; + new_entities jsonb; + new_next_reset_at bigint; + rollover_obj jsonb; + + db_next_reset_at bigint; + updated_row record; + + applied_json jsonb := '{}'::jsonb; + skipped_ids jsonb := '[]'::jsonb; +BEGIN + IF resets_param IS NULL OR jsonb_array_length(resets_param) = 0 THEN + RETURN jsonb_build_object('applied', '{}'::jsonb, 'skipped', '[]'::jsonb); + END IF; + + FOR reset_obj IN SELECT * FROM jsonb_array_elements(resets_param) + LOOP + ent_id := reset_obj->>'cus_ent_id'; + new_balance := (reset_obj->>'balance')::numeric; + new_additional_balance := (reset_obj->>'additional_balance')::numeric; + new_adjustment := (reset_obj->>'adjustment')::numeric; + new_entities := reset_obj->'entities'; + new_next_reset_at := (reset_obj->>'next_reset_at')::bigint; + rollover_obj := reset_obj->'rollover_insert'; + + -- Lock and read the single row + SELECT ce.next_reset_at INTO db_next_reset_at + FROM customer_entitlements ce + WHERE ce.id = ent_id + FOR UPDATE; + + -- Optimistic lock: skip if next_reset_at already equals the new value + IF db_next_reset_at IS NOT DISTINCT FROM new_next_reset_at THEN + skipped_ids := skipped_ids || to_jsonb(ent_id); + CONTINUE; + END IF; + + -- Apply the reset update and capture the updated row + UPDATE customer_entitlements ce + SET + balance = COALESCE(new_balance, ce.balance), + additional_balance = COALESCE(new_additional_balance, ce.additional_balance), + adjustment = COALESCE(new_adjustment, ce.adjustment), + entities = COALESCE(new_entities, ce.entities), + next_reset_at = new_next_reset_at, + cache_version = COALESCE(ce.cache_version, 0) + 1 + WHERE ce.id = ent_id + RETURNING ce.balance, ce.additional_balance, ce.adjustment, ce.entities, + ce.next_reset_at, ce.cache_version + INTO updated_row; + + -- Insert rollover row if provided + IF rollover_obj IS NOT NULL AND rollover_obj != 'null'::jsonb THEN + INSERT INTO rollovers (id, cus_ent_id, balance, usage, expires_at, entities) + VALUES ( + rollover_obj->>'id', + rollover_obj->>'cus_ent_id', + (rollover_obj->>'balance')::numeric, + (rollover_obj->>'usage')::numeric, + (rollover_obj->>'expires_at')::numeric, + COALESCE(rollover_obj->'entities', '{}'::jsonb) + ); + END IF; + + -- Record the latest state of the updated cusEnt + applied_json := jsonb_set( + applied_json, + ARRAY[ent_id], + jsonb_build_object( + 'balance', updated_row.balance, + 'additional_balance', updated_row.additional_balance, + 'adjustment', updated_row.adjustment, + 'entities', updated_row.entities, + 'next_reset_at', updated_row.next_reset_at, + 'cache_version', updated_row.cache_version, + 'rollover', CASE + WHEN rollover_obj IS NOT NULL AND rollover_obj != 'null'::jsonb THEN rollover_obj + ELSE NULL + END + ) + ); + END LOOP; + + RETURN jsonb_build_object( + 'applied', applied_json, + 'skipped', skipped_ids + ); +END; +$$; diff --git a/server/src/internal/balances/utils/sync/syncItemV3.ts b/server/src/internal/balances/utils/sync/syncItemV3.ts index 4eb3bce39..301e006a1 100644 --- a/server/src/internal/balances/utils/sync/syncItemV3.ts +++ b/server/src/internal/balances/utils/sync/syncItemV3.ts @@ -196,14 +196,13 @@ export const syncItemV3 = async ({ ctx: AutumnContext; payload: SyncItemV3; }): Promise => { - const { customerId, orgId, env, region, cusEntIds, rolloverIds } = payload; + const { customerId, region, cusEntIds, rolloverIds } = payload; const { db, logger } = ctx; const redisInstance = region ? getRegionalRedis(region) : undefined; const fullCustomer = await getCachedFullCustomer({ - orgId, - env, + ctx, customerId, redisInstance, }); diff --git a/server/src/internal/billing/v2/workflows/verifyCacheConsistency/verifyCacheConsistency.ts b/server/src/internal/billing/v2/workflows/verifyCacheConsistency/verifyCacheConsistency.ts index ba1ee37ef..2255210c3 100644 --- a/server/src/internal/billing/v2/workflows/verifyCacheConsistency/verifyCacheConsistency.ts +++ b/server/src/internal/billing/v2/workflows/verifyCacheConsistency/verifyCacheConsistency.ts @@ -79,8 +79,7 @@ verifyCacheConsistency?.task({ // Get from cache (now using full customer cache) const cachedFullCustomer = await getCachedFullCustomer({ - orgId: autumnContext.org.id, - env: autumnContext.env, + ctx: autumnContext, customerId, }); diff --git a/server/src/internal/customers/CusBatchService.ts b/server/src/internal/customers/CusBatchService.ts index e6b7ee753..1b2457496 100644 --- a/server/src/internal/customers/CusBatchService.ts +++ b/server/src/internal/customers/CusBatchService.ts @@ -1,37 +1,32 @@ import { AffectedResource, type ApiCustomerV5, - type AppEnv, applyResponseVersionChanges, type CusProductStatus, CustomerExpand, type CustomerLegacyData, type FullCustomer, type ListCustomersV2Params, - type Organization, RELEVANT_STATUSES, } from "@autumn/shared"; -import type { DrizzleCli } from "@/db/initDrizzle.js"; -import type { RequestContext } from "@/honoUtils/HonoEnv.js"; -import { batchResetCustomerEntitlements } from "./actions/resetCustomerEntitlements/batchResetCustomerEntitlements.js"; +import * as Sentry from "@sentry/bun"; +import type { AutumnContext, RequestContext } from "@/honoUtils/HonoEnv.js"; +import { triggerBatchResetCustomerEntitlements } from "./actions/resetCustomerEntitlements/triggerBatchResetCustomerEntitlements.js"; import { getApiCustomerBase } from "./cusUtils/apiCusUtils/getApiCustomerBase.js"; import { getPaginatedFullCusQuery } from "./getFullCusQuery.js"; export class CusBatchService { static async getByInternalIds({ - db, - org, - env, + ctx, internalCustomerIds, }: { - db: DrizzleCli; - org: Organization; - env: AppEnv; + ctx: AutumnContext; internalCustomerIds: string[]; }) { + const { org, env, db } = ctx; const query = getPaginatedFullCusQuery({ - orgId: org.id, - env, + orgId: ctx.org.id, + env: ctx.env, includeInvoices: true, withEntities: true, withTrialsUsed: false, @@ -44,16 +39,15 @@ export class CusBatchService { const fullCustomers = results as unknown as FullCustomer[]; // Fire-and-forget: queue SQS job for any stale entitlement resets - batchResetCustomerEntitlements({ + triggerBatchResetCustomerEntitlements({ + ctx, fullCustomers, - orgId: org.id, - env, - }).catch((err) => - console.error( - "[CusBatchService.getByInternalIds] batch reset failed:", - err, - ), - ); + }).catch((err) => { + ctx.logger.error( + `[CusBatchService.getByInternalIds] batch reset failed: ${err}`, + ); + Sentry.captureException(err); + }); return fullCustomers; } @@ -120,18 +114,18 @@ export class CusBatchService { finals.push(versionedCustomer); } catch (error) { - console.error(`Failed to process customer ${result.id}:`, error); + ctx.logger.error(`Failed to process customer ${result.id}: ${error}`); } } // Fire-and-forget: queue SQS job for any stale entitlement resets - batchResetCustomerEntitlements({ + triggerBatchResetCustomerEntitlements({ + ctx, fullCustomers, - orgId: ctx.org.id, - env: ctx.env, - }).catch((err) => - console.error("[CusBatchService.getPage] batch reset failed:", err), - ); + }).catch((err) => { + ctx.logger.error("[CusBatchService.getPage] batch reset failed:", err); + Sentry.captureException(err); + }); return finals; } diff --git a/server/src/internal/customers/CusService.ts b/server/src/internal/customers/CusService.ts index 3719d4f45..dc696852b 100644 --- a/server/src/internal/customers/CusService.ts +++ b/server/src/internal/customers/CusService.ts @@ -25,6 +25,7 @@ import { import type { DrizzleCli } from "@/db/initDrizzle.js"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { withSpan } from "../analytics/tracer/spanUtils.js"; +import { resetCustomerEntitlements } from "./actions/resetCustomerEntitlements/resetCustomerEntitlements.js"; import { RELEVANT_STATUSES } from "./cusProducts/CusProductService.js"; import { getFullCusQuery } from "./getFullCusQuery.js"; @@ -108,7 +109,15 @@ export class CusService { } } - return data as FullCustomer; + const fullCus = data as FullCustomer; + + // Lazy reset stale entitlements (mutates fullCus in-memory + writes DB) + await resetCustomerEntitlements({ + fullCus, + ctx, + }); + + return fullCus; }, }); } diff --git a/server/src/internal/customers/actions/resetCustomerEntitlements/applyResetResults.ts b/server/src/internal/customers/actions/resetCustomerEntitlements/applyResetResults.ts new file mode 100644 index 000000000..670fa6c0c --- /dev/null +++ b/server/src/internal/customers/actions/resetCustomerEntitlements/applyResetResults.ts @@ -0,0 +1,76 @@ +import type { + FullCustomer, + FullCustomerEntitlement, + Rollover, +} from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { RolloverService } from "@/internal/customers/cusProducts/cusEnts/cusRollovers/RolloverService.js"; +import type { ProcessResetResult } from "./processReset.js"; + +/** Find a cusEnt on the FullCustomer by ID. */ +const findCusEnt = ({ + fullCus, + cusEntId, +}: { + fullCus: FullCustomer; + cusEntId: string; +}): FullCustomerEntitlement | null => { + for (const cusProduct of fullCus.customer_products) { + for (const cusEnt of cusProduct.customer_entitlements) { + if (cusEnt.id === cusEntId) return cusEnt; + } + } + for (const cusEnt of fullCus.extra_customer_entitlements || []) { + if (cusEnt.id === cusEntId) return cusEnt; + } + return null; +}; + +/** + * Applies computed reset values to in-memory FullCustomer for all cusEnts, + * and runs rollover max-clearing only for DB-applied (non-skipped) ones. + */ +export const applyResetResults = async ({ + ctx, + fullCus, + computed, + skipped, +}: { + ctx: AutumnContext; + fullCus: FullCustomer; + computed: Array<{ cusEntId: string; result: ProcessResetResult }>; + skipped: string[]; +}): Promise => { + const { db } = ctx; + const skippedSet = new Set(skipped); + const clearingPromises: Promise[] = []; + + for (const { cusEntId, result } of computed) { + const original = findCusEnt({ fullCus, cusEntId }); + if (!original) continue; + + const { updates } = result; + if (updates.balance !== null) original.balance = updates.balance; + if (updates.additional_balance !== null) + original.additional_balance = updates.additional_balance; + original.adjustment = updates.adjustment; + if (updates.entities !== null) original.entities = updates.entities; + original.next_reset_at = updates.next_reset_at; + + // Only run rollover clearing for DB-applied entries. + // Skipped entries were already cleared by the winning request. + if (!skippedSet.has(cusEntId) && result.rolloverInsert) { + clearingPromises.push( + RolloverService.clearExcessRollovers({ + db, + newRows: result.rolloverInsert.rows, + fullCusEnt: original, + }), + ); + } + } + + if (clearingPromises.length > 0) { + await Promise.all(clearingPromises); + } +}; diff --git a/server/src/internal/customers/actions/resetCustomerEntitlements/batchResetCustomerEntitlements.ts b/server/src/internal/customers/actions/resetCustomerEntitlements/batchResetCustomerEntitlements.ts index 949c3ee45..86262b901 100644 --- a/server/src/internal/customers/actions/resetCustomerEntitlements/batchResetCustomerEntitlements.ts +++ b/server/src/internal/customers/actions/resetCustomerEntitlements/batchResetCustomerEntitlements.ts @@ -1,40 +1,42 @@ -import { - CusProductStatus, - type FullCustomer, - fullCustomerToCustomerEntitlements, -} from "@autumn/shared"; -import { workflows } from "@/queue/workflows.js"; +import { CusProductStatus } from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import type { BatchResetCusEntsPayload } from "@/queue/workflows.js"; +import { CusService } from "../../CusService.js"; +import { resetCustomerEntitlements } from "./resetCustomerEntitlements.js"; /** - * Checks a list of FullCustomers for entitlements needing reset, - * and queues an SQS job with the cusEnt IDs if any are found. + * SQS worker handler: fetches cusEnts by ID, groups by customer, + * fetches each FullCustomer, and runs the lazy reset logic. */ export const batchResetCustomerEntitlements = async ({ - fullCustomers, - orgId, - env, + ctx, + payload, }: { - fullCustomers: FullCustomer[]; - orgId: string; - env: string; + ctx: AutumnContext; + payload: BatchResetCusEntsPayload; }): Promise => { - const now = Date.now(); - const cusEntIds: string[] = []; + const { resets } = payload; - for (const fullCus of fullCustomers) { - const cusEnts = fullCustomerToCustomerEntitlements({ - fullCustomer: fullCus, - inStatuses: [CusProductStatus.Active, CusProductStatus.PastDue], - }); + if (resets.length === 0) return; - for (const cusEnt of cusEnts) { - if (cusEnt.next_reset_at && cusEnt.next_reset_at < now) { - cusEntIds.push(cusEnt.id); - } - } + const BATCH_SIZE = 100; + + for (let i = 0; i < resets.length; i += BATCH_SIZE) { + const batch = resets.slice(i, i + BATCH_SIZE); + + await Promise.all( + batch.map(async (reset) => { + const fullCus = await CusService.getFull({ + ctx, + idOrInternalId: reset.internalCustomerId, + inStatuses: [CusProductStatus.Active, CusProductStatus.PastDue], + }); + + await resetCustomerEntitlements({ + ctx, + fullCus, + }); + }), + ); } - - if (cusEntIds.length === 0) return; - - await workflows.triggerBatchResetCusEnts({ orgId, env, cusEntIds }); }; diff --git a/server/src/internal/customers/actions/resetCustomerEntitlements/executeResetCache.ts b/server/src/internal/customers/actions/resetCustomerEntitlements/executeResetCache.ts new file mode 100644 index 000000000..fbb9f1395 --- /dev/null +++ b/server/src/internal/customers/actions/resetCustomerEntitlements/executeResetCache.ts @@ -0,0 +1,34 @@ +import { redis } from "@/external/redis/initRedis.js"; +import type { AutumnContext } from "@/honoUtils/HonoEnv"; +import type { ResetCusEntParam } from "@/internal/balances/utils/sql/client.js"; +import { buildFullCustomerCacheKey } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/fullCustomerCacheConfig.js"; +import { tryRedisWrite } from "@/utils/cacheUtils/cacheUtils.js"; + +/** + * Atomically resets cusEnt fields in the cached FullCustomer blob. + * Skips gracefully if the cache doesn't exist or the cusEnt was already reset. + * Fire-and-forget — failures are logged but don't propagate. + */ +export const executeResetCache = async ({ + ctx, + customerId, + resets, +}: { + ctx: AutumnContext; + customerId: string; + resets: ResetCusEntParam[]; +}): Promise => { + if (resets.length === 0) return; + + const { org, env, logger } = ctx; + + const cacheKey = buildFullCustomerCacheKey({ + orgId: org.id, + env, + customerId, + }); + + await tryRedisWrite(() => + redis.resetCustomerEntitlements(cacheKey, JSON.stringify({ resets })), + ); +}; diff --git a/server/src/internal/customers/actions/resetCustomerEntitlements/processReset.ts b/server/src/internal/customers/actions/resetCustomerEntitlements/processReset.ts index 03262199e..390524e75 100644 --- a/server/src/internal/customers/actions/resetCustomerEntitlements/processReset.ts +++ b/server/src/internal/customers/actions/resetCustomerEntitlements/processReset.ts @@ -1,44 +1,51 @@ import { - type AppEnv, cusEntToOptions, type EntInterval, + type EntityBalance, type FullCusEntWithFullCusProduct, type FullCustomerEntitlement, getStartingBalance, isLifetimeEntitlement, isUnlimitedEntitlement, - type Organization, type Rollover, } from "@autumn/shared"; import { logger } from "better-auth"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { getRolloverUpdates } from "@/internal/customers/cusProducts/cusEnts/cusRollovers/rolloverUtils.js"; import { getResetBalancesUpdate } from "@/internal/customers/cusProducts/cusEnts/groupByUtils.js"; import { getResetAtUpdate } from "./getResetAtUpdate.js"; +export type ResetUpdates = { + balance: number | null; + additional_balance: number | null; + adjustment: number; + entities: Record | null; + next_reset_at: number; +}; + export type ProcessResetResult = { - updates: Record; + updates: ResetUpdates; rolloverInsert?: { rows: Rollover[]; fullCusEnt: FullCustomerEntitlement }; }; /** Processes a single cusEnt reset. Returns updates + optional rollover insert, or null if skipped. */ export const processReset = async ({ cusEnt, - org, - env, + ctx, }: { cusEnt: FullCusEntWithFullCusProduct; - org: Organization; - env: AppEnv; + ctx: AutumnContext; }): Promise => { const ent = cusEnt.entitlement; const cusProduct = cusEnt.customer_product; - // Handle unlimited entitlements + // Unlimited / lifetime cusEnts should never reach here + // (getCusEntsNeedingReset filters them out), but guard defensively if ( isUnlimitedEntitlement({ entitlement: ent }) || isLifetimeEntitlement({ entitlement: ent }) ) { - return { updates: { next_reset_at: null } }; + return null; } const options = cusEntToOptions({ cusEnt }); @@ -56,6 +63,8 @@ export const processReset = async ({ return null; } + const { org, env } = ctx; + // Compute next reset time (with Stripe anchor adjustment on edge dates) const nextResetAt = await getResetAtUpdate({ curResetAt: cusEnt.next_reset_at, @@ -78,11 +87,22 @@ export const processReset = async ({ allowance: resetBalance, }); - const updates = { - ...resetBalanceUpdate, - next_reset_at: nextResetAt, - adjustment: 0, - }; + const updates: ResetUpdates = + "entities" in resetBalanceUpdate + ? { + balance: null, + additional_balance: null, + adjustment: 0, + entities: resetBalanceUpdate.entities, + next_reset_at: nextResetAt, + } + : { + balance: resetBalanceUpdate.balance, + additional_balance: resetBalanceUpdate.additional_balance, + adjustment: 0, + entities: null, + next_reset_at: nextResetAt, + }; let rolloverInsert: | { rows: Rollover[]; fullCusEnt: FullCustomerEntitlement } diff --git a/server/src/internal/customers/actions/resetCustomerEntitlements/resetCustomerEntitlements.ts b/server/src/internal/customers/actions/resetCustomerEntitlements/resetCustomerEntitlements.ts index 85692fc8d..823e78355 100644 --- a/server/src/internal/customers/actions/resetCustomerEntitlements/resetCustomerEntitlements.ts +++ b/server/src/internal/customers/actions/resetCustomerEntitlements/resetCustomerEntitlements.ts @@ -1,101 +1,107 @@ -import type { - AppEnv, - FullCustomer, - FullCustomerEntitlement, - Organization, - Rollover, -} from "@autumn/shared"; -import type { DrizzleCli } from "@/db/initDrizzle.js"; -import { CusEntService } from "@/internal/customers/cusProducts/cusEnts/CusEntitlementService.js"; -import { RolloverService } from "@/internal/customers/cusProducts/cusEnts/cusRollovers/RolloverService.js"; +import type { FullCustomer } from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { + type ResetCusEntParam, + resetCusEnts, +} from "@/internal/balances/utils/sql/client.js"; +import { applyResetResults } from "./applyResetResults.js"; +import { executeResetCache } from "./executeResetCache.js"; import { getCusEntsNeedingReset } from "./getCusEntsNeedingReset.js"; -import { processReset } from "./processReset.js"; +import { type ProcessResetResult, processReset } from "./processReset.js"; -/** Find the original cusEnt reference on the FullCustomer by ID. */ -const findOriginalCusEnt = ({ - fullCus, +/** Maps a processReset result into the JSONB shape for the SQL function. */ +const toResetParam = ({ cusEntId, + result, }: { - fullCus: FullCustomer; cusEntId: string; -}): FullCustomerEntitlement | null => { - for (const cusProduct of fullCus.customer_products) { - for (const cusEnt of cusProduct.customer_entitlements) { - if (cusEnt.id === cusEntId) return cusEnt; - } - } + result: ProcessResetResult; +}): ResetCusEntParam => { + const { updates } = result; + const firstRollover = result.rolloverInsert?.rows[0] ?? null; - for (const cusEnt of fullCus.extra_customer_entitlements || []) { - if (cusEnt.id === cusEntId) return cusEnt; - } - - return null; + return { + cus_ent_id: cusEntId, + balance: updates.balance, + additional_balance: updates.additional_balance, + adjustment: updates.adjustment, + entities: updates.entities, + next_reset_at: updates.next_reset_at, + rollover_insert: firstRollover, + }; }; /** * Lazily resets customer entitlements that have passed their next_reset_at. - * Mutates the FullCustomer in-memory and awaits DB + rollover writes. + * Uses an atomic Postgres function with per-row locking to prevent double-resets. + * Mutates the FullCustomer in-memory using the latest DB state from applied resets. * Returns true if any entitlements were reset. */ export const resetCustomerEntitlements = async ({ + ctx, fullCus, - db, - org, - env, }: { + ctx: AutumnContext; fullCus: FullCustomer; - db: DrizzleCli; - org: Organization; - env: AppEnv; }): Promise => { const now = Date.now(); + const { logger } = ctx; + const customerId = fullCus.id || fullCus.internal_id; + const cusEntsNeedingReset = getCusEntsNeedingReset({ fullCus, now }); if (cusEntsNeedingReset.length === 0) return false; - const dbUpdates: Array<{ - id: string; - updates: Record; - }> = []; - const rolloverInserts: Array<{ - rows: Rollover[]; - fullCusEnt: FullCustomerEntitlement; + logger.info( + `[resetCustomerEntitlements] customer=${customerId}, cusEnts needing reset: ${cusEntsNeedingReset.length}`, + ); + + // 1. Compute all resets (pure computation, no DB writes) + const computed: Array<{ + cusEntId: string; + result: ProcessResetResult; }> = []; for (const cusEnt of cusEntsNeedingReset) { - const result = await processReset({ - cusEnt, - org, - env, - }); - + const result = await processReset({ cusEnt, ctx }); if (!result) continue; - - // Mutate the original cusEnt on fullCus (not the spread copy) - const original = findOriginalCusEnt({ fullCus, cusEntId: cusEnt.id }); - if (original) { - Object.assign(original, result.updates); - } - - dbUpdates.push({ id: cusEnt.id, updates: result.updates }); - - if (result.rolloverInsert) { - rolloverInserts.push(result.rolloverInsert); - } + computed.push({ cusEntId: cusEnt.id, result }); } - if (dbUpdates.length === 0) return false; + if (computed.length === 0) return false; - // Await all DB writes - await Promise.all([ - ...dbUpdates.map(({ id, updates }) => - CusEntService.update({ db, id, updates }), - ), - ...rolloverInserts.map(({ rows, fullCusEnt }) => - RolloverService.insert({ db, rows, fullCusEnt }), - ), - ]); + // 2. Execute atomic DB writes via Postgres function + const resets = computed.map(({ cusEntId, result }) => + toResetParam({ cusEntId, result }), + ); + + const { applied, skipped } = await resetCusEnts({ ctx, resets }); + + logger.info( + `[resetCustomerEntitlements] customer=${customerId}, applied: ${Object.keys(applied).length}, skipped: ${skipped.length}`, + ); + + // 3. Apply computed reset values to in-memory FullCustomer. + // Both DB-applied and DB-skipped cusEnts get their in-memory state updated + // (skipped means another request already wrote the same values to DB). + // Rollover clearing only runs for DB-applied entries. + await applyResetResults({ ctx, fullCus, computed, skipped }); + + // 4. Update Redis cache atomically (fire-and-forget) + // Only needed when we actually wrote to DB — skipped means cache was + // already updated by the winning request. + if (Object.keys(applied).length > 0) { + await executeResetCache({ + ctx, + customerId, + resets, + }); + + logger.info( + `[resetCustomerEntitlements] customer=${customerId}, Redis cache updated`, + ); + } return true; }; diff --git a/server/src/internal/customers/actions/resetCustomerEntitlements/runBatchResetCusEntsTask.ts b/server/src/internal/customers/actions/resetCustomerEntitlements/runBatchResetCusEntsTask.ts deleted file mode 100644 index 56fd5ad12..000000000 --- a/server/src/internal/customers/actions/resetCustomerEntitlements/runBatchResetCusEntsTask.ts +++ /dev/null @@ -1,59 +0,0 @@ -import { CusProductStatus } from "@autumn/shared"; -import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; -import type { BatchResetCusEntsPayload } from "@/queue/workflows.js"; -import { CusService } from "../../CusService.js"; -import { CusEntService } from "../../cusProducts/cusEnts/CusEntitlementService.js"; -import { resetCustomerEntitlements } from "./resetCustomerEntitlements.js"; - -/** - * SQS worker handler: fetches cusEnts by ID, groups by customer, - * fetches each FullCustomer, and runs the lazy reset logic. - */ -export const runBatchResetCusEntsTask = async ({ - ctx, - payload, -}: { - ctx: AutumnContext; - payload: BatchResetCusEntsPayload; -}): Promise => { - const { db, org, env } = ctx; - const { cusEntIds } = payload; - - if (cusEntIds.length === 0) return; - - // Fetch cusEnts to get their internal_customer_id for grouping - const cusEnts = await CusEntService.getByIds({ db, ids: cusEntIds }); - - // Group by internal_customer_id - const byCustomer = new Map(); - for (const cusEnt of cusEnts) { - const cusId = cusEnt.internal_customer_id; - if (!byCustomer.has(cusId)) { - byCustomer.set(cusId, []); - } - byCustomer.get(cusId)!.push(cusEnt.id); - } - - // For each customer, fetch full customer and run reset - for (const [internalCustomerId] of byCustomer) { - try { - const fullCus = await CusService.getFull({ - ctx, - idOrInternalId: internalCustomerId, - inStatuses: [CusProductStatus.Active, CusProductStatus.PastDue], - }); - - await resetCustomerEntitlements({ - fullCus, - db, - org, - env, - }); - } catch (error) { - console.error( - `[BatchResetCusEnts] Failed to reset for customer ${internalCustomerId}:`, - error, - ); - } - } -}; diff --git a/server/src/internal/customers/actions/resetCustomerEntitlements/triggerBatchResetCustomerEntitlements.ts b/server/src/internal/customers/actions/resetCustomerEntitlements/triggerBatchResetCustomerEntitlements.ts new file mode 100644 index 000000000..13a7b7f39 --- /dev/null +++ b/server/src/internal/customers/actions/resetCustomerEntitlements/triggerBatchResetCustomerEntitlements.ts @@ -0,0 +1,40 @@ +import type { FullCustomer } from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv"; +import { type BatchResetCusEntsPayload, workflows } from "@/queue/workflows.js"; +import { getCusEntsNeedingReset } from "./getCusEntsNeedingReset"; + +/** + * Checks a list of FullCustomers for entitlements needing reset, + * and queues an SQS job with the cusEnt IDs if any are found. + */ +export const triggerBatchResetCustomerEntitlements = async ({ + ctx, + fullCustomers, +}: { + ctx: AutumnContext; + fullCustomers: FullCustomer[]; +}): Promise => { + const now = Date.now(); + + const resets: BatchResetCusEntsPayload["resets"] = []; + for (const fullCus of fullCustomers) { + const cusEntsNeedingReset = getCusEntsNeedingReset({ + fullCus, + now, + }); + + resets.push({ + internalCustomerId: fullCus.internal_id, + customerId: fullCus.id ?? "", + cusEntIds: cusEntsNeedingReset.map((cusEnt) => cusEnt.id), + }); + } + + if (resets.length === 0) return; + + await workflows.triggerBatchResetCusEnts({ + orgId: ctx.org.id, + env: ctx.env, + resets, + }); +}; diff --git a/server/src/internal/customers/cusProducts/cusEnts/cusRollovers/RolloverService.ts b/server/src/internal/customers/cusProducts/cusEnts/cusRollovers/RolloverService.ts index 632b9cb01..baebd26f9 100644 --- a/server/src/internal/customers/cusProducts/cusEnts/cusRollovers/RolloverService.ts +++ b/server/src/internal/customers/cusProducts/cusEnts/cusRollovers/RolloverService.ts @@ -88,7 +88,24 @@ export class RolloverService { await db.insert(rollovers).values(rows).returning(); - let curRollovers = [...fullCusEnt.rollovers, ...rows]; + return RolloverService.clearExcessRollovers({ + db, + newRows: rows, + fullCusEnt, + }); + } + + /** Enforces the rollover max cap after new rollovers have been inserted into the DB. */ + static async clearExcessRollovers({ + db, + newRows, + fullCusEnt, + }: { + db: DrizzleCli; + newRows: Rollover[]; + fullCusEnt: FullCustomerEntitlement; + }): Promise { + const curRollovers = [...fullCusEnt.rollovers, ...newRows]; const { toDelete, toUpdate } = performMaximumClearing({ rows: curRollovers as Rollover[], @@ -103,17 +120,9 @@ export class RolloverService { await RolloverService.upsert({ db, rows: toUpdate }); } - // Return latest rollovers...? - curRollovers = curRollovers.filter((r) => toDelete.includes(r.id)); - curRollovers = curRollovers.map((r) => { - const updatedRow = toUpdate.find((u) => u.id === r.id); - if (updatedRow) { - return updatedRow; - } - return r; - }); - - return curRollovers; + return curRollovers + .filter((r) => !toDelete.includes(r.id)) + .map((r) => toUpdate.find((u) => u.id === r.id) ?? r); } static async delete({ db, ids }: { db: DrizzleCli; ids: string[] }) { diff --git a/server/src/internal/customers/cusProducts/cusEnts/groupByUtils.ts b/server/src/internal/customers/cusProducts/cusEnts/groupByUtils.ts index 1f32300e1..8b90c7de7 100644 --- a/server/src/internal/customers/cusProducts/cusEnts/groupByUtils.ts +++ b/server/src/internal/customers/cusProducts/cusEnts/groupByUtils.ts @@ -1,14 +1,17 @@ -import type { FullCustomerEntitlement } from "@autumn/shared"; +import type { EntityBalance, FullCustomerEntitlement } from "@autumn/shared"; import { notNullish } from "@/utils/genUtils.js"; +export type ResetBalancesUpdate = + | { entities: Record } + | { balance: number; additional_balance: number; adjustment: number }; + export const getResetBalancesUpdate = ({ cusEnt, allowance, }: { cusEnt: FullCustomerEntitlement; allowance?: number; -}) => { - let update = {}; +}): ResetBalancesUpdate => { const newBalance = notNullish(allowance) ? allowance! : cusEnt.entitlement.allowance || 0; @@ -21,14 +24,12 @@ export const getResetBalancesUpdate = ({ newEntities[entityId].balance = newBalance; newEntities[entityId].adjustment = 0; } - update = { entities: newEntities }; - } else { - update = { - balance: newBalance, - additional_balance: 0, - adjustment: 0, - }; + return { entities: newEntities }; } - return update; + return { + balance: newBalance, + additional_balance: 0, + adjustment: 0, + }; }; diff --git a/server/src/internal/customers/cusUtils/fullCustomerCacheUtils/getCachedFullCustomer.ts b/server/src/internal/customers/cusUtils/fullCustomerCacheUtils/getCachedFullCustomer.ts index 4d6d0dc85..975894008 100644 --- a/server/src/internal/customers/cusUtils/fullCustomerCacheUtils/getCachedFullCustomer.ts +++ b/server/src/internal/customers/cusUtils/fullCustomerCacheUtils/getCachedFullCustomer.ts @@ -2,7 +2,9 @@ import type { FullCustomer } from "@autumn/shared"; import { Decimal } from "decimal.js"; import type { Redis } from "ioredis"; import { redis } from "@/external/redis/initRedis.js"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { tryRedisRead } from "@/utils/cacheUtils/cacheUtils.js"; +import { resetCustomerEntitlements } from "../../actions/resetCustomerEntitlements/resetCustomerEntitlements.js"; import { buildFullCustomerCacheKey } from "./fullCustomerCacheConfig.js"; /** @@ -80,23 +82,26 @@ const roundFullCustomerBalances = ( }; /** - * Get FullCustomer from Redis cache + * Get FullCustomer from Redis cache. Lazily resets stale entitlements. * @returns FullCustomer if found, null if not in cache */ export const getCachedFullCustomer = async ({ - orgId, - env, + ctx, customerId, entityId, redisInstance, }: { - orgId: string; - env: string; + ctx: AutumnContext; customerId: string; entityId?: string; redisInstance?: Redis; }): Promise => { - const cacheKey = buildFullCustomerCacheKey({ orgId, env, customerId }); + const { org, env } = ctx; + const cacheKey = buildFullCustomerCacheKey({ + orgId: org.id, + env, + customerId, + }); const redisClient = redisInstance || redis; const cached = await tryRedisRead( @@ -123,6 +128,9 @@ export const getCachedFullCustomer = async ({ fullCustomer.send_email_receipts = false; } + // Lazy reset stale entitlements (DB + in-memory + cache via Lua) + await resetCustomerEntitlements({ ctx, fullCus: fullCustomer }); + // Round balance fields to handle floating-point precision from JSON.NUMINCRBY return roundFullCustomerBalances(fullCustomer); }; diff --git a/server/src/internal/customers/cusUtils/fullCustomerCacheUtils/getOrCreateCachedFullCustomer.ts b/server/src/internal/customers/cusUtils/fullCustomerCacheUtils/getOrCreateCachedFullCustomer.ts index fc127daad..732788f08 100644 --- a/server/src/internal/customers/cusUtils/fullCustomerCacheUtils/getOrCreateCachedFullCustomer.ts +++ b/server/src/internal/customers/cusUtils/fullCustomerCacheUtils/getOrCreateCachedFullCustomer.ts @@ -1,5 +1,4 @@ import { - type AppEnv, type CheckParams, CustomerExpand, type Entity, @@ -9,7 +8,6 @@ import { import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { customerActions } from "@/internal/customers/actions/index.js"; import { autoCreateEntity } from "@/internal/entities/handlers/handleCreateEntity/autoCreateEntity.js"; -import { resetCustomerEntitlements } from "../../actions/resetCustomerEntitlements/resetCustomerEntitlements.js"; import { CusService } from "../../CusService.js"; import { updateCustomerDetails } from "../cusUtils.js"; import { deleteCachedFullCustomer } from "./deleteCachedFullCustomer.js"; @@ -30,7 +28,7 @@ export const getOrCreateCachedFullCustomer = async ({ }; source?: string; }): Promise => { - const { org, env, db, skipCache, logger } = ctx; + const { skipCache, logger } = ctx; const { customer_id: customerId, customer_data: customerData, @@ -41,37 +39,23 @@ export const getOrCreateCachedFullCustomer = async ({ let fullCustomer: FullCustomer | undefined; const fetchTimeMs = Date.now(); - // 1. Try cache first + // 1. Try cache first (getCachedFullCustomer handles lazy reset internally) let setCache = true; if (customerId && !skipCache) { fullCustomer = (await getCachedFullCustomer({ - orgId: org.id, - env, + ctx, customerId, entityId, })) ?? undefined; if (fullCustomer) { logger.debug(`[getOrCreateCachedFullCustomer] Cache hit: ${customerId}`); - - // Lazy reset stale entitlements - const didReset = await resetCustomerEntitlements({ - fullCus: fullCustomer, - db, - org, - env: env as AppEnv, - }); - - if (didReset) { - setCache = true; - } else { - setCache = false; - } + setCache = false; } } - // 2. Try DB if not in cache + // 2. Try DB if not in cache (CusService.getFull handles lazy reset internally) if (!fullCustomer && customerId) { fullCustomer = await CusService.getFull({ ctx, @@ -81,16 +65,6 @@ export const getOrCreateCachedFullCustomer = async ({ expand: [CustomerExpand.Invoices], allowNotFound: true, }); - - // Lazy reset stale entitlements (DB path) - if (fullCustomer) { - await resetCustomerEntitlements({ - fullCus: fullCustomer, - db, - org, - env: env as AppEnv, - }); - } } // 3. Create if not found diff --git a/server/src/internal/customers/cusUtils/fullCustomerCacheUtils/getOrSetCachedFullCustomer.ts b/server/src/internal/customers/cusUtils/fullCustomerCacheUtils/getOrSetCachedFullCustomer.ts index d089a954f..47611e915 100644 --- a/server/src/internal/customers/cusUtils/fullCustomerCacheUtils/getOrSetCachedFullCustomer.ts +++ b/server/src/internal/customers/cusUtils/fullCustomerCacheUtils/getOrSetCachedFullCustomer.ts @@ -1,12 +1,10 @@ import { - type AppEnv, CustomerExpand, CustomerNotFoundError, EntityNotFoundError, type FullCustomer, } from "@autumn/shared"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; -import { resetCustomerEntitlements } from "../../actions/resetCustomerEntitlements/resetCustomerEntitlements.js"; import { CusService } from "../../CusService.js"; import { getCachedFullCustomer } from "./getCachedFullCustomer.js"; import { setCachedFullCustomer } from "./setCachedFullCustomer.js"; @@ -26,13 +24,12 @@ export const getOrSetCachedFullCustomer = async ({ entityId?: string; source?: string; }): Promise => { - const { org, env, db, skipCache, logger } = ctx; + const { skipCache, logger } = ctx; - // 1. Try cache first + // 1. Try cache first (getCachedFullCustomer handles lazy reset internally) if (!skipCache) { const cached = await getCachedFullCustomer({ - orgId: org.id, - env, + ctx, customerId, }); @@ -41,27 +38,6 @@ export const getOrSetCachedFullCustomer = async ({ `[getOrSetCachedFullCustomer] Cache hit for ${customerId}, source: ${source}`, ); - // // Lazy reset stale entitlements - // const didReset = await resetCustomerEntitlements({ - // fullCus: cached, - // db, - // org, - // env: env as AppEnv, - // }); - - // if (didReset) { - // await setCachedFullCustomer({ - // ctx, - // fullCustomer: cached, - // customerId, - // fetchTimeMs: Date.now(), - // source: "lazyReset", - // overwrite: true, - // }).catch((err) => - // logger.error(`[lazyReset] Failed to update cache: ${err}`), - // ); - // } - // Set entity if entityId is provided, otherwise clear it if (entityId) { cached.entity = cached.entities?.find((e) => e.id === entityId); @@ -96,14 +72,6 @@ export const getOrSetCachedFullCustomer = async ({ throw new CustomerNotFoundError({ customerId }); } - // Lazy reset stale entitlements (DB path) - await resetCustomerEntitlements({ - fullCus: fullCustomer, - db, - org, - env: env as AppEnv, - }); - if (entityId) { fullCustomer.entity = fullCustomer.entities?.find((e) => e.id === entityId); if (!fullCustomer.entity) { diff --git a/server/src/internal/customers/internalHandlers/handleGetCustomerEvents.ts b/server/src/internal/customers/internalHandlers/handleGetCustomerEvents.ts index 802adafe1..89ae36de3 100644 --- a/server/src/internal/customers/internalHandlers/handleGetCustomerEvents.ts +++ b/server/src/internal/customers/internalHandlers/handleGetCustomerEvents.ts @@ -19,13 +19,11 @@ export const handleGetCustomerEvents = createRoute({ query: QuerySchema, handler: async (c) => { const ctx = c.get("ctx"); - const { db, org, env } = ctx; const { customer_id } = c.req.param(); const { interval, limit } = c.req.valid("query"); const customer = await getCachedFullCustomer({ - orgId: org.id, - env, + ctx, customerId: customer_id, }); diff --git a/server/src/internal/customers/internalHandlers/handleGetFullCustomers.ts b/server/src/internal/customers/internalHandlers/handleGetFullCustomers.ts index 54c8fc8ef..3b57e7bd1 100644 --- a/server/src/internal/customers/internalHandlers/handleGetFullCustomers.ts +++ b/server/src/internal/customers/internalHandlers/handleGetFullCustomers.ts @@ -19,9 +19,11 @@ export const handleGetFullCustomers = createRoute({ filters: z.any().optional(), }), handler: async (c) => { - const { db, org, env } = c.get("ctx"); + const ctx = c.get("ctx"); const { search, page_size, page, last_item, filters } = c.req.valid("json"); + const { org, env, db } = ctx; + const { data: customers } = await CusSearchService.search({ db, orgId: org.id, @@ -34,9 +36,7 @@ export const handleGetFullCustomers = createRoute({ }); const fullCustomers = await CusBatchService.getByInternalIds({ - db, - org, - env, + ctx, internalCustomerIds: customers.map( (customer: Customer) => customer.internal_id, ), diff --git a/server/src/queue/processMessage.ts b/server/src/queue/processMessage.ts index c7a6a6213..9e606b603 100644 --- a/server/src/queue/processMessage.ts +++ b/server/src/queue/processMessage.ts @@ -8,7 +8,7 @@ import { runInsertEventBatch } from "@/internal/balances/events/runInsertEventBa import { syncItemV3 } from "@/internal/balances/utils/sync/syncItemV3.js"; import { grantCheckoutReward } from "@/internal/billing/v2/workflows/grantCheckoutReward/grantCheckoutReward.js"; import { sendProductsUpdated } from "@/internal/billing/v2/workflows/sendProductsUpdated/sendProductsUpdated.js"; -import { runBatchResetCusEntsTask } from "@/internal/customers/actions/resetCustomerEntitlements/runBatchResetCusEntsTask.js"; +import { batchResetCustomerEntitlements } from "@/internal/customers/actions/resetCustomerEntitlements/batchResetCustomerEntitlements.js"; import { runClearCreditSystemCacheTask } from "@/internal/features/featureActions/runClearCreditSystemCacheTask.js"; import { generateFeatureDisplay } from "@/internal/features/workflows/generateFeatureDisplay.js"; import { runMigrationTask } from "@/internal/migrations/runMigrationTask.js"; @@ -192,7 +192,7 @@ export const processMessage = async ({ workerLogger.error("No context found for batch reset cus ents job"); return; } - await runBatchResetCusEntsTask({ + await batchResetCustomerEntitlements({ ctx, payload: job.data, }); diff --git a/server/src/queue/workflows.ts b/server/src/queue/workflows.ts index bef984cb3..be2309ce2 100644 --- a/server/src/queue/workflows.ts +++ b/server/src/queue/workflows.ts @@ -38,7 +38,11 @@ export type GrantCheckoutRewardPayload = { export type BatchResetCusEntsPayload = { orgId: string; env: string; - cusEntIds: string[]; + resets: { + internalCustomerId: string; + customerId: string; + cusEntIds: string[]; + }[]; }; // ============ Workflow Registry ============ diff --git a/server/tests/integration/crud/customers/reset-customer-entitlements.test.ts b/server/tests/integration/crud/customers/reset-customer-entitlements.test.ts deleted file mode 100644 index 3fdef73fe..000000000 --- a/server/tests/integration/crud/customers/reset-customer-entitlements.test.ts +++ /dev/null @@ -1,21 +0,0 @@ -import { test } from "bun:test"; -import { items } from "@tests/utils/fixtures/items"; -import { products } from "@tests/utils/fixtures/products"; -import { initScenario, s } from "@tests/utils/testInitUtils/initScenario"; -import chalk from "chalk"; - -test.concurrent(`${chalk.yellowBright("create: basic with ID")}`, async () => { - const customerId = "create-basic-id"; - - const proProduct = products.pro({ - items: [ - items.monthlyMessages(), - items.consumableWords({ includedUsage: 100 }), - ], - }); - const { autumnV1 } = await initScenario({ - customerId, - setup: [s.customer({ testClock: false })], - actions: [], - }); -}); diff --git a/server/tests/integration/crud/customers/reset-customer-entitlements/get-customer-reset-concurrent.test.ts b/server/tests/integration/crud/customers/reset-customer-entitlements/get-customer-reset-concurrent.test.ts new file mode 100644 index 000000000..6469b8f25 --- /dev/null +++ b/server/tests/integration/crud/customers/reset-customer-entitlements/get-customer-reset-concurrent.test.ts @@ -0,0 +1,251 @@ +import { expect, test } from "bun:test"; +import type { ApiCustomer, CheckResponseV2 } from "@autumn/shared"; +import { findCustomerEntitlement } from "@tests/balances/utils/findCustomerEntitlement.js"; +import { TestFeature } from "@tests/setup/v2Features.js"; +import { expireCusEntForReset } from "@tests/utils/cusProductUtils/resetTestUtils.js"; +import { items } from "@tests/utils/fixtures/items.js"; +import { products } from "@tests/utils/fixtures/products.js"; +import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js"; +import chalk from "chalk"; + +// ───────────────────────────────────────────────────────────────── +// Concurrent GET /customers — multiple reads trigger reset exactly once +// ───────────────────────────────────────────────────────────────── + +test.concurrent(`${chalk.yellowBright("concurrent reset: multiple GET customers all return reset balance")}`, async () => { + const messagesItem = items.monthlyMessages({ includedUsage: 100 }); + const pro = products.pro({ items: [messagesItem] }); + + const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ + customerId: "reset-concurrent-get", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [pro] }), + ], + actions: [s.attach({ productId: pro.id })], + }); + + await autumnV1.track({ + customer_id: customerId, + feature_id: TestFeature.Messages, + value: 40, + }); + await new Promise((resolve) => setTimeout(resolve, 2000)); + + await expireCusEntForReset({ + ctx, + customerId, + featureId: TestFeature.Messages, + }); + + // Fire 5 concurrent GET requests — all should see reset balance + const results = await Promise.all( + Array.from({ length: 5 }, () => + autumnV2.customers.get(customerId), + ), + ); + + for (const customer of results) { + expect(customer.balances[TestFeature.Messages].current_balance).toBe(100); + expect(customer.balances[TestFeature.Messages].usage).toBe(0); + } + + // DB should also reflect the reset (only applied once) + const cusEntAfter = await findCustomerEntitlement({ + ctx, + customerId, + featureId: TestFeature.Messages, + }); + expect(cusEntAfter).toBeDefined(); + expect(cusEntAfter!.next_reset_at).toBeGreaterThan(Date.now()); +}); + +// ───────────────────────────────────────────────────────────────── +// Concurrent checks — all return reset balance +// ───────────────────────────────────────────────────────────────── + +test.concurrent(`${chalk.yellowBright("concurrent reset: multiple checks all return reset balance")}`, async () => { + const messagesItem = items.monthlyMessages({ includedUsage: 100 }); + const pro = products.pro({ items: [messagesItem] }); + + const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ + customerId: "reset-concurrent-check", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [pro] }), + ], + actions: [s.attach({ productId: pro.id })], + }); + + await autumnV1.track({ + customer_id: customerId, + feature_id: TestFeature.Messages, + value: 70, + }); + await new Promise((resolve) => setTimeout(resolve, 2000)); + + await expireCusEntForReset({ + ctx, + customerId, + featureId: TestFeature.Messages, + }); + + // Fire 5 concurrent check requests + const results = await Promise.all( + Array.from({ length: 5 }, () => + autumnV2.check({ + customer_id: customerId, + feature_id: TestFeature.Messages, + }), + ), + ); + + for (const res of results) { + const check = res as unknown as CheckResponseV2; + expect(check.allowed).toBe(true); + expect(check.balance?.current_balance).toBe(100); + expect(check.balance?.usage).toBe(0); + } +}); + +// ───────────────────────────────────────────────────────────────── +// Concurrent tracks — reset once, all deductions applied atomically +// ───────────────────────────────────────────────────────────────── + +test.concurrent(`${chalk.yellowBright("concurrent reset: multiple tracks reset once then deduct atomically")}`, async () => { + const messagesItem = items.monthlyMessages({ includedUsage: 100 }); + const pro = products.pro({ items: [messagesItem] }); + + const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ + customerId: "reset-concurrent-track", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [pro] }), + ], + actions: [s.attach({ productId: pro.id })], + }); + + await autumnV1.track({ + customer_id: customerId, + feature_id: TestFeature.Messages, + value: 50, + }); + await new Promise((resolve) => setTimeout(resolve, 2000)); + + await expireCusEntForReset({ + ctx, + customerId, + featureId: TestFeature.Messages, + }); + + // Fire 5 concurrent tracks of 10 each — should reset to 100, then deduct 50 total + await Promise.all( + Array.from({ length: 5 }, () => + autumnV2.track({ + customer_id: customerId, + feature_id: TestFeature.Messages, + value: 10, + }), + ), + ); + + // Verify final balance: 100 (reset) - 50 (5 * 10) = 50 + const customer = await autumnV2.customers.get(customerId); + expect(customer.balances[TestFeature.Messages].current_balance).toBe(50); + expect(customer.balances[TestFeature.Messages].usage).toBe(50); + + // Wait for DB sync and verify DB agrees + await new Promise((resolve) => setTimeout(resolve, 2000)); + + const customerDb = await autumnV2.customers.get(customerId, { + skip_cache: "true", + }); + expect(customerDb.balances[TestFeature.Messages].current_balance).toBe(50); + expect(customerDb.balances[TestFeature.Messages].usage).toBe(50); +}); + +// ───────────────────────────────────────────────────────────────── +// Mixed concurrent: GET + check + track all hit a stale cusEnt +// ───────────────────────────────────────────────────────────────── + +test.concurrent(`${chalk.yellowBright("concurrent reset: mixed GET/check/track all handle stale cusEnt correctly")}`, async () => { + const messagesItem = items.monthlyMessages({ includedUsage: 100 }); + const pro = products.pro({ items: [messagesItem] }); + + const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ + customerId: "reset-concurrent-mixed", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [pro] }), + ], + actions: [s.attach({ productId: pro.id })], + }); + + await autumnV1.track({ + customer_id: customerId, + feature_id: TestFeature.Messages, + value: 80, + }); + await new Promise((resolve) => setTimeout(resolve, 2000)); + + await expireCusEntForReset({ + ctx, + customerId, + featureId: TestFeature.Messages, + }); + + // Fire mixed concurrent requests: 2 GETs, 2 checks, 1 track(15) + const [get1, get2, check1, check2, _trackRes] = await Promise.all([ + autumnV2.customers.get(customerId), + autumnV2.customers.get(customerId), + autumnV2.check({ + customer_id: customerId, + feature_id: TestFeature.Messages, + }) as unknown as Promise, + autumnV2.check({ + customer_id: customerId, + feature_id: TestFeature.Messages, + }) as unknown as Promise, + autumnV2.track({ + customer_id: customerId, + feature_id: TestFeature.Messages, + value: 15, + }), + ]); + + // GETs should show reset balance (may or may not include the track deduction depending on ordering) + for (const customer of [get1, get2]) { + // Balance should be >= 85 (reset 100 minus at most 15 from track) + expect( + customer.balances[TestFeature.Messages].current_balance, + ).toBeGreaterThanOrEqual(85); + expect( + customer.balances[TestFeature.Messages].current_balance, + ).toBeLessThanOrEqual(100); + } + + // Checks should show reset balance + for (const check of [check1, check2]) { + expect(check.allowed).toBe(true); + expect(check.balance?.current_balance).toBeGreaterThanOrEqual(85); + expect(check.balance?.current_balance).toBeLessThanOrEqual(100); + } + + // Wait for dust to settle, verify final state + await new Promise((resolve) => setTimeout(resolve, 2000)); + + const finalDb = await autumnV2.customers.get(customerId, { + skip_cache: "true", + }); + // Final balance must be exactly 85: reset to 100, one track of 15 + expect(finalDb.balances[TestFeature.Messages].current_balance).toBe(85); + expect(finalDb.balances[TestFeature.Messages].usage).toBe(15); + + const cusEntAfter = await findCustomerEntitlement({ + ctx, + customerId, + featureId: TestFeature.Messages, + }); + expect(cusEntAfter).toBeDefined(); + expect(cusEntAfter!.next_reset_at).toBeGreaterThan(Date.now()); +}); diff --git a/server/tests/integration/crud/customers/reset-customer-entitlements/get-customer-reset.test.ts b/server/tests/integration/crud/customers/reset-customer-entitlements/get-customer-reset.test.ts new file mode 100644 index 000000000..c4e954327 --- /dev/null +++ b/server/tests/integration/crud/customers/reset-customer-entitlements/get-customer-reset.test.ts @@ -0,0 +1,265 @@ +import { expect, test } from "bun:test"; +import type { + ApiCustomer, + CheckResponseV2, + TrackResponseV2, +} from "@autumn/shared"; +import { findCustomerEntitlement } from "@tests/balances/utils/findCustomerEntitlement.js"; +import { TestFeature } from "@tests/setup/v2Features.js"; +import { expireCusEntForReset } from "@tests/utils/cusProductUtils/resetTestUtils.js"; +import { items } from "@tests/utils/fixtures/items.js"; +import { products } from "@tests/utils/fixtures/products.js"; +import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js"; +import chalk from "chalk"; + +// ───────────────────────────────────────────────────────────────── +// GET /customers (skip_cache) — DB path lazy reset +// ───────────────────────────────────────────────────────────────── + +test.concurrent(`${chalk.yellowBright("lazy reset (DB): GET customer resets balance after next_reset_at passes")}`, async () => { + const messagesItem = items.monthlyMessages({ includedUsage: 100 }); + const pro = products.pro({ items: [messagesItem] }); + + const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ + customerId: "reset-get-db", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [pro] }), + ], + actions: [s.attach({ productId: pro.id })], + }); + + await autumnV1.track({ + customer_id: customerId, + feature_id: TestFeature.Messages, + value: 30, + }); + await new Promise((resolve) => setTimeout(resolve, 2000)); + + const before = await autumnV2.customers.get(customerId, { + skip_cache: "true", + }); + expect(before.balances[TestFeature.Messages].current_balance).toBe(70); + expect(before.balances[TestFeature.Messages].usage).toBe(30); + + await expireCusEntForReset({ + ctx, + customerId, + featureId: TestFeature.Messages, + }); + + const after = await autumnV2.customers.get(customerId, { + skip_cache: "true", + }); + expect(after.balances[TestFeature.Messages].current_balance).toBe(100); + expect(after.balances[TestFeature.Messages].usage).toBe(0); + + const cusEntAfter = await findCustomerEntitlement({ + ctx, + customerId, + featureId: TestFeature.Messages, + }); + expect(cusEntAfter).toBeDefined(); + expect(cusEntAfter!.next_reset_at).toBeGreaterThan(Date.now()); +}); + +// ───────────────────────────────────────────────────────────────── +// GET /customers (cached) — cache path lazy reset +// ───────────────────────────────────────────────────────────────── + +test.concurrent(`${chalk.yellowBright("lazy reset (cache): GET customer resets balance from cache after next_reset_at passes")}`, async () => { + const messagesItem = items.monthlyMessages({ includedUsage: 100 }); + const pro = products.pro({ items: [messagesItem] }); + + const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ + customerId: "reset-get-cache", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [pro] }), + ], + actions: [s.attach({ productId: pro.id })], + }); + + await autumnV1.track({ + customer_id: customerId, + feature_id: TestFeature.Messages, + value: 30, + }); + await new Promise((resolve) => setTimeout(resolve, 2000)); + + // Warm the cache + const before = await autumnV2.customers.get(customerId); + expect(before.balances[TestFeature.Messages].current_balance).toBe(70); + + await expireCusEntForReset({ + ctx, + customerId, + featureId: TestFeature.Messages, + }); + + const after = await autumnV2.customers.get(customerId); + expect(after.balances[TestFeature.Messages].current_balance).toBe(100); + expect(after.balances[TestFeature.Messages].usage).toBe(0); + + const cusEntAfter = await findCustomerEntitlement({ + ctx, + customerId, + featureId: TestFeature.Messages, + }); + expect(cusEntAfter).toBeDefined(); + expect(cusEntAfter!.next_reset_at).toBeGreaterThan(Date.now()); +}); + +// ───────────────────────────────────────────────────────────────── +// POST /check — lazy reset before check +// ───────────────────────────────────────────────────────────────── + +test.concurrent(`${chalk.yellowBright("lazy reset: check returns reset balance after next_reset_at passes")}`, async () => { + const messagesItem = items.monthlyMessages({ includedUsage: 100 }); + const pro = products.pro({ items: [messagesItem] }); + + const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ + customerId: "reset-check", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [pro] }), + ], + actions: [s.attach({ productId: pro.id })], + }); + + await autumnV1.track({ + customer_id: customerId, + feature_id: TestFeature.Messages, + value: 60, + }); + await new Promise((resolve) => setTimeout(resolve, 2000)); + + // Verify pre-reset state + const checkBefore = (await autumnV2.check({ + customer_id: customerId, + feature_id: TestFeature.Messages, + })) as unknown as CheckResponseV2; + expect(checkBefore.allowed).toBe(true); + expect(checkBefore.balance?.current_balance).toBe(40); + expect(checkBefore.balance?.usage).toBe(60); + + await expireCusEntForReset({ + ctx, + customerId, + featureId: TestFeature.Messages, + }); + + // Check should trigger lazy reset and return reset balance + const checkAfter = (await autumnV2.check({ + customer_id: customerId, + feature_id: TestFeature.Messages, + })) as unknown as CheckResponseV2; + expect(checkAfter.allowed).toBe(true); + expect(checkAfter.balance?.current_balance).toBe(100); + expect(checkAfter.balance?.usage).toBe(0); +}); + +// ───────────────────────────────────────────────────────────────── +// POST /track — lazy reset then deduction +// ───────────────────────────────────────────────────────────────── + +test.concurrent(`${chalk.yellowBright("lazy reset: track resets balance then deducts correctly")}`, async () => { + const messagesItem = items.monthlyMessages({ includedUsage: 100 }); + const pro = products.pro({ items: [messagesItem] }); + + const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ + customerId: "reset-track", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [pro] }), + ], + actions: [s.attach({ productId: pro.id })], + }); + + await autumnV1.track({ + customer_id: customerId, + feature_id: TestFeature.Messages, + value: 50, + }); + await new Promise((resolve) => setTimeout(resolve, 2000)); + + await expireCusEntForReset({ + ctx, + customerId, + featureId: TestFeature.Messages, + }); + + // Track 20 — should reset (100) then deduct (100 - 20 = 80) + const trackRes: TrackResponseV2 = await autumnV2.track({ + customer_id: customerId, + feature_id: TestFeature.Messages, + value: 20, + }); + expect(trackRes.balance?.current_balance).toBe(80); + expect(trackRes.balance?.usage).toBe(20); + + // Verify cache reflects reset + deduction + const customer = await autumnV2.customers.get(customerId); + expect(customer.balances[TestFeature.Messages].current_balance).toBe(80); + expect(customer.balances[TestFeature.Messages].usage).toBe(20); + + // Wait for DB sync and verify DB state + await new Promise((resolve) => setTimeout(resolve, 2000)); + + const customerDb = await autumnV2.customers.get(customerId, { + skip_cache: "true", + }); + expect(customerDb.balances[TestFeature.Messages].current_balance).toBe(80); + expect(customerDb.balances[TestFeature.Messages].usage).toBe(20); + + const cusEntAfter = await findCustomerEntitlement({ + ctx, + customerId, + featureId: TestFeature.Messages, + }); + expect(cusEntAfter).toBeDefined(); + expect(cusEntAfter!.next_reset_at).toBeGreaterThan(Date.now()); +}); + +// ───────────────────────────────────────────────────────────────── +// POST /customers (create-or-get) — lazy reset on existing customer +// ───────────────────────────────────────────────────────────────── + +test.concurrent(`${chalk.yellowBright("lazy reset: POST /check on existing customer triggers reset")}`, async () => { + const messagesItem = items.monthlyMessages({ includedUsage: 100 }); + const pro = products.pro({ items: [messagesItem] }); + + const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ + customerId: "reset-post-check", + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.products({ list: [pro] }), + ], + actions: [s.attach({ productId: pro.id })], + }); + + await autumnV1.track({ + customer_id: customerId, + feature_id: TestFeature.Messages, + value: 40, + }); + await new Promise((resolve) => setTimeout(resolve, 2000)); + + await expireCusEntForReset({ + ctx, + customerId, + featureId: TestFeature.Messages, + }); + + // POST /check fetches the customer (triggering reset) before checking + const checkRes = (await autumnV2.check({ + customer_id: customerId, + feature_id: TestFeature.Messages, + required_balance: 90, + })) as unknown as CheckResponseV2; + + // Should be allowed since balance was reset to 100 + expect(checkRes.allowed).toBe(true); + expect(checkRes.balance?.current_balance).toBe(100); + expect(checkRes.balance?.usage).toBe(0); +}); diff --git a/server/tests/integration/crud/customers/reset-customer-entitlements/list-customers-reset.test.ts b/server/tests/integration/crud/customers/reset-customer-entitlements/list-customers-reset.test.ts new file mode 100644 index 000000000..361bd54c8 --- /dev/null +++ b/server/tests/integration/crud/customers/reset-customer-entitlements/list-customers-reset.test.ts @@ -0,0 +1,131 @@ +import { expect, test } from "bun:test"; +import type { ApiCustomer } from "@autumn/shared"; +import { findCustomerEntitlement } from "@tests/balances/utils/findCustomerEntitlement.js"; +import { TestFeature } from "@tests/setup/v2Features.js"; +import { expireCusEntForReset } from "@tests/utils/cusProductUtils/resetTestUtils.js"; +import { items } from "@tests/utils/fixtures/items.js"; +import { products } from "@tests/utils/fixtures/products.js"; +import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js"; +import chalk from "chalk"; + +// ───────────────────────────────────────────────────────────────── +// POST /customers/list — batch reset: only stale customers reset +// +// 6 customers with shared prefix. Customers 2, 4, 6 have their +// next_reset_at expired. After listV2 triggers the async SQS batch +// reset, only those 3 should be reset; the other 3 keep their usage. +// ───────────────────────────────────────────────────────────────── + +const PREFIX = "reset-list-cohort"; +const OTHER_IDS = Array.from({ length: 5 }, (_, i) => `${PREFIX}-${i + 2}`); + +// Indices 2, 4, 6 are stale (0-indexed: 1, 3, 5 in the all-customers array) +const STALE_IDS = [`${PREFIX}-2`, `${PREFIX}-4`, `${PREFIX}-6`]; +const FRESH_IDS = [`${PREFIX}-1`, `${PREFIX}-3`, `${PREFIX}-5`]; + +// Each customer tracks a different amount so we can verify individually +const USAGE: Record = { + [`${PREFIX}-1`]: 20, + [`${PREFIX}-2`]: 35, + [`${PREFIX}-3`]: 50, + [`${PREFIX}-4`]: 15, + [`${PREFIX}-5`]: 70, + [`${PREFIX}-6`]: 40, +}; + +test.concurrent(`${chalk.yellowBright("list customers reset: only stale customers are reset, fresh customers keep usage")}`, async () => { + const messagesItem = items.monthlyMessages({ includedUsage: 100 }); + const freePlan = products.base({ + id: "free", + items: [messagesItem], + }); + + const primaryId = `${PREFIX}-1`; + + const { autumnV1, autumnV2, ctx } = await initScenario({ + customerId: primaryId, + setup: [ + s.customer({ paymentMethod: "success", testClock: false }), + s.otherCustomers( + OTHER_IDS.map((id) => ({ + id, + testClock: false, + })), + ), + s.products({ list: [freePlan] }), + ], + actions: [ + // Attach free plan to all 6 customers + s.attach({ productId: freePlan.id }), + ...OTHER_IDS.map((id) => + s.attach({ productId: freePlan.id, customerId: id }), + ), + ], + }); + + // 1. Track different usage on each customer + const allIds = [primaryId, ...OTHER_IDS]; + await Promise.all( + allIds.map((id) => + autumnV1.track({ + customer_id: id, + feature_id: TestFeature.Messages, + value: USAGE[id], + }), + ), + ); + await new Promise((resolve) => setTimeout(resolve, 2000)); + + // 2. Expire only the stale customers (2, 4, 6) + await Promise.all( + STALE_IDS.map((id) => + expireCusEntForReset({ + ctx, + customerId: id, + featureId: TestFeature.Messages, + }), + ), + ); + + // 3. Call listV2 with the shared prefix — triggers batch reset via SQS + const listRes = (await autumnV2.customers.listV2({ + search: PREFIX, + })) as { list: ApiCustomer[] }; + + // All 6 should appear in the list + for (const id of allIds) { + const found = listRes.list.find((c) => c.id === id); + expect(found).toBeDefined(); + } + + // 4. Wait for the SQS batch reset worker to process + await new Promise((resolve) => setTimeout(resolve, 5000)); + + // 5. Verify stale customers (2, 4, 6) were reset to full balance + for (const id of STALE_IDS) { + const customer = await autumnV2.customers.get(id, { + skip_cache: "true", + }); + expect(customer.balances[TestFeature.Messages].current_balance).toBe(100); + expect(customer.balances[TestFeature.Messages].usage).toBe(0); + + const cusEnt = await findCustomerEntitlement({ + ctx, + customerId: id, + featureId: TestFeature.Messages, + }); + expect(cusEnt).toBeDefined(); + expect(cusEnt!.next_reset_at).toBeGreaterThan(Date.now()); + } + + // 6. Verify fresh customers (1, 3, 5) kept their usage — no reset + for (const id of FRESH_IDS) { + const customer = await autumnV2.customers.get(id, { + skip_cache: "true", + }); + expect(customer.balances[TestFeature.Messages].current_balance).toBe( + 100 - USAGE[id], + ); + expect(customer.balances[TestFeature.Messages].usage).toBe(USAGE[id]); + } +}); diff --git a/server/tests/utils/cusProductUtils/resetTestUtils.ts b/server/tests/utils/cusProductUtils/resetTestUtils.ts new file mode 100644 index 000000000..6c72380ac --- /dev/null +++ b/server/tests/utils/cusProductUtils/resetTestUtils.ts @@ -0,0 +1,110 @@ +import { customerEntitlements, type FullCustomer } from "@autumn/shared"; +import { findCustomerEntitlement } from "@tests/balances/utils/findCustomerEntitlement.js"; +import type { TestContext } from "@tests/utils/testInitUtils/createTestContext.js"; +import { eq } from "drizzle-orm"; +import { redis } from "@/external/redis/initRedis.js"; +import { buildFullCustomerCacheKey } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/fullCustomerCacheConfig.js"; + +/** + * Update next_reset_at for a specific cusEnt in the Redis FullCustomer cache. + * Reads the cached blob, finds the cusEnt by ID, then uses JSON.SET on the exact path. + */ +export const setCachedCusEntField = async ({ + orgId, + env, + customerId, + cusEntId, + field, + value, +}: { + orgId: string; + env: string; + customerId: string; + cusEntId: string; + field: string; + value: number | string | null; +}): Promise => { + const cacheKey = buildFullCustomerCacheKey({ orgId, env, customerId }); + + const raw = (await redis.call("JSON.GET", cacheKey)) as string | null; + if (!raw) return; + + const fullCustomer = JSON.parse(raw) as FullCustomer; + const serializedValue = value === null ? "null" : JSON.stringify(value); + + for (let cpIdx = 0; cpIdx < fullCustomer.customer_products.length; cpIdx++) { + const cusEnts = fullCustomer.customer_products[cpIdx].customer_entitlements; + for (let ceIdx = 0; ceIdx < cusEnts.length; ceIdx++) { + if (cusEnts[ceIdx].id === cusEntId) { + await redis.call( + "JSON.SET", + cacheKey, + `$.customer_products[${cpIdx}].customer_entitlements[${ceIdx}].${field}`, + serializedValue, + ); + return; + } + } + } + + const extras = fullCustomer.extra_customer_entitlements || []; + for (let eIdx = 0; eIdx < extras.length; eIdx++) { + if (extras[eIdx].id === cusEntId) { + await redis.call( + "JSON.SET", + cacheKey, + `$.extra_customer_entitlements[${eIdx}].${field}`, + serializedValue, + ); + return; + } + } +}; + +/** + * Expire a cusEnt's next_reset_at in both Postgres and Redis cache, + * so the next read triggers a lazy reset. Returns the cusEnt for assertions. + */ +export const expireCusEntForReset = async ({ + ctx, + customerId, + featureId, + pastTimeMs, +}: { + ctx: TestContext; + customerId: string; + featureId: string; + pastTimeMs?: number; +}) => { + const cusEnt = await findCustomerEntitlement({ + ctx, + customerId, + featureId, + }); + + if (!cusEnt) { + throw new Error( + `cusEnt not found for customer=${customerId} feature=${featureId}`, + ); + } + + const pastTime = pastTimeMs ?? Date.now() - 1000; + + // Update Postgres + await ctx.db + .update(customerEntitlements) + .set({ next_reset_at: pastTime }) + .where(eq(customerEntitlements.id, cusEnt.id)); + + // Update Redis cache + await setCachedCusEntField({ + orgId: ctx.org.id, + env: ctx.env, + customerId, + cusEntId: cusEnt.id, + field: "next_reset_at", + value: pastTime, + }); + + return cusEnt; +}; diff --git a/server/tests/utils/testInitUtils/initScenario.ts b/server/tests/utils/testInitUtils/initScenario.ts index 36cfbacbc..3e89f3eea 100644 --- a/server/tests/utils/testInitUtils/initScenario.ts +++ b/server/tests/utils/testInitUtils/initScenario.ts @@ -1114,7 +1114,7 @@ export async function initScenario({ customerData: otherCusConfig.data, attachPm: otherCusConfig.paymentMethod, withTestClock: false, // Don't create a new test clock - existingTestClockId: testClockId, // Reuse primary customer's test clock + ...(testClockId ? { existingTestClockId: testClockId } : {}), withDefault: false, defaultGroup: productPrefix, skipWebhooks: config.skipWebhooks,