From cc1c527e17a0facd0fdfca89eff3a36742f09bb4 Mon Sep 17 00:00:00 2001 From: John Yeo Date: Mon, 26 Jan 2026 21:22:36 +0000 Subject: [PATCH] fix: one time product race condition --- .../add-product/createOneTimeCusProduct.ts | 28 +++++++++++++++ .../handleQuantityDowngrade.ts | 19 +++++++++++ .../handleQuantityUpgrade.ts | 19 +++++++++++ .../customers/cusProducts/AttachParams.ts | 3 +- .../attach/migrations/migration3.test.ts | 2 ++ server/tests/attach/others/others2.test.ts | 1 + .../track-race-condition2.test.ts | 34 ++++++++----------- .../attach/attach-update-quantity.test.ts | 20 +++++------ 8 files changed, 95 insertions(+), 31 deletions(-) diff --git a/server/src/internal/customers/add-product/createOneTimeCusProduct.ts b/server/src/internal/customers/add-product/createOneTimeCusProduct.ts index 8f7bead32..537a9fee0 100644 --- a/server/src/internal/customers/add-product/createOneTimeCusProduct.ts +++ b/server/src/internal/customers/add-product/createOneTimeCusProduct.ts @@ -13,6 +13,8 @@ import { } from "@autumn/shared"; import type { DrizzleCli } from "@/db/initDrizzle.js"; import { addProductsUpdatedWebhookTask } from "@/internal/analytics/handlers/handleProductsUpdated.js"; +import { executeRedisDeduction } from "@/internal/balances/utils/deduction/executeRedisDeduction.js"; +import type { FeatureDeduction } from "@/internal/balances/utils/types/featureDeduction.js"; import { triggerVerifyCacheConsistency } from "@/internal/billing/v2/workflows/verifyCacheConsistency/triggerVerifyCacheConsistency.js"; import { getEntRelatedPrice } from "@/internal/products/entitlements/entitlementUtils.js"; import { getEntOptions } from "@/internal/products/prices/priceUtils.js"; @@ -32,6 +34,7 @@ const updateOneOffExistingEntitlement = async ({ options, relatedPrice, logger, + attachParams, }: { db: DrizzleCli; cusEnt: FullCustomerEntitlement; @@ -41,6 +44,7 @@ const updateOneOffExistingEntitlement = async ({ options?: FeatureOptions; relatedPrice?: Price; logger: any; + attachParams: InsertCusProductParams; }) => { if (entitlement.allowance_type === AllowanceType.Unlimited) { return; @@ -76,6 +80,29 @@ const updateOneOffExistingEntitlement = async ({ }, }); + const context = attachParams.req; + if (context) { + const featureDeductions: FeatureDeduction[] = [ + { + feature: entitlement.feature, + deduction: -resetBalance, + }, + ]; + + try { + await executeRedisDeduction({ + ctx: context, + fullCustomer: attachParams.customer, + deductions: featureDeductions, + deductionOptions: { + overageBehaviour: "allow", + }, + }); + } catch (error) { + logger.warn(`Failed to execute Redis deduction: ${error}`); + } + } + return; }; @@ -128,6 +155,7 @@ export const updateOneTimeCusProduct = async ({ options: options || undefined, relatedPrice, logger, + attachParams, }); } else { const newCusEnt = initCusEntitlement({ diff --git a/server/src/internal/customers/attach/attachFunctions/updateQuantityFlow/handleQuantityDowngrade.ts b/server/src/internal/customers/attach/attachFunctions/updateQuantityFlow/handleQuantityDowngrade.ts index 89d435333..0a7251350 100644 --- a/server/src/internal/customers/attach/attachFunctions/updateQuantityFlow/handleQuantityDowngrade.ts +++ b/server/src/internal/customers/attach/attachFunctions/updateQuantityFlow/handleQuantityDowngrade.ts @@ -17,6 +17,7 @@ import { import { Decimal } from "decimal.js"; import type { Stripe } from "stripe"; import { subToPeriodStartEnd } from "@/external/stripe/stripeSubUtils/convertSubUtils.js"; +import { executeRedisDeduction } from "@/internal/balances/utils/deduction/executeRedisDeduction"; import type { AttachParams } from "@/internal/customers/cusProducts/AttachParams.js"; import { CusEntService } from "@/internal/customers/cusProducts/cusEnts/CusEntitlementService.js"; import { InvoiceService } from "@/internal/invoices/InvoiceService.js"; @@ -193,5 +194,23 @@ export const handleQuantityDowngrade = async ({ id: cusEnt.id, amount: decrementBy, }); + + try { + await executeRedisDeduction({ + ctx, + fullCustomer: attachParams.customer, + deductions: [ + { + feature: cusEnt.entitlement.feature, + deduction: decrementBy, + }, + ], + deductionOptions: { + overageBehaviour: "allow", + }, + }); + } catch (error) { + logger.warn(`Failed to execute Redis deduction: ${error}`); + } } }; diff --git a/server/src/internal/customers/attach/attachFunctions/updateQuantityFlow/handleQuantityUpgrade.ts b/server/src/internal/customers/attach/attachFunctions/updateQuantityFlow/handleQuantityUpgrade.ts index 01d7497bc..572091708 100644 --- a/server/src/internal/customers/attach/attachFunctions/updateQuantityFlow/handleQuantityUpgrade.ts +++ b/server/src/internal/customers/attach/attachFunctions/updateQuantityFlow/handleQuantityUpgrade.ts @@ -17,6 +17,7 @@ import { import { Decimal } from "decimal.js"; import type { Stripe } from "stripe"; import { subToPeriodStartEnd } from "@/external/stripe/stripeSubUtils/convertSubUtils.js"; +import { executeRedisDeduction } from "@/internal/balances/utils/deduction/executeRedisDeduction"; import type { AttachParams } from "@/internal/customers/cusProducts/AttachParams.js"; import { CusEntService } from "@/internal/customers/cusProducts/cusEnts/CusEntitlementService.js"; import { InvoiceService } from "@/internal/invoices/InvoiceService.js"; @@ -201,6 +202,24 @@ export const handleQuantityUpgrade = async ({ id: cusEnt.id, amount: incrementBy, }); + + try { + await executeRedisDeduction({ + ctx, + fullCustomer: attachParams.customer, + deductions: [ + { + feature: cusEnt.entitlement.feature, + deduction: -incrementBy, + }, + ], + deductionOptions: { + overageBehaviour: "allow", + }, + }); + } catch (error) { + logger.warn(`Failed to execute Redis deduction: ${error}`); + } } return { invoice }; }; diff --git a/server/src/internal/customers/cusProducts/AttachParams.ts b/server/src/internal/customers/cusProducts/AttachParams.ts index e5ddcdba2..b6d84eea4 100644 --- a/server/src/internal/customers/cusProducts/AttachParams.ts +++ b/server/src/internal/customers/cusProducts/AttachParams.ts @@ -4,7 +4,6 @@ import type { AttachConfig, AttachReplaceable, AttachScenario, - Customer, EntitlementWithFeature, Entity, Feature, @@ -96,7 +95,7 @@ export type InsertCusProductParams = { req?: AutumnContext; now?: number; - customer: Customer; + customer: FullCustomer; org: Organization; product: FullProduct; prices: Price[]; diff --git a/server/tests/attach/migrations/migration3.test.ts b/server/tests/attach/migrations/migration3.test.ts index ad339a7c6..bc9a82746 100644 --- a/server/tests/attach/migrations/migration3.test.ts +++ b/server/tests/attach/migrations/migration3.test.ts @@ -7,6 +7,7 @@ import { import { defaultApiVersion } from "@tests/constants.js"; import { TestFeature } from "@tests/setup/v2Features.js"; import { attachAndExpectCorrect } from "@tests/utils/expectUtils/expectAttach.js"; +import { timeout } from "@tests/utils/genUtils.js"; import { advanceTestClock } from "@tests/utils/stripeUtils.js"; import ctx from "@tests/utils/testInitUtils/createTestContext.js"; import chalk from "chalk"; @@ -100,6 +101,7 @@ describe(`${chalk.yellowBright(`${testCase}: Testing migration for pro with tria test("should attach track usage and get correct balance", async () => { const wordsUsage = 120000; + await timeout(2000); await autumn.track({ customer_id: customerId, value: wordsUsage, diff --git a/server/tests/attach/others/others2.test.ts b/server/tests/attach/others/others2.test.ts index 623e0bb04..1778824f5 100644 --- a/server/tests/attach/others/others2.test.ts +++ b/server/tests/attach/others/others2.test.ts @@ -38,6 +38,7 @@ describe(`${chalk.yellowBright(`${testCase}: Testing one-off`)}`, () => { ctx, products: [oneOff], prefix: testCase, + customerId, }); const { testClockId: testClockId1 } = await initCustomerV3({ diff --git a/server/tests/balances/track/race-condition/track-race-condition2.test.ts b/server/tests/balances/track/race-condition/track-race-condition2.test.ts index adbed1601..254cecd0d 100644 --- a/server/tests/balances/track/race-condition/track-race-condition2.test.ts +++ b/server/tests/balances/track/race-condition/track-race-condition2.test.ts @@ -8,11 +8,7 @@ import { currentRegion } from "@/external/redis/initRedis.js"; import { executeRedisDeduction } from "@/internal/balances/utils/deduction/executeRedisDeduction.js"; import { syncItemV3 } from "@/internal/balances/utils/sync/syncItemV3.js"; import { getOrSetCachedFullCustomer } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/getOrSetCachedFullCustomer.js"; -import { constructFeatureItem } from "@/utils/scriptUtils/constructItem.js"; -import { - constructProduct, - constructRawProduct, -} from "@/utils/scriptUtils/createTestProducts.js"; +import { constructRawProduct } from "@/utils/scriptUtils/createTestProducts.js"; import { initCustomerV3 } from "@/utils/scriptUtils/testUtils/initCustomerV3.js"; import { initProductsV0 } from "@/utils/scriptUtils/testUtils/initProductsV0.js"; import { deleteCachedFullCustomer } from "../../../../src/internal/customers/cusUtils/fullCustomerCacheUtils/deleteCachedFullCustomer.js"; @@ -23,15 +19,15 @@ import { import { constructPrepaidItem } from "../../../../src/utils/scriptUtils/constructItem.js"; import { timeout } from "../../../utils/genUtils"; -const pro = constructProduct({ - type: "pro", - items: [ - constructFeatureItem({ - featureId: TestFeature.Messages, - includedUsage: 100, - }), - ], -}); +// const pro = constructProduct({ +// type: "pro", +// items: [ +// constructFeatureItem({ +// featureId: TestFeature.Messages, +// includedUsage: 100, +// }), +// ], +// }); const oneOffCredits = constructRawProduct({ id: "one_off_messages", @@ -72,13 +68,13 @@ describe(`${chalk.yellowBright("track-race-condition2: sync should not wipe out await initProductsV0({ ctx, - products: [pro, oneOffCredits], + products: [oneOffCredits], prefix: testCase, }); await autumnV2.attach({ customer_id: customerId, - product_ids: [pro.id, oneOffCredits.id], + product_ids: [oneOffCredits.id], options: [ { feature_id: TestFeature.Messages, @@ -122,7 +118,7 @@ describe(`${chalk.yellowBright("track-race-condition2: sync should not wipe out deductions: [ { feature: messagesFeature, - deduction: 5, + deduction: 95, // use up a BIT of one-off credits }, ], fullCustomer, @@ -209,7 +205,7 @@ describe(`${chalk.yellowBright("track-race-condition2: sync should not wipe out // Expected: 100 (pro) + 100 (initial one-off) - 5 (tracked) + 100 (attached one-off) = 295 expect(cachedCustomer.balances[TestFeature.Messages].current_balance).toBe( - 295, + 105, ); const customerAfterSync = await autumnV2.customers.get( @@ -220,6 +216,6 @@ describe(`${chalk.yellowBright("track-race-condition2: sync should not wipe out ); expect( customerAfterSync.balances[TestFeature.Messages].current_balance, - ).toBe(295); + ).toBe(105); }); }); diff --git a/server/tests/integration/billing/legacy/attach/attach-update-quantity.test.ts b/server/tests/integration/billing/legacy/attach/attach-update-quantity.test.ts index 34aeec083..1ae08d734 100644 --- a/server/tests/integration/billing/legacy/attach/attach-update-quantity.test.ts +++ b/server/tests/integration/billing/legacy/attach/attach-update-quantity.test.ts @@ -175,7 +175,7 @@ test.concurrent(`${chalk.yellowBright("attach: quantity decrease → increase // Decrease to 200 (on_decrease: none - sets upcoming_quantity, no immediate change) await autumnV1.attach({ customer_id: customerId, - product_id: `${pro.id}_${customerId}`, + product_id: pro.id, options: [{ feature_id: TestFeature.Messages, quantity: 200 }], }); @@ -190,14 +190,14 @@ test.concurrent(`${chalk.yellowBright("attach: quantity decrease → increase customer: customerAfterDecrease, productId: pro.id, featureId: TestFeature.Messages, - quantity: 3, // Still 300 / 100 - upcomingQuantity: 2, // 200 / 100 + quantity: 300, // Still 300 / 100 + upcomingQuantity: 200, // 200 / 100 }); // Increase to 400 (prorate_immediately - creates invoice, immediate change) await autumnV1.attach({ customer_id: customerId, - product_id: `${pro.id}_${customerId}`, + product_id: pro.id, options: [{ feature_id: TestFeature.Messages, quantity: 400 }], }); @@ -219,7 +219,7 @@ test.concurrent(`${chalk.yellowBright("attach: quantity decrease → increase // Decrease back to 200 (sets upcoming_quantity again) await autumnV1.attach({ customer_id: customerId, - product_id: `${pro.id}_${customerId}`, + product_id: pro.id, options: [{ feature_id: TestFeature.Messages, quantity: 200 }], }); @@ -231,8 +231,8 @@ test.concurrent(`${chalk.yellowBright("attach: quantity decrease → increase customer: customerFinal, productId: pro.id, featureId: TestFeature.Messages, - quantity: 4, // Still 400 / 100 - upcomingQuantity: 2, // 200 / 100 + quantity: 400, // Still 400 / 100 + upcomingQuantity: 200, // 200 / 100 }); }); @@ -358,8 +358,8 @@ test.concurrent(`${chalk.yellowBright("attach: prepaid add-on with entities - up customer: entity2, productId: prepaidAddOn.id, featureId: TestFeature.Messages, - quantity: entity2OriginalQuantity / 100, // billingUnits = 100 - upcomingQuantity: entity2DowngradedQuantity / 100, + quantity: entity2OriginalQuantity, // billingUnits = 100 + upcomingQuantity: entity2DowngradedQuantity, }); }); @@ -438,7 +438,7 @@ test.concurrent(`${chalk.yellowBright("attach: quantity upgrade with prorate-nex customer: customerFinal, productId: pro.id, featureId: TestFeature.Messages, - quantity: 4, // 400 / 100 billingUnits + quantity: 400, upcomingQuantity: "undefined", // No upcoming_quantity since it's an upgrade });