refactor: reset cron

2026-02-22 12:03:14 +00:00
parent 4a3c05714d
commit 653690a3f0
6 changed files with 359 additions and 63 deletions
--- a/server/src/cron/cronInit.ts
+++ b/server/src/cron/cronInit.ts
@@ -1,74 +1,14 @@
 import "../sentry.ts";
-import type { CustomerEntitlement, ResetCusEnt } from "@autumn/shared";
-import { UTCDate } from "@date-fns/utc";
 import { CronJob } from "cron";
-import { format } from "date-fns";
 import { initDrizzle } from "../db/initDrizzle.js";
 import { logger } from "../external/logtail/logtailUtils.js";
-import { CusEntService } from "../internal/customers/cusProducts/cusEnts/CusEntitlementService.js";
-import { notNullish } from "../utils/genUtils.js";
-import {
-	clearCusEntsFromCache,
-	resetCustomerEntitlement,
-} from "./cronUtils.js";
 import { runInvoiceCron } from "./invoiceCron/runInvoiceCron.js";
 import { runProductCron } from "./productCron/runProductCron.js";
+import { runResetCron } from "./resetCron/runResetCron.js";
 import type { CronContext } from "./utils/CronContext.js";

 const { db, client } = initDrizzle();

-const cronTask = async () => {
-	try {
-		// const [productCusEnts, looseCusEnts] = await Promise.all([
-		// 	CusEntService.getActiveResetPassed({ db, batchSize: 500 }),
-		// 	CusEntService.getLooseResetPassed({ db, batchSize: 500 }),
-		// ]);
-		// const cusEnts: ResetCusEnt[] = [...productCusEnts, ...looseCusEnts];
-		const cusEnts = await CusEntService.getActiveResetPassed({
-			db,
-			batchSize: 500,
-		});
-
-		const batchSize = 100;
-		for (let i = 0; i < cusEnts.length; i += batchSize) {
-			const batch = cusEnts.slice(i, i + batchSize);
-			const batchResets = [];
-			const updatedCusEnts: ResetCusEnt[] = [];
-			for (const cusEnt of batch) {
-				batchResets.push(
-					resetCustomerEntitlement({
-						db,
-						cusEnt: cusEnt,
-						updatedCusEnts,
-					}),
-				);
-			}
-
-			const results = await Promise.all(batchResets);
-
-			const toUpsert = results.filter(notNullish);
-			await CusEntService.upsert({
-				db,
-				data: toUpsert as CustomerEntitlement[],
-			});
-			console.log(`Upserted ${toUpsert.length} short entitlements`);
-
-			await clearCusEntsFromCache({ cusEnts: updatedCusEnts });
-		}
-
-		console.log(
-			"FINISHED RESET CRON:",
-			format(new UTCDate(), "yyyy-MM-dd HH:mm:ss"),
-		);
-		console.log("----------------------------------\n");
-	} catch (error) {
-		console.error("Error getting entitlements for reset:", error);
-		return;
-	}
-
-	// await client.end();
-};
-
 const main = async () => {
 	if (process.env.DISABLE_CRON === "true") {
 		console.log(`Cron disabled!`);
@@ -79,7 +19,11 @@ const main = async () => {
 		db,
 		logger,
 	};
-	await Promise.all([cronTask(), runProductCron(), runInvoiceCron({ ctx })]);
+	await Promise.all([
+		runResetCron({ ctx }),
+		runProductCron(),
+		runInvoiceCron({ ctx }),
+	]);
 };

 new CronJob(
--- a/server/src/cron/cronUtils.ts
+++ b/server/src/cron/cronUtils.ts
@@ -233,7 +233,6 @@ export const resetCustomerEntitlement = async ({
 		});

 		// 1. Check if should reset
-
 		let nextResetAt = getNextResetAt({
 			curReset: new UTCDate(cusEnt.next_reset_at!),
 			interval: cusEnt.entitlement.interval as EntInterval,
--- a/server/src/cron/resetCron/getStripeSubscriptionAnchor.ts
+++ b/server/src/cron/resetCron/getStripeSubscriptionAnchor.ts
@@ -0,0 +1,69 @@
+import type { AppEnv, Organization, ResetCusEnt } from "@autumn/shared";
+import { UTCDate } from "@date-fns/utc";
+import { format, getDate, getMonth, setDate } from "date-fns";
+import type { DrizzleCli } from "@/db/initDrizzle";
+import { createStripeCli } from "@/external/connect/createStripeCli";
+import { CusProductService } from "@/internal/customers/cusProducts/CusProductService";
+
+export const getStripeSubscriptionAnchor = async ({
+	db,
+	cusEnt,
+	nextResetAt,
+}: {
+	db: DrizzleCli;
+	cusEnt: ResetCusEnt;
+	nextResetAt: number;
+}) => {
+	let nextResetAtDate = new UTCDate(nextResetAt);
+
+	// If nextResetAt is on the 28th of March, or Day 30, then do this check.
+	const nextResetAtDay = getDate(nextResetAtDate);
+	const nextResetAtMonth = getMonth(nextResetAtDate);
+
+	const shouldCheck =
+		nextResetAtDay === 30 || (nextResetAtDay === 28 && nextResetAtMonth === 2);
+
+	if (!shouldCheck) return nextResetAt;
+
+	// 1. Get the customer product
+	const cusProduct = await CusProductService.getByIdForReset({
+		db,
+		id: cusEnt.customer_product_id ?? "",
+	});
+
+	// Get org and env
+	const env = cusProduct.product.env as AppEnv;
+	const org = cusProduct.product.org as Organization;
+
+	const stripeCli = createStripeCli({ org, env });
+	if (
+		!cusProduct.subscription_ids ||
+		cusProduct.subscription_ids.length === 0
+	) {
+		return nextResetAt;
+	}
+
+	const subId = cusProduct.subscription_ids[0];
+	const sub = await stripeCli.subscriptions.retrieve(subId);
+
+	const billingCycleAnchor = sub.billing_cycle_anchor * 1000;
+	console.log("Checking billing cycle anchor");
+	console.log(
+		"Next reset at       ",
+		format(new UTCDate(nextResetAt), "dd MMM yyyy HH:mm:ss"),
+	);
+	console.log(
+		"Billing cycle anchor",
+		format(new UTCDate(billingCycleAnchor), "dd MMM yyyy HH:mm:ss"),
+	);
+
+	const billingCycleDay = getDate(new UTCDate(billingCycleAnchor));
+	const nextResetDay = getDate(nextResetAtDate);
+
+	if (billingCycleDay > nextResetDay) {
+		nextResetAtDate = setDate(nextResetAtDate, billingCycleDay);
+		return nextResetAtDate.getTime();
+	} else {
+		return nextResetAt;
+	}
+};
--- a/server/src/cron/resetCron/resetCustomerEntitlement.ts
+++ b/server/src/cron/resetCron/resetCustomerEntitlement.ts
@@ -0,0 +1,168 @@
+import {
+	AllowanceType,
+	EntInterval,
+	type FullEntitlement,
+	getStartingBalance,
+	type ResetCusEnt,
+} from "@autumn/shared";
+import { UTCDate } from "@date-fns/utc";
+import { format } from "date-fns";
+import type { DrizzleCli } from "@/db/initDrizzle.js";
+import { CusEntService } from "@/internal/customers/cusProducts/cusEnts/CusEntitlementService";
+import { getRelatedCusPrice } from "@/internal/customers/cusProducts/cusEnts/cusEntUtils.js";
+import { RolloverService } from "@/internal/customers/cusProducts/cusEnts/cusRollovers/RolloverService";
+import { getRolloverUpdates } from "@/internal/customers/cusProducts/cusEnts/cusRollovers/rolloverUtils";
+import { getResetBalancesUpdate } from "@/internal/customers/cusProducts/cusEnts/groupByUtils";
+import { CusPriceService } from "@/internal/customers/cusProducts/cusPrices/CusPriceService.js";
+import { getEntOptions } from "@/internal/products/prices/priceUtils.js";
+import { getNextResetAt } from "@/utils/timeUtils.js";
+import { getStripeSubscriptionAnchor } from "./getStripeSubscriptionAnchor";
+import { resetShortDurationCustomerEntitlement } from "./resetShortDurationCustomerEntitlement";
+
+const shortDurations = [EntInterval.Minute, EntInterval.Hour, EntInterval.Day];
+
+export const resetCustomerEntitlement = async ({
+	db,
+	cusEnt,
+	updatedCusEnts,
+}: {
+	db: DrizzleCli;
+	cusEnt: ResetCusEnt;
+	updatedCusEnts: ResetCusEnt[];
+}) => {
+	try {
+		const ent = cusEnt.entitlement as FullEntitlement;
+
+		if (
+			ent.allowance_type === AllowanceType.Fixed &&
+			shortDurations.includes(ent.interval as EntInterval)
+		) {
+			return await resetShortDurationCustomerEntitlement({
+				db,
+				cusEnt,
+				updatedCusEnts,
+			});
+		}
+
+		// Fetch related price (skip for loose ents)
+		let relatedCusPrice = null;
+		if (cusEnt.customer_product_id) {
+			const cusPrices = await CusPriceService.getByCustomerProductId({
+				db,
+				customerProductId: cusEnt.customer_product_id,
+			});
+			relatedCusPrice = getRelatedCusPrice(cusEnt, cusPrices);
+			if (relatedCusPrice) {
+				return;
+			}
+		}
+
+		const entOptions = getEntOptions(
+			cusEnt.customer_product?.options ?? [],
+			cusEnt.entitlement,
+		);
+
+		// Handle if entitlement changed to unlimited...
+		const entitlement = cusEnt.entitlement;
+		if (entitlement.allowance_type === AllowanceType.Unlimited) {
+			await CusEntService.update({
+				db,
+				id: cusEnt.id,
+				updates: {
+					unlimited: true,
+					next_reset_at: null,
+				},
+			});
+
+			console.log(
+				`Reset ${cusEnt.id} | customer: ${cusEnt.customer_id} | feature: ${cusEnt.feature_id} | new balance: unlimited`,
+			);
+			return;
+		}
+
+		if (entitlement.interval === EntInterval.Lifetime) {
+			await CusEntService.update({
+				db,
+				id: cusEnt.id,
+				updates: {
+					next_reset_at: null,
+				},
+			});
+
+			console.log(
+				`Reset ${cusEnt.id} | customer: ${cusEnt.customer_id} | feature: ${cusEnt.feature_id} | reset to lifetime (next_reset_at: null)`,
+			);
+			return;
+		}
+
+		const resetBalance = getStartingBalance({
+			entitlement: cusEnt.entitlement,
+			options: entOptions || undefined,
+			relatedPrice: undefined,
+			productQuantity: cusEnt.customer_product?.quantity ?? 1,
+		});
+
+		if (!cusEnt.next_reset_at) return;
+
+		// 1. Check if should reset
+		let nextResetAt = getNextResetAt({
+			curReset: new UTCDate(cusEnt.next_reset_at),
+			interval: cusEnt.entitlement.interval as EntInterval,
+			intervalCount: cusEnt.entitlement.interval_count,
+		});
+
+		const rolloverUpdate = getRolloverUpdates({
+			cusEnt,
+			nextResetAt: cusEnt.next_reset_at,
+		});
+
+		const resetBalanceUpdate = getResetBalancesUpdate({
+			cusEnt,
+			allowance: resetBalance || undefined,
+		});
+
+		// Only check sub anchor for product-based ents (loose ents have no subscription)
+		if (cusEnt.customer_product) {
+			try {
+				nextResetAt = await getStripeSubscriptionAnchor({
+					db,
+					cusEnt,
+					nextResetAt,
+				});
+			} catch (error) {
+				console.log(
+					`WARNING: Failed to check sub anchor: ${error}, Org: ${cusEnt.customer.org_id}`,
+				);
+				console.log(error);
+			}
+		}
+
+		await CusEntService.update({
+			db,
+			id: cusEnt.id,
+			updates: {
+				...resetBalanceUpdate,
+				next_reset_at: nextResetAt,
+				adjustment: 0,
+			},
+		});
+
+		if (rolloverUpdate?.toInsert && rolloverUpdate.toInsert.length > 0) {
+			await RolloverService.insert({
+				db,
+				rows: rolloverUpdate.toInsert,
+				fullCusEnt: cusEnt,
+			});
+		}
+
+		updatedCusEnts.push(cusEnt);
+
+		console.log(
+			`Reset ${cusEnt.id} | customer: ${cusEnt.customer_id} | feature: ${cusEnt.feature_id} | new balance: ${resetBalance} | new next_reset_at: ${format(new UTCDate(nextResetAt), "dd MMM yyyy HH:mm:ss")}`,
+		);
+	} catch (error) {
+		console.log(
+			`Failed to reset ${cusEnt.id} | ${cusEnt.customer_id} | ${cusEnt.feature_id}, error: ${error}`,
+		);
+	}
+};
--- a/server/src/cron/resetCron/resetShortDurationCustomerEntitlement.ts
+++ b/server/src/cron/resetCron/resetShortDurationCustomerEntitlement.ts
@@ -0,0 +1,60 @@
+import type { EntInterval, FullEntitlement, ResetCusEnt } from "@autumn/shared";
+import { UTCDate } from "@date-fns/utc";
+import { Decimal } from "decimal.js";
+import type { DrizzleCli } from "@/db/initDrizzle.js";
+import { RolloverService } from "@/internal/customers/cusProducts/cusEnts/cusRollovers/RolloverService";
+import { getRolloverUpdates } from "@/internal/customers/cusProducts/cusEnts/cusRollovers/rolloverUtils";
+import { getResetBalancesUpdate } from "@/internal/customers/cusProducts/cusEnts/groupByUtils";
+import { getNextResetAt } from "@/utils/timeUtils.js";
+
+export const resetShortDurationCustomerEntitlement = async ({
+	db,
+	cusEnt,
+	updatedCusEnts,
+}: {
+	db: DrizzleCli;
+	cusEnt: ResetCusEnt;
+	updatedCusEnts: ResetCusEnt[];
+}) => {
+	const ent = cusEnt.entitlement as FullEntitlement;
+
+	if (!cusEnt.next_reset_at) return;
+
+	const resetCusEnt = {
+		...cusEnt,
+		next_reset_at: getNextResetAt({
+			curReset: new UTCDate(cusEnt.next_reset_at),
+			interval: ent.interval as EntInterval,
+			intervalCount: ent.interval_count,
+		}),
+		adjustment: 0,
+		...getResetBalancesUpdate({
+			cusEnt,
+			allowance: new Decimal(ent.allowance || 0)
+				.mul(cusEnt.customer_product?.quantity ?? 1)
+				.toNumber(),
+		}),
+	};
+	const newCusEnt = resetCusEnt;
+
+	const rolloverUpdate = getRolloverUpdates({
+		cusEnt,
+		nextResetAt: cusEnt.next_reset_at,
+	});
+
+	if (rolloverUpdate?.toInsert && rolloverUpdate.toInsert.length > 0) {
+		await RolloverService.insert({
+			db,
+			rows: rolloverUpdate.toInsert,
+			fullCusEnt: cusEnt,
+		});
+	}
+
+	console.log(
+		`Reseting short cus ent (${cusEnt.feature_id}) [${ent.interval}], customer: ${cusEnt.customer_id}, org: ${cusEnt.customer.org_id}`,
+	);
+
+	updatedCusEnts.push(newCusEnt);
+
+	return newCusEnt;
+};
--- a/server/src/cron/resetCron/runResetCron.ts
+++ b/server/src/cron/resetCron/runResetCron.ts
@@ -0,0 +1,56 @@
+import {
+	type CustomerEntitlement,
+	notNullish,
+	type ResetCusEnt,
+} from "@autumn/shared";
+import { UTCDate } from "@date-fns/utc";
+import { format } from "date-fns";
+import { CusEntService } from "@/internal/customers/cusProducts/cusEnts/CusEntitlementService";
+import { clearCusEntsFromCache, resetCustomerEntitlement } from "../cronUtils";
+import type { CronContext } from "../utils/CronContext";
+
+export const runResetCron = async ({ ctx }: { ctx: CronContext }) => {
+	const { db } = ctx;
+	try {
+		const cusEnts = await CusEntService.getActiveResetPassed({
+			db,
+			batchSize: 500,
+		});
+
+		const batchSize = 100;
+		for (let i = 0; i < cusEnts.length; i += batchSize) {
+			const batch = cusEnts.slice(i, i + batchSize);
+			const batchResets = [];
+			const updatedCusEnts: ResetCusEnt[] = [];
+			for (const cusEnt of batch) {
+				batchResets.push(
+					resetCustomerEntitlement({
+						db,
+						cusEnt: cusEnt,
+						updatedCusEnts,
+					}),
+				);
+			}
+
+			const results = await Promise.all(batchResets);
+
+			const toUpsert = results.filter(notNullish);
+			await CusEntService.upsert({
+				db,
+				data: toUpsert as CustomerEntitlement[],
+			});
+			console.log(`Upserted ${toUpsert.length} short entitlements`);
+
+			await clearCusEntsFromCache({ cusEnts: updatedCusEnts });
+		}
+
+		console.log(
+			"FINISHED RESET CRON:",
+			format(new UTCDate(), "yyyy-MM-dd HH:mm:ss"),
+		);
+		console.log("----------------------------------\n");
+	} catch (error) {
+		console.error("Error getting entitlements for reset:", error);
+		return;
+	}
+};