575 lines
14 KiB
TypeScript
575 lines
14 KiB
TypeScript
import {
|
|
AllowanceType,
|
|
type AppEnv,
|
|
CusProductStatus,
|
|
type Customer,
|
|
ErrCode,
|
|
type Feature,
|
|
FeatureType,
|
|
FeatureUsageType,
|
|
type FullCustomerEntitlement,
|
|
type Organization,
|
|
} from "@autumn/shared";
|
|
import { Decimal } from "decimal.js";
|
|
import { sql } from "drizzle-orm";
|
|
import { StatusCodes } from "http-status-codes";
|
|
import type { DrizzleCli } from "@/db/initDrizzle.js";
|
|
import { CusService } from "@/internal/customers/CusService.js";
|
|
import { refreshCusCache } from "@/internal/customers/cusCache/updateCachedCus.js";
|
|
import { getFeatureBalance } from "@/internal/customers/cusProducts/cusEnts/cusEntUtils.js";
|
|
import { deductFromApiCusRollovers } from "@/internal/customers/cusProducts/cusEnts/cusRollovers/rolloverDeductionUtils.js";
|
|
import { getCusEntsInFeatures } from "@/internal/customers/cusUtils/cusUtils.js";
|
|
import { featureToCreditSystem } from "@/internal/features/creditSystemUtils.js";
|
|
import RecaseError from "@/utils/errorUtils.js";
|
|
import { handleThresholdReached } from "./handleThresholdReached.js";
|
|
import {
|
|
deductAllowanceFromCusEnt,
|
|
deductFromUsageBasedCusEnt,
|
|
} from "./updateBalanceTask.js";
|
|
|
|
// 2. Get deductions for each feature
|
|
const getFeatureDeductions = ({
|
|
cusEnts,
|
|
value,
|
|
features,
|
|
shouldSet,
|
|
entityId,
|
|
}: {
|
|
cusEnts: FullCustomerEntitlement[];
|
|
value: number;
|
|
features: Feature[];
|
|
shouldSet: boolean;
|
|
entityId?: string;
|
|
}) => {
|
|
const meteredFeature =
|
|
features.find((f) => f.type === FeatureType.Metered) || features[0];
|
|
|
|
const featureDeductions = [];
|
|
for (const feature of features) {
|
|
let newValue = value;
|
|
const unlimitedExists = cusEnts.some(
|
|
(cusEnt) =>
|
|
cusEnt.entitlement.allowance_type === AllowanceType.Unlimited &&
|
|
cusEnt.entitlement.internal_feature_id === feature.internal_id,
|
|
);
|
|
|
|
if (unlimitedExists) {
|
|
continue;
|
|
}
|
|
|
|
if (feature.type === FeatureType.CreditSystem) {
|
|
newValue = featureToCreditSystem({
|
|
featureId: meteredFeature.id,
|
|
creditSystem: feature,
|
|
amount: value,
|
|
});
|
|
}
|
|
|
|
// If it's set
|
|
let deduction = newValue;
|
|
|
|
if (shouldSet) {
|
|
const totalAllowance = cusEnts.reduce((acc, curr) => {
|
|
return acc + (curr.entitlement.allowance || 0);
|
|
}, 0);
|
|
|
|
const targetBalance = new Decimal(totalAllowance).sub(value).toNumber();
|
|
|
|
const totalBalance = getFeatureBalance({
|
|
cusEnts,
|
|
internalFeatureId: feature.internal_id!,
|
|
entityId,
|
|
})!;
|
|
|
|
deduction = new Decimal(totalBalance).sub(targetBalance).toNumber();
|
|
}
|
|
|
|
if (deduction === 0) {
|
|
console.log(` - Skipping feature ${feature.id} -- deduction is 0`);
|
|
continue;
|
|
}
|
|
|
|
featureDeductions.push({
|
|
feature,
|
|
deduction,
|
|
});
|
|
}
|
|
|
|
featureDeductions.sort((a, b) => {
|
|
if (
|
|
a.feature.type === FeatureType.CreditSystem &&
|
|
b.feature.type !== FeatureType.CreditSystem
|
|
) {
|
|
return 1;
|
|
}
|
|
|
|
if (
|
|
a.feature.type !== FeatureType.CreditSystem &&
|
|
b.feature.type === FeatureType.CreditSystem
|
|
) {
|
|
return -1;
|
|
}
|
|
|
|
return a.feature.id.localeCompare(b.feature.id);
|
|
});
|
|
|
|
return featureDeductions;
|
|
};
|
|
|
|
/**
|
|
* Calculate total available rollover balance for a feature
|
|
*/
|
|
const calculateAvailableRolloverBalance = ({
|
|
cusEnts,
|
|
feature,
|
|
entityId,
|
|
}: {
|
|
cusEnts: FullCustomerEntitlement[];
|
|
feature: Feature;
|
|
entityId?: string;
|
|
}) => {
|
|
const featureCusEnts = cusEnts.filter(
|
|
(cusEnt) => cusEnt.entitlement.internal_feature_id === feature.internal_id,
|
|
);
|
|
|
|
if (!entityId) {
|
|
// Non-entity: sum rollover.balance
|
|
return featureCusEnts.reduce((sum, cusEnt) => {
|
|
const rolloverSum = cusEnt.rollovers.reduce(
|
|
(rSum, rollover) =>
|
|
new Decimal(rSum).add(rollover.balance || 0).toNumber(),
|
|
0,
|
|
);
|
|
return new Decimal(sum).add(rolloverSum).toNumber();
|
|
}, 0);
|
|
} else {
|
|
// Entity: sum rollover.entities[entityId].balance
|
|
return featureCusEnts.reduce((sum, cusEnt) => {
|
|
const rolloverSum = cusEnt.rollovers.reduce((rSum, rollover) => {
|
|
const entityRollover = rollover.entities?.[entityId];
|
|
if (entityRollover) {
|
|
return new Decimal(rSum).add(entityRollover.balance || 0).toNumber();
|
|
}
|
|
return rSum;
|
|
}, 0);
|
|
return new Decimal(sum).add(rolloverSum).toNumber();
|
|
}, 0);
|
|
}
|
|
};
|
|
|
|
/**
|
|
* Validate that the deduction is possible given the current balance and usage allowed.
|
|
* Constraint 1: Insufficient balance without usage_allowed.
|
|
* Constraint 2: Usage limit exceeded for customer entitlements with usage_allowed.
|
|
*/
|
|
const validateDeductionPossible = ({
|
|
cusEnts,
|
|
featureDeductions,
|
|
entityId,
|
|
}: {
|
|
cusEnts: FullCustomerEntitlement[];
|
|
featureDeductions: { feature: Feature; deduction: number }[];
|
|
entityId?: string;
|
|
}) => {
|
|
for (const { feature, deduction } of featureDeductions) {
|
|
const featureCusEnts = cusEnts.filter(
|
|
(customerEntitlement) =>
|
|
customerEntitlement.entitlement.internal_feature_id ===
|
|
feature.internal_id,
|
|
);
|
|
|
|
// CONSTRAINT 1: Insufficient balance without usage_allowed
|
|
const cusEntBalance = getFeatureBalance({
|
|
cusEnts: featureCusEnts,
|
|
internalFeatureId: feature.internal_id!,
|
|
entityId,
|
|
});
|
|
|
|
// If unlimited, skip validation
|
|
if (cusEntBalance === null) {
|
|
continue;
|
|
}
|
|
const rolloverBalance = calculateAvailableRolloverBalance({
|
|
cusEnts,
|
|
feature,
|
|
entityId,
|
|
});
|
|
const totalBalance = new Decimal(cusEntBalance)
|
|
.add(rolloverBalance)
|
|
.toNumber();
|
|
|
|
const hasUsageAllowed = featureCusEnts.some(
|
|
(customerEntitlement) => customerEntitlement.usage_allowed,
|
|
);
|
|
|
|
// Check if this is a "free" feature (single-use with included_usage but no pricing)
|
|
// Only apply to SingleUse features; ContinuousUse (allocated) features should reject
|
|
const isFreeFeature =
|
|
feature.type === FeatureType.Metered &&
|
|
feature.config?.usage_type === FeatureUsageType.Single &&
|
|
featureCusEnts.some(
|
|
(cusEnt) =>
|
|
cusEnt.entitlement.allowance && cusEnt.entitlement.allowance > 0,
|
|
) &&
|
|
!hasUsageAllowed;
|
|
|
|
// For free SingleUse features, allow tracking beyond balance (will cap at 0 in performDeduction)
|
|
// For prepaid/allocated/other features without usage_allowed, reject insufficient balance
|
|
if (totalBalance < deduction && !hasUsageAllowed && !isFreeFeature) {
|
|
throw new RecaseError({
|
|
message: `Insufficient balance for feature ${feature.id}. Available: ${totalBalance} (${cusEntBalance} + ${rolloverBalance} rollover), Required: ${deduction}`,
|
|
code: ErrCode.InsufficientBalance,
|
|
statusCode: StatusCodes.BAD_REQUEST,
|
|
data: {
|
|
feature_id: feature.id,
|
|
available: totalBalance,
|
|
cus_ent_balance: cusEntBalance,
|
|
rollover_balance: rolloverBalance,
|
|
required: deduction,
|
|
},
|
|
});
|
|
}
|
|
|
|
// CONSTRAINT 2: Usage limit exceeded for customer entitlements with usage_allowed
|
|
const entitlementDeduction =
|
|
new Decimal(deduction).sub(rolloverBalance).toNumber() > 0
|
|
? new Decimal(deduction).sub(rolloverBalance).toNumber()
|
|
: 0;
|
|
|
|
if (entitlementDeduction > 0) {
|
|
const featureCusEntsWithUsageAllowed = featureCusEnts.filter(
|
|
(customerEntitlement) => customerEntitlement.usage_allowed,
|
|
);
|
|
|
|
const totalRemainingLimit = featureCusEntsWithUsageAllowed.reduce(
|
|
(sum, cusEnt) => {
|
|
const usageLimit = cusEnt.entitlement.usage_limit;
|
|
if (!usageLimit) {
|
|
return sum;
|
|
}
|
|
|
|
const featureBalance = getFeatureBalance({
|
|
cusEnts: [cusEnt],
|
|
internalFeatureId: feature.internal_id!,
|
|
entityId,
|
|
});
|
|
|
|
// Skip if unlimited
|
|
if (featureBalance === null) {
|
|
return sum;
|
|
}
|
|
|
|
const allowance = new Decimal(cusEnt.entitlement.allowance || 0);
|
|
const currentBalance = new Decimal(featureBalance);
|
|
const currentUsed = allowance.sub(currentBalance);
|
|
const remainingLimit = new Decimal(usageLimit).sub(currentUsed);
|
|
|
|
return new Decimal(sum)
|
|
.add(Decimal.max(0, remainingLimit))
|
|
.toNumber();
|
|
},
|
|
0,
|
|
);
|
|
|
|
if (
|
|
featureCusEntsWithUsageAllowed.length > 0 &&
|
|
entitlementDeduction > totalRemainingLimit
|
|
) {
|
|
throw new RecaseError({
|
|
message: `Usage limit exceeded for feature ${feature.id}. Total remaining capacity: ${totalRemainingLimit}, Requested from entitlement: ${entitlementDeduction} (${rolloverBalance} covered by rollovers)`,
|
|
code: ErrCode.InsufficientBalance,
|
|
statusCode: StatusCodes.BAD_REQUEST,
|
|
data: {
|
|
feature_id: feature.id,
|
|
total_remaining_capacity: totalRemainingLimit,
|
|
requested_from_entitlement: entitlementDeduction,
|
|
covered_by_rollovers: rolloverBalance,
|
|
total_requested: deduction,
|
|
},
|
|
});
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
const logUsageUpdate = ({
|
|
customer,
|
|
features,
|
|
cusEnts,
|
|
featureDeductions,
|
|
org,
|
|
setUsage,
|
|
entityId,
|
|
}: {
|
|
customer: Customer;
|
|
features: Feature[];
|
|
cusEnts: FullCustomerEntitlement[];
|
|
featureDeductions: any;
|
|
org: Organization;
|
|
setUsage: boolean;
|
|
entityId?: string;
|
|
}) => {
|
|
console.log(
|
|
` - Customer: ${customer.id} (${customer.env}) | Org: ${
|
|
org.slug
|
|
} | Features: ${features.map((f) => f.id).join(", ")} | Set Usage: ${
|
|
setUsage ? "true" : "false"
|
|
}`,
|
|
);
|
|
|
|
console.log(
|
|
" - CusEnts:",
|
|
cusEnts.map((cusEnt: any) => {
|
|
let balanceStr = cusEnt.balance;
|
|
try {
|
|
if (cusEnt.entitlement.allowance_type === AllowanceType.Unlimited) {
|
|
balanceStr = "Unlimited";
|
|
}
|
|
} catch (_error) {
|
|
balanceStr = "failed_to_get_balance";
|
|
}
|
|
|
|
if (entityId && cusEnt.entities) {
|
|
balanceStr = `${cusEnt.entities?.[entityId!]?.balance} [${entityId}]`;
|
|
}
|
|
|
|
return `${cusEnt.feature_id} - ${balanceStr} (${
|
|
cusEnt.customer_product ? cusEnt.customer_product.product_id : ""
|
|
})`;
|
|
}),
|
|
"| Deductions:",
|
|
featureDeductions.map((f: any) => `${f.feature.id}: ${f.deduction}`),
|
|
);
|
|
};
|
|
|
|
// Main function to update customer balance
|
|
export const updateUsage = async ({
|
|
db,
|
|
customerId,
|
|
features,
|
|
org,
|
|
env,
|
|
value,
|
|
properties,
|
|
setUsage,
|
|
logger,
|
|
entityId,
|
|
allFeatures,
|
|
}: {
|
|
db: DrizzleCli;
|
|
customerId: string;
|
|
features: Feature[];
|
|
org: Organization;
|
|
env: AppEnv;
|
|
value: number;
|
|
properties: any;
|
|
setUsage: boolean;
|
|
logger: any;
|
|
entityId?: string;
|
|
allFeatures: Feature[];
|
|
}) => {
|
|
const customer = await CusService.getFull({
|
|
db,
|
|
idOrInternalId: customerId,
|
|
orgId: org.id,
|
|
env,
|
|
inStatuses: [CusProductStatus.Active, CusProductStatus.PastDue],
|
|
entityId,
|
|
withSubs: true,
|
|
});
|
|
|
|
const { cusEnts, cusPrices } = await getCusEntsInFeatures({
|
|
customer,
|
|
internalFeatureIds: features.map((f) => f.internal_id!),
|
|
logger,
|
|
reverseOrder: org.config?.reverse_deduction_order,
|
|
});
|
|
|
|
// 1. Get deductions for each feature
|
|
const featureDeductions = getFeatureDeductions({
|
|
cusEnts,
|
|
value,
|
|
shouldSet: setUsage,
|
|
features,
|
|
entityId,
|
|
});
|
|
|
|
logUsageUpdate({
|
|
customer,
|
|
features,
|
|
cusEnts,
|
|
featureDeductions,
|
|
org,
|
|
setUsage,
|
|
entityId,
|
|
});
|
|
|
|
// 3. Return if no customer entitlements or features found
|
|
if (cusEnts.length === 0 || features.length === 0) {
|
|
console.log(" - No customer entitlements or features found");
|
|
return;
|
|
}
|
|
|
|
validateDeductionPossible({ cusEnts, featureDeductions, entityId });
|
|
|
|
const originalCusEnts = structuredClone(cusEnts);
|
|
for (const obj of featureDeductions) {
|
|
let { feature, deduction: toDeduct } = obj;
|
|
|
|
for (const cusEnt of cusEnts) {
|
|
if (cusEnt.entitlement.internal_feature_id !== feature.internal_id) {
|
|
continue;
|
|
}
|
|
|
|
toDeduct = await deductFromApiCusRollovers({
|
|
toDeduct,
|
|
cusEnt,
|
|
deductParams: {
|
|
db,
|
|
feature,
|
|
env,
|
|
entity: customer.entity ? customer.entity : undefined,
|
|
},
|
|
});
|
|
|
|
if (toDeduct === 0) continue;
|
|
|
|
toDeduct = await deductAllowanceFromCusEnt({
|
|
toDeduct,
|
|
cusEnt,
|
|
deductParams: {
|
|
db,
|
|
feature,
|
|
env,
|
|
org,
|
|
cusPrices: cusPrices as any[],
|
|
customer,
|
|
entity: customer.entity,
|
|
},
|
|
featureDeductions,
|
|
willDeductCredits: true,
|
|
setZeroAdjustment: true,
|
|
});
|
|
}
|
|
|
|
if (toDeduct !== 0) {
|
|
await deductFromUsageBasedCusEnt({
|
|
toDeduct,
|
|
cusEnts,
|
|
deductParams: {
|
|
db,
|
|
feature,
|
|
env,
|
|
org,
|
|
cusPrices: cusPrices as any[],
|
|
customer,
|
|
entity: customer.entity,
|
|
},
|
|
setZeroAdjustment: true,
|
|
});
|
|
}
|
|
|
|
handleThresholdReached({
|
|
org,
|
|
env,
|
|
features: allFeatures,
|
|
db,
|
|
|
|
feature,
|
|
cusEnts: originalCusEnts,
|
|
newCusEnts: cusEnts,
|
|
fullCus: customer,
|
|
logger,
|
|
});
|
|
}
|
|
|
|
return cusEnts;
|
|
};
|
|
|
|
// MAIN FUNCTION
|
|
export const runUpdateUsageTask = async ({
|
|
payload,
|
|
logger,
|
|
db,
|
|
throwError = false,
|
|
}: {
|
|
payload: any;
|
|
logger: any;
|
|
db: DrizzleCli;
|
|
throwError?: boolean;
|
|
}) => {
|
|
try {
|
|
// 1. Update customer balance
|
|
const {
|
|
internalCustomerId,
|
|
customerId,
|
|
eventId,
|
|
features,
|
|
value,
|
|
set_usage,
|
|
properties,
|
|
org,
|
|
env,
|
|
entityId,
|
|
allFeatures,
|
|
} = payload;
|
|
|
|
console.log("--------------------------------");
|
|
console.log(
|
|
`HANDLING USAGE TASK FOR CUSTOMER (${customerId}), ORG: ${org.slug}, EVENT ID: ${eventId}`,
|
|
);
|
|
|
|
const cusEnts = await db.transaction(
|
|
async (tx) => {
|
|
// Acquire advisory lock for this customer (and entity if provided) to serialize concurrent requests
|
|
// Include entity_id in lock key so different entities can update concurrently
|
|
const lockKeyStr = `${internalCustomerId}_${org.id}_${env}${entityId ? `_${entityId}` : ""}`;
|
|
const hash =
|
|
lockKeyStr.split("").reduce((acc, char) => {
|
|
return (acc << 5) - acc + char.charCodeAt(0);
|
|
}, 0) | 0; // Convert to 32-bit integer
|
|
|
|
console.log(
|
|
` 🔒 [${eventId}] Acquiring advisory lock (hash=${hash}) for: ${lockKeyStr}`,
|
|
);
|
|
await tx.execute(sql`SELECT pg_advisory_xact_lock(${hash})`);
|
|
console.log(
|
|
` 🔓 [${eventId}] Advisory lock acquired, proceeding with update`,
|
|
);
|
|
|
|
return await updateUsage({
|
|
db: tx as unknown as DrizzleCli,
|
|
customerId,
|
|
features,
|
|
value,
|
|
properties,
|
|
org,
|
|
env,
|
|
setUsage: set_usage,
|
|
logger,
|
|
entityId,
|
|
allFeatures,
|
|
});
|
|
},
|
|
{
|
|
isolationLevel: "read committed",
|
|
},
|
|
);
|
|
|
|
await refreshCusCache({
|
|
db,
|
|
customerId,
|
|
entityId,
|
|
org,
|
|
env,
|
|
});
|
|
} catch (error) {
|
|
logger.error(`ERROR UPDATING USAGE`);
|
|
logger.error(error);
|
|
|
|
if (throwError) {
|
|
throw error;
|
|
}
|
|
}
|
|
};
|