diff --git a/server/experiments/sanitizeBenchmark.ts b/server/experiments/sanitizeBenchmark.ts new file mode 100644 index 000000000..890fdae13 --- /dev/null +++ b/server/experiments/sanitizeBenchmark.ts @@ -0,0 +1,354 @@ +/** + * Benchmark: sanitizeCachedFullSubject + sanitizeCachedSubjectBalance + * + * Simulates a customer with 1000 customer products, each with 3 customer + * entitlements (3000 SubjectBalance blobs total), and measures how long + * sanitization takes. + * + * Run with: bun run experiments/sanitizeBenchmark.ts + */ + +import { AppEnv } from "@autumn/shared"; +import { CollectionMethod, CusProductStatus } from "@shared/models/cusProductModels/cusProductEnums.js"; +import { FeatureType } from "@shared/models/featureModels/featureEnums.js"; +import { AllowanceType } from "@shared/models/productModels/entModels/entModels.js"; +import { EntInterval } from "@shared/models/productModels/intervals/entitlementInterval.js"; +import type { CachedFullSubject } from "@/internal/customers/cache/fullSubject/fullSubjectCacheModel.js"; +import { sanitizeCachedFullSubject } from "@/internal/customers/cache/fullSubject/sanitize/sanitizeCachedFullSubject.js"; +import { sanitizeCachedSubjectBalance } from "@/internal/customers/cache/fullSubject/sanitize/sanitizeCachedSubjectBalance.js"; + +const NUM_CUSTOMER_PRODUCTS = 1000; +const CUS_ENTS_PER_PRODUCT = 3; +const TOTAL_CUS_ENTS = NUM_CUSTOMER_PRODUCTS * CUS_ENTS_PER_PRODUCT; +const WARMUP_ITERATIONS = 100; +const BENCHMARK_ITERATIONS = 1000; + +const now = Date.now(); + +const buildSubjectBalance = ({ index }: { index: number }) => ({ + id: `cus_ent_${index}`, + customer_product_id: `cp_${Math.floor(index / CUS_ENTS_PER_PRODUCT)}`, + entitlement_id: `ent_${index}`, + internal_customer_id: "cus_int_1", + internal_entity_id: null, + internal_feature_id: `feat_int_${index}`, + feature_id: `feature_${index}`, + unlimited: false, + balance: 100 + index, + adjustment: 0, + additional_balance: 0, + usage_allowed: true, + next_reset_at: now + 86400000, + expires_at: null, + external_id: null, + cache_version: 1, + created_at: now, + customer_id: "cus_1", + rollovers: index % 5 === 0 + ? [ + { + id: `rollover_${index}`, + cus_ent_id: `cus_ent_${index}`, + balance: 50, + usage: 10, + expires_at: now + 172800000, + entities: {}, + }, + ] + : [], + entities: index % 3 === 0 + ? { + entity_a: { id: "entity_a", balance: 50, adjustment: 0 }, + entity_b: { id: "entity_b", balance: 30, adjustment: 0 }, + } + : null, + entitlement: { + id: `ent_${index}`, + created_at: now, + internal_feature_id: `feat_int_${index}`, + internal_product_id: `prod_int_${Math.floor(index / CUS_ENTS_PER_PRODUCT)}`, + is_custom: false, + allowance_type: AllowanceType.Fixed, + allowance: 1000, + interval: EntInterval.Month, + interval_count: 1, + carry_from_previous: false, + entity_feature_id: null, + org_id: "org_1", + feature_id: `feature_${index}`, + usage_limit: null, + rollover: null, + feature: { + internal_id: `feat_int_${index}`, + org_id: "org_1", + created_at: now, + env: AppEnv.Sandbox, + id: `feature_${index}`, + name: `Feature ${index}`, + type: FeatureType.Metered, + config: { usage_type: "single" }, + display: null, + archived: false, + event_names: ["track_event"], + }, + }, + customerPrice: null, + customerProductOptions: null, + customerProductQuantity: 1, +}); + +const buildCachedFullSubject = (): CachedFullSubject => { + const customerProducts = Array.from( + { length: NUM_CUSTOMER_PRODUCTS }, + (_, index) => ({ + id: `cp_${index}`, + internal_product_id: `prod_int_${index}`, + product_id: `prod_${index}`, + internal_customer_id: "cus_int_1", + customer_id: "cus_1", + internal_entity_id: null, + entity_id: null, + created_at: now, + status: CusProductStatus.Active, + canceled: false, + starts_at: now, + trial_ends_at: null, + canceled_at: null, + ended_at: null, + options: [ + { + feature_id: `feature_${index * CUS_ENTS_PER_PRODUCT}`, + quantity: 1, + upcoming_quantity: null, + adjustable_quantity: false, + internal_feature_id: `feat_int_${index * CUS_ENTS_PER_PRODUCT}`, + }, + ], + free_trial_id: null, + collection_method: CollectionMethod.ChargeAutomatically, + subscription_ids: ["sub_1"], + scheduled_ids: null, + processor: null, + quantity: 1, + api_semver: null, + is_custom: false, + billing_version: "v1" as const, + external_id: null, + }), + ); + + const entitlements = Array.from( + { length: TOTAL_CUS_ENTS }, + (_, index) => ({ + id: `ent_${index}`, + created_at: now, + internal_feature_id: `feat_int_${index}`, + internal_product_id: `prod_int_${Math.floor(index / CUS_ENTS_PER_PRODUCT)}`, + is_custom: false, + allowance_type: AllowanceType.Fixed, + allowance: 1000, + interval: EntInterval.Month, + interval_count: 1, + carry_from_previous: false, + entity_feature_id: null, + org_id: "org_1", + feature_id: `feature_${index}`, + usage_limit: null, + rollover: null, + feature: { + internal_id: `feat_int_${index}`, + org_id: "org_1", + created_at: now, + env: AppEnv.Sandbox, + id: `feature_${index}`, + name: `Feature ${index}`, + type: FeatureType.Metered, + config: { usage_type: "single" }, + display: null, + archived: false, + event_names: ["track_event"], + }, + }), + ); + + const products = Array.from( + { length: NUM_CUSTOMER_PRODUCTS }, + (_, index) => ({ + internal_id: `prod_int_${index}`, + id: `prod_${index}`, + name: `Product ${index}`, + description: null, + is_add_on: false, + is_default: false, + version: 1, + group: "", + env: AppEnv.Sandbox, + org_id: "org_1", + created_at: now, + processor: null, + base_variant_id: null, + archived: false, + }), + ); + + const customerEntitlementIdsByFeatureId: Record = {}; + for (let i = 0; i < TOTAL_CUS_ENTS; i++) { + customerEntitlementIdsByFeatureId[`feature_${i}`] = [`cus_ent_${i}`]; + } + + return { + subjectType: "customer", + customerId: "cus_1", + internalCustomerId: "cus_int_1", + customer: { + id: "cus_1", + internal_id: "cus_int_1", + org_id: "org_1", + env: AppEnv.Live, + created_at: now, + name: "Benchmark Customer", + email: "bench@test.com", + fingerprint: null, + processor: null, + processors: null, + metadata: {}, + send_email_receipts: false, + auto_topups: [], + spend_limits: [], + usage_alerts: [], + overage_allowed: [], + }, + customer_products: customerProducts, + products, + entitlements, + prices: [], + free_trials: [], + subscriptions: [], + invoices: [], + flags: {}, + _cachedAt: now, + meteredFeatures: Array.from({ length: TOTAL_CUS_ENTS }, (_, i) => `feature_${i}`), + customerEntitlementIdsByFeatureId, + subjectViewEpoch: 1, + } as unknown as CachedFullSubject; +}; + +const subjectBalances = Array.from( + { length: TOTAL_CUS_ENTS }, + (_, index) => buildSubjectBalance({ index }), +); + +const cachedFullSubject = buildCachedFullSubject(); + +const formatNs = (nanoseconds: bigint): string => { + const microseconds = Number(nanoseconds) / 1000; + if (microseconds < 1000) return `${microseconds.toFixed(1)}us`; + return `${(microseconds / 1000).toFixed(3)}ms`; +}; + +console.log("=== Sanitize Benchmark ==="); +console.log(`Customer products: ${NUM_CUSTOMER_PRODUCTS}`); +console.log(`CusEnts per product: ${CUS_ENTS_PER_PRODUCT}`); +console.log(`Total SubjectBalance blobs: ${TOTAL_CUS_ENTS}`); +console.log(`Warmup iterations: ${WARMUP_ITERATIONS}`); +console.log(`Benchmark iterations: ${BENCHMARK_ITERATIONS}`); +console.log(); + +// --- Benchmark: sanitizeCachedSubjectBalance (single) --- +for (let i = 0; i < WARMUP_ITERATIONS; i++) { + sanitizeCachedSubjectBalance({ subjectBalance: subjectBalances[0] as any }); +} + +let singleTotal = 0n; +for (let i = 0; i < BENCHMARK_ITERATIONS; i++) { + const balance = subjectBalances[i % subjectBalances.length]; + const start = Bun.nanoseconds(); + sanitizeCachedSubjectBalance({ subjectBalance: balance as any }); + singleTotal += BigInt(Math.round(Bun.nanoseconds() - start)); +} +const singleAvg = singleTotal / BigInt(BENCHMARK_ITERATIONS); +console.log(`sanitizeCachedSubjectBalance (1 balance):`); +console.log(` avg: ${formatNs(singleAvg)}`); +console.log(` total for ${BENCHMARK_ITERATIONS} calls: ${formatNs(singleTotal)}`); +console.log(); + +// --- Benchmark: sanitizeCachedSubjectBalance (all 3000) --- +for (let i = 0; i < WARMUP_ITERATIONS; i++) { + for (const balance of subjectBalances) { + sanitizeCachedSubjectBalance({ subjectBalance: balance as any }); + } +} + +let batchTotal = 0n; +for (let i = 0; i < 10; i++) { + const start = Bun.nanoseconds(); + for (const balance of subjectBalances) { + sanitizeCachedSubjectBalance({ subjectBalance: balance as any }); + } + batchTotal += BigInt(Math.round(Bun.nanoseconds() - start)); +} +const batchAvg = batchTotal / 10n; +console.log(`sanitizeCachedSubjectBalance (all ${TOTAL_CUS_ENTS} balances):`); +console.log(` avg per full batch: ${formatNs(batchAvg)}`); +console.log(` avg per balance: ${formatNs(batchAvg / BigInt(TOTAL_CUS_ENTS))}`); +console.log(); + +// --- Benchmark: sanitizeCachedFullSubject --- +for (let i = 0; i < WARMUP_ITERATIONS; i++) { + sanitizeCachedFullSubject({ cachedFullSubject: cachedFullSubject }); +} + +let shellTotal = 0n; +for (let i = 0; i < BENCHMARK_ITERATIONS; i++) { + const start = Bun.nanoseconds(); + sanitizeCachedFullSubject({ cachedFullSubject: cachedFullSubject }); + shellTotal += BigInt(Math.round(Bun.nanoseconds() - start)); +} +const shellAvg = shellTotal / BigInt(BENCHMARK_ITERATIONS); +console.log(`sanitizeCachedFullSubject (${NUM_CUSTOMER_PRODUCTS} products, ${TOTAL_CUS_ENTS} entitlements):`); +console.log(` avg: ${formatNs(shellAvg)}`); +console.log(` total for ${BENCHMARK_ITERATIONS} calls: ${formatNs(shellTotal)}`); +console.log(); + +// --- Combined: full cache read sanitization --- +for (let i = 0; i < WARMUP_ITERATIONS; i++) { + sanitizeCachedFullSubject({ cachedFullSubject: cachedFullSubject }); + for (const balance of subjectBalances) { + sanitizeCachedSubjectBalance({ subjectBalance: balance as any }); + } +} + +let combinedTotal = 0n; +for (let i = 0; i < 10; i++) { + const start = Bun.nanoseconds(); + sanitizeCachedFullSubject({ cachedFullSubject: cachedFullSubject }); + for (const balance of subjectBalances) { + sanitizeCachedSubjectBalance({ subjectBalance: balance as any }); + } + combinedTotal += BigInt(Math.round(Bun.nanoseconds() - start)); +} +const combinedAvg = combinedTotal / 10n; +console.log(`Combined (shell + all ${TOTAL_CUS_ENTS} balances):`); +console.log(` avg: ${formatNs(combinedAvg)}`); +console.log(); + +// --- JSON.parse baseline for comparison --- +const serializedBalances = subjectBalances.map((balance) => JSON.stringify(balance)); +const serializedShell = JSON.stringify(cachedFullSubject); + +let parseTotal = 0n; +for (let i = 0; i < 10; i++) { + const start = Bun.nanoseconds(); + JSON.parse(serializedShell); + for (const json of serializedBalances) { + JSON.parse(json); + } + parseTotal += BigInt(Math.round(Bun.nanoseconds() - start)); +} +const parseAvg = parseTotal / 10n; +console.log(`Baseline: JSON.parse (shell + ${TOTAL_CUS_ENTS} balances):`); +console.log(` avg: ${formatNs(parseAvg)}`); +console.log(); + +const overheadPct = Number(combinedAvg) / Number(parseAvg) * 100; +console.log(`Sanitize overhead vs JSON.parse: ${overheadPct.toFixed(1)}%`); diff --git a/server/src/_luaScriptsV2/fullSubject/updateEntityDataV2.lua b/server/src/_luaScriptsV2/fullSubject/updateEntityDataV2.lua new file mode 100644 index 000000000..8bc42eebd --- /dev/null +++ b/server/src/_luaScriptsV2/fullSubject/updateEntityDataV2.lua @@ -0,0 +1,55 @@ +--[[ + Lua Script: Update Entity Data in FullSubject V2 Redis Cache + + Atomically updates top-level entity fields in the cached FullSubject JSON. + + KEYS[1] = FullSubject cache key (entity-specific) + + ARGV[1] = updates JSON object + ARGV[2] = cache TTL in seconds + ARGV[3] = current timestamp in ms + + Returns JSON: + { "success": true, "updated_fields": ["spend_limits", "usage_alerts"] } + or + { "success": false, "cache_miss": true } + or + { "success": false, "no_entity": true } +]] + +local subject_key = KEYS[1] +local updates = cjson.decode(ARGV[1]) +local cache_ttl = tonumber(ARGV[2]) +local now_ms = tonumber(ARGV[3]) + +local has_updates = false +for _ in pairs(updates) do + has_updates = true + break +end + +if not has_updates then + return cjson.encode({ success = true, updated_fields = {} }) +end + +local current_raw = redis.call("GET", subject_key) +if not current_raw then + return cjson.encode({ success = false, cache_miss = true }) +end + +local cached = cjson.decode(current_raw) + +if not cached.entity then + return cjson.encode({ success = false, no_entity = true }) +end + +local updated_fields = {} + +for field_name, field_value in pairs(updates) do + cached.entity[field_name] = field_value + table.insert(updated_fields, field_name) +end + +redis.call("SET", subject_key, cjson.encode(cached), "EX", cache_ttl) + +return cjson.encode({ success = true, updated_fields = updated_fields }) diff --git a/server/src/_luaScriptsV2/fullSubject/updateSubjectBalances.lua b/server/src/_luaScriptsV2/fullSubject/updateSubjectBalances.lua new file mode 100644 index 000000000..fd375108f --- /dev/null +++ b/server/src/_luaScriptsV2/fullSubject/updateSubjectBalances.lua @@ -0,0 +1,185 @@ +--[[ + Lua Script: Update Subject Balances in V2 Cache (per-feature hash) + + Atomically updates SubjectBalance entries in a single per-feature + balance hash. Each hash field is cusEntId → JSON(SubjectBalance). + + Supports: scalar field updates, rollover operations (insert/overwrite/delete), + replaceable operations (insert/delete), and an expected_next_reset_at guard. + + Helper functions prepended via string interpolation from: + - luaUtils.lua (safe_number, is_nil, safe_table) + + KEYS[1] = balance hash key + e.g. {customerId}:orgId:env:full_subject:shared_balances:{featureId} + + ARGV[1] = JSON params: + { + ttl_seconds: number, + updates: [{ + cus_ent_id: string, + balance: number | null, + additional_balance: number | null, + adjustment: number | null, + entities: object | null, + next_reset_at: number | null, + expected_next_reset_at: number | null, + rollover_insert: { id, cus_ent_id, balance, usage, expires_at, entities } | null, + rollover_overwrites: [{ id, balance, usage, entities }] | null, + rollover_delete_ids: string[] | null, + new_replaceables: Replaceable[] | null, + deleted_replaceable_ids: string[] | null, + }] + } + + Returns JSON: + { "applied": { "": true }, "skipped": ["id1"] } +]] + +local balance_key = KEYS[1] +local params = cjson.decode(ARGV[1]) +local updates = params.updates or {} +local ttl_seconds = params.ttl_seconds + +if #updates == 0 then + return cjson.encode({ applied = {}, skipped = {} }) +end + +local applied = {} +local skipped = {} + +for _, update in ipairs(updates) do + local cus_ent_id = update.cus_ent_id + + -- Read the SubjectBalance from the hash field + local raw = redis.call('HGET', balance_key, cus_ent_id) + if not raw then + table.insert(skipped, cus_ent_id) + else + local subject_balance = cjson.decode(raw) + + -- Optimistic guard: skip if expected_next_reset_at doesn't match + local should_skip = false + if not is_nil(update.expected_next_reset_at) then + local current_reset_at = safe_number(subject_balance.next_reset_at) + if current_reset_at ~= update.expected_next_reset_at then + should_skip = true + end + end + + if should_skip then + table.insert(skipped, cus_ent_id) + else + -- ================================================================ + -- Apply scalar field updates + -- ================================================================ + if not is_nil(update.balance) then + subject_balance.balance = update.balance + end + + if not is_nil(update.additional_balance) then + subject_balance.additional_balance = update.additional_balance + end + + if not is_nil(update.adjustment) then + subject_balance.adjustment = update.adjustment + end + + if not is_nil(update.entities) then + subject_balance.entities = update.entities + end + + if not is_nil(update.next_reset_at) then + subject_balance.next_reset_at = update.next_reset_at + end + + -- ================================================================ + -- Rollover operations + -- ================================================================ + + -- Ensure rollovers array exists + subject_balance.rollovers = safe_table(subject_balance.rollovers) + + -- APPEND a new rollover + if not is_nil(update.rollover_insert) then + table.insert(subject_balance.rollovers, update.rollover_insert) + end + + -- OVERWRITE existing rollovers by ID + if not is_nil(update.rollover_overwrites) then + local overwrite_map = {} + for _, ow in ipairs(update.rollover_overwrites) do + overwrite_map[ow.id] = ow + end + + for i, rollover in ipairs(subject_balance.rollovers) do + local ow = overwrite_map[rollover.id] + if ow then + subject_balance.rollovers[i].balance = ow.balance + subject_balance.rollovers[i].usage = ow.usage + if not is_nil(ow.entities) then + subject_balance.rollovers[i].entities = ow.entities + end + end + end + end + + -- DELETE rollovers by ID + if not is_nil(update.rollover_delete_ids) then + local delete_set = {} + for _, del_id in ipairs(update.rollover_delete_ids) do + delete_set[del_id] = true + end + + local new_rollovers = {} + for _, rollover in ipairs(subject_balance.rollovers) do + if not delete_set[rollover.id] then + table.insert(new_rollovers, rollover) + end + end + subject_balance.rollovers = new_rollovers + end + + -- ================================================================ + -- Replaceable operations + -- ================================================================ + + -- APPEND new replaceables + if not is_nil(update.new_replaceables) then + subject_balance.replaceables = safe_table(subject_balance.replaceables) + for _, replaceable in ipairs(update.new_replaceables) do + table.insert(subject_balance.replaceables, replaceable) + end + end + + -- DELETE replaceables by ID + if not is_nil(update.deleted_replaceable_ids) then + if not is_nil(subject_balance.replaceables) then + local delete_set = {} + for _, del_id in ipairs(update.deleted_replaceable_ids) do + delete_set[del_id] = true + end + + local new_replaceables = {} + for _, replaceable in ipairs(subject_balance.replaceables) do + if not delete_set[replaceable.id] then + table.insert(new_replaceables, replaceable) + end + end + subject_balance.replaceables = new_replaceables + end + end + + -- Write back the updated SubjectBalance + redis.call('HSET', balance_key, cus_ent_id, cjson.encode(subject_balance)) + applied[cus_ent_id] = true + end + end +end + +-- Refresh TTL on the hash key +if ttl_seconds and ttl_seconds > 0 then + redis.call('EXPIRE', balance_key, ttl_seconds) +end + +return cjson.encode({ applied = applied, skipped = skipped }) diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/contextUtilsV2.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/contextUtilsV2.lua index 1eb4e8357..faece34fb 100644 --- a/server/src/_luaScriptsV2/fullSubjectDeduction/contextUtilsV2.lua +++ b/server/src/_luaScriptsV2/fullSubjectDeduction/contextUtilsV2.lua @@ -1,138 +1,93 @@ -- ============================================================================ -- CONTEXT UTILITIES --- Functions for managing in-memory context during deductions +-- Functions for managing in-memory SubjectBalance context during deductions -- ============================================================================ ---[[ - init_context(params) +local function mark_customer_entitlement_for_update(context, customer_entitlement_id) - Initializes context object with current balances for all customer_entitlements - and builds a rollover index for fast lookups. - Reads from Redis once upfront to avoid multiple reads during passes. + if is_nil(customer_entitlement_id) then + return + end - params: - cache_key: string - customer_entitlement_ids: array of customer_entitlement IDs - full_customer: decoded FullCustomer object (nil when path index is available) - pathidx_key: string (path index Redis Hash key) - has_pathidx: boolean (true when path index exists) + if context.pending_write_ids[customer_entitlement_id] then + return + end + + context.pending_write_ids[customer_entitlement_id] = true + table.insert(context.pending_writes, customer_entitlement_id) +end - Returns: context table with: - customer_entitlements: { [cus_ent_id]: { base_path, balance, adjustment, entities } } - rollovers: { [rollover_id]: { base_path, cus_ent_id, balance, usage, entities } } - mutation_logs: {} (ordered mutation items for receipts and replay) - pending_writes: {} (empty array to queue writes) - logs: {} (debug logs) - logger: { log(fmt, ...): function } (logger that appends to logs) -]] local function init_context(params) local logs = {} - local has_pathidx = params.has_pathidx - local pathidx_key = params.pathidx_key + local read_result = read_subject_balances({ + org_id = params.org_id, + env = params.env, + customer_id = params.customer_id, + customer_entitlement_deductions = params.customer_entitlement_deductions, + }) local context = { customer_entitlements = {}, rollovers = {}, - cache_key = params.cache_key, - full_customer = params.full_customer, - pathidx_key = pathidx_key, - has_pathidx = has_pathidx, + org_id = params.org_id, + env = params.env, + customer_id = params.customer_id, mutation_logs = {}, pending_writes = {}, + pending_write_ids = {}, + missing_customer_entitlement_ids = + read_result.missing_customer_entitlement_ids or {}, logs = logs, logger = { log = function(fmt, ...) table.insert(logs, string.format(fmt, ...)) - end + end, }, } - for _, ent_id in ipairs(params.customer_entitlement_ids or {}) do - local base_path - local has_entity_scope - local is_loose - local adjustment - local unlimited - local cus_ent_rollovers - local cus_ent_balance - local cus_ent_entities + for _, ent_obj in ipairs(params.customer_entitlement_deductions or {}) do + local ent_id = ent_obj.customer_entitlement_id + local balance_entry = read_result.balances_by_id[ent_id] - if has_pathidx then - local result = get_customer_entitlement_via_index({ - pathidx_key = pathidx_key, - cache_key = params.cache_key, - cus_ent_id = ent_id, - }) - if result then - base_path = result.base_path - has_entity_scope = result.has_entity_scope - is_loose = result.is_loose - local sub = result.sub - adjustment = safe_number(sub.adjustment or 0) - unlimited = sub.unlimited - cus_ent_rollovers = sub.rollovers - cus_ent_balance = safe_number(sub.balance or 0) - cus_ent_entities = safe_table(sub.entities) - end - else - -- Fallback path: decode full customer + nested loop search - local cus_ent, cus_product, ce_idx, cp_idx = find_entitlement(params.full_customer, ent_id) - - if cus_ent then - is_loose = (cp_idx == nil) - - if is_loose then - local ece_idx_0 = ce_idx - 1 - base_path = '$.extra_customer_entitlements[' .. ece_idx_0 .. ']' - else - local cp_idx_0 = cp_idx - 1 - local ce_idx_0 = ce_idx - 1 - base_path = '$.customer_products[' .. cp_idx_0 .. '].customer_entitlements[' .. ce_idx_0 .. ']' - end - - local entitlement = cus_ent.entitlement - has_entity_scope = not is_nil(entitlement) and not is_nil(entitlement.entity_feature_id) - adjustment = cus_ent.adjustment or 0 - unlimited = cus_ent.unlimited - cus_ent_rollovers = cus_ent.rollovers - cus_ent_balance = safe_number(cus_ent.balance or 0) - cus_ent_entities = safe_table(cus_ent.entities) - end - end - - if base_path then - local ent_data = { - base_path = base_path, - has_entity_scope = has_entity_scope, - adjustment = adjustment, - unlimited = unlimited, - is_loose = is_loose, - } + if balance_entry then + local subject_balance = balance_entry.subject_balance + local has_entity_scope = not is_nil(ent_obj.entity_feature_id) + local entities = nil if has_entity_scope then - ent_data.balance = 0 - ent_data.entities = cus_ent_entities or {} - else - ent_data.balance = cus_ent_balance or 0 - ent_data.entities = nil + entities = safe_table(subject_balance.entities) + subject_balance.entities = entities end + local ent_data = { + base_path = ent_id, + balance_key = balance_entry.balance_key, + subject_balance = subject_balance, + customer_entitlement_id = ent_id, + feature_id = balance_entry.feature_id, + has_entity_scope = has_entity_scope, + adjustment = safe_number(subject_balance.adjustment), + unlimited = subject_balance.unlimited, + is_loose = is_nil(subject_balance.customer_product_id), + balance = has_entity_scope and 0 or safe_number(subject_balance.balance), + entities = has_entity_scope and entities or nil, + } + context.customer_entitlements[ent_id] = ent_data - if cus_ent_rollovers and type(cus_ent_rollovers) == 'table' then - for r_idx, rollover in ipairs(cus_ent_rollovers) do - if rollover and rollover.id then - local r_idx_0 = r_idx - 1 - local rollover_path = base_path .. '.rollovers[' .. r_idx_0 .. ']' + for _, rollover in ipairs(subject_balance.rollovers or {}) do + if rollover and rollover.id then + local rollover_entities = safe_table(rollover.entities) + rollover.entities = rollover_entities - context.rollovers[rollover.id] = { - base_path = rollover_path, - cus_ent_id = ent_id, - balance = safe_number(rollover.balance or 0), - usage = safe_number(rollover.usage or 0), - entities = safe_table(rollover.entities), - } - end + context.rollovers[rollover.id] = { + base_path = rollover.id, + cus_ent_id = ent_id, + rollover_ref = rollover, + balance = safe_number(rollover.balance), + usage = safe_number(rollover.usage), + entities = rollover_entities, + } end end end @@ -141,11 +96,6 @@ local function init_context(params) return context end ---[[ - append_mutation_log(params) - - Appends one ordered mutation log entry for later receipt persistence and replay. -]] local function append_mutation_log(params) local context = params.context table.insert(context.mutation_logs, { @@ -161,12 +111,6 @@ local function append_mutation_log(params) }) end ---[[ - update_in_memory_customer_entitlement_mutation(params) - - Applies an arbitrary balance/adjustment mutation to an in-memory - customer_entitlement target. -]] local function update_in_memory_customer_entitlement_mutation(params) local target = params.target local entity_id = params.entity_id @@ -185,20 +129,23 @@ local function update_in_memory_customer_entitlement_mutation(params) if not target[entity_id] then target[entity_id] = { balance = 0, adjustment = 0 } end - target[entity_id].balance = (target[entity_id].balance or 0) + balance_delta - target[entity_id].adjustment = (target[entity_id].adjustment or 0) + adjustment_delta + + target[entity_id].balance = + (target[entity_id].balance or 0) + balance_delta + target[entity_id].adjustment = + (target[entity_id].adjustment or 0) + adjustment_delta return end target.balance = (target.balance or 0) + balance_delta target.adjustment = (target.adjustment or 0) + adjustment_delta + + if target.subject_balance then + target.subject_balance.balance = target.balance + target.subject_balance.adjustment = target.adjustment + end end ---[[ - update_in_memory_rollover_mutation(params) - - Applies an arbitrary balance/usage mutation to an in-memory rollover target. -]] local function update_in_memory_rollover_mutation(params) local target = params.target local entity_id = params.entity_id @@ -209,6 +156,7 @@ local function update_in_memory_rollover_mutation(params) if not target[entity_id] then target[entity_id] = { balance = 0, usage = 0 } end + target[entity_id].balance = (target[entity_id].balance or 0) + balance_delta target[entity_id].usage = (target[entity_id].usage or 0) + usage_delta return @@ -216,17 +164,15 @@ local function update_in_memory_rollover_mutation(params) target.balance = (target.balance or 0) + balance_delta target.usage = (target.usage or 0) + usage_delta + + if target.rollover_ref then + target.rollover_ref.balance = target.balance + target.rollover_ref.usage = target.usage + end end ---[[ - queue_customer_entitlement_mutation(params) - - Queues a generic customer_entitlement mutation into pending_writes and - mutation_logs. -]] local function queue_customer_entitlement_mutation(params) local context = params.context - local path = params.path local balance_delta = params.balance_delta local adjustment_delta = params.adjustment_delta @@ -238,13 +184,14 @@ local function queue_customer_entitlement_mutation(params) adjustment_delta = params.alter_granted_balance and balance_delta or 0 end - if balance_delta ~= 0 then - table.insert(context.pending_writes, { path = path .. '.balance', delta = balance_delta }) + if balance_delta == 0 and adjustment_delta == 0 then + return end - if adjustment_delta ~= 0 then - table.insert(context.pending_writes, { path = path .. '.adjustment', delta = adjustment_delta }) - end + mark_customer_entitlement_for_update( + context, + params.customer_entitlement_id + ) append_mutation_log({ context = context, @@ -260,27 +207,23 @@ local function queue_customer_entitlement_mutation(params) }) end ---[[ - queue_rollover_mutation(params) - - Queues a generic rollover mutation into pending_writes and mutation_logs. -]] local function queue_rollover_mutation(params) local context = params.context - local path = params.path local balance_delta = params.balance_delta or 0 local usage_delta = params.usage_delta or 0 local rollover_id = params.rollover_id - if balance_delta ~= 0 then - table.insert(context.pending_writes, { path = path .. '.balance', delta = balance_delta }) - end - - if usage_delta ~= 0 then - table.insert(context.pending_writes, { path = path .. '.usage', delta = usage_delta }) + if balance_delta == 0 and usage_delta == 0 then + return end local rollover_data = context.rollovers[rollover_id] + + mark_customer_entitlement_for_update( + context, + rollover_data and rollover_data.cus_ent_id or nil + ) + append_mutation_log({ context = context, target_type = 'rollover', @@ -295,18 +238,11 @@ local function queue_rollover_mutation(params) }) end ---[[ - queue_rollover_update(params) - - Queues a rollover balance/usage update to pending_writes. - Rollovers track both balance (decrements) and usage (increments). -]] local function queue_rollover_update(params) local deduct_amount = params.deduct_amount queue_rollover_mutation({ context = params.context, - path = params.path, rollover_id = params.rollover_id, entity_id = params.entity_id, credit_cost = params.credit_cost, @@ -316,11 +252,6 @@ local function queue_rollover_update(params) }) end ---[[ - update_in_memory_rollover(params) - - Backwards-compatible wrapper for main deduction paths. -]] local function update_in_memory_rollover(params) update_in_memory_rollover_mutation({ target = params.target, @@ -330,14 +261,35 @@ local function update_in_memory_rollover(params) }) end ---[[ - apply_pending_writes(cache_key, context) +local function apply_pending_writes(_, context) + local writes_by_balance_key = {} - Applies all queued writes to Redis. - Called only after validation passes. -]] -local function apply_pending_writes(cache_key, context) - for _, write in ipairs(context.pending_writes) do - redis.call('JSON.NUMINCRBY', cache_key, write.path, write.delta) + for _, customer_entitlement_id in ipairs(context.pending_writes) do + local ent_data = context.customer_entitlements[customer_entitlement_id] + + if ent_data and ent_data.balance_key and ent_data.subject_balance then + if writes_by_balance_key[ent_data.balance_key] == nil then + writes_by_balance_key[ent_data.balance_key] = {} + end + + table.insert( + writes_by_balance_key[ent_data.balance_key], + customer_entitlement_id + ) + table.insert( + writes_by_balance_key[ent_data.balance_key], + cjson.encode(ent_data.subject_balance) + ) + end + end + + for balance_key, write_args in pairs(writes_by_balance_key) do + if #write_args > 0 then + local redis_args = { 'HSET', balance_key } + for _, arg in ipairs(write_args) do + table.insert(redis_args, arg) + end + redis.call(unpack(redis_args)) + end end end diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromMainBalanceV2.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromMainBalanceV2.lua index 7a2cc0b62..41db0a493 100644 --- a/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromMainBalanceV2.lua +++ b/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromMainBalanceV2.lua @@ -78,7 +78,7 @@ end Uses context object for in-memory operations: - Reads balances from context.customer_entitlements - - Queues writes via queue_balance_update (applied later) + - Queues customer_entitlement updates for later shared-hash writeback - Updates context balances in-memory after calculating change - Logs to context.logs @@ -112,8 +112,6 @@ local function deduct_from_main_balance(params) local amount = params.amount * params.credit_cost local deducted = 0 local prefix = params.log_prefix or "" - local base_path = ent_data.base_path - -- Base calc_params (adjustment is set per-case since entities have their own) local base_calc_params = { available_overage = params.available_overage, @@ -147,11 +145,8 @@ local function deduct_from_main_balance(params) logger.log("%s type: entity, entity_id: %s, balance: %s, to_change: %s", prefix, params.target_entity_id, balance, to_change) if to_change ~= 0 then - local entity_path = build_entity_path(base_path, params.target_entity_id) - queue_customer_entitlement_mutation({ context = context, - path = entity_path, delta = -to_change, alter_granted_balance = params.alter_granted_balance, customer_entitlement_id = ent_id, @@ -199,11 +194,8 @@ local function deduct_from_main_balance(params) local to_change = calculate_change(balance, remaining, calc_params) if to_change ~= 0 then - local entity_path = build_entity_path(base_path, entity_key) - queue_customer_entitlement_mutation({ context = context, - path = entity_path, delta = -to_change, alter_granted_balance = params.alter_granted_balance, customer_entitlement_id = ent_id, @@ -256,7 +248,6 @@ local function deduct_from_main_balance(params) queue_customer_entitlement_mutation({ context = context, - path = base_path, delta = delta, alter_granted_balance = params.alter_granted_balance, customer_entitlement_id = ent_id, diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromRolloversV2.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromRolloversV2.lua index 27dbcdf9c..7ee5f7f47 100644 --- a/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromRolloversV2.lua +++ b/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromRolloversV2.lua @@ -82,8 +82,6 @@ local function deduct_from_rollovers(params) if not rollover_data then logger.log(" Rollover %s not found in context", rollover_id) else - local base_path = rollover_data.base_path - -- Convert remaining (feature units) to credits for this rollover local remaining_credits = remaining * credit_cost @@ -101,11 +99,8 @@ local function deduct_from_rollovers(params) rollover_id, target_entity_id, balance, credit_cost, to_change) if to_change > 0 then - local entity_path = base_path .. '["entities"]["' .. target_entity_id .. '"]' - queue_rollover_update({ context = context, - path = entity_path, deduct_amount = to_change, rollover_id = rollover_id, entity_id = target_entity_id, @@ -147,11 +142,8 @@ local function deduct_from_rollovers(params) rollover_id, entity_key, balance, credit_cost, to_change) if to_change > 0 then - local entity_path = base_path .. '["entities"]["' .. entity_key .. '"]' - queue_rollover_update({ context = context, - path = entity_path, deduct_amount = to_change, rollover_id = rollover_id, entity_id = entity_key, @@ -185,7 +177,6 @@ local function deduct_from_rollovers(params) if to_change > 0 then queue_rollover_update({ context = context, - path = base_path, deduct_amount = to_change, rollover_id = rollover_id, entity_id = nil, diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromSubjectBalances.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromSubjectBalances.lua index 18ad670f4..9eab032e7 100644 --- a/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromSubjectBalances.lua +++ b/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromSubjectBalances.lua @@ -1,8 +1,8 @@ --[[ - Lua Script: Deduct from Customer Entitlements in Redis + Lua Script: Deduct from Subject Balances in Redis - Uses JSON.NUMINCRBY for atomic incremental updates. - Reads CURRENT balance from Redis before each calculation to avoid stale reads. + Reads shared SubjectBalance payloads from per-feature Redis hashes and + writes back only the touched customer_entitlement fields. Deduction Order (mirrors SQL performDeduction.sql): 1. Deduct from rollovers first (oldest first by expires_at) @@ -10,23 +10,21 @@ 3. Pass 2: Allow negative if usage_allowed Helper functions are prepended via string interpolation from: - - fullCustomerKeyBuilders.lua (build_path_index_key, etc.) - luaUtils.lua (safe_table, safe_number, sorted_keys, is_nil) - - fullCustomerUtils.lua (find_entitlement, find_entitlement_from_index, build_entity_path, etc.) - - readBalances.lua (read_current_balance, read_current_entity_balance, read_current_entities, read_rollover_data) - - contextUtils.lua (init_context, queue_balance_update, apply_pending_writes) - - deductFromRollovers.lua (deduct_from_rollovers) - - deductFromMainBalance.lua (calculate_change, deduct_from_main_balance) - - getTotalBalance.lua (get_total_balance) + - readSubjectBalances.lua + - contextUtilsV2.lua + - deductFromRolloversV2.lua + - deductFromMainBalanceV2.lua + - getTotalBalance.lua - KEYS[1] = FullCustomer cache key (used for cluster slot routing) + KEYS[1] = shared balance routing key (used for cluster slot routing) ARGV[1] = JSON params: { org_id: string, env: string, customer_id: string, - sorted_entitlements: [{ customer_entitlement_id, credit_cost, feature_id, entity_feature_id, usage_allowed, min_balance, max_balance }], + customer_entitlement_deductions: [{ customer_entitlement_id, credit_cost, feature_id, entity_feature_id, usage_allowed, min_balance, max_balance }], spend_limit_by_feature_id: { [feature_id]: { feature_id, enabled, overage_limit } } | null, usage_based_cus_ent_ids_by_feature_id: { [feature_id]: string[] } | null, amount_to_deduct: number | null, @@ -53,16 +51,16 @@ -- ============================================================================ -- MAIN SCRIPT -- ============================================================================ -local cache_key = KEYS[1] +local routing_key = KEYS[1] local params = cjson.decode(ARGV[1]) --- Extract org/env/customer for path index key construction local org_id = params.org_id local env = params.env local customer_id = params.customer_id -- Extract parameters -local sorted_entitlements = params.sorted_entitlements or {} +local customer_entitlement_deductions = + params.customer_entitlement_deductions or {} local spend_limit_by_feature_id = params.spend_limit_by_feature_id local usage_based_cus_ent_ids_by_feature_id = params.usage_based_cus_ent_ids_by_feature_id local amount_to_deduct = params.amount_to_deduct @@ -77,56 +75,39 @@ local lock = params.lock local unwind_value = params.unwind_value local lock_receipt_key = params.lock_receipt_key --- Compute overage_behavior_is_allow once -local overage_behavior_is_allow = alter_granted_balance or overage_behaviour == 'allow' - local empty_logs = cjson.decode('[]') --- Check if customer exists -local key_exists = redis.call('EXISTS', cache_key) -if key_exists == 0 then - return cjson.encode({ error = 'CUSTOMER_NOT_FOUND', updates = {}, rollover_updates = {}, mutation_logs = empty_logs, remaining = 0 }) -end - --- Build path index key and check existence (fast path vs fallback) -local pathidx_key = build_path_index_key(org_id, env, customer_id) -local has_pathidx = redis.call('EXISTS', pathidx_key) == 1 - --- Only decode full customer if path index is NOT available (fallback) -local full_customer = nil -if not has_pathidx then - local full_customer_json = redis.call('JSON.GET', cache_key, '.') - if not full_customer_json then - return cjson.encode({ error = 'CUSTOMER_NOT_FOUND', updates = {}, rollover_updates = {}, mutation_logs = empty_logs, remaining = 0 }) - end - - full_customer = cjson.decode(full_customer_json) - - if not full_customer.customer_products then - return cjson.encode({ - error = 'NO_CUSTOMER_PRODUCTS', - updates = {}, - rollover_updates = {}, - mutation_logs = empty_logs, - remaining = 0 - }) - end -end - -- Initialize context with in-memory state from Redis -local customer_entitlement_ids = {} -for _, ent_obj in ipairs(sorted_entitlements) do - table.insert(customer_entitlement_ids, ent_obj.customer_entitlement_id) +if #customer_entitlement_deductions == 0 then + return cjson.encode({ + updates = {}, + rollover_updates = {}, + mutation_logs = empty_logs, + remaining = 0, + error = cjson.null, + logs = empty_logs, + }) end local context = init_context({ - cache_key = cache_key, - customer_entitlement_ids = customer_entitlement_ids, - full_customer = full_customer, - pathidx_key = pathidx_key, - has_pathidx = has_pathidx, + org_id = org_id, + env = env, + customer_id = customer_id, + customer_entitlement_deductions = customer_entitlement_deductions, }) +if #(context.missing_customer_entitlement_ids or {}) > 0 then + return cjson.encode({ + error = 'SUBJECT_BALANCE_NOT_FOUND', + updates = {}, + rollover_updates = {}, + mutation_logs = empty_logs, + remaining = 0, + logs = context.logs, + missing_customer_entitlement_ids = context.missing_customer_entitlement_ids, + }) +end + local unwind_modified_cus_ent_ids = {} if not is_nil(unwind_value) and safe_number(unwind_value) > 0 then @@ -168,7 +149,7 @@ logger.log(" target_entity_id: %s", tostring(target_entity_id or "nil")) logger.log(" overage_behaviour: %s", tostring(overage_behaviour or "nil")) local deduction_result = run_deduction_on_context({ context = context, - sorted_entitlements = sorted_entitlements, + customer_entitlement_deductions = customer_entitlement_deductions, spend_limit_by_feature_id = spend_limit_by_feature_id, usage_based_cus_ent_ids_by_feature_id = usage_based_cus_ent_ids_by_feature_id, rollovers = rollovers, @@ -259,7 +240,7 @@ then end -- Apply all pending writes to Redis (only after validation passes) -apply_pending_writes(cache_key, context) +apply_pending_writes(routing_key, context) logger.log("=== LUA DEDUCTION END ===") diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/readSubjectBalances.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/readSubjectBalances.lua index e2592b88e..c6aabcb71 100644 --- a/server/src/_luaScriptsV2/fullSubjectDeduction/readSubjectBalances.lua +++ b/server/src/_luaScriptsV2/fullSubjectDeduction/readSubjectBalances.lua @@ -1,101 +1,139 @@ -- ============================================================================ --- READ BALANCES --- Functions to read current balance state from Redis +-- READ SUBJECT BALANCES +-- Functions to read shared SubjectBalance payloads from Redis hashes -- ============================================================================ --- ============================================================================ --- HELPER: Read rollover data from Redis (fresh read) --- Returns: { balance, usage, entities } or nil if not found --- ============================================================================ -local function read_rollover_data(cache_key, rollover_path) - local result = redis.call('JSON.GET', cache_key, rollover_path) - if not result or result == cjson.null then +local function build_shared_subject_balance_key(params) + return '{' + .. params.customer_id + .. '}:' + .. params.org_id + .. ':' + .. params.env + .. ':full_subject:shared_balances:' + .. params.feature_id +end + +local function decode_subject_balance(raw_value) + if is_nil(raw_value) then return nil end - - local decoded = cjson.decode(result) - -- JSONPath returns an array of matches, extract the first element - if type(decoded) == 'table' and decoded[1] ~= nil then - decoded = decoded[1] - end - + + local decoded = cjson.decode(raw_value) if type(decoded) ~= 'table' then return nil end - - return { - balance = safe_number(decoded.balance), - usage = safe_number(decoded.usage), - entities = safe_table(decoded.entities), - } -end --- ============================================================================ --- HELPER: Read current balance from Redis (fresh read, not from snapshot) --- ============================================================================ -local function read_current_balance(cache_key, base_path) - local result = redis.call('JSON.GET', cache_key, base_path .. '.balance') - if not result or result == cjson.null then - return 0 - end - -- JSON.GET returns a JSON string, need to decode - local decoded = cjson.decode(result) - -- JSONPath returns an array of matches, extract the first element - if type(decoded) == 'table' and decoded[1] ~= nil then - decoded = decoded[1] - end - return safe_number(decoded) -end - --- ============================================================================ --- HELPER: Read current entity balance from Redis (fresh read) --- ============================================================================ -local function read_current_entity_balance(cache_key, base_path, entity_id) - -- Use bracket notation for entity access (handles special characters in entity_id) - local entity_path = build_entity_path(base_path, entity_id) - local result = redis.call('JSON.GET', cache_key, entity_path) - - if not result or result == cjson.null then - return nil -- Entity doesn't exist - end - - local decoded = cjson.decode(result) - - -- JSONPath returns an array of matches, extract the first element - if type(decoded) == 'table' and decoded[1] ~= nil then - decoded = decoded[1] - end - - if type(decoded) ~= 'table' then - return nil - end - - return { - balance = safe_number(decoded.balance), - adjustment = safe_number(decoded.adjustment) - } -end - --- ============================================================================ --- HELPER: Read all entity balances from Redis (fresh read) --- ============================================================================ -local function read_current_entities(cache_key, base_path) - local entities_path = base_path .. '.entities' - local result = redis.call('JSON.GET', cache_key, entities_path) - if not result or result == cjson.null then - return {} - end - local decoded = cjson.decode(result) - -- JSONPath returns an array of matches, extract the first element - if type(decoded) == 'table' and decoded[1] ~= nil and type(decoded[1]) == 'table' then - -- Check if this looks like a JSONPath array wrapper (first element is an object with entity keys) - -- vs an actual entity object (first element would be a number or have 'balance' field directly) - if decoded.balance == nil and decoded.adjustment == nil then - decoded = decoded[1] - end - end - if type(decoded) ~= 'table' then - return {} - end return decoded end + +local function read_subject_balances(params) + local balances_by_id = {} + local missing_customer_entitlement_ids = {} + local entries_by_balance_key = {} + + for _, ent_obj in ipairs(params.customer_entitlement_deductions or {}) do + local customer_entitlement_id = ent_obj.customer_entitlement_id + local feature_id = ent_obj.feature_id + + if customer_entitlement_id and feature_id then + local balance_key = build_shared_subject_balance_key({ + org_id = params.org_id, + env = params.env, + customer_id = params.customer_id, + feature_id = feature_id, + }) + + if entries_by_balance_key[balance_key] == nil then + entries_by_balance_key[balance_key] = { + feature_id = feature_id, + customer_entitlement_ids = {}, + } + end + + table.insert( + entries_by_balance_key[balance_key].customer_entitlement_ids, + customer_entitlement_id + ) + end + end + + for balance_key, entry in pairs(entries_by_balance_key) do + local hmget_args = { 'HMGET', balance_key } + for _, customer_entitlement_id in ipairs(entry.customer_entitlement_ids) do + table.insert(hmget_args, customer_entitlement_id) + end + + local raw_values = redis.call(unpack(hmget_args)) + + for index, customer_entitlement_id in ipairs(entry.customer_entitlement_ids) do + local raw_value = raw_values[index] + local subject_balance = decode_subject_balance(raw_value) + + if subject_balance == nil then + table.insert( + missing_customer_entitlement_ids, + customer_entitlement_id + ) + else + balances_by_id[customer_entitlement_id] = { + balance_key = balance_key, + customer_entitlement_id = customer_entitlement_id, + feature_id = entry.feature_id, + subject_balance = subject_balance, + } + end + end + end + + return { + balances_by_id = balances_by_id, + missing_customer_entitlement_ids = missing_customer_entitlement_ids, + } +end + +local function read_rollover_data(rollover) + if type(rollover) ~= 'table' then + return nil + end + + return { + balance = safe_number(rollover.balance), + usage = safe_number(rollover.usage), + entities = safe_table(rollover.entities), + } +end + +local function read_current_balance(subject_balance) + if type(subject_balance) ~= 'table' then + return 0 + end + + return safe_number(subject_balance.balance) +end + +local function read_current_entity_balance(subject_balance, entity_id) + if type(subject_balance) ~= 'table' or is_nil(entity_id) then + return nil + end + + local entities = safe_table(subject_balance.entities) + local entity_balance = entities[entity_id] + + if type(entity_balance) ~= 'table' then + return nil + end + + return { + balance = safe_number(entity_balance.balance), + adjustment = safe_number(entity_balance.adjustment), + } +end + +local function read_current_entities(subject_balance) + if type(subject_balance) ~= 'table' then + return {} + end + + return safe_table(subject_balance.entities) +end diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/runDeductionOnContextV2.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/runDeductionOnContextV2.lua index 4cf3bc61d..e092ee3a3 100644 --- a/server/src/_luaScriptsV2/fullSubjectDeduction/runDeductionOnContextV2.lua +++ b/server/src/_luaScriptsV2/fullSubjectDeduction/runDeductionOnContextV2.lua @@ -16,7 +16,7 @@ end --[[ process_deduction_pass(params) - Runs one main-balance deduction pass over all sorted customer_entitlements. + Runs one main-balance deduction pass over all customer_entitlement deductions. Returns: { @@ -26,7 +26,8 @@ end ]] local function process_deduction_pass(params) local context = params.context - local sorted_entitlements = params.sorted_entitlements or {} + local customer_entitlement_deductions = + params.customer_entitlement_deductions or {} local target_entity_id = params.target_entity_id local spend_limit_by_feature_id = params.spend_limit_by_feature_id local usage_based_cus_ent_ids_by_feature_id = params.usage_based_cus_ent_ids_by_feature_id @@ -41,7 +42,7 @@ local function process_deduction_pass(params) logger.log("=== %s START ===", pass_name) - for _, ent_obj in ipairs(sorted_entitlements) do + for _, ent_obj in ipairs(customer_entitlement_deductions) do if remaining_amount == 0 then break end @@ -134,7 +135,8 @@ end ]] local function process_rollover_deduction(params) local context = params.context - local sorted_entitlements = params.sorted_entitlements or {} + local customer_entitlement_deductions = + params.customer_entitlement_deductions or {} local rollovers = params.rollovers local target_entity_id = params.target_entity_id local remaining_amount = params.remaining_amount or 0 @@ -144,7 +146,7 @@ local function process_rollover_deduction(params) return 0 end - local first_ent = sorted_entitlements[1] + local first_ent = customer_entitlement_deductions[1] local has_entity_scope = false if first_ent then has_entity_scope = first_ent.entity_feature_id ~= nil and first_ent.entity_feature_id ~= cjson.null @@ -171,7 +173,7 @@ end params: context: initialized context - sorted_entitlements: deduction inputs + customer_entitlement_deductions: deduction inputs rollovers: rollover inputs | nil amount_to_deduct: number | nil target_balance: number | nil @@ -188,7 +190,8 @@ end ]] local function run_deduction_on_context(params) local context = params.context - local sorted_entitlements = params.sorted_entitlements or {} + local customer_entitlement_deductions = + params.customer_entitlement_deductions or {} local rollovers = params.rollovers local target_entity_id = params.target_entity_id local spend_limit_by_feature_id = params.spend_limit_by_feature_id @@ -202,7 +205,7 @@ local function run_deduction_on_context(params) if not is_nil(params.target_balance) then local current_total = get_total_balance({ context = context, - sorted_entitlements = sorted_entitlements, + sorted_entitlements = customer_entitlement_deductions, target_entity_id = target_entity_id, }) remaining_amount = current_total - params.target_balance @@ -215,7 +218,7 @@ local function run_deduction_on_context(params) if not alter_granted_balance then local rollover_deducted = process_rollover_deduction({ context = context, - sorted_entitlements = sorted_entitlements, + customer_entitlement_deductions = customer_entitlement_deductions, rollovers = rollovers, target_entity_id = target_entity_id, remaining_amount = remaining_amount, @@ -225,7 +228,7 @@ local function run_deduction_on_context(params) local pass_one_result = process_deduction_pass({ context = context, - sorted_entitlements = sorted_entitlements, + customer_entitlement_deductions = customer_entitlement_deductions, target_entity_id = target_entity_id, spend_limit_by_feature_id = spend_limit_by_feature_id, usage_based_cus_ent_ids_by_feature_id = usage_based_cus_ent_ids_by_feature_id, @@ -242,7 +245,7 @@ local function run_deduction_on_context(params) if remaining_amount ~= 0 then local pass_two_result = process_deduction_pass({ context = context, - sorted_entitlements = sorted_entitlements, + customer_entitlement_deductions = customer_entitlement_deductions, target_entity_id = target_entity_id, spend_limit_by_feature_id = spend_limit_by_feature_id, usage_based_cus_ent_ids_by_feature_id = usage_based_cus_ent_ids_by_feature_id, diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/spendLimitUtils.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/spendLimitUtils.lua deleted file mode 100644 index e43aec44c..000000000 --- a/server/src/_luaScriptsV2/fullSubjectDeduction/spendLimitUtils.lua +++ /dev/null @@ -1,106 +0,0 @@ --- ============================================================================ --- SPEND LIMIT UTILITIES --- Computes atomic available overage from current in-memory/live entitlement data --- ============================================================================ - -local function build_ent_data_from_full_customer(params) - local full_customer = params.full_customer - local cus_ent_id = params.cus_ent_id - local cache_key = params.cache_key - local pathidx_key = params.pathidx_key - - if is_nil(cus_ent_id) then - return nil - end - - -- Fast path: single JSON.GET on the sub-object via path index - if params.has_pathidx then - local result = get_customer_entitlement_via_index({ - pathidx_key = pathidx_key, - cache_key = cache_key, - cus_ent_id = cus_ent_id, - }) - if is_nil(result) then return nil end - - if result.has_entity_scope then - return { - has_entity_scope = true, - balance = 0, - entities = safe_table(result.sub.entities), - } - else - return { - has_entity_scope = false, - balance = safe_number(result.sub.balance or 0), - entities = {}, - } - end - end - - -- Fallback path: decode from full customer - if is_nil(full_customer) then - return nil - end - - local cus_ent = find_entitlement(full_customer, cus_ent_id) - if is_nil(cus_ent) then - return nil - end - - local entitlement = cus_ent.entitlement - local has_entity_scope = not is_nil(entitlement) - and not is_nil(entitlement.entity_feature_id) - - return { - has_entity_scope = has_entity_scope, - balance = safe_number(cus_ent.balance), - entities = cus_ent.entities or {}, - } -end - -local function get_available_overage_from_spend_limit(params) - local context = params.context - local spend_limit = params.spend_limit - local usage_based_cus_ent_ids = params.usage_based_cus_ent_ids or {} - local target_entity_id = params.target_entity_id - - if is_nil(spend_limit) or is_nil(spend_limit.overage_limit) then - return nil - end - - local total_overage = 0 - - for _, cus_ent_id in ipairs(usage_based_cus_ent_ids) do - local ent_data = context.customer_entitlements[cus_ent_id] - - if is_nil(ent_data) then - ent_data = build_ent_data_from_full_customer({ - full_customer = context.full_customer, - cus_ent_id = cus_ent_id, - cache_key = context.cache_key, - pathidx_key = context.pathidx_key, - has_pathidx = context.has_pathidx, - }) - end - - if ent_data then - if ent_data.has_entity_scope then - if not is_nil(target_entity_id) then - local entity_data = ent_data.entities and ent_data.entities[target_entity_id] - local balance = entity_data and safe_number(entity_data.balance) or 0 - total_overage = total_overage + math.max(-balance, 0) - else - for _, entity_data in pairs(ent_data.entities or {}) do - local balance = safe_number(entity_data.balance) - total_overage = total_overage + math.max(-balance, 0) - end - end - else - local balance = safe_number(ent_data.balance) - total_overage = total_overage + math.max(-balance, 0) - end - end - end - - return math.max(0, safe_number(spend_limit.overage_limit) - total_overage) -end diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/spendLimitUtilsV2.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/spendLimitUtilsV2.lua new file mode 100644 index 000000000..0cfd7da61 --- /dev/null +++ b/server/src/_luaScriptsV2/fullSubjectDeduction/spendLimitUtilsV2.lua @@ -0,0 +1,41 @@ +-- ============================================================================ +-- SPEND LIMIT UTILITIES +-- Computes available overage from the initialized in-memory deduction context +-- ============================================================================ + +local function get_available_overage_from_spend_limit(params) + local context = params.context + local spend_limit = params.spend_limit + local usage_based_cus_ent_ids = params.usage_based_cus_ent_ids or {} + local target_entity_id = params.target_entity_id + + if is_nil(spend_limit) or is_nil(spend_limit.overage_limit) then + return nil + end + + local total_overage = 0 + + for _, cus_ent_id in ipairs(usage_based_cus_ent_ids) do + local ent_data = context.customer_entitlements[cus_ent_id] + + if ent_data then + if ent_data.has_entity_scope then + if not is_nil(target_entity_id) then + local entity_data = ent_data.entities and ent_data.entities[target_entity_id] + local balance = entity_data and safe_number(entity_data.balance) or 0 + total_overage = total_overage + math.max(-balance, 0) + else + for _, entity_data in pairs(ent_data.entities or {}) do + local balance = safe_number(entity_data.balance) + total_overage = total_overage + math.max(-balance, 0) + end + end + else + local balance = safe_number(ent_data.balance) + total_overage = total_overage + math.max(-balance, 0) + end + end + end + + return math.max(0, safe_number(spend_limit.overage_limit) - total_overage) +end diff --git a/server/src/_luaScriptsV2/luaScriptsV2.ts b/server/src/_luaScriptsV2/luaScriptsV2.ts index 869b67292..b7bf889f9 100644 --- a/server/src/_luaScriptsV2/luaScriptsV2.ts +++ b/server/src/_luaScriptsV2/luaScriptsV2.ts @@ -14,6 +14,7 @@ const DELETE_CACHE_DIR = join(__dirname, "deleteFullCustomerCache"); const RESET_DIR = join(__dirname, "resetCustomerEntitlements"); const UPDATE_DIR = join(__dirname, "updateCustomerEntitlements"); const FULL_SUBJECT_DIR = join(__dirname, "fullSubject"); +const FULL_SUBJECT_DEDUCTION_DIR = join(__dirname, "fullSubjectDeduction"); // ============================================================================ // HELPER MODULES @@ -21,7 +22,7 @@ const FULL_SUBJECT_DIR = join(__dirname, "fullSubject"); const FULL_CUSTOMER_DIR = join(__dirname, "fullCustomer"); -const LUA_UTILS = readFileSync(join(DEDUCT_DIR, "luaUtils.lua"), "utf-8"); +const LUA_UTILS = readFileSync(join(__dirname, "luaUtils.lua"), "utf-8"); const FULL_CUSTOMER_UTILS = readFileSync( join(FULL_CUSTOMER_DIR, "fullCustomerUtils.lua"), @@ -191,6 +192,14 @@ const updateCustomerDataV2Script = readFileSync( /** Atomically update top-level customer fields in the cached FullSubject. */ export const UPDATE_CUSTOMER_DATA_V2_SCRIPT = updateCustomerDataV2Script; +const updateEntityDataV2Script = readFileSync( + join(FULL_SUBJECT_DIR, "updateEntityDataV2.lua"), + "utf-8", +); + +/** Atomically update top-level entity fields in the cached FullSubject. */ +export const UPDATE_ENTITY_DATA_V2_SCRIPT = updateEntityDataV2Script; + // ============================================================================ // RESET CUSTOMER ENTITLEMENTS SCRIPT (deprecated — kept for backward compat) // ============================================================================ @@ -300,3 +309,79 @@ export const UPDATE_CUSTOMER_PRODUCT_SCRIPT = readFileSync( join(CUS_PRODUCT_DIR, "updateCustomerProduct.lua"), "utf-8", ); + +// ============================================================================ +// FULL SUBJECT DEDUCTION SCRIPT (V2 cache — per-feature hash balances) +// ============================================================================ + +const READ_SUBJECT_BALANCES = readFileSync( + join(FULL_SUBJECT_DEDUCTION_DIR, "readSubjectBalances.lua"), + "utf-8", +); + +const CONTEXT_UTILS_V2 = readFileSync( + join(FULL_SUBJECT_DEDUCTION_DIR, "contextUtilsV2.lua"), + "utf-8", +); + +const DEDUCT_FROM_ROLLOVERS_V2 = readFileSync( + join(FULL_SUBJECT_DEDUCTION_DIR, "deductFromRolloversV2.lua"), + "utf-8", +); + +const DEDUCT_FROM_MAIN_BALANCE_V2 = readFileSync( + join(FULL_SUBJECT_DEDUCTION_DIR, "deductFromMainBalanceV2.lua"), + "utf-8", +); + +const RUN_DEDUCTION_ON_CONTEXT_V2 = readFileSync( + join(FULL_SUBJECT_DEDUCTION_DIR, "runDeductionOnContextV2.lua"), + "utf-8", +); + +const SPEND_LIMIT_UTILS_V2 = readFileSync( + join(FULL_SUBJECT_DEDUCTION_DIR, "spendLimitUtilsV2.lua"), + "utf-8", +); + +const DEDUCT_FROM_SUBJECT_BALANCES_MAIN = readFileSync( + join(FULL_SUBJECT_DEDUCTION_DIR, "deductFromSubjectBalances.lua"), + "utf-8", +); + +/** + * Lua script for deducting from subject balances in Redis (V2 cache). + * Reads from per-feature hash fields and writes back touched entitlements. + * Composed from shared helper modules + V2-specific storage adapters. + */ +export const DEDUCT_FROM_SUBJECT_BALANCES_SCRIPT = `${LUA_UTILS} +${READ_SUBJECT_BALANCES} +${CONTEXT_UTILS_V2} +${GET_TOTAL_BALANCE} +${DEDUCT_FROM_ROLLOVERS_V2} +${DEDUCT_FROM_MAIN_BALANCE_V2} +${SPEND_LIMIT_UTILS_V2} +${RUN_DEDUCTION_ON_CONTEXT_V2} +${MUTATION_ITEM_UTILS} +${LOCK_RECEIPT_UTILS} +${LOCK_STATE_UTILS} +${LOCK_UNWIND_UTILS} +${DEDUCT_FROM_SUBJECT_BALANCES_MAIN}`; + +// ============================================================================ +// UPDATE SUBJECT BALANCES SCRIPT (V2 cache — per-feature hash updates) +// ============================================================================ + +const UPDATE_SUBJECT_BALANCES_MAIN = readFileSync( + join(FULL_SUBJECT_DIR, "updateSubjectBalances.lua"), + "utf-8", +); + +/** + * Lua script for atomically updating SubjectBalance entries in a single + * per-feature balance hash. Supports scalar updates, rollover ops, + * replaceable ops, and expected_next_reset_at guard. + * Called once per feature via pipeline. + */ +export const UPDATE_SUBJECT_BALANCES_SCRIPT = `${LUA_UTILS} +${UPDATE_SUBJECT_BALANCES_MAIN}`; diff --git a/server/src/_luaScriptsV2/deductFromCustomerEntitlements/luaUtils.lua b/server/src/_luaScriptsV2/luaUtils.lua similarity index 100% rename from server/src/_luaScriptsV2/deductFromCustomerEntitlements/luaUtils.lua rename to server/src/_luaScriptsV2/luaUtils.lua diff --git a/server/src/external/redis/initUtils/redisTypes.ts b/server/src/external/redis/initUtils/redisTypes.ts index 742f705a5..1e561c614 100644 --- a/server/src/external/redis/initUtils/redisTypes.ts +++ b/server/src/external/redis/initUtils/redisTypes.ts @@ -84,6 +84,14 @@ declare module "ioredis" { cacheKey: string, paramsJson: string, ): Promise; + deductFromSubjectBalances( + routingKey: string, + paramsJson: string, + ): Promise; + updateSubjectBalances( + balanceKey: string, + paramsJson: string, + ): Promise; deleteFullCustomerCache( cacheKey: string, orgId: string, @@ -136,6 +144,12 @@ declare module "ioredis" { cacheTtlSeconds: string, nowMs: string, ): Promise; + updateFullSubjectEntityDataV2( + subjectKey: string, + updatesJson: string, + cacheTtlSeconds: string, + nowMs: string, + ): Promise; appendEntityToCustomer( cacheKey: string, entityJson: string, diff --git a/server/src/external/redis/initUtils/registerRedisCommands.ts b/server/src/external/redis/initUtils/registerRedisCommands.ts index 602d0c642..308716fee 100644 --- a/server/src/external/redis/initUtils/registerRedisCommands.ts +++ b/server/src/external/redis/initUtils/registerRedisCommands.ts @@ -18,6 +18,7 @@ import { APPEND_ENTITY_TO_CUSTOMER_SCRIPT, CLAIM_LOCK_RECEIPT_SCRIPT, DEDUCT_FROM_CUSTOMER_ENTITLEMENTS_SCRIPT, + DEDUCT_FROM_SUBJECT_BALANCES_SCRIPT, DELETE_FULL_CUSTOMER_CACHE_SCRIPT, RELEASE_FULL_SUBJECT_RESERVATION_SCRIPT, RESERVE_FULL_SUBJECT_WRITE_SCRIPT, @@ -25,9 +26,11 @@ import { SET_FULL_CUSTOMER_CACHE_SCRIPT, UPDATE_CUSTOMER_DATA_SCRIPT, UPDATE_CUSTOMER_DATA_V2_SCRIPT, + UPDATE_ENTITY_DATA_V2_SCRIPT, UPDATE_CUSTOMER_ENTITLEMENTS_SCRIPT, UPDATE_CUSTOMER_PRODUCT_SCRIPT, UPDATE_ENTITY_IN_CUSTOMER_SCRIPT, + UPDATE_SUBJECT_BALANCES_SCRIPT, UPSERT_INVOICE_IN_CUSTOMER_SCRIPT, } from "../../../_luaScriptsV2/luaScriptsV2.js"; @@ -104,6 +107,16 @@ export const registerRedisCommands = ({ lua: DEDUCT_FROM_CUSTOMER_ENTITLEMENTS_SCRIPT, }); + redisInstance.defineCommand("deductFromSubjectBalances", { + numberOfKeys: 1, + lua: DEDUCT_FROM_SUBJECT_BALANCES_SCRIPT, + }); + + redisInstance.defineCommand("updateSubjectBalances", { + numberOfKeys: 1, + lua: UPDATE_SUBJECT_BALANCES_SCRIPT, + }); + redisInstance.defineCommand("deleteFullCustomerCache", { numberOfKeys: 1, lua: DELETE_FULL_CUSTOMER_CACHE_SCRIPT, @@ -144,6 +157,11 @@ export const registerRedisCommands = ({ lua: UPDATE_CUSTOMER_DATA_V2_SCRIPT, }); + redisInstance.defineCommand("updateFullSubjectEntityDataV2", { + numberOfKeys: 1, + lua: UPDATE_ENTITY_DATA_V2_SCRIPT, + }); + redisInstance.defineCommand("appendEntityToCustomer", { numberOfKeys: 1, lua: APPEND_ENTITY_TO_CUSTOMER_SCRIPT, diff --git a/server/src/honoMiddlewares/refreshCacheConfigs.ts b/server/src/honoMiddlewares/refreshCacheConfigs.ts index 0cbca4c54..aa3c382f9 100644 --- a/server/src/honoMiddlewares/refreshCacheConfigs.ts +++ b/server/src/honoMiddlewares/refreshCacheConfigs.ts @@ -20,6 +20,16 @@ export const REFRESH_CACHE_ROUTE_CONFIGS: RefreshCacheRouteConfig[] = [ url: "/customers/:customer_id", }), + route({ + method: "POST", + url: "/customers/:customer_id", + }), + + route({ + method: "PATCH", + url: "/customers/:customer_id", + }), + route({ method: "POST", url: "/customers/:customer_id/balances", @@ -111,6 +121,16 @@ export const REFRESH_CACHE_ROUTE_CONFIGS: RefreshCacheRouteConfig[] = [ url: "/entities.delete", }), + route({ + method: "POST", + url: "/entities.update", + }), + + route({ + method: "POST", + url: "/customers.update", + }), + route({ method: "POST", url: "/customers.delete", diff --git a/server/src/internal/balances/balancesRouter.ts b/server/src/internal/balances/balancesRouter.ts index 424a80343..56f096be8 100644 --- a/server/src/internal/balances/balancesRouter.ts +++ b/server/src/internal/balances/balancesRouter.ts @@ -5,9 +5,9 @@ import { handleCreateBalance } from "./handlers/handleCreateBalance.js"; import { handleDeleteBalance } from "./handlers/handleDeleteBalance.js"; import { handleFinalizeLock } from "./handlers/handleFinalizeLock.js"; import { handleListBalances } from "./handlers/handleListBalances.js"; +import { handleSetUsage } from "./handlers/handleSetUsage.js"; import { handleTrack } from "./handlers/handleTrack.js"; import { handleUpdateBalance } from "./handlers/handleUpdateBalance.js"; -import { handleSetUsage } from "./setUsage/handleSetUsage.js"; // Create a Hono app for products export const balancesRouter = new Hono(); @@ -26,6 +26,7 @@ balancesRouter.post("/check", ...handleCheck); // Legacy balancesRouter.post("/usage", ...handleSetUsage); + export const balancesRpcRouter = new Hono(); balancesRpcRouter.post("/balances.create", ...handleCreateBalance); balancesRpcRouter.post("/balances.update", ...handleUpdateBalance); diff --git a/server/src/internal/balances/check/getCheckDataV2.ts b/server/src/internal/balances/check/getCheckDataV2.ts index bba8993db..f8f131bfb 100644 --- a/server/src/internal/balances/check/getCheckDataV2.ts +++ b/server/src/internal/balances/check/getCheckDataV2.ts @@ -73,6 +73,8 @@ export const getCheckDataV2 = async ({ source: "getCheckDataV2", }); + console.log("Full subject", fullSubject); + const apiSubject = await getApiSubject({ ctx, fullSubject, diff --git a/server/src/internal/balances/check/runCheckV2.ts b/server/src/internal/balances/check/runCheckV2.ts index f24846926..e38dff937 100644 --- a/server/src/internal/balances/check/runCheckV2.ts +++ b/server/src/internal/balances/check/runCheckV2.ts @@ -3,6 +3,7 @@ import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import type { CheckDataV2 } from "./checkTypes/CheckDataV2.js"; import { getCheckDataV2 } from "./getCheckDataV2.js"; import { getCheckResponseV2 } from "./getCheckResponseV2.js"; +import { runCheckWithTrackV2 } from "./runCheckWithTrackV2.js"; export const runCheckV2 = async ({ ctx, @@ -22,10 +23,18 @@ export const runCheckV2 = async ({ requiredBalance, }); - const response = await getCheckResponseV2({ - checkData, - requiredBalance, - }); + const response = + body.send_event || body.lock?.enabled + ? await runCheckWithTrackV2({ + ctx, + body, + requiredBalance, + checkData, + }) + : await getCheckResponseV2({ + checkData, + requiredBalance, + }); return { checkData, diff --git a/server/src/internal/balances/check/runCheckWithRollout.ts b/server/src/internal/balances/check/runCheckWithRollout.ts index bdc5b2b77..d6eca0482 100644 --- a/server/src/internal/balances/check/runCheckWithRollout.ts +++ b/server/src/internal/balances/check/runCheckWithRollout.ts @@ -17,20 +17,8 @@ export const runCheckWithRollout = async ({ checkData: CheckData; response: CheckResponseV3; }> => { - if ( - isFullSubjectRolloutEnabled({ ctx }) && - !body.send_event && - !body.lock?.enabled - ) { - return runCheckV2({ - ctx, - body, - requiredBalance, - }); - } - if (isFullSubjectRolloutEnabled({ ctx })) { - return runCheckLegacyFlow({ + return runCheckV2({ ctx, body, requiredBalance, diff --git a/server/src/internal/balances/check/runCheckWithTrackV2.ts b/server/src/internal/balances/check/runCheckWithTrackV2.ts new file mode 100644 index 000000000..caed4bfc3 --- /dev/null +++ b/server/src/internal/balances/check/runCheckWithTrackV2.ts @@ -0,0 +1,143 @@ +import { + CheckResponseV3Schema, + ErrCode, + FeatureType, + featureUtils, + InsufficientBalanceError, + InternalError, + type ParsedCheckParams, + RecaseError, + type TrackParams, +} from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { getTrackFeatureDeductions } from "@/internal/balances/track/utils/getFeatureDeductions.js"; +import { runTrackV3 } from "@/internal/balances/track/v3/runTrackV3.js"; +import { buildLockScheduleName } from "@/internal/balances/utils/lock/buildLockScheduleName.js"; +import { featureToCreditSystem } from "@/internal/features/creditSystemUtils.js"; +import { workflows } from "@/queue/workflows.js"; +import type { CheckDataV2 } from "./checkTypes/CheckDataV2.js"; + +export const runCheckWithTrackV2 = async ({ + ctx, + body, + requiredBalance, + checkData, +}: { + ctx: AutumnContext; + body: ParsedCheckParams; + requiredBalance: number; + checkData: CheckDataV2; +}) => { + if (!body.feature_id) { + throw new InternalError({ + message: "ran check with track but no feature ID", + }); + } + + if (ctx.isPublic) { + throw new RecaseError({ + message: + "Can't pass in 'send_event: true' when using publishable key for Autumn", + }); + } + + if (body.lock && featureUtils.isAllocated(checkData.featureToUse)) { + throw new RecaseError({ + message: "Lock is not supported for allocated features", + code: ErrCode.InvalidRequest, + statusCode: 400, + }); + } + + if (checkData.originalFeature.type === FeatureType.Boolean) { + throw new RecaseError({ + message: "Not allowed to pass in send_event: true for a boolean feature", + code: ErrCode.InvalidRequest, + statusCode: 400, + }); + } + + const featureDeductions = getTrackFeatureDeductions({ + ctx, + featureId: body.feature_id, + lock: body.lock, + value: requiredBalance, + }); + + const trackBody: TrackParams = { + customer_id: body.customer_id, + entity_id: body.entity_id, + feature_id: body.feature_id, + value: requiredBalance, + properties: body.properties, + skip_event: body.skip_event, + overage_behavior: "reject", + lock: body.lock, + }; + + let allowed = true; + + try { + const response = await runTrackV3({ + ctx, + body: trackBody, + featureDeductions, + }); + + checkData.apiBalance = response.balance ?? undefined; + checkData.evaluationApiBalance = response.balance ?? undefined; + } catch (error) { + if (error instanceof InsufficientBalanceError) { + allowed = false; + } else { + throw error; + } + } + + const { featureToUse, originalFeature } = checkData; + if ( + featureToUse.type === FeatureType.CreditSystem && + featureToUse.id !== originalFeature.id + ) { + requiredBalance = featureToCreditSystem({ + featureId: originalFeature.id, + creditSystem: featureToUse, + amount: requiredBalance, + }); + } + + if (body.lock?.expires_at && allowed) { + try { + const scheduleName = buildLockScheduleName({ + orgId: ctx.org.id, + env: ctx.env, + hashedKey: body.lock.hashed_key, + }); + + await workflows.triggerExpireLockReceipt( + { + orgId: ctx.org.id, + env: ctx.env, + customerId: body.customer_id, + lockId: body.lock.lock_id, + hashedKey: body.lock.hashed_key, + }, + { + scheduleAt: new Date(body.lock.expires_at), + scheduleName, + }, + ); + } catch (error) { + ctx.logger.error(`Failed to schedule lock expiration: ${error}`); + } + } + + return CheckResponseV3Schema.parse({ + allowed, + customer_id: checkData.customerId || "", + entity_id: checkData.entityId, + required_balance: requiredBalance, + balance: checkData.apiBalance ?? null, + flag: checkData.apiFlag ?? null, + }); +}; diff --git a/server/src/internal/balances/finalizeLock/buildFinalizeLockContext.ts b/server/src/internal/balances/finalizeLock/buildFinalizeLockContext.ts index 60651dc35..2650db1b8 100644 --- a/server/src/internal/balances/finalizeLock/buildFinalizeLockContext.ts +++ b/server/src/internal/balances/finalizeLock/buildFinalizeLockContext.ts @@ -1,5 +1,11 @@ import type { Feature, FullCustomer } from "@autumn/shared"; import { type FinalizeLockParamsV0, findFeatureById } from "@autumn/shared"; +import type { Redis } from "ioredis"; +import { + currentRegion, + getRegionalRedis, + redis, +} from "@/external/redis/initRedis.js"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { fetchLockReceipt, @@ -15,6 +21,7 @@ import { getOrSetCachedFullCustomer } from "@/internal/customers/cusUtils/fullCu export type FinalizeLockContext = { receipt: LockReceipt; lockReceiptKey: string; + redisInstance: Redis; fullCustomer: FullCustomer; feature: Feature; lockValue: number; @@ -60,9 +67,15 @@ export const buildFinalizeLockContext = async ({ errorOnNotFound: true, }); + const redisInstance = + receipt.region && receipt.region !== currentRegion + ? getRegionalRedis(receipt.region) + : redis; + return { receipt, lockReceiptKey, + redisInstance, fullCustomer, feature, lockValue, diff --git a/server/src/internal/balances/finalizeLock/insertFinalizeLockEventV2.ts b/server/src/internal/balances/finalizeLock/insertFinalizeLockEventV2.ts new file mode 100644 index 000000000..a8abc63ce --- /dev/null +++ b/server/src/internal/balances/finalizeLock/insertFinalizeLockEventV2.ts @@ -0,0 +1,31 @@ +import { Decimal } from "decimal.js"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { globalEventBatchingManager } from "@/internal/balances/events/EventBatchingManager.js"; +import { initEvent } from "@/internal/balances/events/initEvent.js"; +import type { FinalizeLockContextV2 } from "@/internal/balances/utils/lockV2/buildFinalizeLockContextV2.js"; + +export const insertFinalizeLockEventV2 = ({ + ctx, + finalizeLockContext, +}: { + ctx: AutumnContext; + finalizeLockContext: FinalizeLockContextV2; +}) => { + const { receipt, fullSubject, finalValue, lockValue, properties } = + finalizeLockContext; + + const event = initEvent({ + ctx, + eventInfo: { + event_name: receipt.feature_id, + value: new Decimal(finalValue).sub(lockValue).toNumber(), + properties, + }, + internalCustomerId: fullSubject.internalCustomerId, + internalEntityId: fullSubject.internalEntityId, + customerId: receipt.customer_id, + entityId: receipt.entity_id ?? undefined, + }); + + globalEventBatchingManager.addEvent(event); +}; diff --git a/server/src/internal/balances/finalizeLock/runFinalizeLock.ts b/server/src/internal/balances/finalizeLock/runFinalizeLock.ts index 81b8fe8cc..d22765d6a 100644 --- a/server/src/internal/balances/finalizeLock/runFinalizeLock.ts +++ b/server/src/internal/balances/finalizeLock/runFinalizeLock.ts @@ -4,7 +4,9 @@ import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { cancelLockExpiry } from "@/internal/balances/utils/lock/cancelLockExpiry.js"; import { claimLockReceipt } from "@/internal/balances/utils/lock/claimLockReceipt.js"; import { deleteLockReceipt } from "@/internal/balances/utils/lock/deleteLockReceipt.js"; +import { fetchLockReceipt } from "@/internal/balances/utils/lock/fetchLockReceipt.js"; import { buildFinalizeLockContext } from "./buildFinalizeLockContext.js"; +import { runFinalizeLockV2 } from "./runFinalizeLockV2.js"; import { runRedisFinalizeLock } from "./runRedisFinalizeLock.js"; export const runFinalizeLock = async ({ @@ -14,14 +16,27 @@ export const runFinalizeLock = async ({ ctx: AutumnContext; params: FinalizeLockParamsV0; }) => { + const fetchedReceipt = await fetchLockReceipt({ + ctx, + lockId: params.lock_id, + }); + + if (fetchedReceipt.source === "redis_v2") { + return runFinalizeLockV2({ + ctx, + params, + receipt: fetchedReceipt.receipt, + lockReceiptKey: fetchedReceipt.lockReceiptKey, + }); + } + const finalizeLockContext = await buildFinalizeLockContext({ ctx, params }); - const { lockReceiptKey, receipt, finalValue, lockValue } = + const { lockReceiptKey, receipt, finalValue, lockValue, redisInstance } = finalizeLockContext; - // Claim on the receipt's origin region to prevent cross-region double-claim - const { redisInstance } = await claimLockReceipt({ + await claimLockReceipt({ lockReceiptKey, - receiptRegion: receipt.region, + redisInstance, }); try { diff --git a/server/src/internal/balances/finalizeLock/runFinalizeLockV2.ts b/server/src/internal/balances/finalizeLock/runFinalizeLockV2.ts new file mode 100644 index 000000000..79a1d6489 --- /dev/null +++ b/server/src/internal/balances/finalizeLock/runFinalizeLockV2.ts @@ -0,0 +1,56 @@ +import { type FinalizeLockParamsV0, notNullish } from "@autumn/shared"; +import { Decimal } from "decimal.js"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { cancelLockExpiry } from "@/internal/balances/utils/lock/cancelLockExpiry.js"; +import { claimLockReceipt } from "@/internal/balances/utils/lock/claimLockReceipt.js"; +import { deleteLockReceipt } from "@/internal/balances/utils/lock/deleteLockReceipt.js"; +import type { LockReceipt } from "@/internal/balances/utils/lock/fetchLockReceipt.js"; +import { buildFinalizeLockContextV2 } from "@/internal/balances/utils/lockV2/buildFinalizeLockContextV2.js"; +import { runRedisFinalizeLockV2 } from "./runRedisFinalizeLockV2.js"; + +export const runFinalizeLockV2 = async ({ + ctx, + params, + receipt, + lockReceiptKey, +}: { + ctx: AutumnContext; + params: FinalizeLockParamsV0; + receipt: LockReceipt; + lockReceiptKey: string; +}) => { + const finalizeLockContext = await buildFinalizeLockContextV2({ + ctx, + params, + receipt, + lockReceiptKey, + }); + const { redisInstance, finalValue, lockValue } = finalizeLockContext; + + await claimLockReceipt({ + lockReceiptKey, + redisInstance, + }); + + try { + if (notNullish(receipt.expires_at)) { + await cancelLockExpiry({ + orgId: ctx.org.id, + env: ctx.env, + hashedKey: Bun.hash(params.lock_id).toString(), + }); + } + } catch (error) { + ctx.logger.error(`Failed to cancel lock expiry: ${error}`); + } + + if (new Decimal(finalValue).equals(lockValue)) { + await deleteLockReceipt({ lockReceiptKey, redisInstance }); + return { success: true }; + } + + await runRedisFinalizeLockV2({ ctx, finalizeLockContext }); + await deleteLockReceipt({ lockReceiptKey, redisInstance }); + + return { success: true }; +}; diff --git a/server/src/internal/balances/finalizeLock/runPostgresFinalizeLockV2.ts b/server/src/internal/balances/finalizeLock/runPostgresFinalizeLockV2.ts new file mode 100644 index 000000000..1e7d894a6 --- /dev/null +++ b/server/src/internal/balances/finalizeLock/runPostgresFinalizeLockV2.ts @@ -0,0 +1,26 @@ +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { executePostgresDeductionV2 } from "@/internal/balances/utils/deductionV2/executePostgresDeductionV2.js"; +import type { FinalizeLockContextV2 } from "@/internal/balances/utils/lockV2/buildFinalizeLockContextV2.js"; +import { insertFinalizeLockEventV2 } from "./insertFinalizeLockEventV2.js"; + +export const runPostgresFinalizeLockV2 = async ({ + ctx, + finalizeLockContext, +}: { + ctx: AutumnContext; + finalizeLockContext: FinalizeLockContextV2; +}) => { + const { receipt, fullSubject, deduction, deductionOptions } = + finalizeLockContext; + + await executePostgresDeductionV2({ + ctx, + fullSubject, + customerId: receipt.customer_id, + entityId: receipt.entity_id ?? undefined, + deductions: [deduction], + options: deductionOptions, + }); + + insertFinalizeLockEventV2({ ctx, finalizeLockContext }); +}; diff --git a/server/src/internal/balances/finalizeLock/runRedisFinalizeLockV2.ts b/server/src/internal/balances/finalizeLock/runRedisFinalizeLockV2.ts new file mode 100644 index 000000000..568c21465 --- /dev/null +++ b/server/src/internal/balances/finalizeLock/runRedisFinalizeLockV2.ts @@ -0,0 +1,63 @@ +import { currentRegion } from "@/external/redis/initRedis.js"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { executeRedisDeductionV2 } from "@/internal/balances/utils/deductionV2/executeRedisDeductionV2.js"; +import type { FinalizeLockContextV2 } from "@/internal/balances/utils/lockV2/buildFinalizeLockContextV2.js"; +import { deductionUpdatesToModifiedIds } from "@/internal/balances/utils/sync/deductionUpdatesToModifiedIds.js"; +import { globalSyncBatchingManagerV3 } from "@/internal/balances/utils/sync/SyncBatchingManagerV3.js"; +import { RedisDeductionError } from "@/internal/balances/utils/types/redisDeductionError.js"; +import { insertFinalizeLockEventV2 } from "./insertFinalizeLockEventV2.js"; +import { runPostgresFinalizeLockV2 } from "./runPostgresFinalizeLockV2.js"; + +export const runRedisFinalizeLockV2 = async ({ + ctx, + finalizeLockContext, +}: { + ctx: AutumnContext; + finalizeLockContext: FinalizeLockContextV2; +}) => { + const { receipt, fullSubject, deduction, deductionOptions } = + finalizeLockContext; + + let redisResult: Awaited>; + + try { + redisResult = await executeRedisDeductionV2({ + ctx, + fullSubject, + entityId: receipt.entity_id ?? undefined, + deductions: [deduction], + deductionOptions, + }); + } catch (error) { + if (error instanceof RedisDeductionError && error.shouldFallback()) { + ctx.logger.warn( + `[FINALIZE LOCK V2] Falling back to Postgres: ${error.code}`, + ); + await runPostgresFinalizeLockV2({ ctx, finalizeLockContext }); + return; + } + + throw error; + } + + const { updates, rolloverUpdates, modifiedCusEntIdsByFeatureId } = + redisResult; + const modifiedCusEntIds = deductionUpdatesToModifiedIds({ updates }); + const rolloverIds = Object.keys(rolloverUpdates); + + if (modifiedCusEntIds.length > 0 || rolloverIds.length > 0) { + ctx.logger.info(`[QUEUE SYNC V4] (${receipt.customer_id})`); + globalSyncBatchingManagerV3.addSyncItem({ + customerId: receipt.customer_id, + orgId: ctx.org.id, + env: ctx.env, + cusEntIds: modifiedCusEntIds, + rolloverIds, + region: currentRegion, + entityId: receipt.entity_id ?? undefined, + modifiedCusEntIdsByFeatureId, + }); + } + + insertFinalizeLockEventV2({ ctx, finalizeLockContext }); +}; diff --git a/server/src/internal/balances/handlers/handleCreateBalance.ts b/server/src/internal/balances/handlers/handleCreateBalance.ts index 03d13a329..391e3bcdb 100644 --- a/server/src/internal/balances/handlers/handleCreateBalance.ts +++ b/server/src/internal/balances/handlers/handleCreateBalance.ts @@ -1,32 +1,39 @@ -import { CreateBalanceParamsV0Schema } from "@autumn/shared"; +import { CreateBalanceParamsV0Schema, fullSubjectToFullCustomer } from "@autumn/shared"; import { FeatureNotFoundError } from "@shared/index"; import { createRoute } from "@/honoMiddlewares/routeHandler"; import { prepareNewBalanceForInsertion } from "@/internal/balances/createBalance/prepareNewBalanceForInsertion"; import { validateCreateBalanceParams } from "@/internal/balances/createBalance/validateCreateBalance"; +import { getOrSetCachedFullSubject } from "@/internal/customers/cache/fullSubject/actions/getOrSetCachedFullSubject.js"; import { CusService } from "@/internal/customers/CusService"; import { CusEntService } from "@/internal/customers/cusProducts/cusEnts/CusEntitlementService"; import { EntitlementService } from "@/internal/products/entitlements/EntitlementService"; +import { isFullSubjectRolloutEnabled } from "@/internal/misc/rollouts/fullSubjectRolloutUtils.js"; export const handleCreateBalance = createRoute({ body: CreateBalanceParamsV0Schema, handler: async (c) => { const ctx = c.get("ctx"); - const { org, env } = ctx; - const createBalanceParams = c.req.valid("json"); const { feature_id, customer_id, entity_id } = createBalanceParams; const feature = ctx.features.find((f) => f.id === feature_id); - if (!feature) { - throw new FeatureNotFoundError({ featureId: feature_id }); - } + if (!feature) throw new FeatureNotFoundError({ featureId: feature_id }); - const fullCustomer = await CusService.getFull({ - ctx, - idOrInternalId: customer_id, - entityId: entity_id, - withEntities: true, - }); + const fullCustomer = isFullSubjectRolloutEnabled({ ctx }) + ? fullSubjectToFullCustomer({ + fullSubject: await getOrSetCachedFullSubject({ + ctx, + customerId: customer_id, + entityId: entity_id, + source: "handleCreateBalance", + }), + }) + : await CusService.getFull({ + ctx, + idOrInternalId: customer_id, + entityId: entity_id, + withEntities: true, + }); await validateCreateBalanceParams({ ctx, diff --git a/server/src/internal/balances/handlers/handleSetUsage.ts b/server/src/internal/balances/handlers/handleSetUsage.ts new file mode 100644 index 000000000..2ea3220bf --- /dev/null +++ b/server/src/internal/balances/handlers/handleSetUsage.ts @@ -0,0 +1,59 @@ +import { SetUsageParamsSchema } from "@autumn/shared"; +import { createRoute } from "@/honoMiddlewares/routeHandler.js"; +import { getOrCreateCachedFullSubject } from "@/internal/customers/cache/fullSubject/actions/getOrCreateCachedFullSubject.js"; +import { getOrCreateCachedFullCustomer } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/getOrCreateCachedFullCustomer.js"; +import { isFullSubjectRolloutEnabled } from "@/internal/misc/rollouts/fullSubjectRolloutUtils.js"; +import { runUpdateUsage } from "../updateBalance/runUpdateUsage.js"; +import { updateUsageV2 } from "../updateBalance/v2/updateUsageV2.js"; + +export const handleSetUsage = createRoute({ + body: SetUsageParamsSchema, + handler: async (c) => { + const body = c.req.valid("json"); + const ctx = c.get("ctx"); + + if (isFullSubjectRolloutEnabled({ ctx })) { + const fullSubject = await getOrCreateCachedFullSubject({ + ctx, + params: { + customer_id: body.customer_id, + entity_id: body.entity_id, + }, + source: "handleSetUsage", + }); + + await updateUsageV2({ + ctx, + fullSubject, + params: { + customer_id: body.customer_id, + feature_id: body.feature_id, + usage: body.value, + entity_id: body.entity_id, + }, + }); + } else { + const fullCustomer = await getOrCreateCachedFullCustomer({ + ctx, + params: { + customer_id: body.customer_id, + entity_id: body.entity_id, + }, + source: "handleSetUsage", + }); + + await runUpdateUsage({ + ctx, + params: { + customer_id: body.customer_id, + feature_id: body.feature_id, + usage: body.value, + entity_id: body.entity_id, + }, + fullCustomer, + }); + } + + return c.json({ success: true }); + }, +}); diff --git a/server/src/internal/balances/handlers/handleUpdateBalance.ts b/server/src/internal/balances/handlers/handleUpdateBalance.ts index e1e7bb00d..0858283b2 100644 --- a/server/src/internal/balances/handlers/handleUpdateBalance.ts +++ b/server/src/internal/balances/handlers/handleUpdateBalance.ts @@ -8,12 +8,9 @@ import { } from "@autumn/shared"; import { StatusCodes } from "http-status-codes"; import { createRoute } from "@/honoMiddlewares/routeHandler"; -import { runUpdateBalanceV2 } from "@/internal/balances/updateBalance/runUpdateBalanceV2"; -import { runUpdateUsage } from "@/internal/balances/updateBalance/runUpdateUsage"; -import { updateGrantedBalance } from "@/internal/balances/updateBalance/updateGrantedBalance"; -import { updateNextResetAt } from "@/internal/balances/updateBalance/updateNextResetAt"; -import { buildCustomerEntitlementFilters } from "@/internal/balances/utils/buildCustomerEntitlementFilters"; -import { getOrSetCachedFullCustomer } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/getOrSetCachedFullCustomer"; +import { isFullSubjectRolloutEnabled } from "@/internal/misc/rollouts/fullSubjectRolloutUtils.js"; +import { updateBalanceV1 } from "@/internal/balances/updateBalance/updateBalanceV1.js"; +import { updateBalanceV2 } from "@/internal/balances/updateBalance/v2/updateBalanceV2.js"; export const handleUpdateBalance = createRoute({ body: UpdateBalanceParamsV0Schema.extend({}), @@ -41,54 +38,10 @@ export const handleUpdateBalance = createRoute({ }); } - let fullCustomer = await getOrSetCachedFullCustomer({ - ctx, - customerId: params.customer_id, - entityId: params.entity_id, - source: "handleUpdateBalance", - }); - - if (notNullish(params.add_to_balance) || notNullish(targetBalance)) { - const result = await runUpdateBalanceV2({ ctx, params, fullCustomer }); - fullCustomer = result?.fullCus ?? fullCustomer; - } - - if (notNullish(params.usage)) { - const result = await runUpdateUsage({ ctx, params, fullCustomer }); - fullCustomer = result?.fullCus ?? fullCustomer; - } - - if (notNullish(params.included_grant)) { - - ctx.logger.info( - `updating granted balance for feature ${params.feature_id} to ${params.included_grant}`, - ); - - const customerEntitlementFilters = buildCustomerEntitlementFilters({ - params, - }); - - await updateGrantedBalance({ - ctx, - fullCustomer, - featureId: params.feature_id, - targetGrantedBalance: params.included_grant, - customerEntitlementFilters, - }); - } - - if (notNullish(params.next_reset_at)) { - const customerEntitlementFilters = buildCustomerEntitlementFilters({ - params, - }); - - await updateNextResetAt({ - ctx, - fullCustomer, - featureId: params.feature_id, - nextResetAt: params.next_reset_at, - customerEntitlementFilters, - }); + if (isFullSubjectRolloutEnabled({ ctx })) { + await updateBalanceV2({ ctx, params, targetBalance }); + } else { + await updateBalanceV1({ ctx, params, targetBalance }); } return c.json({ success: true }); diff --git a/server/src/internal/balances/setUsage/handleSetUsage.ts b/server/src/internal/balances/setUsage/handleSetUsage.ts deleted file mode 100644 index 8dea1c72f..000000000 --- a/server/src/internal/balances/setUsage/handleSetUsage.ts +++ /dev/null @@ -1,34 +0,0 @@ -import { SetUsageParamsSchema } from "@autumn/shared"; -import { getOrCreateCachedFullCustomer } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/getOrCreateCachedFullCustomer.js"; -import { createRoute } from "../../../honoMiddlewares/routeHandler.js"; -import { runUpdateUsage } from "../updateBalance/runUpdateUsage.js"; - -export const handleSetUsage = createRoute({ - body: SetUsageParamsSchema, - handler: async (c) => { - const body = c.req.valid("json"); - const ctx = c.get("ctx"); - - const fullCustomer = await getOrCreateCachedFullCustomer({ - ctx, - params: { - customer_id: body.customer_id, - entity_id: body.entity_id, - }, - source: "handleSetUsage", - }); - - await runUpdateUsage({ - ctx, - params: { - customer_id: body.customer_id, - feature_id: body.feature_id, - usage: body.value, - entity_id: body.entity_id, - }, - fullCustomer, - }); - - return c.json({ success: true }); - }, -}); diff --git a/server/src/internal/balances/track/runTrackWithRollout.ts b/server/src/internal/balances/track/runTrackWithRollout.ts index bff55068e..bd26099e4 100644 --- a/server/src/internal/balances/track/runTrackWithRollout.ts +++ b/server/src/internal/balances/track/runTrackWithRollout.ts @@ -5,7 +5,7 @@ import type { FeatureDeduction } from "../utils/types/featureDeduction.js"; import { runTrackV2 } from "./runTrackV2.js"; import { runTrackV3 } from "./v3/runTrackV3.js"; -const TRACK_V3_ENABLED = false; +const TRACK_V3_ENABLED = true; export const shouldUseTrackV3 = ({ ctx }: { ctx: AutumnContext }): boolean => TRACK_V3_ENABLED && isFullSubjectRolloutEnabled({ ctx }); diff --git a/server/src/internal/balances/track/v3/handleRedisTrackErrorV3.ts b/server/src/internal/balances/track/v3/handleRedisTrackErrorV3.ts new file mode 100644 index 000000000..e3f6eac26 --- /dev/null +++ b/server/src/internal/balances/track/v3/handleRedisTrackErrorV3.ts @@ -0,0 +1,63 @@ +import { + ErrCode, + type FullSubject, + InsufficientBalanceError, + RecaseError, + type TrackParams, + type TrackResponseV3, +} from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; +import { + RedisDeductionError, + RedisDeductionErrorCode, +} from "../../utils/types/redisDeductionError.js"; +import { runPostgresTrackV3 } from "./runPostgresTrackV3.js"; + +/** Handles errors from V2 Redis deduction. Falls back to Postgres V3 path. */ +export const handleRedisTrackErrorV3 = async ({ + ctx, + error, + body, + fullSubject, + featureDeductions, +}: { + ctx: AutumnContext; + error: Error; + body: TrackParams; + fullSubject: FullSubject; + featureDeductions: FeatureDeduction[]; +}): Promise => { + if (!(error instanceof RedisDeductionError)) throw error; + + if (error.code === RedisDeductionErrorCode.InsufficientBalance) { + throw new InsufficientBalanceError({ + value: body.value ?? 1, + featureId: body.feature_id, + eventName: body.event_name, + }); + } + + if (error.code === RedisDeductionErrorCode.LockAlreadyExists) { + throw new RecaseError({ + message: "A lock with this ID already exists", + code: ErrCode.LockAlreadyExists, + statusCode: 409, + }); + } + + if (error.shouldFallback()) { + ctx.logger.warn( + `Falling back to Postgres V3 for track operation: ${error.code}`, + ); + + return runPostgresTrackV3({ + ctx, + fullSubject, + body, + featureDeductions, + }); + } + + throw error; +}; diff --git a/server/src/internal/balances/track/v3/runPostgresTrackV3.ts b/server/src/internal/balances/track/v3/runPostgresTrackV3.ts new file mode 100644 index 000000000..d2f5acfa6 --- /dev/null +++ b/server/src/internal/balances/track/v3/runPostgresTrackV3.ts @@ -0,0 +1,79 @@ +import type { FullSubject, TrackParams, TrackResponseV3 } from "@autumn/shared"; +import { tryCatch } from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { globalEventBatchingManager } from "@/internal/balances/events/EventBatchingManager.js"; +import { + buildEventInfo, + initEvent, +} from "@/internal/balances/events/initEvent.js"; +import { + deductionToTrackResponseV2, + executePostgresDeductionV2, +} from "@/internal/balances/utils/deductionV2/index.js"; +import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; +import { handlePostgresTrackError } from "../utils/handlePostgresTrackError.js"; + +export const runPostgresTrackV3 = async ({ + ctx, + fullSubject, + body, + featureDeductions, +}: { + ctx: AutumnContext; + fullSubject: FullSubject; + body: TrackParams; + featureDeductions: FeatureDeduction[]; +}): Promise => { + const { data: result, error } = await tryCatch( + executePostgresDeductionV2({ + ctx, + fullSubject, + customerId: body.customer_id, + entityId: body.entity_id, + deductions: featureDeductions, + options: { + overageBehaviour: body.overage_behavior || "cap", + triggerAutoTopUp: true, + }, + }), + ); + + if (error || !result) { + return handlePostgresTrackError({ + error: error ?? new Error("Unknown error"), + body, + }); + } + + const { fullSubject: updatedFullSubject, updates } = result; + + if (!body.skip_event && !body.idempotency_key) { + const eventInfo = buildEventInfo(body); + const event = initEvent({ + ctx, + eventInfo, + internalCustomerId: updatedFullSubject.internalCustomerId, + internalEntityId: updatedFullSubject.internalEntityId, + customerId: body.customer_id, + entityId: body.entity_id, + }); + + globalEventBatchingManager.addEvent(event); + } + + const { balance, balances } = await deductionToTrackResponseV2({ + ctx, + fullSubject: updatedFullSubject, + featureDeductions, + updates, + }); + + return { + customer_id: body.customer_id, + entity_id: body.entity_id, + event_name: body.event_name, + value: body.value ?? 1, + balance, + balances, + }; +}; diff --git a/server/src/internal/balances/track/v3/runRedisTrackV3.ts b/server/src/internal/balances/track/v3/runRedisTrackV3.ts index ccd02fc2c..f0895de7b 100644 --- a/server/src/internal/balances/track/v3/runRedisTrackV3.ts +++ b/server/src/internal/balances/track/v3/runRedisTrackV3.ts @@ -1,10 +1,74 @@ import type { FullSubject, TrackParams, TrackResponseV3 } from "@autumn/shared"; +import { tryCatch } from "@autumn/shared"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { globalEventBatchingManager } from "@/internal/balances/events/EventBatchingManager.js"; +import { + buildEventInfo, + initEvent, +} from "@/internal/balances/events/initEvent.js"; import { deductionToTrackResponseV2, executeRedisDeductionV2, } from "@/internal/balances/utils/deductionV2/index.js"; +import { globalSyncBatchingManagerV3 } from "@/internal/balances/utils/sync/SyncBatchingManagerV3.js"; import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; +import type { RolloverUpdate } from "../../utils/types/rolloverUpdate.js"; +import { handleRedisTrackErrorV3 } from "./handleRedisTrackErrorV3.js"; + +const queueSyncItem = ({ + ctx, + body, + fullSubject, + rolloverUpdates, + modifiedCusEntIdsByFeatureId, +}: { + ctx: AutumnContext; + body: TrackParams; + fullSubject: FullSubject; + rolloverUpdates: Record; + modifiedCusEntIdsByFeatureId: Record; +}): void => { + const cusEntIds = Object.values(modifiedCusEntIdsByFeatureId).flat(); + const rolloverIds = Object.keys(rolloverUpdates); + + if (cusEntIds.length === 0 && rolloverIds.length === 0) return; + + ctx.logger.info(`[QUEUE SYNC V4] (${body.customer_id})`); + globalSyncBatchingManagerV3.addSyncItem({ + customerId: body.customer_id, + orgId: ctx.org.id, + env: ctx.env, + cusEntIds, + rolloverIds, + entityId: fullSubject.entityId, + modifiedCusEntIdsByFeatureId, + }); +}; + +const queueEvent = ({ + ctx, + body, + fullSubject, +}: { + ctx: AutumnContext; + body: TrackParams; + fullSubject: FullSubject; +}): void => { + if (body.skip_event) return; + + const eventInfo = buildEventInfo(body); + + globalEventBatchingManager.addEvent( + initEvent({ + ctx, + eventInfo, + internalCustomerId: fullSubject.internalCustomerId, + internalEntityId: fullSubject.internalEntityId, + customerId: body.customer_id, + entityId: body.entity_id, + }), + ); +}; export const runRedisTrackV3 = async ({ ctx, @@ -19,8 +83,8 @@ export const runRedisTrackV3 = async ({ overageBehavior: "cap" | "reject"; body: TrackParams; }): Promise => { - const { fullSubject: updatedFullSubject, updates } = - await executeRedisDeductionV2({ + const { data: result, error } = await tryCatch( + executeRedisDeductionV2({ ctx, fullSubject, entityId: fullSubject.entity?.id ?? undefined, @@ -29,7 +93,35 @@ export const runRedisTrackV3 = async ({ overageBehaviour: overageBehavior, triggerAutoTopUp: true, }, + }), + ); + + if (error) { + return handleRedisTrackErrorV3({ + ctx, + error, + body, + fullSubject, + featureDeductions, }); + } + + const { + updates, + fullSubject: updatedFullSubject, + rolloverUpdates, + modifiedCusEntIdsByFeatureId, + } = result; + + queueSyncItem({ + ctx, + body, + fullSubject: updatedFullSubject, + rolloverUpdates, + modifiedCusEntIdsByFeatureId, + }); + + queueEvent({ ctx, body, fullSubject }); const { balance, balances } = await deductionToTrackResponseV2({ ctx, diff --git a/server/src/internal/balances/updateBalance/updateBalanceV1.ts b/server/src/internal/balances/updateBalance/updateBalanceV1.ts new file mode 100644 index 000000000..0b144ac68 --- /dev/null +++ b/server/src/internal/balances/updateBalance/updateBalanceV1.ts @@ -0,0 +1,67 @@ +import { notNullish, type UpdateBalanceParamsV0 } from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { getOrSetCachedFullCustomer } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/getOrSetCachedFullCustomer.js"; +import { buildCustomerEntitlementFilters } from "../utils/buildCustomerEntitlementFilters.js"; +import { runUpdateUsage } from "./runUpdateUsage.js"; +import { updateGrantedBalance } from "./updateGrantedBalance.js"; +import { updateNextResetAt } from "./updateNextResetAt.js"; +import { updateRemainingV1 } from "./updateRemainingV1.js"; + +export const updateBalanceV1 = async ({ + ctx, + params, + targetBalance, +}: { + ctx: AutumnContext; + params: UpdateBalanceParamsV0; + targetBalance?: number; +}) => { + let fullCustomer = await getOrSetCachedFullCustomer({ + ctx, + customerId: params.customer_id, + entityId: params.entity_id, + source: "handleUpdateBalance", + }); + + if (notNullish(params.add_to_balance) || notNullish(targetBalance)) { + const result = await updateRemainingV1({ ctx, params, fullCustomer }); + fullCustomer = result?.fullCus ?? fullCustomer; + } + + if (notNullish(params.usage)) { + const result = await runUpdateUsage({ ctx, params, fullCustomer }); + fullCustomer = result?.fullCus ?? fullCustomer; + } + + if (notNullish(params.included_grant)) { + ctx.logger.info( + `updating granted balance for feature ${params.feature_id} to ${params.included_grant}`, + ); + + const customerEntitlementFilters = buildCustomerEntitlementFilters({ + params, + }); + + await updateGrantedBalance({ + ctx, + fullCustomer, + featureId: params.feature_id, + targetGrantedBalance: params.included_grant, + customerEntitlementFilters, + }); + } + + if (notNullish(params.next_reset_at)) { + const customerEntitlementFilters = buildCustomerEntitlementFilters({ + params, + }); + + await updateNextResetAt({ + ctx, + fullCustomer, + featureId: params.feature_id, + nextResetAt: params.next_reset_at, + customerEntitlementFilters, + }); + } +}; diff --git a/server/src/internal/balances/updateBalance/runUpdateBalanceV2.ts b/server/src/internal/balances/updateBalance/updateRemainingV1.ts similarity index 97% rename from server/src/internal/balances/updateBalance/runUpdateBalanceV2.ts rename to server/src/internal/balances/updateBalance/updateRemainingV1.ts index f366304cd..6b033fdb2 100644 --- a/server/src/internal/balances/updateBalance/runUpdateBalanceV2.ts +++ b/server/src/internal/balances/updateBalance/updateRemainingV1.ts @@ -18,7 +18,7 @@ import { runRedisUpdateBalanceV2 } from "./runRedisUpdateBalanceV2.js"; * 3. Call runRedisUpdateBalanceV2 to update Redis * 4. Returns the result (caller handles sync/events if needed) */ -export const runUpdateBalanceV2 = async ({ +export const updateRemainingV1 = async ({ ctx, params, fullCustomer, diff --git a/server/src/internal/balances/updateBalance/v2/handleUpdateBalanceDeductionErrorV2.ts b/server/src/internal/balances/updateBalance/v2/handleUpdateBalanceDeductionErrorV2.ts new file mode 100644 index 000000000..2ca628f65 --- /dev/null +++ b/server/src/internal/balances/updateBalance/v2/handleUpdateBalanceDeductionErrorV2.ts @@ -0,0 +1,38 @@ +import type { CustomerEntitlementFilters, FullSubject } from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { executePostgresDeductionV2 } from "@/internal/balances/utils/deductionV2/executePostgresDeductionV2.js"; +import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; +import { RedisDeductionError } from "../../utils/types/redisDeductionError.js"; + +/** Handles Redis deduction errors for update balance V2. Falls back to Postgres when recoverable. */ +export const handleUpdateBalanceDeductionErrorV2 = async ({ + ctx, + error, + fullSubject, + featureDeductions, + customerEntitlementFilters, +}: { + ctx: AutumnContext; + error: Error; + fullSubject: FullSubject; + featureDeductions: FeatureDeduction[]; + customerEntitlementFilters?: CustomerEntitlementFilters; +}) => { + if (!(error instanceof RedisDeductionError) || !error.shouldFallback()) + throw error; + + ctx.logger.info(`[updateBalanceV2] Falling back to Postgres (${error.code})`); + + await executePostgresDeductionV2({ + ctx, + fullSubject, + customerId: fullSubject.customerId, + entityId: fullSubject.entityId, + deductions: featureDeductions, + options: { + overageBehaviour: "allow", + customerEntitlementFilters, + alterGrantedBalance: false, + }, + }); +}; diff --git a/server/src/internal/balances/updateBalance/v2/updateBalanceV2.ts b/server/src/internal/balances/updateBalance/v2/updateBalanceV2.ts new file mode 100644 index 000000000..67a7b569f --- /dev/null +++ b/server/src/internal/balances/updateBalance/v2/updateBalanceV2.ts @@ -0,0 +1,66 @@ +import { notNullish, type UpdateBalanceParamsV0 } from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { getOrSetCachedFullSubject } from "@/internal/customers/cache/fullSubject/actions/getOrSetCachedFullSubject.js"; +import { buildCustomerEntitlementFilters } from "../../utils/buildCustomerEntitlementFilters.js"; +import { updateIncludedGrantV2 } from "./updateIncludedGrantV2.js"; +import { updateNextResetAtV2 } from "./updateNextResetAtV2.js"; +import { updateRemainingV2 } from "./updateRemainingV2.js"; +import { updateUsageV2 } from "./updateUsageV2.js"; + +/** Update balance using the FullSubject cache path. */ +export const updateBalanceV2 = async ({ + ctx, + params, + targetBalance, +}: { + ctx: AutumnContext; + params: UpdateBalanceParamsV0; + targetBalance?: number; +}) => { + const fullSubject = await getOrSetCachedFullSubject({ + ctx, + customerId: params.customer_id, + entityId: params.entity_id, + source: "handleUpdateBalance", + }); + + if (notNullish(params.add_to_balance) || notNullish(targetBalance)) { + await updateRemainingV2({ ctx, fullSubject, params }); + } + + if (notNullish(params.usage)) { + await updateUsageV2({ ctx, fullSubject, params }); + } + + if (notNullish(params.included_grant)) { + ctx.logger.info( + `updating granted balance for feature ${params.feature_id} to ${params.included_grant}`, + ); + + const customerEntitlementFilters = buildCustomerEntitlementFilters({ + params, + }); + + await updateIncludedGrantV2({ + ctx, + fullSubject, + featureId: params.feature_id, + targetGrantedBalance: params.included_grant, + customerEntitlementFilters, + }); + } + + if (notNullish(params.next_reset_at)) { + const customerEntitlementFilters = buildCustomerEntitlementFilters({ + params, + }); + + await updateNextResetAtV2({ + ctx, + fullSubject, + featureId: params.feature_id, + nextResetAt: params.next_reset_at, + customerEntitlementFilters, + }); + } +}; diff --git a/server/src/internal/balances/updateBalance/v2/updateIncludedGrantV2.ts b/server/src/internal/balances/updateBalance/v2/updateIncludedGrantV2.ts new file mode 100644 index 000000000..54131e30a --- /dev/null +++ b/server/src/internal/balances/updateBalance/v2/updateIncludedGrantV2.ts @@ -0,0 +1,111 @@ +import { + type CustomerEntitlementFilters, + cusEntsToAllowance, + type FullSubject, + fullSubjectToCustomerEntitlements, + InternalError, + isEntityScopedCusEnt, + notNullish, + nullish, + orgToInStatuses, + RecaseError, +} from "@autumn/shared"; +import { Decimal } from "decimal.js"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { CusEntService } from "@/internal/customers/cusProducts/cusEnts/CusEntitlementService.js"; +import { updateSubjectBalanceCache } from "@/internal/customers/cusProducts/cusEnts/actions/cache/updateSubjectBalanceCache.js"; + +export const updateIncludedGrantV2 = async ({ + ctx, + fullSubject, + featureId, + targetGrantedBalance, + customerEntitlementFilters = {}, +}: { + ctx: AutumnContext; + fullSubject: FullSubject; + featureId: string | undefined; + targetGrantedBalance: number; + customerEntitlementFilters?: CustomerEntitlementFilters; +}) => { + const cusEnts = fullSubjectToCustomerEntitlements({ + fullSubject, + featureIds: featureId ? [featureId] : undefined, + inStatuses: orgToInStatuses({ org: ctx.org }), + customerEntitlementFilters, + }); + + if (cusEnts.length === 0) { + throw new RecaseError({ + message: `[updateIncludedGrantV2] No balances to update for feature ${featureId}, customer ${fullSubject.customerId}`, + }); + } + + const currentAllowance = cusEntsToAllowance({ + cusEnts, + entityId: fullSubject.entityId, + withRollovers: false, + }); + + const requiredAdjustment = new Decimal(targetGrantedBalance) + .sub(currentAllowance) + .toNumber(); + + const targetCusEnt = cusEnts[0]; + const targetFeatureId = featureId ?? targetCusEnt.entitlement.feature.id; + const isEntityScoped = isEntityScopedCusEnt(targetCusEnt); + const entityId = fullSubject.entityId; + + if (isEntityScoped) { + const entityKeys = Object.keys(targetCusEnt.entities ?? {}); + const targetEntityId = notNullish(entityId) ? entityId : entityKeys[0]; + + if ( + nullish(targetEntityId) || + nullish(targetCusEnt.entities?.[targetEntityId]) + ) { + throw new InternalError({ + message: `[updateIncludedGrantV2] No entity balance found for feature ${featureId}, customer ${fullSubject.customerId}`, + }); + } + + const currentEntity = targetCusEnt.entities[targetEntityId]; + const newEntities = { + ...targetCusEnt.entities, + [targetEntityId]: { + id: targetEntityId, + balance: currentEntity.balance, + adjustment: requiredAdjustment, + additional_balance: currentEntity.additional_balance, + }, + }; + + await CusEntService.update({ + ctx, + id: targetCusEnt.id, + updates: { entities: newEntities }, + }); + + await updateSubjectBalanceCache({ + ctx, + customerId: fullSubject.customerId, + featureId: targetFeatureId, + customerEntitlementId: targetCusEnt.id, + updates: { entities: newEntities }, + }); + } else { + await CusEntService.update({ + ctx, + id: targetCusEnt.id, + updates: { adjustment: requiredAdjustment }, + }); + + await updateSubjectBalanceCache({ + ctx, + customerId: fullSubject.customerId, + featureId: targetFeatureId, + customerEntitlementId: targetCusEnt.id, + updates: { adjustment: requiredAdjustment }, + }); + } +}; diff --git a/server/src/internal/balances/updateBalance/v2/updateNextResetAtV2.ts b/server/src/internal/balances/updateBalance/v2/updateNextResetAtV2.ts new file mode 100644 index 000000000..125586ac1 --- /dev/null +++ b/server/src/internal/balances/updateBalance/v2/updateNextResetAtV2.ts @@ -0,0 +1,68 @@ +import { + type CustomerEntitlementFilters, + EntInterval, + type FullSubject, + fullSubjectToCustomerEntitlements, + orgToInStatuses, + RecaseError, +} from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { CusEntService } from "@/internal/customers/cusProducts/cusEnts/CusEntitlementService.js"; +import { updateSubjectBalanceCache } from "@/internal/customers/cusProducts/cusEnts/actions/cache/updateSubjectBalanceCache.js"; + +export const updateNextResetAtV2 = async ({ + ctx, + fullSubject, + featureId, + nextResetAt, + customerEntitlementFilters, +}: { + ctx: AutumnContext; + fullSubject: FullSubject; + featureId: string | undefined; + nextResetAt: number; + customerEntitlementFilters?: CustomerEntitlementFilters; +}) => { + const cusEnts = fullSubjectToCustomerEntitlements({ + fullSubject, + featureIds: featureId ? [featureId] : undefined, + inStatuses: orgToInStatuses({ org: ctx.org }), + customerEntitlementFilters, + }); + + if (cusEnts.length === 0) { + throw new RecaseError({ + message: `[updateNextResetAtV2] No balances found for feature ${featureId}, customer ${fullSubject.customerId}`, + }); + } + + const sorted = [...cusEnts].sort((a, b) => { + const aReset = a.next_reset_at ?? Number.POSITIVE_INFINITY; + const bReset = b.next_reset_at ?? Number.POSITIVE_INFINITY; + return aReset - bReset; + }); + + const targetCusEnt = sorted[0]; + + if (targetCusEnt.entitlement.interval === EntInterval.Lifetime) { + throw new RecaseError({ + message: `Cannot update next reset at for lifetime balance (feature ${featureId}, customer ${fullSubject.customerId})`, + }); + } + + const targetFeatureId = featureId ?? targetCusEnt.entitlement.feature.id; + + await CusEntService.update({ + ctx, + id: targetCusEnt.id, + updates: { next_reset_at: nextResetAt }, + }); + + await updateSubjectBalanceCache({ + ctx, + customerId: fullSubject.customerId, + featureId: targetFeatureId, + customerEntitlementId: targetCusEnt.id, + updates: { next_reset_at: nextResetAt }, + }); +}; diff --git a/server/src/internal/balances/updateBalance/v2/updateRemainingV2.ts b/server/src/internal/balances/updateBalance/v2/updateRemainingV2.ts new file mode 100644 index 000000000..7c1fa1787 --- /dev/null +++ b/server/src/internal/balances/updateBalance/v2/updateRemainingV2.ts @@ -0,0 +1,87 @@ +import { + FeatureNotFoundError, + type FullSubject, + notNullish, + tryCatch, + type UpdateBalanceParamsV0, +} from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { executeRedisDeductionV2 } from "@/internal/balances/utils/deductionV2/executeRedisDeductionV2.js"; +import { globalSyncBatchingManagerV3 } from "@/internal/balances/utils/sync/SyncBatchingManagerV3.js"; +import { buildCustomerEntitlementFilters } from "../../utils/buildCustomerEntitlementFilters.js"; +import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; +import { handleUpdateBalanceDeductionErrorV2 } from "./handleUpdateBalanceDeductionErrorV2.js"; + +/** Updates remaining balance using the FullSubject cache path. */ +export const updateRemainingV2 = async ({ + ctx, + fullSubject, + params, +}: { + ctx: AutumnContext; + fullSubject: FullSubject; + params: UpdateBalanceParamsV0; +}) => { + const { features } = ctx; + const { feature_id: featureId, add_to_balance: addToBalance } = params; + const targetBalance = params.remaining ?? params.current_balance; + + const feature = features.find((f) => f.id === featureId); + if (!feature) throw new FeatureNotFoundError({ featureId }); + + const customerEntitlementFilters = buildCustomerEntitlementFilters({ + params, + }); + + const featureDeductions: FeatureDeduction[] = [ + { + feature, + deduction: notNullish(addToBalance) ? -addToBalance : 0, + targetBalance: notNullish(targetBalance) ? targetBalance : undefined, + }, + ]; + + const entityId = fullSubject.entityId; + + const { data: result, error } = await tryCatch( + executeRedisDeductionV2({ + ctx, + fullSubject, + entityId, + deductions: featureDeductions, + deductionOptions: { + overageBehaviour: "allow", + customerEntitlementFilters, + alterGrantedBalance: false, + }, + }), + ); + + if (error) { + return handleUpdateBalanceDeductionErrorV2({ + ctx, + error, + fullSubject, + featureDeductions, + customerEntitlementFilters, + }); + } + + const { rolloverUpdates, modifiedCusEntIdsByFeatureId } = result; + const cusEntIds = Object.values(modifiedCusEntIdsByFeatureId).flat(); + const rolloverIds = Object.keys(rolloverUpdates); + + if (cusEntIds.length > 0 || rolloverIds.length > 0) { + globalSyncBatchingManagerV3.addSyncItem({ + customerId: fullSubject.customerId, + orgId: ctx.org.id, + env: ctx.env, + cusEntIds, + rolloverIds, + entityId: fullSubject.entityId, + modifiedCusEntIdsByFeatureId, + }); + } + + return result; +}; diff --git a/server/src/internal/balances/updateBalance/v2/updateUsageV2.ts b/server/src/internal/balances/updateBalance/v2/updateUsageV2.ts new file mode 100644 index 000000000..d64ba1578 --- /dev/null +++ b/server/src/internal/balances/updateBalance/v2/updateUsageV2.ts @@ -0,0 +1,131 @@ +import { + type CustomerEntitlementFilters, + cusEntsToGrantedBalance, + cusEntsToPrepaidQuantity, + FeatureNotFoundError, + type FullSubject, + fullSubjectToCustomerEntitlements, + nullish, + tryCatch, + type UpdateBalanceParamsV0, +} from "@autumn/shared"; +import { Decimal } from "decimal.js"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { executeRedisDeductionV2 } from "@/internal/balances/utils/deductionV2/executeRedisDeductionV2.js"; +import { globalSyncBatchingManagerV3 } from "@/internal/balances/utils/sync/SyncBatchingManagerV3.js"; +import { buildCustomerEntitlementFilters } from "../../utils/buildCustomerEntitlementFilters.js"; +import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; +import { handleUpdateBalanceDeductionErrorV2 } from "./handleUpdateBalanceDeductionErrorV2.js"; + +const getUpdateUsageTargetBalance = ({ + fullSubject, + featureId, + entityId, + usage, + customerEntitlementFilters, +}: { + fullSubject: FullSubject; + featureId: string; + entityId?: string; + usage: number; + customerEntitlementFilters?: CustomerEntitlementFilters; +}) => { + const cusEnts = fullSubjectToCustomerEntitlements({ + fullSubject, + featureIds: [featureId], + customerEntitlementFilters, + }); + + const grantedBalance = cusEntsToGrantedBalance({ + cusEnts, + entityId, + }); + + const prepaidQuantity = cusEntsToPrepaidQuantity({ + cusEnts, + sumAcrossEntities: nullish(entityId), + }); + + return new Decimal(grantedBalance).add(prepaidQuantity).sub(usage).toNumber(); +}; + +/** Updates balance by setting usage to an exact value, using the FullSubject cache path. */ +export const updateUsageV2 = async ({ + ctx, + fullSubject, + params, +}: { + ctx: AutumnContext; + fullSubject: FullSubject; + params: UpdateBalanceParamsV0; +}) => { + const { features } = ctx; + const { feature_id: featureId, usage } = params; + + const feature = features.find((f) => f.id === featureId); + if (!feature) throw new FeatureNotFoundError({ featureId }); + + const customerEntitlementFilters = buildCustomerEntitlementFilters({ + params, + }); + + const entityId = fullSubject.entityId; + + const targetBalance = getUpdateUsageTargetBalance({ + fullSubject, + featureId, + entityId, + usage: usage!, + customerEntitlementFilters, + }); + + const featureDeductions: FeatureDeduction[] = [ + { + feature, + deduction: 0, + targetBalance, + }, + ]; + + const { data: result, error } = await tryCatch( + executeRedisDeductionV2({ + ctx, + fullSubject, + entityId, + deductions: featureDeductions, + deductionOptions: { + overageBehaviour: "allow", + customerEntitlementFilters, + alterGrantedBalance: false, + }, + }), + ); + + if (error) { + return handleUpdateBalanceDeductionErrorV2({ + ctx, + error, + fullSubject, + featureDeductions, + customerEntitlementFilters, + }); + } + + const { rolloverUpdates, modifiedCusEntIdsByFeatureId } = result; + const cusEntIds = Object.values(modifiedCusEntIdsByFeatureId).flat(); + const rolloverIds = Object.keys(rolloverUpdates); + + if (cusEntIds.length > 0 || rolloverIds.length > 0) { + globalSyncBatchingManagerV3.addSyncItem({ + customerId: fullSubject.customerId, + orgId: ctx.org.id, + env: ctx.env, + cusEntIds, + rolloverIds, + entityId: fullSubject.entityId, + modifiedCusEntIdsByFeatureId, + }); + } + + return result; +}; diff --git a/server/src/internal/balances/utils/deductionV2/executePostgresDeductionV2.ts b/server/src/internal/balances/utils/deductionV2/executePostgresDeductionV2.ts new file mode 100644 index 000000000..b69ca469e --- /dev/null +++ b/server/src/internal/balances/utils/deductionV2/executePostgresDeductionV2.ts @@ -0,0 +1,305 @@ +import { + type EntityRolloverBalance, + type FullCusEntWithFullCusProduct, + type FullSubject, + fullSubjectToFullCustomer, + InternalError, +} from "@autumn/shared"; +import { sql } from "drizzle-orm"; +import { withLock } from "@/external/redis/redisUtils.js"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { triggerAutoTopUp } from "@/internal/balances/autoTopUp/triggerAutoTopUp.js"; +import { fireTrackWebhooks } from "@/internal/balances/trackWebhooks/fireTrackWebhooks.js"; +import { createAllocatedInvoice } from "@/internal/balances/utils/allocatedInvoice/createAllocatedInvoice.js"; +import { saveLockReceipt } from "@/internal/balances/utils/lock/saveLockReceipt.js"; +import type { DeductionOptions } from "../types/deductionTypes.js"; +import type { DeductionUpdate } from "../types/deductionUpdate.js"; +import type { FeatureDeduction } from "../types/featureDeduction.js"; +import type { MutationLogItem } from "../types/mutationLogItem.js"; +import { applyDeductionUpdateToFullSubject } from "./applyDeductionUpdateToFullSubject.js"; +import { applyRolloverUpdatesToFullSubject } from "./applyRolloverUpdatesToFullSubject.js"; +import { logDeductionUpdatesV2 } from "./logDeductionUpdatesV2.js"; +import { mutationLogsToFeaturesV2 } from "./mutationLogsToFeaturesV2.js"; +import { prepareDeductionOptionsV2 } from "./prepareDeductionOptionsV2.js"; +import { prepareFeatureDeductionV2 } from "./prepareFeatureDeductionV2.js"; +import { rollbackDeductionV2 } from "./rollbackDeductionV2.js"; +import { syncDeductionUpdatesToFullSubjectCache } from "./syncDeductionUpdatesToFullSubjectCache.js"; + +interface RolloverOverwrite { + id: string; + cus_ent_id: string; + balance: number; + usage: number; + entities: Record; +} + +export const executePostgresDeductionV2 = async ({ + ctx, + fullSubject, + customerId, + entityId, + deductions, + options = {}, +}: { + ctx: AutumnContext; + customerId: string; + entityId?: string; + fullSubject: FullSubject; + deductions: FeatureDeduction[]; + options?: DeductionOptions; +}): Promise<{ + oldFullSubject: FullSubject; + fullSubject: FullSubject; + updates: Record; + mutationLogs: MutationLogItem[]; + modifiedCusEntIdsByFeatureId: Record; +}> => { + const { db, org, env } = ctx; + + ctx.logger.info( + `executing postgres deduction v2, deductions: ${JSON.stringify( + deductions.map((d) => ({ + featureId: d.feature.id, + deduction: d.deduction, + targetBalance: d.targetBalance, + })), + )}`, + ); + + const oldFullSubject = structuredClone(fullSubject); + + const resolvedOptions = prepareDeductionOptionsV2({ + ctx, + fullSubject, + options, + deductions, + }); + + if (resolvedOptions.paidAllocated && deductions.some((d) => d.lock)) { + throw new InternalError({ + message: "Locks are not supported for paid allocated features", + }); + } + + const executeDeduction = async (): Promise<{ + updates: Record; + mutationLogs: MutationLogItem[]; + modifiedCusEntIdsByFeatureId: Record; + }> => { + let allUpdates: Record = {}; + let allRolloverOverwrites: RolloverOverwrite[] = []; + let allMutationLogs: MutationLogItem[] = []; + const allModifiedCusEntIdsByFeatureId: Record = {}; + + for (const deduction of deductions) { + const { + feature, + deduction: toDeduct, + targetBalance, + lockReceipt, + unwindValue, + } = deduction; + + const { + customerEntitlementDeductions, + spendLimitByFeatureId, + usageBasedCusEntIdsByFeatureId, + rollovers, + customerEntitlements, + unlimitedFeatureIds, + lock: preparedLock, + } = prepareFeatureDeductionV2({ + ctx, + fullSubject, + deduction, + options: resolvedOptions, + }); + + if (customerEntitlements.length === 0 || unlimitedFeatureIds.length > 0) + continue; + + const result = await db.execute( + sql`SELECT * FROM deduct_from_cus_ents( + ${JSON.stringify({ + sorted_entitlements: customerEntitlementDeductions, + spend_limit_by_feature_id: spendLimitByFeatureId ?? null, + usage_based_cus_ent_ids_by_feature_id: + usageBasedCusEntIdsByFeatureId ?? null, + amount_to_deduct: toDeduct ?? null, + target_balance: targetBalance ?? null, + lock_receipt: lockReceipt ?? null, + unwind_value: unwindValue ?? null, + target_entity_id: entityId || null, + rollovers: rollovers.length > 0 ? rollovers : null, + cus_ent_ids: customerEntitlements.map((ce) => ce.id), + skip_additional_balance: resolvedOptions.skipAdditionalBalance, + alter_granted_balance: resolvedOptions.alterGrantedBalance, + overage_behaviour: resolvedOptions.overageBehaviour, + feature_id: feature.id, + })}::jsonb + )`, + ); + + const resultJson = result[0]?.deduct_from_cus_ents as { + updates: Record; + remaining: number; + rollover_updates: RolloverOverwrite[]; + mutation_logs: MutationLogItem[]; + }; + + if (!resultJson) { + throw new InternalError({ + message: "Failed to deduct from entitlements", + }); + } + + const { updates, rollover_updates, mutation_logs } = resultJson; + logDeductionUpdatesV2({ + ctx, + fullSubject, + updates, + source: "executePostgresDeductionV2", + }); + allUpdates = { ...allUpdates, ...updates }; + allMutationLogs = [...allMutationLogs, ...(mutation_logs ?? [])]; + if (rollover_updates?.length > 0) { + allRolloverOverwrites = [...allRolloverOverwrites, ...rollover_updates]; + } + + for (const ced of customerEntitlementDeductions) { + if (!updates[ced.customer_entitlement_id]) continue; + if (!allModifiedCusEntIdsByFeatureId[ced.feature_id]) { + allModifiedCusEntIdsByFeatureId[ced.feature_id] = []; + } + allModifiedCusEntIdsByFeatureId[ced.feature_id].push( + ced.customer_entitlement_id, + ); + } + + const oldFullCustomer = fullSubjectToFullCustomer({ + fullSubject: oldFullSubject, + }); + + try { + applyRolloverUpdatesToFullSubject({ + fullSubject, + rolloverUpdates: Object.fromEntries( + (rollover_updates ?? []).map((rollover) => [ + rollover.id, + { + balance: rollover.balance, + usage: rollover.usage, + entities: rollover.entities, + }, + ]), + ), + }); + + for (const customerEntitlementId of Object.keys(updates)) { + const update = updates[customerEntitlementId]; + const customerEntitlement = customerEntitlements.find( + (ce: FullCusEntWithFullCusProduct) => + ce.id === customerEntitlementId, + ); + + if (!customerEntitlement) continue; + + await createAllocatedInvoice({ + ctx, + customerEntitlement, + oldFullCustomer, + update, + }); + + applyDeductionUpdateToFullSubject({ + fullSubject, + customerEntitlementId, + update, + }); + } + + if (preparedLock?.enabled) { + await saveLockReceipt({ + lock: preparedLock, + customerId: fullSubject.customerId || customerId, + featureId: feature.id, + entityId, + items: mutation_logs ?? [], + }); + } + } catch (error) { + if (error instanceof Error && !error?.message?.includes("declined")) { + ctx.logger.error( + `[executePostgresDeductionV2] Attempting rollback due to error: ${error}`, + ); + } + await rollbackDeductionV2({ + ctx, + oldFullSubject, + updates, + }); + throw error; + } + + const featuresFromMutationLogs = mutationLogsToFeaturesV2({ + fullSubject, + mutationLogs: mutation_logs ?? [], + }); + + const newFullCustomer = fullSubjectToFullCustomer({ fullSubject }); + + fireTrackWebhooks({ + ctx, + oldFullCus: oldFullCustomer, + newFullCus: newFullCustomer, + feature: deduction.feature, + entityId, + featuresFromMutationLogs, + }); + + if (resolvedOptions.triggerAutoTopUp) { + triggerAutoTopUp({ + ctx, + newFullCus: newFullCustomer, + feature: deduction.feature, + }).catch((error) => { + ctx.logger.error( + `[executePostgresDeductionV2] Failed to trigger auto top-up: ${error}`, + ); + }); + } + } + + await syncDeductionUpdatesToFullSubjectCache({ + ctx, + customerId, + fullSubject: oldFullSubject, + cusEntUpdates: allUpdates, + rolloverOverwrites: allRolloverOverwrites, + modifiedCusEntIdsByFeatureId: allModifiedCusEntIdsByFeatureId, + }); + + return { + updates: allUpdates, + mutationLogs: allMutationLogs, + modifiedCusEntIdsByFeatureId: allModifiedCusEntIdsByFeatureId, + }; + }; + + const deductionResult = resolvedOptions.paidAllocated + ? await withLock({ + lockKey: `lock:deduction:${org.id}:${env}:${customerId}`, + ttlMs: 60000, + errorMessage: `Deduction for paid feature ${deductions[0]?.feature?.name} already in progress for customer ${customerId}.`, + fn: executeDeduction, + }) + : await executeDeduction(); + + return { + oldFullSubject, + fullSubject, + updates: deductionResult.updates, + mutationLogs: deductionResult.mutationLogs, + modifiedCusEntIdsByFeatureId: deductionResult.modifiedCusEntIdsByFeatureId, + }; +}; diff --git a/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts b/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts index e074be0c4..fb6265263 100644 --- a/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts +++ b/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts @@ -1,16 +1,33 @@ import { - type FullCustomer, + type FullCusEntWithFullCusProduct, type FullSubject, - InternalError, + fullSubjectToFullCustomer, } from "@autumn/shared"; +import type { Redis } from "ioredis"; +import { currentRegion, redis } from "@/external/redis/initRedis.js"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; -import { prepareDeductionOptions } from "../deduction/prepareDeductionOptions.js"; +import { triggerAutoTopUp } from "@/internal/balances/autoTopUp/triggerAutoTopUp.js"; +import { fireTrackWebhooks } from "@/internal/balances/trackWebhooks/fireTrackWebhooks.js"; +import { createAllocatedInvoice } from "@/internal/balances/utils/allocatedInvoice/createAllocatedInvoice.js"; +import { buildFullSubjectKey } from "@/internal/customers/cache/fullSubject/builders/buildFullSubjectKey.js"; +import { tryRedisWrite } from "@/utils/cacheUtils/cacheUtils.js"; import type { DeductionOptions } from "../types/deductionTypes.js"; import type { DeductionUpdate } from "../types/deductionUpdate.js"; import type { FeatureDeduction } from "../types/featureDeduction.js"; import type { MutationLogItem } from "../types/mutationLogItem.js"; +import { + RedisDeductionError, + RedisDeductionErrorCode, +} from "../types/redisDeductionError.js"; +import type { LuaDeductionResult } from "../types/redisDeductionResult.js"; import type { RolloverUpdate } from "../types/rolloverUpdate.js"; +import { applyDeductionUpdateToFullSubject } from "./applyDeductionUpdateToFullSubject.js"; +import { applyRolloverUpdatesToFullSubject } from "./applyRolloverUpdatesToFullSubject.js"; +import { logDeductionUpdatesV2 } from "./logDeductionUpdatesV2.js"; +import { mutationLogsToFeaturesV2 } from "./mutationLogsToFeaturesV2.js"; +import { prepareDeductionOptionsV2 } from "./prepareDeductionOptionsV2.js"; import { prepareFeatureDeductionV2 } from "./prepareFeatureDeductionV2.js"; +import { rollbackDeductionV2 } from "./rollbackDeductionV2.js"; export const executeRedisDeductionV2 = async ({ ctx, @@ -18,42 +35,258 @@ export const executeRedisDeductionV2 = async ({ entityId, deductions, deductionOptions = {}, + redisInstance, }: { ctx: AutumnContext; fullSubject: FullSubject; entityId?: string; deductions: FeatureDeduction[]; deductionOptions?: DeductionOptions; + redisInstance?: Redis; }): Promise<{ oldFullSubject: FullSubject; fullSubject: FullSubject; updates: Record; rolloverUpdates: Record; mutationLogs: MutationLogItem[]; + modifiedCusEntIdsByFeatureId: Record; }> => { - const _oldFullSubject = structuredClone(fullSubject); - const resolvedOptions = prepareDeductionOptions({ + const { org, env } = ctx; + const oldFullSubject = structuredClone(fullSubject); + + const options = prepareDeductionOptionsV2({ + ctx, + fullSubject, options: deductionOptions, - fullCustomer: fullSubject.customer as unknown as FullCustomer, deductions, }); - const preparedDeductions = deductions.map((deduction) => - prepareFeatureDeductionV2({ + + if (options.paidAllocated) { + throw new RedisDeductionError({ + message: "Paid allocated deductions are not supported for Redis", + code: RedisDeductionErrorCode.PaidAllocated, + }); + } + + if (options.paidAllocated && deductions.some((d) => d.lock)) { + throw new RedisDeductionError({ + message: "Locks are not supported for paid allocated features", + code: RedisDeductionErrorCode.PaidAllocated, + }); + } + + if (ctx.skipCache) { + throw new RedisDeductionError({ + message: "Skipping cache is not supported for Redis", + code: RedisDeductionErrorCode.SkipCache, + }); + } + + let allUpdates: Record = {}; + let allRolloverUpdates: Record = {}; + let allMutationLogs: MutationLogItem[] = []; + const allModifiedCusEntIdsByFeatureId: Record = {}; + + const customerId = fullSubject.customerId; + const routingKey = buildFullSubjectKey({ + orgId: org.id, + env, + customerId, + entityId: fullSubject.entityId, + }); + + for (const deduction of deductions) { + const { + feature, + deduction: toDeduct, + targetBalance, + unwindValue, + lockReceiptKey, + } = deduction; + + const { + customerEntitlementDeductions, + spendLimitByFeatureId, + usageBasedCusEntIdsByFeatureId, + rollovers, + customerEntitlements, + unlimitedFeatureIds, + lock: preparedLock, + } = prepareFeatureDeductionV2({ ctx, fullSubject, deduction, - options: resolvedOptions, - }), - ); + options, + }); - throw new InternalError({ - message: "FullSubject Redis deduction is not implemented", - code: "full_subject_redis_deduction_not_implemented", - data: { - entityId: entityId ?? null, - deductionsCount: deductions.length, - preparedDeductionsCount: preparedDeductions.length, - subjectType: _oldFullSubject.subjectType, - }, - }); + if (unlimitedFeatureIds.length > 0) { + continue; + } + + const luaParams = { + org_id: org.id, + env, + customer_id: customerId, + customer_entitlement_deductions: customerEntitlementDeductions, + spend_limit_by_feature_id: spendLimitByFeatureId ?? null, + usage_based_cus_ent_ids_by_feature_id: + usageBasedCusEntIdsByFeatureId ?? null, + amount_to_deduct: toDeduct ?? null, + target_balance: targetBalance ?? null, + target_entity_id: entityId || null, + rollovers: rollovers.length > 0 ? rollovers : null, + skip_additional_balance: options.skipAdditionalBalance, + alter_granted_balance: options.alterGrantedBalance, + overage_behaviour: options.overageBehaviour, + feature_id: feature.id, + lock: preparedLock + ? { + ...preparedLock, + region: currentRegion, + } + : null, + unwind_value: unwindValue ?? null, + lock_receipt_key: lockReceiptKey ?? null, + }; + + const targetRedis = redisInstance ?? redis; + const result = await tryRedisWrite( + () => + targetRedis.deductFromSubjectBalances( + routingKey, + JSON.stringify(luaParams), + ), + redisInstance, + ); + + if (!result) { + throw new RedisDeductionError({ + message: "Redis not ready for deduction", + code: RedisDeductionErrorCode.SubjectBalanceNotFound, + }); + } + + const resultJson = JSON.parse(result) as LuaDeductionResult; + + if (resultJson.logs && resultJson.logs.length > 0) { + ctx.logger.debug( + `[executeRedisDeductionV2] Logs: ${resultJson.logs.join("\n")}`, + ); + } + + if (resultJson.error) { + throw new RedisDeductionError({ + message: `Redis deduction failed: ${resultJson.error}`, + code: resultJson.error as RedisDeductionErrorCode, + }); + } + + const { updates, rollover_updates } = resultJson; + const mutationLogs = Array.isArray(resultJson.mutation_logs) + ? resultJson.mutation_logs + : []; + + logDeductionUpdatesV2({ + ctx, + fullSubject, + updates, + source: "executeRedisDeductionV2", + }); + + allUpdates = { ...allUpdates, ...updates }; + allRolloverUpdates = { ...allRolloverUpdates, ...rollover_updates }; + allMutationLogs = [...allMutationLogs, ...mutationLogs]; + + for (const ced of customerEntitlementDeductions) { + if (!updates[ced.customer_entitlement_id]) continue; + if (!allModifiedCusEntIdsByFeatureId[ced.feature_id]) { + allModifiedCusEntIdsByFeatureId[ced.feature_id] = []; + } + allModifiedCusEntIdsByFeatureId[ced.feature_id].push( + ced.customer_entitlement_id, + ); + } + + const oldFullCustomer = fullSubjectToFullCustomer({ + fullSubject: oldFullSubject, + }); + + try { + applyRolloverUpdatesToFullSubject({ + fullSubject, + rolloverUpdates: rollover_updates, + }); + + for (const customerEntitlementId of Object.keys(updates)) { + const update = updates[customerEntitlementId]; + const customerEntitlement = customerEntitlements.find( + (ce: FullCusEntWithFullCusProduct) => ce.id === customerEntitlementId, + ); + + if (!customerEntitlement) continue; + + await createAllocatedInvoice({ + ctx, + customerEntitlement, + oldFullCustomer, + update, + }); + + applyDeductionUpdateToFullSubject({ + fullSubject, + customerEntitlementId, + update, + }); + } + } catch (error) { + if (error instanceof Error && !error?.message?.includes("declined")) { + ctx.logger.error( + `[executeRedisDeductionV2] Attempting rollback due to error: ${error}`, + ); + } + await rollbackDeductionV2({ + ctx, + oldFullSubject, + updates, + }); + throw error; + } + + const featuresFromMutationLogs = mutationLogsToFeaturesV2({ + fullSubject, + mutationLogs, + }); + + const newFullCustomer = fullSubjectToFullCustomer({ fullSubject }); + + fireTrackWebhooks({ + ctx, + oldFullCus: oldFullCustomer, + newFullCus: newFullCustomer, + feature: deduction.feature, + entityId, + featuresFromMutationLogs, + }); + + if (options.triggerAutoTopUp) { + triggerAutoTopUp({ + ctx, + newFullCus: newFullCustomer, + feature: deduction.feature, + }).catch((error) => { + ctx.logger.error( + `[executeRedisDeductionV2] Failed to trigger auto top-up: ${error}`, + ); + }); + } + } + + return { + oldFullSubject, + fullSubject, + updates: allUpdates, + rolloverUpdates: allRolloverUpdates, + mutationLogs: allMutationLogs, + modifiedCusEntIdsByFeatureId: allModifiedCusEntIdsByFeatureId, + }; }; diff --git a/server/src/internal/balances/utils/deductionV2/index.ts b/server/src/internal/balances/utils/deductionV2/index.ts index 1eca353c9..4195d4102 100644 --- a/server/src/internal/balances/utils/deductionV2/index.ts +++ b/server/src/internal/balances/utils/deductionV2/index.ts @@ -1,5 +1,10 @@ export { applyDeductionUpdateToFullSubject } from "./applyDeductionUpdateToFullSubject.js"; export { applyRolloverUpdatesToFullSubject } from "./applyRolloverUpdatesToFullSubject.js"; export { deductionToTrackResponseV2 } from "./deductionToTrackResponseV2.js"; +export { executePostgresDeductionV2 } from "./executePostgresDeductionV2.js"; export { executeRedisDeductionV2 } from "./executeRedisDeductionV2.js"; +export { logDeductionUpdatesV2 } from "./logDeductionUpdatesV2.js"; +export { mutationLogsToFeaturesV2 } from "./mutationLogsToFeaturesV2.js"; +export { prepareDeductionOptionsV2 } from "./prepareDeductionOptionsV2.js"; export { prepareFeatureDeductionV2 } from "./prepareFeatureDeductionV2.js"; +export { rollbackDeductionV2 } from "./rollbackDeductionV2.js"; diff --git a/server/src/internal/balances/utils/deductionV2/logDeductionUpdatesV2.ts b/server/src/internal/balances/utils/deductionV2/logDeductionUpdatesV2.ts new file mode 100644 index 000000000..be886ddbe --- /dev/null +++ b/server/src/internal/balances/utils/deductionV2/logDeductionUpdatesV2.ts @@ -0,0 +1,49 @@ +import { + type FullSubject, + findCustomerEntitlementById, + fullSubjectToCustomerEntitlements, +} from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import type { DeductionUpdate } from "../types/deductionUpdate.js"; + +/** Logs deduction updates with customer entitlement details (FullSubject version). */ +export const logDeductionUpdatesV2 = ({ + ctx, + fullSubject, + updates, + source, +}: { + ctx: AutumnContext; + fullSubject: FullSubject; + updates: Record; + source?: string; +}): void => { + if (Object.keys(updates).length === 0) return; + + const customerEntitlements = fullSubjectToCustomerEntitlements({ + fullSubject, + }); + + for (const [customerEntitlementId, update] of Object.entries(updates)) { + const customerEntitlement = findCustomerEntitlementById({ + cusEnts: customerEntitlements, + id: customerEntitlementId, + }); + + const featureId = customerEntitlement?.entitlement.feature.id ?? "unknown"; + const entityScope = customerEntitlement?.entitlement.entity_feature_id + ? "entity" + : "customer"; + + ctx.logger.info(`[${source}] Deduction updates:`, { + data2: { + cusEntId: customerEntitlementId, + featureId, + entityScope, + balance: update.balance, + adjustment: update.adjustment, + entities: update.entities, + }, + }); + } +}; diff --git a/server/src/internal/balances/utils/deductionV2/mutationLogsToFeaturesV2.ts b/server/src/internal/balances/utils/deductionV2/mutationLogsToFeaturesV2.ts new file mode 100644 index 000000000..7f839d167 --- /dev/null +++ b/server/src/internal/balances/utils/deductionV2/mutationLogsToFeaturesV2.ts @@ -0,0 +1,46 @@ +import type { Feature, FullSubject } from "@autumn/shared"; +import { fullSubjectToCustomerEntitlements } from "@autumn/shared"; +import type { MutationLogItem } from "../types/mutationLogItem.js"; + +/** Maps customer entitlement and rollover mutation targets to their features (FullSubject version). */ +export const mutationLogsToFeaturesV2 = ({ + fullSubject, + mutationLogs, +}: { + fullSubject: FullSubject; + mutationLogs: MutationLogItem[]; +}): Feature[] => { + const customerEntitlements = fullSubjectToCustomerEntitlements({ + fullSubject, + }); + + const customerEntitlementIdToFeature = new Map(); + const rolloverIdToFeature = new Map(); + + for (const customerEntitlement of customerEntitlements) { + const feature = customerEntitlement.entitlement.feature; + customerEntitlementIdToFeature.set(customerEntitlement.id, feature); + for (const rollover of customerEntitlement.rollovers ?? []) { + rolloverIdToFeature.set(rollover.id, feature); + } + } + + const featuresById = new Map(); + + for (const log of mutationLogs) { + if ( + log.target_type === "customer_entitlement" && + log.customer_entitlement_id + ) { + const resolved = customerEntitlementIdToFeature.get( + log.customer_entitlement_id, + ); + if (resolved) featuresById.set(resolved.id, resolved); + } else if (log.target_type === "rollover" && log.rollover_id) { + const resolved = rolloverIdToFeature.get(log.rollover_id); + if (resolved) featuresById.set(resolved.id, resolved); + } + } + + return [...featuresById.values()]; +}; diff --git a/server/src/internal/balances/utils/deductionV2/prepareDeductionOptionsV2.ts b/server/src/internal/balances/utils/deductionV2/prepareDeductionOptionsV2.ts new file mode 100644 index 000000000..a9aca1354 --- /dev/null +++ b/server/src/internal/balances/utils/deductionV2/prepareDeductionOptionsV2.ts @@ -0,0 +1,48 @@ +import { + type FullSubject, + fullSubjectHasUsageBasedAllocated, + orgToInStatuses, +} from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import type { DeductionOptions } from "../types/deductionTypes.js"; +import type { FeatureDeduction } from "../types/featureDeduction.js"; + +/** Deduction options with all defaults resolved (no optional fields) */ +type ResolvedDeductionOptions = Required< + Omit +> & { + customerEntitlementFilters: DeductionOptions["customerEntitlementFilters"]; +}; + +/** + * Prepares deduction options with defaults and paidAllocated overrides. + * FullSubject version of prepareDeductionOptions. + */ +export const prepareDeductionOptionsV2 = ({ + ctx, + fullSubject, + options = {}, + deductions, +}: { + ctx: AutumnContext; + fullSubject: FullSubject; + options?: DeductionOptions; + deductions: FeatureDeduction[]; +}): ResolvedDeductionOptions => { + const isPaidAllocated = fullSubjectHasUsageBasedAllocated({ + fullSubject, + features: deductions.map((d) => d.feature), + inStatuses: orgToInStatuses({ org: ctx.org }), + }); + + return { + overageBehaviour: isPaidAllocated + ? "reject" + : (options.overageBehaviour ?? "cap"), + skipAdditionalBalance: true, + alterGrantedBalance: options.alterGrantedBalance ?? false, + customerEntitlementFilters: options.customerEntitlementFilters, + paidAllocated: isPaidAllocated, + triggerAutoTopUp: options.triggerAutoTopUp ?? false, + }; +}; diff --git a/server/src/internal/balances/utils/deductionV2/rollbackDeductionV2.ts b/server/src/internal/balances/utils/deductionV2/rollbackDeductionV2.ts new file mode 100644 index 000000000..85b55e8ac --- /dev/null +++ b/server/src/internal/balances/utils/deductionV2/rollbackDeductionV2.ts @@ -0,0 +1,65 @@ +import { + type FullSubject, + fullSubjectToCustomerEntitlements, +} from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import type { DeductionUpdate } from "@/internal/balances/utils/types/deductionUpdate.js"; +import { CusEntService } from "@/internal/customers/cusProducts/cusEnts/CusEntitlementService.js"; + +/** Rolls back deduction updates by restoring original entitlement values (FullSubject version). */ +export const rollbackDeductionV2 = async ({ + ctx, + oldFullSubject, + updates, +}: { + ctx: AutumnContext; + oldFullSubject: FullSubject; + updates: Record; +}) => { + const { logger } = ctx; + + logger.warn( + `[ROLLBACK] Starting rollback for ${Object.keys(updates).length} entitlements`, + ); + + const customerEntitlements = fullSubjectToCustomerEntitlements({ + fullSubject: oldFullSubject, + }); + + for (const customerEntitlementId of Object.keys(updates)) { + try { + const originalCustomerEntitlement = customerEntitlements.find( + (customerEntitlement) => + customerEntitlement.id === customerEntitlementId, + ); + + if (!originalCustomerEntitlement) { + logger.error( + `[ROLLBACK] Could not find original cusEnt ${customerEntitlementId} in oldFullSubject`, + ); + continue; + } + + await CusEntService.update({ + ctx, + id: customerEntitlementId, + updates: { + balance: originalCustomerEntitlement.balance ?? 0, + additional_balance: originalCustomerEntitlement.additional_balance, + adjustment: originalCustomerEntitlement.adjustment, + entities: originalCustomerEntitlement.entities, + }, + }); + + logger.info( + `[ROLLBACK] Successfully restored cusEnt ${customerEntitlementId} to original state`, + ); + } catch (error) { + logger.error( + `[ROLLBACK] Failed to rollback cusEnt ${customerEntitlementId}: ${error}`, + ); + } + } + + logger.warn("[ROLLBACK] Rollback completed"); +}; diff --git a/server/src/internal/balances/utils/deductionV2/syncDeductionUpdatesToFullSubjectCache.ts b/server/src/internal/balances/utils/deductionV2/syncDeductionUpdatesToFullSubjectCache.ts new file mode 100644 index 000000000..62ea93442 --- /dev/null +++ b/server/src/internal/balances/utils/deductionV2/syncDeductionUpdatesToFullSubjectCache.ts @@ -0,0 +1,138 @@ +import type { EntityRolloverBalance, FullSubject } from "@autumn/shared"; +import { redisV2 } from "@/external/redis/initRedisV2.js"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { buildSharedFullSubjectBalanceKey } from "@/internal/customers/cache/fullSubject/builders/buildSharedFullSubjectBalanceKey.js"; +import { FULL_SUBJECT_CACHE_TTL_SECONDS } from "@/internal/customers/cache/fullSubject/config/fullSubjectCacheConfig.js"; +import { tryRedisWrite } from "@/utils/cacheUtils/cacheUtils.js"; +import type { DeductionUpdate } from "../types/deductionUpdate.js"; + +interface RolloverOverwrite { + id: string; + cus_ent_id: string; + balance: number; + usage: number; + entities: Record; +} + +interface SubjectBalanceUpdate { + cus_ent_id: string; + balance: number | null; + additional_balance: number | null; + adjustment: number | null; + entities: Record | null; + next_reset_at: number | null; + expected_next_reset_at: number | null; + rollover_insert: unknown | null; + rollover_overwrites: RolloverOverwrite[] | null; + rollover_delete_ids: string[] | null; + new_replaceables: unknown[] | null; + deleted_replaceable_ids: string[] | null; +} + +/** + * Syncs deduction updates to the V2 FullSubject balance hashes. + * Groups updates by featureId and pipelines one Lua call per feature. + * Fire-and-forget — failures are logged but don't propagate. + */ +export const syncDeductionUpdatesToFullSubjectCache = async ({ + ctx, + customerId, + fullSubject, + cusEntUpdates, + rolloverOverwrites, + modifiedCusEntIdsByFeatureId, +}: { + ctx: AutumnContext; + customerId: string; + fullSubject: FullSubject; + cusEntUpdates: Record; + rolloverOverwrites: RolloverOverwrite[]; + modifiedCusEntIdsByFeatureId: Record; +}): Promise => { + try { + const rolloverOverwritesByCusEnt: Record = {}; + for (const rolloverOverwrite of rolloverOverwrites) { + if (!rolloverOverwritesByCusEnt[rolloverOverwrite.cus_ent_id]) { + rolloverOverwritesByCusEnt[rolloverOverwrite.cus_ent_id] = []; + } + rolloverOverwritesByCusEnt[rolloverOverwrite.cus_ent_id].push( + rolloverOverwrite, + ); + } + + // Build a lookup of cusEntId -> next_reset_at from the fullSubject + const cusEntNextResetAts: Record = {}; + for (const customerProduct of fullSubject.customer_products) { + for (const customerEntitlement of customerProduct.customer_entitlements) { + cusEntNextResetAts[customerEntitlement.id] = + customerEntitlement.next_reset_at ?? null; + } + } + for (const customerEntitlement of fullSubject.extra_customer_entitlements ?? + []) { + cusEntNextResetAts[customerEntitlement.id] = + customerEntitlement.next_reset_at ?? null; + } + + const { org, env } = ctx; + + // Group updates by featureId and build per-feature update arrays + const updatesByFeatureId: Record = {}; + + for (const [featureId, cusEntIds] of Object.entries( + modifiedCusEntIdsByFeatureId, + )) { + const featureUpdates: SubjectBalanceUpdate[] = []; + + for (const cusEntId of cusEntIds) { + const update = cusEntUpdates[cusEntId]; + if (!update) continue; + + featureUpdates.push({ + cus_ent_id: cusEntId, + balance: update.balance ?? null, + additional_balance: update.additional_balance ?? null, + adjustment: update.adjustment ?? null, + entities: update.entities ?? null, + next_reset_at: null, + expected_next_reset_at: cusEntNextResetAts[cusEntId] ?? null, + rollover_insert: null, + rollover_overwrites: rolloverOverwritesByCusEnt[cusEntId] ?? null, + rollover_delete_ids: null, + new_replaceables: update.newReplaceables ?? null, + deleted_replaceable_ids: + update.deletedReplaceables?.map((r) => r.id) ?? null, + }); + } + + if (featureUpdates.length > 0) { + updatesByFeatureId[featureId] = featureUpdates; + } + } + + if (Object.keys(updatesByFeatureId).length === 0) return; + + const pipeline = redisV2.pipeline(); + for (const [featureId, updates] of Object.entries(updatesByFeatureId)) { + const balanceKey = buildSharedFullSubjectBalanceKey({ + orgId: org.id, + env, + customerId, + featureId, + }); + pipeline.updateSubjectBalances( + balanceKey, + JSON.stringify({ + ttl_seconds: FULL_SUBJECT_CACHE_TTL_SECONDS, + updates, + }), + ); + } + + await tryRedisWrite(() => pipeline.exec(), redisV2); + } catch (error) { + ctx.logger.error( + `[syncDeductionUpdatesToFullSubjectCache] Failed to sync updates to cache: ${error}`, + ); + } +}; diff --git a/server/src/internal/balances/utils/lock/claimLockReceipt.ts b/server/src/internal/balances/utils/lock/claimLockReceipt.ts index 6840f4a31..e236fe9c4 100644 --- a/server/src/internal/balances/utils/lock/claimLockReceipt.ts +++ b/server/src/internal/balances/utils/lock/claimLockReceipt.ts @@ -1,37 +1,20 @@ import { ErrCode, InternalError, RecaseError } from "@autumn/shared"; import type { Redis } from "ioredis"; -import { - currentRegion, - getRegionalRedis, - redis, -} from "@/external/redis/initRedis.js"; import { tryRedisWrite } from "@/utils/cacheUtils/cacheUtils.js"; /** * Atomically claims a lock receipt: pending → processing. * - * Routes the claim to the Redis instance the receipt was originally written to - * (identified by receipt.region) so that Active-Active replication cannot allow - * two concurrent claims on separate regional instances. - * - * Returns the Redis instance that was used — callers must use it for all - * subsequent operations (unwind deduction, delete) to stay on the same instance. - * * Throws RecaseError for terminal/already-processing statuses. * Throws InternalError when Redis is unavailable. */ export const claimLockReceipt = async ({ lockReceiptKey, - receiptRegion, + redisInstance, }: { lockReceiptKey: string; - receiptRegion?: string | null; -}): Promise<{ redisInstance: Redis }> => { - const redisInstance = - receiptRegion && receiptRegion !== currentRegion - ? getRegionalRedis(receiptRegion) - : redis; - + redisInstance: Redis; +}): Promise => { const result = await tryRedisWrite( () => redisInstance.claimLockReceipt(lockReceiptKey), redisInstance, @@ -44,7 +27,7 @@ export const claimLockReceipt = async ({ } if (result === "OK") { - return { redisInstance }; + return; } throw new RecaseError({ diff --git a/server/src/internal/balances/utils/lock/deleteLockReceipt.ts b/server/src/internal/balances/utils/lock/deleteLockReceipt.ts index 64fb0d314..4df82353b 100644 --- a/server/src/internal/balances/utils/lock/deleteLockReceipt.ts +++ b/server/src/internal/balances/utils/lock/deleteLockReceipt.ts @@ -1,5 +1,4 @@ import type { Redis } from "ioredis"; -import { redis } from "@/external/redis/initRedis.js"; import { tryRedisWrite } from "@/utils/cacheUtils/cacheUtils.js"; /** Removes a lock receipt from Redis after successful finalize or expiry. */ @@ -8,11 +7,10 @@ export const deleteLockReceipt = async ({ redisInstance, }: { lockReceiptKey: string; - redisInstance?: Redis; + redisInstance: Redis; }): Promise => { - const targetRedis = redisInstance ?? redis; await tryRedisWrite( - () => targetRedis.del(lockReceiptKey) as Promise, + () => redisInstance.del(lockReceiptKey) as Promise, redisInstance, ); }; diff --git a/server/src/internal/balances/utils/lock/fetchLockReceipt.ts b/server/src/internal/balances/utils/lock/fetchLockReceipt.ts index 6b5829070..884e349f5 100644 --- a/server/src/internal/balances/utils/lock/fetchLockReceipt.ts +++ b/server/src/internal/balances/utils/lock/fetchLockReceipt.ts @@ -1,5 +1,6 @@ -import { ErrCode, RecaseError } from "@autumn/shared"; +import { ErrCode, InternalError, RecaseError } from "@autumn/shared"; import { redis } from "@/external/redis/initRedis.js"; +import { redisV2 } from "@/external/redis/initRedisV2.js"; import type { AutumnContext } from "@/honoUtils/HonoEnv"; import type { MutationLogItem } from "@/internal/balances/utils/types/mutationLogItem.js"; import { tryRedisRead } from "@/utils/cacheUtils/cacheUtils.js"; @@ -15,6 +16,8 @@ export type LockReceipt = { items: MutationLogItem[]; }; +export type LockReceiptSource = "redis_v1" | "redis_v2"; + const normalizeLockReceiptItems = ({ items, lockId, @@ -50,9 +53,28 @@ export const fetchLockReceipt = async ({ lockKey: hashedKey, }); - const rawReceipt = await tryRedisRead( - () => redis.call("JSON.GET", lockReceiptKey, "$") as Promise, - ); + const [rawReceiptV1, rawReceiptV2] = await Promise.all([ + tryRedisRead( + () => + redis.call("JSON.GET", lockReceiptKey, "$") as Promise, + redis, + ), + tryRedisRead( + () => + redisV2.call("JSON.GET", lockReceiptKey, "$") as Promise, + redisV2, + ), + ]); + + if (rawReceiptV1 && rawReceiptV2) { + throw new InternalError({ + message: `Lock receipt found in both Redis stores for ID: ${lockId}`, + code: "lock_receipt_found_in_both_stores", + }); + } + + const rawReceipt = rawReceiptV2 ?? rawReceiptV1; + const source: LockReceiptSource = rawReceiptV2 ? "redis_v2" : "redis_v1"; if (!rawReceipt) { throw new RecaseError({ @@ -91,5 +113,6 @@ export const fetchLockReceipt = async ({ return { receipt, lockReceiptKey, + source, }; }; diff --git a/server/src/internal/balances/utils/lockV2/buildFinalizeLockContextV2.ts b/server/src/internal/balances/utils/lockV2/buildFinalizeLockContextV2.ts new file mode 100644 index 000000000..66ca88824 --- /dev/null +++ b/server/src/internal/balances/utils/lockV2/buildFinalizeLockContextV2.ts @@ -0,0 +1,82 @@ +import type { Feature, FullSubject } from "@autumn/shared"; +import { type FinalizeLockParamsV0, findFeatureById } from "@autumn/shared"; +import type { Redis } from "ioredis"; +import { redisV2 } from "@/external/redis/initRedisV2.js"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import type { LockReceipt } from "@/internal/balances/utils/lock/fetchLockReceipt.js"; +import { + calculateLockValue, + calculateUnwindValue, +} from "@/internal/balances/utils/lock/unwindLockUtils.js"; +import type { FeatureDeduction } from "@/internal/balances/utils/types/featureDeduction.js"; +import { getOrSetCachedFullSubject } from "@/internal/customers/cache/fullSubject/actions/getOrSetCachedFullSubject.js"; + +export type FinalizeLockContextV2 = { + receipt: LockReceipt; + lockReceiptKey: string; + redisInstance: Redis; + fullSubject: FullSubject; + feature: Feature; + lockValue: number; + finalValue: number; + unwindValue: number; + additionalValue: number; + properties?: Record; + deduction: FeatureDeduction; + deductionOptions: { triggerAutoTopUp: boolean }; +}; + +export const buildFinalizeLockContextV2 = async ({ + ctx, + params, + receipt, + lockReceiptKey, +}: { + ctx: AutumnContext; + params: FinalizeLockParamsV0; + receipt: LockReceipt; + lockReceiptKey: string; +}): Promise => { + const fullSubject = await getOrSetCachedFullSubject({ + ctx, + customerId: receipt.customer_id, + entityId: receipt.entity_id ?? undefined, + source: "runFinalizeLockV2", + }); + + const lockValue = calculateLockValue({ items: receipt.items }); + const finalValue = + params.action === "release" ? 0 : (params.override_value ?? lockValue); + + const { unwindValue, additionalValue } = calculateUnwindValue({ + receipt, + finalValue, + }); + + const feature = findFeatureById({ + features: ctx.features, + featureId: receipt.feature_id, + errorOnNotFound: true, + }); + + return { + receipt, + lockReceiptKey, + redisInstance: redisV2, + fullSubject, + feature, + lockValue, + finalValue, + unwindValue, + additionalValue, + properties: params.properties, + deduction: { + feature, + deduction: additionalValue, + lockReceipt: receipt, + unwindValue, + lockReceiptKey, + }, + deductionOptions: { triggerAutoTopUp: true }, + }; +}; diff --git a/server/src/internal/balances/utils/sync/SyncBatchingManagerV3.ts b/server/src/internal/balances/utils/sync/SyncBatchingManagerV3.ts new file mode 100644 index 000000000..7c55d7790 --- /dev/null +++ b/server/src/internal/balances/utils/sync/SyncBatchingManagerV3.ts @@ -0,0 +1,310 @@ +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"; + +interface CustomerBatchContext { + customerId: string; + orgId: string; + env: AppEnv; + region: string; + timestamp: number; + cusEntIds: Set; + rolloverIds: Set; + entityId?: string; + modifiedCusEntIdsByFeatureId: Record; +} + +interface CustomerBatch { + context: CustomerBatchContext; + timer: NodeJS.Timeout | null; +} + +export type QueueSyncV4Payload = { + jobName: string; + payload: { + customerId: string; + orgId: string; + env: AppEnv; + region: string; + timestamp: number; + cusEntIds: string[]; + rolloverIds: string[]; + entityId?: string; + modifiedCusEntIdsByFeatureId: Record; + }; + messageGroupId?: string; + messageDeduplicationId: string; +}; + +/** + * Batches v4 sync jobs per customer using a fixed tumbling window. + * Always requires modifiedCusEntIdsByFeatureId — never mixes with v3 sync items. + */ +export class SyncBatchingManagerV3 { + private customerBatches: Map = new Map(); + + private readonly BATCH_WINDOW_MS: number = 1000; + private readonly MAX_BATCH_SIZE = 1000; + private readonly DEDUP_BUCKET_MS: number = 2500; + + private readonly _addTaskToQueue: (args: QueueSyncV4Payload) => Promise; + + constructor({ + addTaskToQueueFn, + batchWindowMs, + dedupBucketMs, + }: { + addTaskToQueueFn?: (args: QueueSyncV4Payload) => Promise; + batchWindowMs?: number; + dedupBucketMs?: number; + } = {}) { + this._addTaskToQueue = + addTaskToQueueFn ?? + (addTaskToQueue as unknown as ( + args: QueueSyncV4Payload, + ) => Promise); + this.BATCH_WINDOW_MS = batchWindowMs ?? 1000; + this.DEDUP_BUCKET_MS = dedupBucketMs ?? 2500; + } + + addSyncItem({ + customerId, + orgId, + env, + cusEntIds, + rolloverIds, + region, + entityId, + modifiedCusEntIdsByFeatureId, + }: { + customerId: string; + orgId: string; + env: AppEnv; + cusEntIds: string[]; + rolloverIds?: string[]; + region?: string; + entityId?: string; + modifiedCusEntIdsByFeatureId: Record; + }): void { + const batchKey = this.buildBatchKey({ orgId, env, customerId }); + let batch = this.customerBatches.get(batchKey); + + if (!batch) { + batch = this.createBatch({ customerId, orgId, env, region }); + this.customerBatches.set(batchKey, batch); + this.scheduleCustomerBatch({ batchKey }); + } + + this.mergeCusEntIds({ batch, cusEntIds }); + this.mergeRolloverIds({ batch, rolloverIds: rolloverIds ?? [] }); + + if (region) batch.context.region = region; + if (entityId) batch.context.entityId = entityId; + + for (const [featureId, ids] of Object.entries( + modifiedCusEntIdsByFeatureId, + )) { + if (!batch.context.modifiedCusEntIdsByFeatureId[featureId]) { + batch.context.modifiedCusEntIdsByFeatureId[featureId] = []; + } + batch.context.modifiedCusEntIdsByFeatureId[featureId].push(...ids); + } + + const totalSize = + batch.context.cusEntIds.size + batch.context.rolloverIds.size; + if (totalSize >= this.MAX_BATCH_SIZE) { + this.executeCustomerBatch({ batchKey }); + } + } + + getStats(): { + totalCustomers: number; + totalPendingEntitlements: number; + totalPendingRollovers: number; + } { + let totalEntitlements = 0; + let totalRollovers = 0; + for (const batch of this.customerBatches.values()) { + totalEntitlements += batch.context.cusEntIds.size; + totalRollovers += batch.context.rolloverIds.size; + } + return { + totalCustomers: this.customerBatches.size, + totalPendingEntitlements: totalEntitlements, + totalPendingRollovers: totalRollovers, + }; + } + + async flush(): Promise { + const batchKeys = Array.from(this.customerBatches.keys()); + await Promise.all( + batchKeys.map((batchKey) => this.executeCustomerBatch({ batchKey })), + ); + } + + private buildBatchKey({ + orgId, + env, + customerId, + }: { + orgId: string; + env: AppEnv; + customerId: string; + }): string { + return `${orgId}:${env}:${customerId}`; + } + + private createBatch({ + customerId, + orgId, + env, + region, + }: { + customerId: string; + orgId: string; + env: AppEnv; + region?: string; + }): CustomerBatch { + return { + context: { + customerId, + orgId, + env, + region: region || currentRegion, + timestamp: Date.now(), + cusEntIds: new Set(), + rolloverIds: new Set(), + modifiedCusEntIdsByFeatureId: {}, + }, + timer: null, + }; + } + + private mergeCusEntIds({ + batch, + cusEntIds, + }: { + batch: CustomerBatch; + cusEntIds: string[]; + }): void { + for (const id of cusEntIds) { + batch.context.cusEntIds.add(id); + } + } + + private mergeRolloverIds({ + batch, + rolloverIds, + }: { + batch: CustomerBatch; + rolloverIds: string[]; + }): void { + for (const id of rolloverIds) { + batch.context.rolloverIds.add(id); + } + } + + private scheduleCustomerBatch({ batchKey }: { batchKey: string }): void { + const batch = this.customerBatches.get(batchKey); + if (!batch) return; + + batch.timer = setTimeout(() => { + this.executeCustomerBatch({ batchKey }); + }, this.BATCH_WINDOW_MS); + + if (batch.timer.unref) { + batch.timer.unref(); + } + } + + private async executeCustomerBatch({ + batchKey, + }: { + batchKey: string; + }): Promise { + const batch = this.customerBatches.get(batchKey); + if (!batch) return; + + this.clearBatchTimer({ batch }); + this.customerBatches.delete(batchKey); + + const { context } = batch; + if (context.cusEntIds.size === 0 && context.rolloverIds.size === 0) return; + + await this.queueSyncJob({ context }); + } + + private clearBatchTimer({ batch }: { batch: CustomerBatch }): void { + if (batch.timer) { + clearTimeout(batch.timer); + batch.timer = null; + } + } + + private buildDeduplicationId({ + context, + cusEntIds, + rolloverIds, + }: { + context: CustomerBatchContext; + cusEntIds: string[]; + rolloverIds: string[]; + }): string { + const dedupBucket = Math.floor(Date.now() / this.DEDUP_BUCKET_MS); + const dedupKey = JSON.stringify({ + jobName: JobName.SyncBalanceBatchV4, + orgId: context.orgId, + env: context.env, + customerId: context.customerId, + cusEntIds, + rolloverIds, + dedupBucket, + }); + + return Bun.hash(dedupKey).toString(); + } + + private async queueSyncJob({ + context, + }: { + context: CustomerBatchContext; + }): Promise { + const cusEntIds = Array.from(context.cusEntIds).sort(); + const rolloverIds = Array.from(context.rolloverIds).sort(); + const messageDeduplicationId = this.buildDeduplicationId({ + context, + cusEntIds, + rolloverIds, + }); + + try { + await this._addTaskToQueue({ + jobName: JobName.SyncBalanceBatchV4, + payload: { + customerId: context.customerId, + orgId: context.orgId, + env: context.env, + region: context.region, + timestamp: Date.now(), + cusEntIds, + rolloverIds, + entityId: context.entityId, + modifiedCusEntIdsByFeatureId: context.modifiedCusEntIdsByFeatureId, + }, + messageDeduplicationId, + }); + + logger.info( + `[SyncV4] Queued sync for ${context.customerId}, ${cusEntIds.length} entitlements, ${rolloverIds.length} rollovers`, + ); + } catch (error) { + logger.error( + `[SyncV4] Failed to queue sync for ${context.customerId}: ${error}`, + ); + } + } +} + +export const globalSyncBatchingManagerV3 = new SyncBatchingManagerV3(); diff --git a/server/src/internal/balances/utils/sync/syncItemV4.ts b/server/src/internal/balances/utils/sync/syncItemV4.ts new file mode 100644 index 000000000..715efee33 --- /dev/null +++ b/server/src/internal/balances/utils/sync/syncItemV4.ts @@ -0,0 +1,240 @@ +import { + type EntityBalance, + type EntityRolloverBalance, + type SubjectBalance, + tryCatch, +} from "@autumn/shared"; +import { sql } from "drizzle-orm"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { invalidateCachedFullSubject } from "@/internal/customers/cache/fullSubject/actions/invalidate/invalidateFullSubject.js"; +import { getCachedFeatureBalance } from "@/internal/customers/cache/fullSubject/balances/getCachedFeatureBalances.js"; + +const SYNC_CONFLICT_CODES = { + ResetAtMismatch: "RESET_AT_MISMATCH", + EntityCountMismatch: "ENTITY_COUNT_MISMATCH", + CacheVersionMismatch: "CACHE_VERSION_MISMATCH", +} as const; + +const handleSyncPostgresError = async ({ + error, + customerId, + entityId, + ctx, +}: { + error: Error; + customerId: string; + entityId?: string; + ctx: AutumnContext; +}): Promise => { + const message = error.message || ""; + const isConflict = + message.includes(SYNC_CONFLICT_CODES.ResetAtMismatch) || + message.includes(SYNC_CONFLICT_CODES.EntityCountMismatch) || + message.includes(SYNC_CONFLICT_CODES.CacheVersionMismatch); + + if (!isConflict) throw error; + + let code: string = SYNC_CONFLICT_CODES.EntityCountMismatch; + if (message.includes(SYNC_CONFLICT_CODES.ResetAtMismatch)) { + code = SYNC_CONFLICT_CODES.ResetAtMismatch; + } else if (message.includes(SYNC_CONFLICT_CODES.CacheVersionMismatch)) { + code = SYNC_CONFLICT_CODES.CacheVersionMismatch; + } + const cusEntMatch = message.match(/cus_ent_id:(\S+)/); + const cusEntId = cusEntMatch?.[1]; + + ctx.logger.warn( + `[SYNC V4] (${customerId}) Sync conflict detected: ${code}, cus_ent: ${cusEntId}. Invalidating cache.`, + ); + + await invalidateCachedFullSubject({ + ctx, + customerId, + entityId, + source: `sync-conflict-${code}`, + }); + + return true; +}; + +interface SyncItemV4 { + customerId: string; + entityId?: string; + orgId: string; + env: string; + timestamp: number; + rolloverIds?: string[]; + modifiedCusEntIdsByFeatureId: Record; +} + +interface SyncEntry { + customer_entitlement_id: string; + feature_id: string; + balance: number; + adjustment: number; + entities: Record | null; + next_reset_at: number | null; + entity_count: number; + cache_version: number | null; +} + +interface RolloverSyncEntry { + rollover_id: string; + balance: number; + usage: number; + entities: Record | null; +} + +const subjectBalanceToSyncEntry = ({ + subjectBalance, +}: { + subjectBalance: SubjectBalance; +}): SyncEntry => ({ + customer_entitlement_id: subjectBalance.id, + feature_id: subjectBalance.feature_id, + balance: subjectBalance.balance ?? 0, + adjustment: subjectBalance.adjustment ?? 0, + entities: subjectBalance.entities ?? null, + next_reset_at: subjectBalance.next_reset_at ?? null, + entity_count: subjectBalance.entities + ? Object.keys(subjectBalance.entities).length + : 0, + cache_version: subjectBalance.cache_version ?? 0, +}); + +const formatSyncEntry = ({ entry }: { entry: SyncEntry }): string => { + const hasEntities = entry.entities && Object.keys(entry.entities).length > 0; + const entitiesStr = hasEntities + ? `, entities= ${Object.keys(entry.entities!).length}` + : ""; + return `${entry.feature_id} (${entry.customer_entitlement_id}): bal= ${entry.balance}, adj= ${entry.adjustment}${entitiesStr}`; +}; + +const formatRolloverSyncEntry = ({ + entry, +}: { + entry: RolloverSyncEntry; +}): string => { + const hasEntities = entry.entities && Object.keys(entry.entities).length > 0; + const entitiesStr = hasEntities + ? `, entities= ${Object.keys(entry.entities!).length}` + : ""; + return `rollover ${entry.rollover_id}: bal= ${entry.balance}, usage= ${entry.usage}${entitiesStr}`; +}; + +/** Sync cached subject balances to Postgres using targeted hash reads. */ +export const syncItemV4 = async ({ + ctx, + payload, +}: { + ctx: AutumnContext; + payload: SyncItemV4; +}): Promise => { + const { + customerId, + entityId, + orgId, + env, + rolloverIds, + modifiedCusEntIdsByFeatureId, + } = payload; + const { db, logger } = ctx; + + // Read targeted balance hashes + const allSubjectBalances: SubjectBalance[] = []; + for (const [featureId, customerEntitlementIds] of Object.entries( + modifiedCusEntIdsByFeatureId, + )) { + const result = await getCachedFeatureBalance({ + orgId, + env, + customerId, + featureId, + customerEntitlementIds, + }); + + if (!result) { + logger.info( + `[SYNC V4] (${customerId}) Cache miss for feature=${featureId}, skipping`, + ); + return; + } + + allSubjectBalances.push(...result.balances); + } + + // Build sync entries + const entries: SyncEntry[] = allSubjectBalances.map((subjectBalance) => + subjectBalanceToSyncEntry({ subjectBalance }), + ); + + // Build rollover sync entries + const rolloverEntries: RolloverSyncEntry[] = []; + if (rolloverIds && rolloverIds.length > 0) { + const rolloverIdSet = new Set(rolloverIds); + for (const subjectBalance of allSubjectBalances) { + if (!subjectBalance.rollovers) continue; + for (const rollover of subjectBalance.rollovers) { + if (rolloverIdSet.has(rollover.id)) { + rolloverEntries.push({ + rollover_id: rollover.id, + balance: rollover.balance ?? 0, + usage: rollover.usage ?? 0, + entities: rollover.entities ?? null, + }); + } + } + } + } + + if (entries.length === 0 && rolloverEntries.length === 0) { + logger.info(`[SYNC V4] (${customerId}) No entries to sync`); + return; + } + + for (const entry of entries) { + logger.info(`[SYNC V4] (${customerId}) ${formatSyncEntry({ entry })}`); + } + for (const entry of rolloverEntries) { + logger.info( + `[SYNC V4] (${customerId}) ${formatRolloverSyncEntry({ entry })}`, + ); + } + + const { data: result, error } = await tryCatch( + db.execute( + sql`SELECT * FROM sync_balances_v2(${JSON.stringify({ + customer_entitlement_updates: entries, + rollover_updates: rolloverEntries, + })}::jsonb)`, + ), + ); + + if (error) { + await handleSyncPostgresError({ + error, + customerId, + entityId, + ctx, + }); + return; + } + + const syncResult = result[0]?.sync_balances_v2 as + | { + updates?: Record; + rollover_updates?: Record; + } + | undefined; + + const updateCount = syncResult?.updates + ? Object.keys(syncResult.updates).length + : 0; + const rolloverUpdateCount = syncResult?.rollover_updates + ? Object.keys(syncResult.rollover_updates).length + : 0; + + logger.info( + `[SYNC V4] (${customerId}) Done: ${updateCount} cus_ents, ${rolloverUpdateCount} rollovers updated`, + ); +}; diff --git a/server/src/internal/balances/utils/types/redisDeductionError.ts b/server/src/internal/balances/utils/types/redisDeductionError.ts index a39c2f606..5339ace26 100644 --- a/server/src/internal/balances/utils/types/redisDeductionError.ts +++ b/server/src/internal/balances/utils/types/redisDeductionError.ts @@ -2,6 +2,7 @@ export enum RedisDeductionErrorCode { CustomerNotFound = "CUSTOMER_NOT_FOUND", NoCustomerProducts = "NO_CUSTOMER_PRODUCTS", + SubjectBalanceNotFound = "SUBJECT_BALANCE_NOT_FOUND", InsufficientBalance = "INSUFFICIENT_BALANCE", PaidAllocated = "PAID_ALLOCATED", SkipCache = "SKIP_CACHE", @@ -12,6 +13,7 @@ export enum RedisDeductionErrorCode { export const FALLBACK_ERROR_CODES = [ RedisDeductionErrorCode.CustomerNotFound, RedisDeductionErrorCode.NoCustomerProducts, + RedisDeductionErrorCode.SubjectBalanceNotFound, RedisDeductionErrorCode.PaidAllocated, RedisDeductionErrorCode.SkipCache, ] as const; diff --git a/server/src/internal/customers/actions/getApiCustomerByRollout.ts b/server/src/internal/customers/actions/getApiCustomerByRollout.ts index 32599fc23..2e097f824 100644 --- a/server/src/internal/customers/actions/getApiCustomerByRollout.ts +++ b/server/src/internal/customers/actions/getApiCustomerByRollout.ts @@ -1,28 +1,28 @@ -import type { CheckParams, TrackParams } from "@autumn/shared"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; -import { getOrCreateCachedFullSubject } from "@/internal/customers/cache/fullSubject/index.js"; +import { getOrSetCachedFullSubject } from "@/internal/customers/cache/fullSubject/index.js"; import { isFullSubjectRolloutEnabled } from "@/internal/misc/rollouts/fullSubjectRolloutUtils.js"; import { getApiCustomer } from "../cusUtils/apiCusUtils/getApiCustomer.js"; -import { getOrCreateCachedFullCustomer } from "../cusUtils/fullCustomerCacheUtils/getOrCreateCachedFullCustomer.js"; +import { getOrSetCachedFullCustomer } from "../cusUtils/fullCustomerCacheUtils/getOrSetCachedFullCustomer.js"; import { getApiCustomerV2 } from "../cusUtils/getApiCustomerV2/index.js"; export const getApiCustomerByRollout = async ({ ctx, - params, + customerId, + entityId, source, withAutumnId, }: { ctx: AutumnContext; - params: Omit & { - customer_id: string | null; - }; + customerId: string; + entityId?: string; source?: string; withAutumnId?: boolean; }) => { if (isFullSubjectRolloutEnabled({ ctx })) { - const fullSubject = await getOrCreateCachedFullSubject({ + const fullSubject = await getOrSetCachedFullSubject({ ctx, - params, + customerId, + entityId, source, }); @@ -33,9 +33,10 @@ export const getApiCustomerByRollout = async ({ }); } - const fullCustomer = await getOrCreateCachedFullCustomer({ + const fullCustomer = await getOrSetCachedFullCustomer({ ctx, - params, + customerId, + entityId, source, }); diff --git a/server/src/internal/customers/actions/getOrCreateApiCustomerByRollout.ts b/server/src/internal/customers/actions/getOrCreateApiCustomerByRollout.ts new file mode 100644 index 000000000..27572a7d8 --- /dev/null +++ b/server/src/internal/customers/actions/getOrCreateApiCustomerByRollout.ts @@ -0,0 +1,47 @@ +import type { CheckParams, TrackParams } from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { getOrCreateCachedFullSubject } from "@/internal/customers/cache/fullSubject/index.js"; +import { isFullSubjectRolloutEnabled } from "@/internal/misc/rollouts/fullSubjectRolloutUtils.js"; +import { getApiCustomer } from "../cusUtils/apiCusUtils/getApiCustomer.js"; +import { getOrCreateCachedFullCustomer } from "../cusUtils/fullCustomerCacheUtils/getOrCreateCachedFullCustomer.js"; +import { getApiCustomerV2 } from "../cusUtils/getApiCustomerV2/index.js"; + +export const getOrCreateApiCustomerByRollout = async ({ + ctx, + params, + source, + withAutumnId, +}: { + ctx: AutumnContext; + params: Omit & { + customer_id: string | null; + }; + source?: string; + withAutumnId?: boolean; +}) => { + if (isFullSubjectRolloutEnabled({ ctx })) { + const fullSubject = await getOrCreateCachedFullSubject({ + ctx, + params, + source, + }); + + return getApiCustomerV2({ + ctx, + fullSubject, + withAutumnId, + }); + } + + const fullCustomer = await getOrCreateCachedFullCustomer({ + ctx, + params, + source, + }); + + return getApiCustomer({ + ctx, + fullCustomer, + withAutumnId, + }); +}; diff --git a/server/src/internal/customers/actions/update/updateCustomer.ts b/server/src/internal/customers/actions/update/updateCustomer.ts index ecb264052..f13df5b8d 100644 --- a/server/src/internal/customers/actions/update/updateCustomer.ts +++ b/server/src/internal/customers/actions/update/updateCustomer.ts @@ -11,7 +11,6 @@ import type Stripe from "stripe"; import { createStripeCli } from "@/external/connect/createStripeCli"; import type { AutumnContext } from "@/honoUtils/HonoEnv"; import { CusService } from "@/internal/customers/CusService"; -import { updateCachedCustomerData } from "../../cusUtils/fullCustomerCacheUtils/updateCachedCustomerData"; export const updateCustomer = async ({ ctx, @@ -135,12 +134,5 @@ export const updateCustomer = async ({ update: updateData, }); - await updateCachedCustomerData({ - ctx, - customerId: originalCustomer.id || originalCustomer.internal_id, - newCustomerId: newCustomerId ?? undefined, - updates: updateData, - }); - return newCustomerId ?? customerId; }; diff --git a/server/src/internal/customers/cache/fullSubject/actions/getCachedFullSubject.ts b/server/src/internal/customers/cache/fullSubject/actions/getCachedFullSubject.ts index dd2f311f5..75553994d 100644 --- a/server/src/internal/customers/cache/fullSubject/actions/getCachedFullSubject.ts +++ b/server/src/internal/customers/cache/fullSubject/actions/getCachedFullSubject.ts @@ -10,6 +10,7 @@ import { type CachedFullSubject, cachedFullSubjectToNormalized, } from "../fullSubjectCacheModel.js"; +import { sanitizeCachedFullSubject } from "../sanitize/index.js"; import { getOrInitFullSubjectViewEpoch } from "./invalidate/getOrInitFullSubjectViewEpoch.js"; import { invalidateCachedFullSubject } from "./invalidate/invalidateFullSubject.js"; import { invalidateCachedFullSubjectExact } from "./invalidate/invalidateFullSubjectExact.js"; @@ -38,11 +39,20 @@ export const getCachedFullSubject = async ({ let cached: CachedFullSubject; try { - cached = JSON.parse(cachedRaw) as CachedFullSubject; + const parsedCached = JSON.parse(cachedRaw) as CachedFullSubject; + cached = sanitizeCachedFullSubject({ + cachedFullSubject: parsedCached, + }); } catch (error) { logger.warn( `[getCachedFullSubject] Failed to parse cached subject for ${customerId}${entityId ? `:${entityId}` : ""}, source: ${source}, error: ${error}`, ); + await invalidateCachedFullSubject({ + ctx, + customerId, + entityId, + source: "parse-failed", + }); return undefined; } @@ -104,10 +114,23 @@ export const getCachedFullSubject = async ({ return undefined; } - const normalized = cachedFullSubjectToNormalized({ - cached, - customerEntitlements: balances.flatMap((balance) => balance.balances), - }); + try { + const normalized = cachedFullSubjectToNormalized({ + cached, + customerEntitlements: balances.flatMap((balance) => balance.balances), + }); - return normalizedToFullSubject({ normalized }); + return normalizedToFullSubject({ normalized }); + } catch (error) { + logger.warn( + `[getCachedFullSubject] Failed to hydrate cached subject for ${customerId}${entityId ? `:${entityId}` : ""}, source: ${source}, error: ${error}`, + ); + await invalidateCachedFullSubjectExact({ + ctx, + customerId, + entityId, + source: "hydrate-failed", + }); + return undefined; + } }; diff --git a/server/src/internal/customers/cache/fullSubject/actions/invalidate/getOrInitFullSubjectViewEpoch.ts b/server/src/internal/customers/cache/fullSubject/actions/invalidate/getOrInitFullSubjectViewEpoch.ts index ffdba782a..260fcb976 100644 --- a/server/src/internal/customers/cache/fullSubject/actions/invalidate/getOrInitFullSubjectViewEpoch.ts +++ b/server/src/internal/customers/cache/fullSubject/actions/invalidate/getOrInitFullSubjectViewEpoch.ts @@ -2,6 +2,7 @@ import { redisV2 } from "@/external/redis/initRedisV2.js"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { tryRedisRead, tryRedisWrite } from "@/utils/cacheUtils/cacheUtils.js"; import { buildFullSubjectViewEpochKey } from "../../builders/buildFullSubjectViewEpochKey.js"; +import { FULL_SUBJECT_EPOCH_TTL_SECONDS } from "../../config/fullSubjectCacheConfig.js"; export const getOrInitFullSubjectViewEpoch = async ({ ctx, @@ -18,11 +19,18 @@ export const getOrInitFullSubjectViewEpoch = async ({ const currentEpoch = await tryRedisRead(() => redisV2.get(epochKey), redisV2); if (currentEpoch !== null && currentEpoch !== undefined) { + await tryRedisWrite( + () => redisV2.expire(epochKey, FULL_SUBJECT_EPOCH_TTL_SECONDS), + redisV2, + ); const parsedEpoch = Number.parseInt(currentEpoch, 10); return Number.isNaN(parsedEpoch) ? 0 : parsedEpoch; } - await tryRedisWrite(() => redisV2.setnx(epochKey, "0"), redisV2); + await tryRedisWrite( + () => redisV2.set(epochKey, "0", "EX", FULL_SUBJECT_EPOCH_TTL_SECONDS), + redisV2, + ); const initializedEpoch = await tryRedisRead( () => redisV2.get(epochKey), redisV2, diff --git a/server/src/internal/customers/cache/fullSubject/actions/invalidate/incrementFullSubjectViewEpoch.ts b/server/src/internal/customers/cache/fullSubject/actions/invalidate/incrementFullSubjectViewEpoch.ts index 79c31c0f1..2a617d0ce 100644 --- a/server/src/internal/customers/cache/fullSubject/actions/invalidate/incrementFullSubjectViewEpoch.ts +++ b/server/src/internal/customers/cache/fullSubject/actions/invalidate/incrementFullSubjectViewEpoch.ts @@ -2,6 +2,7 @@ import { redisV2 } from "@/external/redis/initRedisV2.js"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { tryRedisWrite } from "@/utils/cacheUtils/cacheUtils.js"; import { buildFullSubjectViewEpochKey } from "../../builders/buildFullSubjectViewEpochKey.js"; +import { FULL_SUBJECT_EPOCH_TTL_SECONDS } from "../../config/fullSubjectCacheConfig.js"; export const incrementFullSubjectViewEpoch = async ({ ctx, @@ -18,5 +19,9 @@ export const incrementFullSubjectViewEpoch = async ({ const nextEpoch = await tryRedisWrite(() => redisV2.incr(epochKey), redisV2); if (nextEpoch === null || nextEpoch === undefined) return null; + await tryRedisWrite( + () => redisV2.expire(epochKey, FULL_SUBJECT_EPOCH_TTL_SECONDS), + redisV2, + ); return nextEpoch; }; diff --git a/server/src/internal/customers/cache/fullSubject/actions/partial/getCachedPartialFullSubject.ts b/server/src/internal/customers/cache/fullSubject/actions/partial/getCachedPartialFullSubject.ts index aee520f9a..714fe80af 100644 --- a/server/src/internal/customers/cache/fullSubject/actions/partial/getCachedPartialFullSubject.ts +++ b/server/src/internal/customers/cache/fullSubject/actions/partial/getCachedPartialFullSubject.ts @@ -12,6 +12,7 @@ import { type CachedFullSubject, cachedFullSubjectToNormalized, } from "../../fullSubjectCacheModel.js"; +import { sanitizeCachedFullSubject } from "../../sanitize/index.js"; import { getOrInitFullSubjectViewEpoch } from "../invalidate/getOrInitFullSubjectViewEpoch.js"; import { invalidateCachedFullSubject } from "../invalidate/invalidateFullSubject.js"; import { invalidateCachedFullSubjectExact } from "../invalidate/invalidateFullSubjectExact.js"; @@ -61,7 +62,10 @@ export const getCachedPartialFullSubject = async ({ let cached: CachedFullSubject; try { - cached = JSON.parse(cachedRaw) as CachedFullSubject; + const parsedCached = JSON.parse(cachedRaw) as CachedFullSubject; + cached = sanitizeCachedFullSubject({ + cachedFullSubject: parsedCached, + }); } catch (error) { logger.warn( `[getCachedPartialFullSubject] Failed to parse cached subject for ${customerId}${entityId ? `:${entityId}` : ""}, source: ${source}, error: ${error}`, @@ -144,13 +148,26 @@ export const getCachedPartialFullSubject = async ({ (featureBalance) => featureBalance.balances, ); - const normalized = filterNormalizedFullSubjectByFeatureIds({ - normalized: cachedFullSubjectToNormalized({ - cached, - customerEntitlements, - }), - featureIds, - }); + try { + const normalized = filterNormalizedFullSubjectByFeatureIds({ + normalized: cachedFullSubjectToNormalized({ + cached, + customerEntitlements, + }), + featureIds, + }); - return normalizedToFullSubject({ normalized }); + return normalizedToFullSubject({ normalized }); + } catch (error) { + logger.warn( + `[getCachedPartialFullSubject] Failed to hydrate cached subject for ${customerId}${entityId ? `:${entityId}` : ""}, source: ${source}, error: ${error}`, + ); + await invalidateCachedFullSubjectExact({ + ctx, + customerId, + entityId, + source: "partial-hydrate-failed", + }); + return undefined; + } }; diff --git a/server/src/internal/customers/cache/fullSubject/actions/setCachedFullSubject/setSharedFullSubjectBalances.ts b/server/src/internal/customers/cache/fullSubject/actions/setCachedFullSubject/setSharedFullSubjectBalances.ts index d23b963e7..d62a8b0a1 100644 --- a/server/src/internal/customers/cache/fullSubject/actions/setCachedFullSubject/setSharedFullSubjectBalances.ts +++ b/server/src/internal/customers/cache/fullSubject/actions/setCachedFullSubject/setSharedFullSubjectBalances.ts @@ -49,7 +49,7 @@ export const appendSharedFullSubjectBalanceWrite = async ({ multi, normalized, meteredFeatures: _meteredFeatures, - overwrite: _overwrite, + overwrite, ttlSeconds, }: { ctx: AutumnContext; @@ -70,7 +70,13 @@ export const appendSharedFullSubjectBalanceWrite = async ({ for (const { balanceKey, fields } of balanceWrites) { if (Object.keys(fields).length > 0) { - multi.hset(balanceKey, fields); + if (overwrite) { + multi.hset(balanceKey, fields); + } else { + for (const [field, value] of Object.entries(fields)) { + multi.hsetnx(balanceKey, field, value); + } + } } multi.expire(balanceKey, ttlSeconds); diff --git a/server/src/internal/customers/cache/fullSubject/actions/updateCachedEntityData.ts b/server/src/internal/customers/cache/fullSubject/actions/updateCachedEntityData.ts new file mode 100644 index 000000000..73aa47000 --- /dev/null +++ b/server/src/internal/customers/cache/fullSubject/actions/updateCachedEntityData.ts @@ -0,0 +1,135 @@ +import type { Entity } from "@autumn/shared"; +import { redisV2 } from "@/external/redis/initRedisV2.js"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { tryRedisRead, tryRedisWrite } from "@/utils/cacheUtils/cacheUtils.js"; +import { logAlertEvent } from "@/utils/logging/logAlertEvent.js"; +import { buildFullSubjectKey } from "../builders/buildFullSubjectKey.js"; +import { FULL_SUBJECT_CACHE_TTL_SECONDS } from "../config/fullSubjectCacheConfig.js"; +import { invalidateCachedFullSubject } from "./invalidate/invalidateFullSubject.js"; + +const FULL_SUBJECT_ALERT_BYTES_THRESHOLD = 1024 * 1024; +const FULL_SUBJECT_UPDATE_SLOW_THRESHOLD_MS = 100; + +export const updateCachedEntityData = async ({ + ctx, + customerId, + entityId, + updates, +}: { + ctx: AutumnContext; + customerId: string; + entityId: string; + updates: Partial>; +}): Promise => { + if (Object.keys(updates).length === 0) return; + + const { org, env, logger } = ctx; + const subjectKey = buildFullSubjectKey({ + orgId: org.id, + env, + customerId, + entityId, + }); + + try { + const currentRaw = await tryRedisRead( + () => redisV2.get(subjectKey), + redisV2, + ); + if (!currentRaw) return; + + const payloadBytes = Buffer.byteLength(currentRaw, "utf8"); + if (payloadBytes > FULL_SUBJECT_ALERT_BYTES_THRESHOLD) { + logAlertEvent({ + ctx, + severity: "warning", + category: "redis", + alertKey: "redis_full_subject_payload_large", + message: `FullSubject payload exceeded soft limit during entity cache update for ${customerId}:${entityId}`, + source: "updateCachedEntityData", + component: "full_subject_cache", + data: { + subjectKey, + payload_bytes: payloadBytes, + threshold_bytes: FULL_SUBJECT_ALERT_BYTES_THRESHOLD, + redis_command: "updateFullSubjectEntityDataV2", + }, + }); + } + + const updatesJson = JSON.stringify(updates); + const startTime = Date.now(); + const result = await tryRedisWrite( + () => + redisV2.updateFullSubjectEntityDataV2( + subjectKey, + updatesJson, + String(FULL_SUBJECT_CACHE_TTL_SECONDS), + String(Date.now()), + ), + redisV2, + ); + const durationMs = Date.now() - startTime; + + if (durationMs > FULL_SUBJECT_UPDATE_SLOW_THRESHOLD_MS) { + logAlertEvent({ + ctx, + severity: "warning", + category: "redis", + alertKey: "redis_full_subject_entity_update_slow", + message: `FullSubject entity cache update was slow for ${customerId}:${entityId}`, + source: "updateCachedEntityData", + component: "full_subject_cache", + data: { + subjectKey, + duration_ms: durationMs, + threshold_ms: FULL_SUBJECT_UPDATE_SLOW_THRESHOLD_MS, + redis_command: "updateFullSubjectEntityDataV2", + payload_bytes: payloadBytes, + }, + }); + } + + if (result === null) { + logger.warn( + `[updateCachedEntityData] Redis write failed for ${customerId}:${entityId}, invalidating cache`, + ); + await invalidateCachedFullSubject({ + ctx, + customerId, + entityId, + source: "updateCachedEntityData:redis_write_failed", + }); + return; + } + + const parsed = JSON.parse(result) as { + success: boolean; + updated_fields?: string[]; + cache_miss?: boolean; + no_entity?: boolean; + }; + if (parsed.cache_miss || parsed.no_entity) return; + if (parsed.success) return; + + logger.warn( + `[updateCachedEntityData] Lua update returned unsuccessful result for ${customerId}:${entityId}, invalidating cache`, + ); + await invalidateCachedFullSubject({ + ctx, + customerId, + entityId, + source: "updateCachedEntityData:lua_unsuccessful", + }); + } catch (error) { + logger.error( + `[updateCachedEntityData] Failed to update entity subject for ${customerId}:${entityId}: ${error}`, + ); + await invalidateCachedFullSubject({ + ctx, + customerId, + entityId, + source: "updateCachedEntityData:error", + }); + } +}; diff --git a/server/src/internal/customers/cache/fullSubject/balances/getCachedFeatureBalances.ts b/server/src/internal/customers/cache/fullSubject/balances/getCachedFeatureBalances.ts index 021d6649a..fc57e2918 100644 --- a/server/src/internal/customers/cache/fullSubject/balances/getCachedFeatureBalances.ts +++ b/server/src/internal/customers/cache/fullSubject/balances/getCachedFeatureBalances.ts @@ -2,6 +2,7 @@ import type { SubjectBalance } from "@autumn/shared"; import { redisV2 } from "@/external/redis/initRedisV2.js"; import { tryRedisRead } from "@/utils/cacheUtils/cacheUtils.js"; import { buildSharedFullSubjectBalanceKey } from "../builders/buildSharedFullSubjectBalanceKey.js"; +import { sanitizeCachedSubjectBalance } from "../sanitize/index.js"; export type FeatureBalanceResult = { featureId: string; @@ -43,7 +44,12 @@ export const getCachedFeatureBalance = async ({ const entryJson = results[i]; if (!entryJson) return undefined; try { - balances.push(JSON.parse(entryJson) as SubjectBalance); + const parsedBalance = JSON.parse(entryJson) as SubjectBalance; + balances.push( + sanitizeCachedSubjectBalance({ + subjectBalance: parsedBalance, + }), + ); } catch { return undefined; } @@ -99,7 +105,12 @@ export const getCachedFeatureBalancesBatch = async ({ for (const entryJson of values) { if (!entryJson) return undefined; try { - balances.push(JSON.parse(entryJson) as SubjectBalance); + const parsedBalance = JSON.parse(entryJson) as SubjectBalance; + balances.push( + sanitizeCachedSubjectBalance({ + subjectBalance: parsedBalance, + }), + ); } catch { return undefined; } diff --git a/server/src/internal/customers/cache/fullSubject/config/fullSubjectCacheConfig.ts b/server/src/internal/customers/cache/fullSubject/config/fullSubjectCacheConfig.ts index f210a7a39..790c68015 100644 --- a/server/src/internal/customers/cache/fullSubject/config/fullSubjectCacheConfig.ts +++ b/server/src/internal/customers/cache/fullSubject/config/fullSubjectCacheConfig.ts @@ -3,3 +3,4 @@ import { seconds } from "@autumn/shared"; export const FULL_SUBJECT_CACHE_TTL_SECONDS = seconds.days(3); export const FULL_SUBJECT_CACHE_RESERVE_TTL_SECONDS = 60; export const FULL_SUBJECT_CACHE_GUARD_TTL_SECONDS = 1; +export const FULL_SUBJECT_EPOCH_TTL_SECONDS = seconds.days(5); diff --git a/server/src/internal/customers/cache/fullSubject/sanitize/index.ts b/server/src/internal/customers/cache/fullSubject/sanitize/index.ts new file mode 100644 index 000000000..07fa5e86a --- /dev/null +++ b/server/src/internal/customers/cache/fullSubject/sanitize/index.ts @@ -0,0 +1,2 @@ +export { sanitizeCachedFullSubject } from "./sanitizeCachedFullSubject.js"; +export { sanitizeCachedSubjectBalance } from "./sanitizeCachedSubjectBalance.js"; diff --git a/server/src/internal/customers/cache/fullSubject/sanitize/sanitizeCacheShapeUtils.ts b/server/src/internal/customers/cache/fullSubject/sanitize/sanitizeCacheShapeUtils.ts new file mode 100644 index 000000000..de3b80611 --- /dev/null +++ b/server/src/internal/customers/cache/fullSubject/sanitize/sanitizeCacheShapeUtils.ts @@ -0,0 +1,71 @@ +/** + * Shape spec describes which fields should be arrays, records, or nested objects + * so the recursive sanitizer can coerce malformed Redis/Lua JSON payloads. + * + * "array" -> coerce non-arrays to [] + * "record" -> coerce non-objects to {} + * "nullable_record" -> coerce non-objects to null (for optional record fields) + * ShapeSpec -> recurse into object fields + * { items: ShapeSpec } -> coerce field to array, then recurse each element + */ +export interface ShapeSpec { + [key: string]: FieldRule; +} + +export type FieldRule = + | "array" + | "record" + | "nullable_record" + | ShapeSpec + | { items: ShapeSpec }; + +const isPlainObject = (value: unknown): value is Record => + !!value && typeof value === "object" && !Array.isArray(value); + +const coerceArray = (value: unknown): unknown[] => + Array.isArray(value) ? value : []; + +const coerceRecord = (value: unknown): Record => + isPlainObject(value) ? value : {}; + +const coerceNullableRecord = ( + value: unknown, +): Record | null => (isPlainObject(value) ? value : null); + +/** + * Recursively sanitizes an object against a shape spec. + * Only touches fields that appear in the spec; all other fields pass through. + */ +export const sanitizeShape = ({ + value, + spec, +}: { + value: unknown; + spec: ShapeSpec; +}): T => { + if (!isPlainObject(value)) return {} as T; + + const result = { ...value } as Record; + + for (const [key, rule] of Object.entries(spec)) { + const fieldValue = result[key]; + + if (rule === "array") { + result[key] = coerceArray(fieldValue); + } else if (rule === "record") { + result[key] = coerceRecord(fieldValue); + } else if (rule === "nullable_record") { + result[key] = coerceNullableRecord(fieldValue); + } else if (isPlainObject(rule) && "items" in rule) { + const itemsSpec = (rule as { items: ShapeSpec }).items; + const arr = coerceArray(fieldValue); + result[key] = arr.map((item) => + sanitizeShape({ value: item, spec: itemsSpec }), + ); + } else if (isPlainObject(rule) && isPlainObject(fieldValue)) { + result[key] = sanitizeShape({ value: fieldValue, spec: rule as ShapeSpec }); + } + } + + return result as T; +}; diff --git a/server/src/internal/customers/cache/fullSubject/sanitize/sanitizeCachedFullSubject.ts b/server/src/internal/customers/cache/fullSubject/sanitize/sanitizeCachedFullSubject.ts new file mode 100644 index 000000000..3c0712806 --- /dev/null +++ b/server/src/internal/customers/cache/fullSubject/sanitize/sanitizeCachedFullSubject.ts @@ -0,0 +1,79 @@ +import type { CachedFullSubject } from "../fullSubjectCacheModel.js"; +import { type ShapeSpec, sanitizeShape } from "./sanitizeCacheShapeUtils.js"; + +const featureShapeSpec: ShapeSpec = { + event_names: "array", +}; + +const entitlementCatalogShapeSpec: ShapeSpec = { + feature: featureShapeSpec, +}; + +const priceConfigShapeSpec: ShapeSpec = { + usage_tiers: "array", +}; + +const priceShapeSpec: ShapeSpec = { + config: priceConfigShapeSpec, +}; + +const customerShapeSpec: ShapeSpec = { + auto_topups: "array", + spend_limits: "array", + usage_alerts: "array", + overage_allowed: "array", +}; + +const entityShapeSpec: ShapeSpec = { + spend_limits: "array", + usage_alerts: "array", + overage_allowed: "array", +}; + +const customerProductShapeSpec: ShapeSpec = { + options: "array", + subscription_ids: "array", + scheduled_ids: "array", +}; + +const subscriptionShapeSpec: ShapeSpec = { + usage_features: "array", +}; + +const invoiceShapeSpec: ShapeSpec = { + product_ids: "array", + internal_product_ids: "array", + discounts: "array", + items: "array", +}; + +const entityAggregationsShapeSpec: ShapeSpec = { + aggregated_customer_products: { items: customerProductShapeSpec }, + aggregated_customer_entitlements: "array", +}; + +const cachedFullSubjectShapeSpec: ShapeSpec = { + customer: customerShapeSpec, + entity: entityShapeSpec, + customer_products: { items: customerProductShapeSpec }, + products: "array", + entitlements: { items: entitlementCatalogShapeSpec }, + prices: { items: priceShapeSpec }, + free_trials: "array", + subscriptions: { items: subscriptionShapeSpec }, + invoices: { items: invoiceShapeSpec }, + flags: "record", + meteredFeatures: "array", + customerEntitlementIdsByFeatureId: "record", + entity_aggregations: entityAggregationsShapeSpec, +}; + +export const sanitizeCachedFullSubject = ({ + cachedFullSubject, +}: { + cachedFullSubject: CachedFullSubject; +}): CachedFullSubject => + sanitizeShape({ + value: cachedFullSubject, + spec: cachedFullSubjectShapeSpec, + }); diff --git a/server/src/internal/customers/cache/fullSubject/sanitize/sanitizeCachedSubjectBalance.ts b/server/src/internal/customers/cache/fullSubject/sanitize/sanitizeCachedSubjectBalance.ts new file mode 100644 index 000000000..ad6f300f6 --- /dev/null +++ b/server/src/internal/customers/cache/fullSubject/sanitize/sanitizeCachedSubjectBalance.ts @@ -0,0 +1,43 @@ +import type { SubjectBalance } from "@autumn/shared"; +import { type ShapeSpec, sanitizeShape } from "./sanitizeCacheShapeUtils.js"; + +const featureShapeSpec: ShapeSpec = { + event_names: "array", +}; + +const entitlementShapeSpec: ShapeSpec = { + feature: featureShapeSpec, +}; + +const priceConfigShapeSpec: ShapeSpec = { + usage_tiers: "array", +}; + +const priceShapeSpec: ShapeSpec = { + config: priceConfigShapeSpec, +}; + +const customerPriceShapeSpec: ShapeSpec = { + price: priceShapeSpec, +}; + +const rolloverShapeSpec: ShapeSpec = { + entities: "record", +}; + +const subjectBalanceShapeSpec: ShapeSpec = { + rollovers: { items: rolloverShapeSpec }, + entities: "nullable_record", + entitlement: entitlementShapeSpec, + customerPrice: customerPriceShapeSpec, +}; + +export const sanitizeCachedSubjectBalance = ({ + subjectBalance, +}: { + subjectBalance: SubjectBalance; +}): SubjectBalance => + sanitizeShape({ + value: subjectBalance, + spec: subjectBalanceShapeSpec, + }); diff --git a/server/src/internal/customers/cusProducts/cusEnts/actions/cache/updateSubjectBalanceCache.ts b/server/src/internal/customers/cusProducts/cusEnts/actions/cache/updateSubjectBalanceCache.ts new file mode 100644 index 000000000..fe67fa15f --- /dev/null +++ b/server/src/internal/customers/cusProducts/cusEnts/actions/cache/updateSubjectBalanceCache.ts @@ -0,0 +1,59 @@ +import { redisV2 } from "@/external/redis/initRedisV2.js"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { buildSharedFullSubjectBalanceKey } from "@/internal/customers/cache/fullSubject/builders/buildSharedFullSubjectBalanceKey.js"; +import { FULL_SUBJECT_CACHE_TTL_SECONDS } from "@/internal/customers/cache/fullSubject/config/fullSubjectCacheConfig.js"; +import { tryRedisWrite } from "@/utils/cacheUtils/cacheUtils.js"; + +export const updateSubjectBalanceCache = async ({ + ctx, + customerId, + featureId, + customerEntitlementId, + updates, +}: { + ctx: AutumnContext; + customerId: string; + featureId: string; + customerEntitlementId: string; + updates: { + balance?: number | null; + additional_balance?: number | null; + adjustment?: number | null; + entities?: Record | null; + next_reset_at?: number | null; + }; +}) => { + const balanceKey = buildSharedFullSubjectBalanceKey({ + orgId: ctx.org.id, + env: ctx.env, + customerId, + featureId, + }); + + await tryRedisWrite( + () => + redisV2.updateSubjectBalances( + balanceKey, + JSON.stringify({ + ttl_seconds: FULL_SUBJECT_CACHE_TTL_SECONDS, + updates: [ + { + cus_ent_id: customerEntitlementId, + balance: updates.balance ?? null, + additional_balance: updates.additional_balance ?? null, + adjustment: updates.adjustment ?? null, + entities: updates.entities ?? null, + next_reset_at: updates.next_reset_at ?? null, + expected_next_reset_at: null, + rollover_insert: null, + rollover_overwrites: null, + rollover_delete_ids: null, + new_replaceables: null, + deleted_replaceable_ids: null, + }, + ], + }), + ), + redisV2, + ); +}; diff --git a/server/src/internal/customers/handlers/handleGetCustomerV2.ts b/server/src/internal/customers/handlers/handleGetCustomerV2.ts index 7fdf4e633..195fba839 100644 --- a/server/src/internal/customers/handlers/handleGetCustomerV2.ts +++ b/server/src/internal/customers/handlers/handleGetCustomerV2.ts @@ -9,8 +9,7 @@ import { V0_2_InvoicesAlwaysExpanded, } from "@autumn/shared"; import { createRoute } from "@/honoMiddlewares/routeHandler.js"; -import { getApiCustomer } from "../cusUtils/apiCusUtils/getApiCustomer.js"; -import { getOrSetCachedFullCustomer } from "../cusUtils/fullCustomerCacheUtils/getOrSetCachedFullCustomer.js"; +import { getApiCustomerByRollout } from "../actions/getApiCustomerByRollout.js"; export const handleGetCustomerV2 = createRoute({ versionedQuery: { @@ -32,8 +31,6 @@ export const handleGetCustomerV2 = createRoute({ }); } - // SIDE EFFECT - // !ctx.org.config.disable_v1_invoices && if ( backwardsChangeActive({ apiVersion: ctx.apiVersion, @@ -45,17 +42,10 @@ export const handleGetCustomerV2 = createRoute({ const start = Date.now(); - // Get FullCustomer from cache or DB - const fullCustomer = await getOrSetCachedFullCustomer({ + const customer = await getApiCustomerByRollout({ ctx, customerId, source: "handleGetCustomerV2", - }); - - // Transform to ApiCustomer with version changes - const customer = await getApiCustomer({ - ctx, - fullCustomer, withAutumnId: with_autumn_id, }); diff --git a/server/src/internal/customers/handlers/handleGetOrCreateCustomer/handleGetOrCreateCustomer.ts b/server/src/internal/customers/handlers/handleGetOrCreateCustomer/handleGetOrCreateCustomer.ts index 7ea6670e4..558168b5d 100644 --- a/server/src/internal/customers/handlers/handleGetOrCreateCustomer/handleGetOrCreateCustomer.ts +++ b/server/src/internal/customers/handlers/handleGetOrCreateCustomer/handleGetOrCreateCustomer.ts @@ -10,8 +10,7 @@ import { V0_2_InvoicesAlwaysExpanded, } from "@autumn/shared"; import { createRoute } from "@/honoMiddlewares/routeHandler.js"; -import { getApiCustomer } from "../../cusUtils/apiCusUtils/getApiCustomer.js"; -import { getOrCreateCachedFullCustomer } from "../../cusUtils/fullCustomerCacheUtils/getOrCreateCachedFullCustomer.js"; +import { getOrCreateApiCustomerByRollout } from "@/internal/customers/actions/getOrCreateApiCustomerByRollout.js"; export const handlePostCustomer = createRoute({ versionedQuery: { @@ -53,7 +52,7 @@ export const handlePostCustomer = createRoute({ const customerData = CustomerDataSchema.parse(createCusParams); - const fullCustomer = await getOrCreateCachedFullCustomer({ + const apiCustomer = await getOrCreateApiCustomerByRollout({ ctx, params: { customer_id: createCusParams.id, @@ -62,11 +61,6 @@ export const handlePostCustomer = createRoute({ entity_data: createCusParams.entity_data, }, source: "handlePostCustomer", - }); - - const apiCustomer = await getApiCustomer({ - ctx, - fullCustomer, withAutumnId: with_autumn_id, }); diff --git a/server/src/internal/customers/handlers/handleGetOrCreateCustomer/handleGetOrCreateCustomerV2.ts b/server/src/internal/customers/handlers/handleGetOrCreateCustomer/handleGetOrCreateCustomerV2.ts index a726d843b..99e8d74e5 100644 --- a/server/src/internal/customers/handlers/handleGetOrCreateCustomer/handleGetOrCreateCustomerV2.ts +++ b/server/src/internal/customers/handlers/handleGetOrCreateCustomer/handleGetOrCreateCustomerV2.ts @@ -4,7 +4,7 @@ import { CustomerDataSchema, } from "@autumn/shared"; import { createRoute } from "@/honoMiddlewares/routeHandler.js"; -import { getApiCustomerByRollout } from "@/internal/customers/actions/getApiCustomerByRollout.js"; +import { getOrCreateApiCustomerByRollout } from "@/internal/customers/actions/getOrCreateApiCustomerByRollout.js"; import { isFullSubjectRolloutEnabled } from "@/internal/misc/rollouts/fullSubjectRolloutUtils.js"; export const handleGetOrCreateCustomerV2 = createRoute({ @@ -19,7 +19,7 @@ export const handleGetOrCreateCustomerV2 = createRoute({ const customerData = CustomerDataSchema.parse(createCusParams); const customerId = createCusParams.customer_id; - const apiCustomer = await getApiCustomerByRollout({ + const apiCustomer = await getOrCreateApiCustomerByRollout({ ctx, params: { customer_id: customerId, diff --git a/server/src/internal/customers/handlers/handleUpdateCustomer/handleUpdateCustomer.ts b/server/src/internal/customers/handlers/handleUpdateCustomer/handleUpdateCustomer.ts index 479553112..ccf8af33f 100644 --- a/server/src/internal/customers/handlers/handleUpdateCustomer/handleUpdateCustomer.ts +++ b/server/src/internal/customers/handlers/handleUpdateCustomer/handleUpdateCustomer.ts @@ -6,8 +6,7 @@ import { } from "@autumn/shared"; import { createRoute } from "@/honoMiddlewares/routeHandler.js"; import { customerActions } from "@/internal/customers/actions"; -import { getApiCustomer } from "@/internal/customers/cusUtils/apiCusUtils/getApiCustomer"; -import { getOrSetCachedFullCustomer } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/getOrSetCachedFullCustomer"; +import { getApiCustomerByRollout } from "@/internal/customers/actions/getApiCustomerByRollout.js"; export const handleUpdateCustomer = createRoute({ body: UpdateCustomerParamsV0Schema, @@ -32,15 +31,10 @@ export const handleUpdateCustomer = createRoute({ }); ctx.skipCache = true; - const fullCustomer = await getOrSetCachedFullCustomer({ + const customerDetails = await getApiCustomerByRollout({ ctx, customerId: newCustomerId, - source: "handleUpdateCustomerV2", - }); - - const customerDetails = await getApiCustomer({ - ctx, - fullCustomer, + source: "handleUpdateCustomer", }); return c.json(customerDetails); diff --git a/server/src/internal/customers/handlers/handleUpdateCustomer/handleUpdateCustomerV2.ts b/server/src/internal/customers/handlers/handleUpdateCustomer/handleUpdateCustomerV2.ts index 05e4797ba..3b5a1d94f 100644 --- a/server/src/internal/customers/handlers/handleUpdateCustomer/handleUpdateCustomerV2.ts +++ b/server/src/internal/customers/handlers/handleUpdateCustomer/handleUpdateCustomerV2.ts @@ -1,8 +1,7 @@ import { AffectedResource, UpdateCustomerParamsV1Schema } from "@autumn/shared"; import { createRoute } from "@/honoMiddlewares/routeHandler.js"; +import { getApiCustomerByRollout } from "@/internal/customers/actions/getApiCustomerByRollout.js"; import { customerActions } from "@/internal/customers/actions/index.js"; -import { getApiCustomer } from "@/internal/customers/cusUtils/apiCusUtils/getApiCustomer.js"; -import { getOrSetCachedFullCustomer } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/getOrSetCachedFullCustomer.js"; export const handleUpdateCustomerV2 = createRoute({ body: UpdateCustomerParamsV1Schema, @@ -16,17 +15,12 @@ export const handleUpdateCustomerV2 = createRoute({ }); ctx.skipCache = true; - const fullCustomer = await getOrSetCachedFullCustomer({ + const customerDetails = await getApiCustomerByRollout({ ctx, customerId, source: "handleUpdateCustomerV2", }); - const customerDetails = await getApiCustomer({ - ctx, - fullCustomer, - }); - return c.json(customerDetails); }, }); diff --git a/server/src/internal/entities/actions/updateEntity.ts b/server/src/internal/entities/actions/updateEntity.ts index 2d2fcc3bc..3e47cf504 100644 --- a/server/src/internal/entities/actions/updateEntity.ts +++ b/server/src/internal/entities/actions/updateEntity.ts @@ -1,12 +1,11 @@ import { - CustomerExpand, CustomerNotFoundError, EntityNotFoundError, type UpdateEntityParams, } from "@autumn/shared"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; -import { CusService } from "@/internal/customers/CusService.js"; -import { updateEntityDbAndCache } from "./updateEntityDbAndCache.js"; +import { EntityService } from "@/internal/api/entities/EntityService.js"; +import { getFullSubject } from "@/internal/customers/repos/getFullSubject/getFullSubject.js"; export const updateEntity = async ({ ctx, @@ -24,30 +23,37 @@ export const updateEntity = async ({ throw new CustomerNotFoundError({ customerId: "" }); } - const fullCustomer = await CusService.getFull({ + const fullSubject = await getFullSubject({ ctx, - idOrInternalId: customerId, - // withEntities: true, - entityId: entityId, - expand: [CustomerExpand.Invoices], + customerId, + entityId, }); - const entity = fullCustomer.entity; + if (!fullSubject) { + throw new CustomerNotFoundError({ customerId }); + } + + const entity = fullSubject.entity; if (!entity) { throw new EntityNotFoundError({ entityId }); } - await updateEntityDbAndCache({ - ctx, - customerId, - entity, - updates: { + const filteredUpdates = Object.fromEntries( + Object.entries({ spend_limits: billing_controls?.spend_limits, usage_alerts: billing_controls?.usage_alerts, overage_allowed: billing_controls?.overage_allowed, - }, - }); + }).filter(([, value]) => value !== undefined), + ); + + if (Object.keys(filteredUpdates).length > 0) { + await EntityService.update({ + db: ctx.db, + internalId: entity.internal_id, + update: filteredUpdates, + }); + } return entity.id ?? entity.internal_id; }; diff --git a/server/src/internal/entities/actions/updateEntityDbAndCache.ts b/server/src/internal/entities/actions/updateEntityDbAndCache.ts index d0a508f01..a6c5a412e 100644 --- a/server/src/internal/entities/actions/updateEntityDbAndCache.ts +++ b/server/src/internal/entities/actions/updateEntityDbAndCache.ts @@ -1,16 +1,13 @@ import type { Entity } from "@autumn/shared"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { EntityService } from "@/internal/api/entities/EntityService.js"; -import { updateEntityInCache } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/updateEntityInCache.js"; export const updateEntityDbAndCache = async ({ ctx, - customerId, entity, updates, }: { ctx: AutumnContext; - customerId: string; entity: Entity; updates: Partial< Pick @@ -26,18 +23,9 @@ export const updateEntityDbAndCache = async ({ return entity; } - const updatedEntity = await EntityService.update({ + return EntityService.update({ db: ctx.db, internalId: entity.internal_id, update: filteredUpdates, }); - - await updateEntityInCache({ - ctx, - customerId, - idOrInternalId: entity.id ?? entity.internal_id, - updates: filteredUpdates, - }); - - return updatedEntity; }; diff --git a/server/src/internal/entities/handlers/handleUpdateEntity/handleUpdateEntity.ts b/server/src/internal/entities/handlers/handleUpdateEntity/handleUpdateEntity.ts index ea6603c40..83b68aa3c 100644 --- a/server/src/internal/entities/handlers/handleUpdateEntity/handleUpdateEntity.ts +++ b/server/src/internal/entities/handlers/handleUpdateEntity/handleUpdateEntity.ts @@ -6,7 +6,7 @@ import { import { createRoute } from "@/honoMiddlewares/routeHandler.js"; import { findCustomerForEntity } from "../../actions/findCustomer.js"; import { entityActions } from "../../actions/index.js"; -import { getApiEntity } from "../../entityUtils/apiEntityUtils/getApiEntity.js"; +import { getApiEntityByRollout } from "../../actions/getApiEntityByRollout.js"; export const handleUpdateEntity = createRoute({ body: UpdateEntityParamsSchema, @@ -37,10 +37,11 @@ export const handleUpdateEntity = createRoute({ }, }); - const apiEntity = await getApiEntity({ + const apiEntity = await getApiEntityByRollout({ ctx, customerId, entityId: body.entity_id, + source: "handleUpdateEntity", }); return c.json(apiEntity); diff --git a/server/src/queue/JobName.ts b/server/src/queue/JobName.ts index 877d86672..b09affd28 100644 --- a/server/src/queue/JobName.ts +++ b/server/src/queue/JobName.ts @@ -18,6 +18,7 @@ export enum JobName { SyncBalanceBatch = "sync-balance-batch", SyncBalanceBatchV2 = "sync-balance-batch-v2", SyncBalanceBatchV3 = "sync-balance-batch-v3", + SyncBalanceBatchV4 = "sync-balance-batch-v4", InsertEventBatch = "insert-event-batch", ClearCreditSystemCustomerCache = "clear-credit-system-customer-cache", diff --git a/server/src/queue/bullmq/initBullMqWorkers.ts b/server/src/queue/bullmq/initBullMqWorkers.ts index 478b5ca2f..b5561ccf7 100644 --- a/server/src/queue/bullmq/initBullMqWorkers.ts +++ b/server/src/queue/bullmq/initBullMqWorkers.ts @@ -6,6 +6,7 @@ import { runActionHandlerTask } from "@/internal/analytics/runActionHandlerTask. import { autoTopup } from "@/internal/balances/autoTopUp/autoTopup.js"; import { runInsertEventBatch } from "@/internal/balances/events/runInsertEventBatch.js"; import { syncItemV3 } from "@/internal/balances/utils/sync/syncItemV3.js"; +import { syncItemV4 } from "@/internal/balances/utils/sync/syncItemV4.js"; import { generateFeatureDisplay } from "@/internal/features/workflows/generateFeatureDisplay.js"; import { runMigrationTask } from "@/internal/migrations/runMigrationTask.js"; import { runRewardMigrationTask } from "@/internal/migrations/runRewardMigrationTask.js"; @@ -107,10 +108,11 @@ const initWorker = ({ id, db }: { id: number; db: DrizzleCli }) => { ); return; } - await syncItemV3({ - ctx, - payload: job.data, - }); + if (job.data.syncVersion === "v4") { + await syncItemV4({ ctx, payload: job.data }); + } else { + await syncItemV3({ ctx, payload: job.data }); + } return; } diff --git a/server/src/queue/processMessage.ts b/server/src/queue/processMessage.ts index f8dc4dcd7..a5cc6432f 100644 --- a/server/src/queue/processMessage.ts +++ b/server/src/queue/processMessage.ts @@ -11,6 +11,7 @@ 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 { 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"; import { sendProductsUpdated } from "@/internal/billing/v2/workflows/sendProductsUpdated/sendProductsUpdated.js"; import { storeDeferredInvoiceLineItems } from "@/internal/billing/v2/workflows/storeDeferredInvoiceLineItems/storeDeferredInvoiceLineItems.js"; @@ -159,10 +160,17 @@ export const processMessage = async ({ return; } - await syncItemV3({ - ctx, - payload: job.data, - }); + await syncItemV3({ ctx, payload: job.data }); + return; + } + + if (job.name === JobName.SyncBalanceBatchV4) { + if (!ctx) { + workerLogger.error("No context found for sync balance batch v4 job"); + return; + } + + await syncItemV4({ ctx, payload: job.data }); return; } @@ -268,7 +276,8 @@ export const processMessage = async ({ // Application errors (RecaseError, InternalError) are swallowed — they // won't fix on retry. DB errors (connection, timeout) will. if ( - job.name === JobName.SyncBalanceBatchV3 && + (job.name === JobName.SyncBalanceBatchV3 || + job.name === JobName.SyncBalanceBatchV4) && isRetryableDbError({ error }) ) { Sentry.captureException(error); diff --git a/server/src/queue/queueUtils.ts b/server/src/queue/queueUtils.ts index c922d4c05..a7f4e27d6 100644 --- a/server/src/queue/queueUtils.ts +++ b/server/src/queue/queueUtils.ts @@ -36,6 +36,18 @@ export interface Payloads { timestamp: number; cusEntIds: string[]; rolloverIds?: string[]; + entityId?: string; + }; + [JobName.SyncBalanceBatchV4]: { + customerId: string; + orgId: string; + env: AppEnv; + region?: string; + timestamp: number; + cusEntIds: string[]; + rolloverIds?: string[]; + entityId?: string; + modifiedCusEntIdsByFeatureId: Record; }; [JobName.InsertEventBatch]: { events: EventInsert[]; diff --git a/server/tests/unit/balances/check-v2/runCheckWithTrackV2.test.ts b/server/tests/unit/balances/check-v2/runCheckWithTrackV2.test.ts new file mode 100644 index 000000000..8012f12a9 --- /dev/null +++ b/server/tests/unit/balances/check-v2/runCheckWithTrackV2.test.ts @@ -0,0 +1,147 @@ +import { describe, expect, test } from "bun:test"; +import { + AppEnv, + ErrCode, + type Feature, + FeatureUsageType, + type FullSubject, + type Organization, + SubjectType, +} from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import type { CheckDataV2 } from "@/internal/balances/check/checkTypes/CheckDataV2.js"; +import { runCheckWithTrackV2 } from "@/internal/balances/check/runCheckWithTrackV2.js"; + +const meteredFeature = { + id: "messages", + internal_id: "feat_messages", + org_id: "org_1", + env: AppEnv.Live, + name: "Messages", + type: "metered", + config: null, + display: null, + created_at: 1, + archived: false, + event_names: [], +} as Feature; + +const buildCtx = ({ isPublic = false }: { isPublic?: boolean } = {}) => + ({ + org: { + id: "org_1", + config: {}, + } as Organization, + env: AppEnv.Live, + isPublic, + logger: { + error: () => {}, + info: () => {}, + warn: () => {}, + debug: () => {}, + }, + features: [meteredFeature], + }) as AutumnContext; + +const buildCheckData = ({ + originalFeature = meteredFeature, + featureToUse = meteredFeature, +}: { + originalFeature?: Feature; + featureToUse?: Feature; +} = {}) => + ({ + customerId: "cus_1", + entityId: undefined, + apiBalance: undefined, + apiFlag: undefined, + apiSubject: {}, + originalFeature, + featureToUse, + fullSubject: { + subjectType: SubjectType.Customer, + customerId: "cus_1", + internalCustomerId: "cus_int_1", + customer: { + id: "cus_1", + internal_id: "cus_int_1", + }, + customer_products: [], + extra_customer_entitlements: [], + subscriptions: [], + invoices: [], + } as unknown as FullSubject, + evaluationApiSubject: {}, + evaluationApiBalance: undefined, + evaluationApiFlag: undefined, + }) as CheckDataV2; + +describe("runCheckWithTrackV2", () => { + test("rejects send_event for public requests", async () => { + await expect( + runCheckWithTrackV2({ + ctx: buildCtx({ isPublic: true }), + body: { + customer_id: "cus_1", + feature_id: "messages", + send_event: true, + } as never, + requiredBalance: 1, + checkData: buildCheckData(), + }), + ).rejects.toMatchObject({ + message: + "Can't pass in 'send_event: true' when using publishable key for Autumn", + }); + }); + + test("rejects boolean features", async () => { + await expect( + runCheckWithTrackV2({ + ctx: buildCtx(), + body: { + customer_id: "cus_1", + feature_id: "messages", + send_event: true, + } as never, + requiredBalance: 1, + checkData: buildCheckData({ + originalFeature: { + ...meteredFeature, + type: "boolean", + } as Feature, + }), + }), + ).rejects.toMatchObject({ + code: ErrCode.InvalidRequest, + message: "Not allowed to pass in send_event: true for a boolean feature", + }); + }); + + test("rejects locks for allocated features", async () => { + await expect( + runCheckWithTrackV2({ + ctx: buildCtx(), + body: { + customer_id: "cus_1", + feature_id: "messages", + lock: { + enabled: true, + }, + } as never, + requiredBalance: 1, + checkData: buildCheckData({ + featureToUse: { + ...meteredFeature, + config: { + usage_type: FeatureUsageType.Continuous, + }, + } as Feature, + }), + }), + ).rejects.toMatchObject({ + code: ErrCode.InvalidRequest, + message: "Lock is not supported for allocated features", + }); + }); +}); diff --git a/server/tests/unit/full-subject-cache/sanitizeCachedPayloads.test.ts b/server/tests/unit/full-subject-cache/sanitizeCachedPayloads.test.ts new file mode 100644 index 000000000..5f31c33d9 --- /dev/null +++ b/server/tests/unit/full-subject-cache/sanitizeCachedPayloads.test.ts @@ -0,0 +1,549 @@ +import { describe, expect, test } from "bun:test"; +import { AppEnv, type SubjectBalance } from "@autumn/shared"; +import type { CachedFullSubject } from "@/internal/customers/cache/fullSubject/fullSubjectCacheModel.js"; +import { sanitizeCachedFullSubject } from "@/internal/customers/cache/fullSubject/sanitize/sanitizeCachedFullSubject.js"; +import { sanitizeCachedSubjectBalance } from "@/internal/customers/cache/fullSubject/sanitize/sanitizeCachedSubjectBalance.js"; +import { sanitizeShape } from "@/internal/customers/cache/fullSubject/sanitize/sanitizeCacheShapeUtils.js"; + +describe("sanitizeShape (recursive core)", () => { + test("should coerce non-array to [] for 'array' rule", () => { + const result = sanitizeShape({ + value: { items: {} }, + spec: { items: "array" }, + }); + expect(result).toEqual({ items: [] }); + }); + + test("should preserve valid arrays", () => { + const result = sanitizeShape({ + value: { items: [1, 2, 3] }, + spec: { items: "array" }, + }); + expect(result).toEqual({ items: [1, 2, 3] }); + }); + + test("should coerce non-object to {} for 'record' rule", () => { + const result = sanitizeShape({ + value: { flags: [] }, + spec: { flags: "record" }, + }); + expect(result).toEqual({ flags: {} }); + }); + + test("should coerce non-object to null for 'nullable_record' rule", () => { + const result = sanitizeShape({ + value: { entities: [] }, + spec: { entities: "nullable_record" }, + }); + expect(result).toEqual({ entities: null }); + }); + + test("should preserve null for 'nullable_record' rule", () => { + const result = sanitizeShape({ + value: { entities: null }, + spec: { entities: "nullable_record" }, + }); + expect(result).toEqual({ entities: null }); + }); + + test("should recurse into nested object specs", () => { + const result = sanitizeShape({ + value: { feature: { event_names: {} } }, + spec: { feature: { event_names: "array" } }, + }); + expect(result).toEqual({ feature: { event_names: [] } }); + }); + + test("should handle { items: spec } for array-of-objects", () => { + const result = sanitizeShape({ + value: { rollovers: [{}, { entities: "bad" }] }, + spec: { rollovers: { items: { entities: "record" } } }, + }); + expect(result).toEqual({ + rollovers: [{ entities: {} }, { entities: {} }], + }); + }); + + test("should coerce then recurse for { items: spec } when field is not an array", () => { + const result = sanitizeShape({ + value: { rollovers: {} }, + spec: { rollovers: { items: { entities: "record" } } }, + }); + expect(result).toEqual({ rollovers: [] }); + }); + + test("should leave unspecified fields untouched", () => { + const result = sanitizeShape({ + value: { name: "test", items: {} }, + spec: { items: "array" }, + }); + expect(result).toEqual({ name: "test", items: [] }); + }); + + test("should return {} for non-object input", () => { + const result = sanitizeShape({ value: "not_an_object", spec: {} }); + expect(result).toEqual({}); + }); +}); + +describe("sanitizeCachedSubjectBalance", () => { + const buildMalformedSubjectBalance = (): unknown => ({ + id: "cus_ent_1", + customer_product_id: "cp_1", + entitlement_id: "ent_1", + internal_customer_id: "cus_int_1", + internal_entity_id: null, + internal_feature_id: "feat_int_1", + feature_id: "messages", + unlimited: false, + balance: 100, + adjustment: 0, + additional_balance: 0, + usage_allowed: true, + next_reset_at: null, + expires_at: null, + external_id: null, + cache_version: 1, + created_at: 1000, + customer_id: "cus_1", + rollovers: {}, + entities: [], + entitlement: { + id: "ent_1", + created_at: 1, + internal_feature_id: "feat_int_1", + internal_product_id: "prod_int_1", + is_custom: false, + interval_count: 1, + feature: { + internal_id: "feat_int_1", + org_id: "org_1", + created_at: 1, + env: AppEnv.Sandbox, + id: "messages", + name: "Messages", + type: "metered", + config: { usage_type: "single", schema: {} }, + archived: false, + event_names: {}, + display: null, + }, + }, + customerPrice: { + id: "cp_1", + internal_customer_id: "cus_int_1", + customer_product_id: "cp_1", + created_at: 1, + price_id: "price_1", + price: { + id: "price_1", + internal_product_id: "prod_int_1", + billing_type: null, + tier_behavior: null, + config: { + type: "usage", + bill_when: "end_of_period", + internal_feature_id: "feat_int_1", + feature_id: "messages", + usage_tiers: {}, + interval: "month", + }, + entitlement_id: null, + proration_config: null, + }, + }, + customerProductOptions: null, + customerProductQuantity: 1, + }); + + test("should coerce rollovers from {} to []", () => { + const malformed = buildMalformedSubjectBalance() as SubjectBalance; + const result = sanitizeCachedSubjectBalance({ + subjectBalance: malformed, + }); + expect(Array.isArray(result.rollovers)).toBe(true); + expect(result.rollovers).toEqual([]); + }); + + test("should coerce entities from [] to null", () => { + const malformed = buildMalformedSubjectBalance() as SubjectBalance; + const result = sanitizeCachedSubjectBalance({ + subjectBalance: malformed, + }); + expect(result.entities).toBeNull(); + }); + + test("should coerce entitlement.feature.event_names from {} to []", () => { + const malformed = buildMalformedSubjectBalance() as SubjectBalance; + const result = sanitizeCachedSubjectBalance({ + subjectBalance: malformed, + }); + expect(Array.isArray(result.entitlement.feature.event_names)).toBe(true); + expect(result.entitlement.feature.event_names).toEqual([]); + }); + + test("should coerce customerPrice.price.config.usage_tiers from {} to []", () => { + const malformed = buildMalformedSubjectBalance() as SubjectBalance; + const result = sanitizeCachedSubjectBalance({ + subjectBalance: malformed, + }); + const config = result.customerPrice?.price?.config as Record< + string, + unknown + >; + expect(Array.isArray(config?.usage_tiers)).toBe(true); + }); + + test("should preserve valid fields untouched", () => { + const malformed = buildMalformedSubjectBalance() as SubjectBalance; + const result = sanitizeCachedSubjectBalance({ + subjectBalance: malformed, + }); + expect(result.id).toBe("cus_ent_1"); + expect(result.balance).toBe(100); + expect(result.feature_id).toBe("messages"); + expect(result.entitlement.feature.name).toBe("Messages"); + }); + + test("should handle rollovers array with nested entities coercion", () => { + const malformed = buildMalformedSubjectBalance() as SubjectBalance; + (malformed as unknown as Record).rollovers = [ + { id: "r1", cus_ent_id: "ce1", balance: 50, usage: 0, entities: "bad" }, + ]; + const result = sanitizeCachedSubjectBalance({ + subjectBalance: malformed, + }); + expect(Array.isArray(result.rollovers)).toBe(true); + expect(result.rollovers.length).toBe(1); + expect(result.rollovers[0].entities).toEqual({}); + }); +}); + +describe("sanitizeCachedFullSubject", () => { + const buildMalformedCachedFullSubject = (): unknown => ({ + subjectType: "customer", + customerId: "cus_1", + internalCustomerId: "cus_int_1", + _cachedAt: Date.now(), + subjectViewEpoch: 1, + meteredFeatures: {}, + customerEntitlementIdsByFeatureId: {}, + customer: { + id: "cus_1", + internal_id: "cus_int_1", + org_id: "org_1", + env: AppEnv.Live, + created_at: 1, + name: "Test", + email: null, + fingerprint: null, + processor: null, + processors: null, + metadata: {}, + send_email_receipts: false, + auto_topups: {}, + spend_limits: {}, + usage_alerts: {}, + overage_allowed: {}, + }, + entity: { + id: "ent_1", + org_id: "org_1", + created_at: 1, + internal_id: "ent_int_1", + internal_customer_id: "cus_int_1", + env: "live", + name: null, + deleted: false, + feature_id: "messages", + internal_feature_id: "feat_int_1", + spend_limits: {}, + usage_alerts: {}, + overage_allowed: {}, + }, + customer_products: {}, + products: {}, + entitlements: [ + { + id: "ent_1", + created_at: 1, + internal_feature_id: "feat_int_1", + internal_product_id: "prod_int_1", + is_custom: false, + interval_count: 1, + feature: { + internal_id: "feat_int_1", + org_id: "org_1", + created_at: 1, + env: "sandbox", + id: "messages", + name: "Messages", + type: "metered", + config: {}, + archived: false, + event_names: {}, + display: null, + }, + }, + ], + prices: [ + { + id: "price_1", + internal_product_id: "prod_int_1", + billing_type: null, + tier_behavior: null, + config: { + type: "usage", + bill_when: "end_of_period", + internal_feature_id: "feat_int_1", + feature_id: "messages", + usage_tiers: {}, + interval: "month", + }, + entitlement_id: null, + proration_config: null, + }, + ], + free_trials: {}, + subscriptions: {}, + invoices: [ + { + id: "inv_1", + created_at: 1, + internal_customer_id: "cus_int_1", + internal_entity_id: null, + product_ids: {}, + internal_product_ids: {}, + stripe_id: "in_1", + status: "paid", + hosted_invoice_url: null, + total: 100, + currency: "usd", + discounts: {}, + items: {}, + }, + ], + flags: [], + entity_aggregations: { + aggregated_customer_products: [ + { + id: "acp_1", + internal_product_id: "prod_int_1", + product_id: "prod_1", + internal_customer_id: "cus_int_1", + created_at: 1, + status: "active", + canceled: false, + starts_at: 1, + options: {}, + collection_method: "charge_automatically", + quantity: 1, + api_semver: null, + is_custom: false, + billing_version: "v1", + external_id: null, + subscription_ids: {}, + scheduled_ids: {}, + }, + ], + aggregated_customer_entitlements: {}, + }, + }); + + test("should coerce top-level arrays from {} to []", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + expect(Array.isArray(result.customer_products)).toBe(true); + expect(result.customer_products).toEqual([]); + expect(Array.isArray(result.products)).toBe(true); + expect(result.products).toEqual([]); + expect(Array.isArray(result.free_trials)).toBe(true); + expect(result.free_trials).toEqual([]); + }); + + test("should coerce meteredFeatures from {} to []", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + expect(Array.isArray(result.meteredFeatures)).toBe(true); + expect(result.meteredFeatures).toEqual([]); + }); + + test("should coerce flags from [] to {}", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + expect(Array.isArray(result.flags)).toBe(false); + expect(typeof result.flags).toBe("object"); + expect(result.flags).toEqual({}); + }); + + test("should coerce customer.auto_topups from {} to []", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + expect(Array.isArray(result.customer.auto_topups)).toBe(true); + }); + + test("should coerce customer.spend_limits from {} to []", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + expect(Array.isArray(result.customer.spend_limits)).toBe(true); + }); + + test("should coerce customer.usage_alerts from {} to []", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + expect(Array.isArray(result.customer.usage_alerts)).toBe(true); + }); + + test("should coerce customer.overage_allowed from {} to []", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + expect(Array.isArray(result.customer.overage_allowed)).toBe(true); + }); + + test("should coerce entity.spend_limits from {} to []", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + expect(Array.isArray(result.entity?.spend_limits)).toBe(true); + }); + + test("should coerce entity.usage_alerts from {} to []", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + expect(Array.isArray(result.entity?.usage_alerts)).toBe(true); + }); + + test("should coerce entity.overage_allowed from {} to []", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + expect(Array.isArray(result.entity?.overage_allowed)).toBe(true); + }); + + test("should coerce entitlements[].feature.event_names from {} to []", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + const entitlements = result.entitlements as Array<{ + feature: { event_names: unknown }; + }>; + expect(entitlements.length).toBe(1); + expect(Array.isArray(entitlements[0].feature.event_names)).toBe(true); + expect(entitlements[0].feature.event_names).toEqual([]); + }); + + test("should coerce prices[].config.usage_tiers from {} to []", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + const prices = result.prices as Array<{ + config: { usage_tiers: unknown }; + }>; + expect(prices.length).toBe(1); + expect(Array.isArray(prices[0].config.usage_tiers)).toBe(true); + }); + + test("should coerce invoices[].product_ids from {} to []", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + expect(result.invoices.length).toBe(1); + expect(Array.isArray(result.invoices[0].product_ids)).toBe(true); + }); + + test("should coerce invoices[].internal_product_ids from {} to []", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + expect(Array.isArray(result.invoices[0].internal_product_ids)).toBe(true); + }); + + test("should coerce invoices[].discounts from {} to []", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + expect(Array.isArray(result.invoices[0].discounts)).toBe(true); + }); + + test("should coerce invoices[].items from {} to []", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + expect(Array.isArray(result.invoices[0].items)).toBe(true); + }); + + test("should coerce entity_aggregations nested arrays", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + const entityAgg = result.entity_aggregations; + expect(entityAgg).toBeDefined(); + expect(Array.isArray(entityAgg?.aggregated_customer_entitlements)).toBe( + true, + ); + const products = entityAgg?.aggregated_customer_products ?? []; + expect(Array.isArray(products)).toBe(true); + expect(products.length).toBe(1); + expect(Array.isArray(products[0].options)).toBe(true); + expect(Array.isArray(products[0].subscription_ids)).toBe(true); + expect(Array.isArray(products[0].scheduled_ids)).toBe(true); + }); + + test("should coerce subscriptions from {} to [] and recurse usage_features", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + (malformed as unknown as Record).subscriptions = [ + { + id: "sub_1", + stripe_id: null, + stripe_schedule_id: null, + created_at: 1, + usage_features: {}, + org_id: "org_1", + current_period_start: null, + current_period_end: null, + env: "live", + }, + ]; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + expect(Array.isArray(result.subscriptions)).toBe(true); + expect(result.subscriptions.length).toBe(1); + expect(Array.isArray(result.subscriptions[0].usage_features)).toBe(true); + }); + + test("should preserve scalar fields untouched", () => { + const malformed = buildMalformedCachedFullSubject() as CachedFullSubject; + const result = sanitizeCachedFullSubject({ + cachedFullSubject: malformed, + }); + expect(result.customerId).toBe("cus_1"); + expect(result.subjectViewEpoch).toBe(1); + expect(result.customer.name).toBe("Test"); + }); +}); diff --git a/shared/utils/featureUtils/apiFeatureToDbFeature.ts b/shared/utils/featureUtils/apiFeatureToDbFeature.ts index 8da0627f7..aa1c5bf1a 100644 --- a/shared/utils/featureUtils/apiFeatureToDbFeature.ts +++ b/shared/utils/featureUtils/apiFeatureToDbFeature.ts @@ -194,13 +194,15 @@ export const dbToApiFeatureV1 = ({ dbFeature.type === FeatureType.CreditSystem || dbFeature.config?.usage_type === FeatureUsageType.Single, - credit_schema: dbFeature.config?.schema?.map( - (schema: CreditSchemaItem) => ({ - metered_feature_id: schema.metered_feature_id, - credit_cost: schema.credit_amount, - }), - ), - event_names: dbFeature.event_names, + credit_schema: Array.isArray(dbFeature.config?.schema) + ? dbFeature.config.schema.map((schema: CreditSchemaItem) => ({ + metered_feature_id: schema.metered_feature_id, + credit_cost: schema.credit_amount, + })) + : undefined, + event_names: Array.isArray(dbFeature.event_names) + ? dbFeature.event_names + : [], archived: dbFeature.archived, display: dbFeature.display diff --git a/shared/utils/fullSubjectUtils/classifyFullSubject.ts b/shared/utils/fullSubjectUtils/classifyFullSubject.ts new file mode 100644 index 000000000..412b3cdff --- /dev/null +++ b/shared/utils/fullSubjectUtils/classifyFullSubject.ts @@ -0,0 +1,28 @@ +import { isUsageBasedAllocatedCustomerEntitlement } from "@utils/cusEntUtils/classifyCusEntUtils.js"; +import { isAllocatedFeature } from "@utils/featureUtils/classifyFeature/isAllocatedFeature.js"; +import type { FullSubject } from "../../models/cusModels/fullSubject/fullSubjectModel.js"; +import type { CusProductStatus } from "../../models/cusProductModels/cusProductEnums.js"; +import type { Feature } from "../../models/featureModels/featureModels.js"; +import { fullSubjectToCustomerEntitlements } from "./fullSubjectToCustomerEntitlements.js"; + +/** Checks if a FullSubject has any usage-based allocated (continuous + allocated) pricing for the given features. */ +export const fullSubjectHasUsageBasedAllocated = ({ + fullSubject, + features, + inStatuses, +}: { + fullSubject: FullSubject; + features: Feature[]; + inStatuses?: CusProductStatus[]; +}): boolean => { + const allocatedFeatures = features.filter(isAllocatedFeature); + const customerEntitlements = fullSubjectToCustomerEntitlements({ + fullSubject, + featureIds: allocatedFeatures.map((feature) => feature.id), + inStatuses, + }); + + return customerEntitlements.some((customerEntitlement) => + isUsageBasedAllocatedCustomerEntitlement(customerEntitlement), + ); +}; diff --git a/shared/utils/fullSubjectUtils/fullSubjectToFullCustomer.ts b/shared/utils/fullSubjectUtils/fullSubjectToFullCustomer.ts new file mode 100644 index 000000000..2407730d6 --- /dev/null +++ b/shared/utils/fullSubjectUtils/fullSubjectToFullCustomer.ts @@ -0,0 +1,16 @@ +import type { FullCustomer, FullSubject } from "../../index.js"; + +/** Converts a FullSubject back to a FullCustomer for legacy helper compatibility. */ +export const fullSubjectToFullCustomer = ({ + fullSubject, +}: { + fullSubject: FullSubject; +}): FullCustomer => ({ + ...fullSubject.customer, + customer_products: fullSubject.customer_products, + entities: fullSubject.entity ? [fullSubject.entity] : [], + entity: fullSubject.entity, + extra_customer_entitlements: fullSubject.extra_customer_entitlements, + subscriptions: fullSubject.subscriptions, + invoices: fullSubject.invoices, +}); diff --git a/shared/utils/fullSubjectUtils/index.ts b/shared/utils/fullSubjectUtils/index.ts index 2d9a53d73..5cf17d8d4 100644 --- a/shared/utils/fullSubjectUtils/index.ts +++ b/shared/utils/fullSubjectUtils/index.ts @@ -1,7 +1,9 @@ export * from "./aggregatedUtils/index.js"; +export { fullSubjectHasUsageBasedAllocated } from "./classifyFullSubject.js"; export { fullCustomerToFullSubject } from "./fullCustomerToFullSubject.js"; export { fullSubjectToApiCustomerProducts } from "./fullSubjectToApiCustomerProducts.js"; export { fullSubjectToCustomerEntitlements } from "./fullSubjectToCustomerEntitlements.js"; +export { fullSubjectToFullCustomer } from "./fullSubjectToFullCustomer.js"; export { fullSubjectToOverageAllowedByFeatureId } from "./fullSubjectToOverageAllowed.js"; export { fullSubjectToSpendLimitByFeatureId, diff --git a/shared/utils/fullSubjectUtils/normalizedToFullSubject.ts b/shared/utils/fullSubjectUtils/normalizedToFullSubject.ts index 705e8e11a..1dfd456be 100644 --- a/shared/utils/fullSubjectUtils/normalizedToFullSubject.ts +++ b/shared/utils/fullSubjectUtils/normalizedToFullSubject.ts @@ -13,6 +13,23 @@ import type { Replaceable } from "../../models/cusProductModels/cusEntModels/rep import type { FullCustomerPrice } from "../../models/cusProductModels/cusPriceModels/cusPriceModels.js"; import type { FullCusProduct } from "../../models/cusProductModels/cusProductModels.js"; +const getArrayEntries = ({ value }: { value: unknown }): T[] => + Array.isArray(value) ? (value as T[]) : []; + +const getObjectEntries = >({ + value, + fallback, +}: { + value: unknown; + fallback: T; +}): T => { + if (!value || Array.isArray(value) || typeof value !== "object") { + return fallback; + } + + return value as T; +}; + const getRolloverSortValue = ({ rollover, }: { @@ -24,6 +41,10 @@ const subjectBalanceToFullCustomerEntitlement = ({ }: { subjectBalance: SubjectBalance; }): FullCustomerEntitlement => { + const rollovers = getArrayEntries({ + value: subjectBalance.rollovers, + }); + return { id: subjectBalance.id, internal_customer_id: subjectBalance.internal_customer_id, @@ -46,7 +67,7 @@ const subjectBalanceToFullCustomerEntitlement = ({ customer_id: subjectBalance.customer_id, entitlement: subjectBalance.entitlement, replaceables: [] as Replaceable[], - rollovers: [...subjectBalance.rollovers].sort( + rollovers: [...rollovers].sort( (left, right) => getRolloverSortValue({ rollover: left }) - getRolloverSortValue({ rollover: right }), @@ -96,30 +117,66 @@ export const normalizedToFullSubject = ({ }: { normalized: NormalizedFullSubject; }): FullSubject => { + const entitlements = getArrayEntries< + NormalizedFullSubject["entitlements"][number] + >({ + value: normalized.entitlements, + }); + const products = getArrayEntries({ + value: normalized.products, + }); + const prices = getArrayEntries({ + value: normalized.prices, + }); + const freeTrials = getArrayEntries< + NormalizedFullSubject["free_trials"][number] + >({ + value: normalized.free_trials, + }); + const customerPrices = getArrayEntries< + NormalizedFullSubject["customer_prices"][number] + >({ + value: normalized.customer_prices, + }); + const customerEntitlements = getArrayEntries< + NormalizedFullSubject["customer_entitlements"][number] + >({ + value: normalized.customer_entitlements, + }); + const customerProductsInput = getArrayEntries< + NormalizedFullSubject["customer_products"][number] + >({ + value: normalized.customer_products, + }); + const subscriptions = getArrayEntries< + NormalizedFullSubject["subscriptions"][number] + >({ + value: normalized.subscriptions, + }); + const invoices = getArrayEntries({ + value: normalized.invoices, + }); + const flags = getObjectEntries({ + value: normalized.flags, + fallback: {}, + }); + const entitlementsById = new Map( - normalized.entitlements.map( - (entitlement) => [entitlement.id, entitlement] as const, - ), + entitlements.map((entitlement) => [entitlement.id, entitlement] as const), ); const productsByInternalId = new Map( - normalized.products.map( - (product) => [product.internal_id, product] as const, - ), - ); - const pricesById = new Map( - normalized.prices.map((price) => [price.id, price] as const), + products.map((product) => [product.internal_id, product] as const), ); + const pricesById = new Map(prices.map((price) => [price.id, price] as const)); const freeTrialsById = new Map( - normalized.free_trials.map( - (freeTrial) => [freeTrial.id, freeTrial] as const, - ), + freeTrials.map((freeTrial) => [freeTrial.id, freeTrial] as const), ); const customerPricesByCustomerProductId = new Map< string, FullCustomerPrice[] >(); - for (const customerPrice of normalized.customer_prices) { + for (const customerPrice of customerPrices) { if (!customerPrice.customer_product_id) continue; const price = customerPrice.price_id @@ -139,7 +196,7 @@ export const normalizedToFullSubject = ({ ); } - for (const customerEntitlement of normalized.customer_entitlements) { + for (const customerEntitlement of customerEntitlements) { if (!customerEntitlement.customer_product_id) continue; if (!customerEntitlement.customerPrice) continue; @@ -165,7 +222,7 @@ export const normalizedToFullSubject = ({ >(); const extraMeteredCes: FullCustomerEntitlement[] = []; - for (const customerEntitlement of normalized.customer_entitlements) { + for (const customerEntitlement of customerEntitlements) { const fullCustomerEntitlement = subjectBalanceToFullCustomerEntitlement({ subjectBalance: customerEntitlement, }); @@ -191,7 +248,7 @@ export const normalizedToFullSubject = ({ >(); const extraBooleanCes: FullCustomerEntitlement[] = []; - for (const flag of Object.values(normalized.flags)) { + for (const flag of Object.values(flags)) { const entitlement = entitlementsById.get(flag.entitlementId); if (!entitlement) continue; @@ -211,7 +268,7 @@ export const normalizedToFullSubject = ({ } const customerProducts: FullCusProduct[] = []; - for (const customerProduct of normalized.customer_products) { + for (const customerProduct of customerProductsInput) { const product = productsByInternalId.get( customerProduct.internal_product_id, ); @@ -243,9 +300,23 @@ export const normalizedToFullSubject = ({ if (normalized.entity_aggregations) { const entityAgg = normalized.entity_aggregations; + const aggregatedCustomerProductsInput = getArrayEntries< + NonNullable< + NormalizedFullSubject["entity_aggregations"] + >["aggregated_customer_products"][number] + >({ + value: entityAgg.aggregated_customer_products, + }); + const aggregatedCustomerEntitlementsInput = getArrayEntries< + NonNullable< + NormalizedFullSubject["entity_aggregations"] + >["aggregated_customer_entitlements"][number] + >({ + value: entityAgg.aggregated_customer_entitlements, + }); aggregatedCustomerProducts = []; - for (const entityCusProduct of entityAgg.aggregated_customer_products) { + for (const entityCusProduct of aggregatedCustomerProductsInput) { const product = productsByInternalId.get( entityCusProduct.internal_product_id, ); @@ -262,11 +333,9 @@ export const normalizedToFullSubject = ({ } as FullCusProduct); } - aggregatedCustomerEntitlements = ( - entityAgg.aggregated_customer_entitlements ?? [] - ) + aggregatedCustomerEntitlements = aggregatedCustomerEntitlementsInput .map((aggregatedCusEnt) => { - const feature = normalized.entitlements.find( + const feature = entitlements.find( (entitlement) => entitlement.internal_feature_id === aggregatedCusEnt.internal_feature_id, @@ -294,8 +363,8 @@ export const normalizedToFullSubject = ({ customer: normalized.customer, customer_products: customerProducts, extra_customer_entitlements: extraCustomerEntitlements, - subscriptions: normalized.subscriptions ?? [], - invoices: normalized.invoices ?? [], + subscriptions, + invoices, ...(aggregatedCustomerProducts ? { aggregated_customer_products: aggregatedCustomerProducts } : {}),