From 92c4c9c4c0bb866bcdc25ac68f4e54b08ef78a34 Mon Sep 17 00:00:00 2001 From: John Yeo Date: Thu, 21 Aug 2025 16:10:35 -0700 Subject: [PATCH] fix: track for cont use race condition --- .../external/autumn/autumnWebhookRouter.ts | 36 ++++-- .../handleCusProductDeleted.ts | 5 - .../handlers/handleProductsUpdated.ts | 34 ++--- server/src/internal/api/events/usageRouter.ts | 1 + .../customers/cusCache/updateCachedCus.ts | 46 +++---- server/src/queue/workersInit.ts | 2 +- server/src/trigger/updateBalanceTask.ts | 15 ++- server/src/trigger/updateUsageTask.ts | 3 +- server/tests/contUse/track/track6.ts | 116 ++++++++++++++++++ shared/models/orgModels/orgConfig.ts | 2 + 10 files changed, 204 insertions(+), 56 deletions(-) create mode 100644 server/tests/contUse/track/track6.ts diff --git a/server/src/external/autumn/autumnWebhookRouter.ts b/server/src/external/autumn/autumnWebhookRouter.ts index ddd1f3226..7a7d612b8 100644 --- a/server/src/external/autumn/autumnWebhookRouter.ts +++ b/server/src/external/autumn/autumnWebhookRouter.ts @@ -44,15 +44,29 @@ autumnWebhookRouter.post( "", express.raw({ type: "application/json" }), async (req, res) => { - const evt = await verifyAutumnWebhook(req, res); - console.log("Received webhook from autumn"); - const { type, data } = evt; - console.log("Type", type, "Scenario:", data?.scenario); - // console.log("Data", data); - // console.log("JSON", JSON.stringify(evt, null, 2)); - res.status(200).json({ - success: true, - message: "Webhook received", - }); - }, + try { + const evt = await verifyAutumnWebhook(req, res); + console.log("Received webhook from autumn"); + const { type, data } = evt; + console.log( + "Type", + type, + "Scenario:", + data?.scenario, + "Product:", + data?.updated_product?.id + ); + + res.status(200).json({ + success: true, + message: "Webhook received", + }); + } catch (error) { + res.status(200).json({ + success: false, + message: "Error: Could not verify webhook", + }); + return; + } + } ); diff --git a/server/src/external/stripe/webhookHandlers/handleSubDeleted/handleCusProductDeleted.ts b/server/src/external/stripe/webhookHandlers/handleSubDeleted/handleCusProductDeleted.ts index 977ba3dbc..8d46dd8d0 100644 --- a/server/src/external/stripe/webhookHandlers/handleSubDeleted/handleCusProductDeleted.ts +++ b/server/src/external/stripe/webhookHandlers/handleSubDeleted/handleCusProductDeleted.ts @@ -82,11 +82,6 @@ export const handleCusProductDeleted = async ({ } } - // if (cusProduct.status === CusProductStatus.Expired) { - // // When attaching eg. main is trial, canceled in attach function, don't handle... - // return; - // } - if (scheduled_ids && scheduled_ids.length > 0 && !prematurelyCanceled) { logger.info( `sub.deleted: removing sub_id from cus product ${cusProduct.id}` diff --git a/server/src/internal/analytics/handlers/handleProductsUpdated.ts b/server/src/internal/analytics/handlers/handleProductsUpdated.ts index c8e927b98..cf6f6599b 100644 --- a/server/src/internal/analytics/handlers/handleProductsUpdated.ts +++ b/server/src/internal/analytics/handlers/handleProductsUpdated.ts @@ -22,6 +22,12 @@ import { ExtendedRequest } from "@/utils/models/Request.js"; import { ActionService } from "@/internal/analytics/ActionService.js"; import { constructAction } from "@/internal/analytics/actionUtils.js"; import { parseReqForAction } from "@/internal/analytics/actionUtils.js"; +import { RELEVANT_STATUSES } from "@/internal/customers/cusProducts/CusProductService.js"; +import { + cusProductToEnts, + cusProductToPrices, + cusProductToProduct, +} from "@/internal/customers/cusProducts/cusProductUtils/convertCusProduct.js"; interface ActionDetails { request_id: string; @@ -119,29 +125,25 @@ export const handleProductsUpdated = async ({ // Product: let product = cusProduct.product; - let prices = cusProduct.customer_prices.map((cp) => cp.price); - let entitlements = cusProduct.customer_entitlements.map( - (ce) => ce.entitlement - ); - let freeTrial = cusProduct.free_trial; + // const prices = cusProductToPrices({ cusProduct }); + // const ents = cusProductToEnts({ cusProduct }); + // let freeTrial = cusProduct.free_trial; - let fullProduct: FullProduct = { - ...product, - prices, - entitlements, - free_trial: freeTrial || null, - }; + let fullProduct: FullProduct = cusProductToProduct({ cusProduct }); + + // { + // ...product, + // prices, + // entitlements: ents, + // free_trial: freeTrial || null, + // }; let customer = await CusService.getFull({ db, idOrInternalId: data.customerId || data.internalCustomerId, orgId: data.org.id, env: data.env, - inStatuses: [ - CusProductStatus.Active, - CusProductStatus.Scheduled, - CusProductStatus.Expired, - ], + inStatuses: RELEVANT_STATUSES, entityId: cusProduct.internal_entity_id || undefined, }); diff --git a/server/src/internal/api/events/usageRouter.ts b/server/src/internal/api/events/usageRouter.ts index d5d046922..2770c4cb3 100644 --- a/server/src/internal/api/events/usageRouter.ts +++ b/server/src/internal/api/events/usageRouter.ts @@ -193,6 +193,7 @@ export const handleUsageEvent = async ({ const payload = { customerId: customer.id, internalCustomerId: customer.internal_id, + eventId: newEvent.id, features, org, env: req.env, diff --git a/server/src/internal/customers/cusCache/updateCachedCus.ts b/server/src/internal/customers/cusCache/updateCachedCus.ts index 40e713b67..87dba4e45 100644 --- a/server/src/internal/customers/cusCache/updateCachedCus.ts +++ b/server/src/internal/customers/cusCache/updateCachedCus.ts @@ -39,31 +39,35 @@ export const refreshCusCache = async ({ const list = await upstash.keys(`${baseKey}*`); + const promises = []; for (const key of list) { - const keyName = key; - let params = keyName.split(":"); - let expandParam = params.find((p) => p.startsWith("expand_")); - let expand = expandParam - ? expandParam.replace("expand_", "").split(",") - : []; + const refresh = async () => { + const keyName = key; + let params = keyName.split(":"); + let expandParam = params.find((p) => p.startsWith("expand_")); + let expand = expandParam + ? expandParam.replace("expand_", "").split(",") + : []; - let entityIdParam = params.find((p) => p.startsWith("entity_")); - let entityId = entityIdParam - ? entityIdParam.replace("entity_", "") - : undefined; + let entityIdParam = params.find((p) => p.startsWith("entity_")); + let entityId = entityIdParam + ? entityIdParam.replace("entity_", "") + : undefined; - await getCusWithCache({ - db, - idOrInternalId: customerId, - org, - env, - expand: expand as CusExpand[], - entityId, - skipGet: true, - logger: console, - }); - // console.log(`updated cache key: ${keyName}`); + await getCusWithCache({ + db, + idOrInternalId: customerId, + org, + env, + expand: expand as CusExpand[], + entityId, + skipGet: true, + logger: console, + }); + }; + promises.push(refresh()); } + await Promise.all(promises); } catch (error) { logger.error("Failed to update cache:", { error }); } diff --git a/server/src/queue/workersInit.ts b/server/src/queue/workersInit.ts index 9cd3dfc27..e7537d50d 100644 --- a/server/src/queue/workersInit.ts +++ b/server/src/queue/workersInit.ts @@ -141,7 +141,7 @@ const initWorker = ({ })) ) { await queue.add(job.name, job.data, { - delay: 50, + delay: 200, }); return; } diff --git a/server/src/trigger/updateBalanceTask.ts b/server/src/trigger/updateBalanceTask.ts index eb0b3575a..256ee068d 100644 --- a/server/src/trigger/updateBalanceTask.ts +++ b/server/src/trigger/updateBalanceTask.ts @@ -10,6 +10,7 @@ import { Organization, FullCusEntWithFullCusProduct, BillingType, + FeatureUsageType, } from "@autumn/shared"; import { CusEntService } from "@/internal/customers/cusProducts/cusEnts/CusEntitlementService.js"; import { Customer, FeatureType } from "@autumn/shared"; @@ -461,13 +462,25 @@ export const deductFromUsageBasedCusEnt = async ({ const { db, feature, env, org, cusPrices, customer, entity } = deductParams; // Deduct from usage-based price - const usageBasedEnt = findCusEnt({ + let usageBasedEnt = findCusEnt({ cusEnts, feature, entity, onlyUsageAllowed: true, }) as FullCusEntWithFullCusProduct; + if ( + !usageBasedEnt && + feature.config?.usage_type == FeatureUsageType.Continuous + ) { + console.log(`FALLING BACK TO REGULAR CUS ENT, FEATURE: ${feature.id}`); + usageBasedEnt = findCusEnt({ + cusEnts, + feature, + entity, + }) as FullCusEntWithFullCusProduct; // fallback to regular cus ent if allowed... + } + if (!usageBasedEnt) { console.log( ` - Feature ${feature.id}, To deduct: ${toDeduct} -> no usage-based entitlement found` diff --git a/server/src/trigger/updateUsageTask.ts b/server/src/trigger/updateUsageTask.ts index 942643f64..746e59268 100644 --- a/server/src/trigger/updateUsageTask.ts +++ b/server/src/trigger/updateUsageTask.ts @@ -305,6 +305,7 @@ export const runUpdateUsageTask = async ({ const { internalCustomerId, customerId, + eventId, features, value, set_usage, @@ -316,7 +317,7 @@ export const runUpdateUsageTask = async ({ console.log("--------------------------------"); console.log( - `HANDLING USAGE TASK FOR CUSTOMER (${customerId}), ORG: ${org.slug}` + `HANDLING USAGE TASK FOR CUSTOMER (${customerId}), ORG: ${org.slug}, EVENT ID: ${eventId}` ); const cusEnts: any = await updateUsage({ diff --git a/server/tests/contUse/track/track6.ts b/server/tests/contUse/track/track6.ts new file mode 100644 index 000000000..722947a01 --- /dev/null +++ b/server/tests/contUse/track/track6.ts @@ -0,0 +1,116 @@ +import { AutumnInt } from "@/external/autumn/autumnCli.js"; +import { initCustomer } from "@/utils/scriptUtils/initCustomer.js"; +import { APIVersion, AppEnv, LimitedItem, Organization } from "@autumn/shared"; +import chalk from "chalk"; +import Stripe from "stripe"; +import { DrizzleCli } from "@/db/initDrizzle.js"; +import { setupBefore } from "tests/before.js"; +import { createProducts } from "tests/utils/productUtils.js"; +import { addPrefixToProducts } from "tests/attach/utils.js"; +import { constructProduct } from "@/utils/scriptUtils/createTestProducts.js"; +import { constructFeatureItem } from "@/utils/scriptUtils/constructItem.js"; +import { TestFeature } from "tests/setup/v2Features.js"; +import { timeout } from "@/utils/genUtils.js"; +import { expect } from "chai"; + +const userItem = constructFeatureItem({ + featureId: TestFeature.Users, + includedUsage: 5, +}) as LimitedItem; + +export let free = constructProduct({ + items: [userItem], + type: "free", + isDefault: false, +}); + +const testCase = "track6"; + +describe(`${chalk.yellowBright(`${testCase}: Testing track cont use, race condition`)}`, () => { + let customerId = testCase; + let autumn: AutumnInt = new AutumnInt({ version: APIVersion.v1_4 }); + let testClockId: string; + let db: DrizzleCli, org: Organization, env: AppEnv; + let stripeCli: Stripe; + + let curUnix = new Date().getTime(); + + before(async function () { + await setupBefore(this); + const { autumnJs } = this; + db = this.db; + org = this.org; + env = this.env; + + stripeCli = this.stripeCli; + + addPrefixToProducts({ + products: [free], + prefix: testCase, + }); + + await createProducts({ + autumn, + products: [free], + customerId, + db, + orgId: org.id, + env, + }); + + const { testClockId: testClockId1 } = await initCustomer({ + autumn: autumnJs, + customerId, + db, + org, + env, + attachPm: "success", + }); + + testClockId = testClockId1!; + }); + + it("should track 5 events in a row and have correct balance", async function () { + let startingBalance = userItem.included_usage; + await autumn.attach({ + customer_id: customerId, + product_id: free.id, + }); + + const promises = []; + for (let i = 0; i < 2; i++) { + console.log("--------------------------------"); + console.log(`Cycle ${i}`); + console.log(`Starting balance: ${startingBalance}`); + const values = []; + for (let i = 0; i < 10; i++) { + const randomVal = + Math.floor(Math.random() * 5) * (Math.random() < 0.3 ? -1 : 1); + promises.push( + autumn.track({ + customer_id: customerId, + feature_id: TestFeature.Users, + value: randomVal, + }) + ); + startingBalance -= randomVal; + values.push(randomVal); + } + + console.log(`New balance: ${startingBalance}`); + + const results = await Promise.all(promises); + + await timeout(10000); + + let customer = await autumn.customers.get(customerId); + let userFeature = customer.features[TestFeature.Users]; + if (userFeature.balance != startingBalance) { + for (let i = 0; i < values.length; i++) { + console.log(`Value: ${values[i]}, Event ID: ${results[i].id}`); + } + } + expect(userFeature.balance).to.equal(startingBalance); + } + }); +}); diff --git a/shared/models/orgModels/orgConfig.ts b/shared/models/orgModels/orgConfig.ts index eb30e1dfc..6efc19b68 100644 --- a/shared/models/orgModels/orgConfig.ts +++ b/shared/models/orgModels/orgConfig.ts @@ -13,6 +13,8 @@ export const OrgConfigSchema = z.object({ reverse_deduction_order: z.boolean().default(false), include_past_due: z.boolean().default(true), + // include_check_past_due: z.boolean().default(false), + sync_status: z.boolean().default(true), merge_billing_cycles: z.boolean().default(true), multiple_trials: z.boolean().default(false),