Merge pull request #2339 from useautumn/feat/migration-capacity-minimal

Migrations: shared queue

Author: og2701

server/src/external/autumn/autumnCli.ts

@@ -1123,6 +1123,7 @@ export class AutumnInt {
 			dry_run?: boolean;
 			only?: string[];
 			limit?: number;
+			/** @deprecated Migration concurrency is fleet-managed. */
 			concurrency?: number;
 			lazy_run?: boolean;
 			retry_item_statuses?: ("failed" | "skipped")[];

server/src/internal/migrations/v2/filters/customers/filterCustomers.ts

@@ -76,6 +76,7 @@ export const filterCustomers = ({
 	includeProcessed,
 	batchSize,
 	limit,
+	afterInternalId,
 }: {
 	ctx: AutumnContext;
 	filter: CustomerFilter;
@@ -85,6 +86,7 @@ export const filterCustomers = ({
 	includeProcessed?: IncludeProcessed;
 	batchSize?: number;
 	limit?: number;
+	afterInternalId?: string;
 }): AsyncGenerator<CustomerRow[]> => {
 	const args = buildArgs({ ctx, filter, checkpoint, search, customerFilters });
 	const source = iterateOverFilterResults<CustomerRow>({
@@ -93,6 +95,7 @@ export const filterCustomers = ({
 			buildRowsSelect({ args, includeProcessed, limit, afterInternalId }),
 		batchSize:
 			limit === undefined ? batchSize : Math.min(batchSize ?? limit, limit),
+		afterInternalId,
 	});
 	return limit === undefined ? source : takeRows(source, limit);
 };

server/src/internal/migrations/v2/filters/iterateOverFilterResults.ts

@@ -13,12 +13,14 @@ export async function* iterateOverFilterResults<
 	db,
 	buildSelect,
 	batchSize = DEFAULT_BATCH_SIZE,
+	afterInternalId,
 }: {
 	db: { execute: (query: SQL) => Promise<unknown> };
 	buildSelect: (args: { limit: number; afterInternalId?: string }) => SQL;
 	batchSize?: number;
+	afterInternalId?: string;
 }): AsyncGenerator<TRow[]> {
-	let cursor: string | undefined;
+	let cursor = afterInternalId;
 	while (true) {
 		const query = buildSelect({ limit: batchSize, afterInternalId: cursor });
 		const rows = (await db.execute(query)) as unknown as TRow[];

server/src/internal/migrations/v2/filters/runFilter.ts

@@ -28,16 +28,22 @@ export const runFilter = async ({
 	dryRun,
 	kind,
 	controls,
+	includeCount = true,
+	afterInternalId,
+	batchSize,
 }: {
 	ctx: AutumnContext;
 	migration: MigrationRuntimeWithEventId;
 	migrationRunId: string;
 	dryRun: boolean;
 	kind: RunScopeKind;
 	controls?: MigrationRunControls;
+	includeCount?: boolean;
+	afterInternalId?: string;
+	batchSize?: number;
 }): Promise<{
 	kind: RunScopeKind;
-	count: number;
+	count: number | null;
 	iterate: () => AsyncGenerator<RunScopeItem[]>;
 }> => {
 	if (kind !== "customer")
@@ -56,12 +62,14 @@ export const runFilter = async ({
 		controls,
 	});
 	const limit = controls?.limit ?? undefined;
-	const count = await countCustomers({
-		ctx,
-		filter,
-		checkpoint,
-		limit,
-	});
+	const count = includeCount
+		? await countCustomers({
+				ctx,
+				filter,
+				checkpoint,
+				limit,
+			})
+		: null;
 
 	ctx.logger.info("runFilter: customer scope resolved", {
 		data: {
@@ -94,6 +102,8 @@ export const runFilter = async ({
 			filter,
 			checkpoint,
 			limit,
+			afterInternalId,
+			batchSize,
 		})) {
 			yield batch.map(
 				(row): RunScopeItem => ({

server/src/internal/migrations/v2/handlers/handleRunMigration.ts

@@ -1,4 +1,4 @@
-import { ErrCode, makeScopeChecker, RecaseError, Scopes } from "@autumn/shared";
+import { ErrCode, RecaseError, Scopes } from "@autumn/shared";
 import { auth } from "@trigger.dev/sdk/v3";
 import { z } from "zod/v4";
 import { createRoute } from "@/honoMiddlewares/routeHandler";
@@ -7,49 +7,39 @@ import { prepare } from "@/internal/migrations/v2/prepare/index.js";
 import { migrationRepo } from "@/internal/migrations/v2/repos/index.js";
 import { RETRYABLE_MIGRATION_ITEM_RUN_STATUSES } from "@/internal/migrations/v2/run/utils/retryItemStatuses.js";
 import { shouldRunMigrationInline } from "@/internal/migrations/v2/utils/shouldRunMigrationInline.js";
+import {
+	getMigrationTriggerOptions,
+	MIGRATION_RUN_CUSTOMER_CONCURRENCY,
+} from "@/trigger/migrations/migrationTaskQueue.js";
 import {
 	executeRunMigration,
 	type RunMigrationPayload,
 	runMigrationTask,
 } from "@/trigger/migrations/runMigrationTask.js";
 
-const MAX_CONCURRENCY = 5;
-const SUPERUSER_MAX_CONCURRENCY = 100;
+const LEGACY_MAX_REQUESTED_CONCURRENCY = 100;
 
 const RunMigrationBody = z.object({
 	id: z.string(),
 	dry_run: z.boolean().default(false),
 	limit: z.number().int().min(1).optional(),
 	only: z.array(z.string()).optional(),
+	/** Accepted for backward compatibility; migration concurrency is fleet-managed. */
 	concurrency: z
 		.number()
 		.int()
 		.min(1)
-		.max(SUPERUSER_MAX_CONCURRENCY)
-		.optional(),
+		.max(LEGACY_MAX_REQUESTED_CONCURRENCY)
+		.optional()
+		.describe("Deprecated: migration concurrency is fleet-managed"),
 	retry_item_statuses: z
 		.array(z.enum(RETRYABLE_MIGRATION_ITEM_RUN_STATUSES))
 		.optional(),
-	/** When true, claim a lazy run alongside the background sweeper. Customers
-	 *  hit on the request path get migrated lazily via `runMigrationCustomerTask`
-	 *  before the sweeper reaches them. Background and lazy run on the same
-	 *  migration_run row — the claim is shared. */
+	/** Lazy runs share one run row with the sweeper and enqueue request-path customer work.
+	 * Targeted `only` is incompatible because lazy matching happens on customer reads. */
 	lazy_run: z.boolean().default(false),
 });
 
-const getRunMigrationTriggerOptions = ({
-	orgId,
-	migrationId,
-	isDev,
-}: {
-	orgId: string;
-	migrationId: string;
-	isDev: boolean;
-}) => ({
-	...(isDev ? { region: "eu-central-1" } : {}),
-	concurrencyKey: `${orgId}:${migrationId}`,
-});
-
 export const handleRunMigration = createRoute({
 	scopes: [Scopes.Migrations.Write],
 	body: RunMigrationBody,
@@ -60,23 +50,10 @@ export const handleRunMigration = createRoute({
 			dry_run: dryRun,
 			limit,
 			only,
-			concurrency,
 			retry_item_statuses: retryItemStatuses,
 			lazy_run: lazyRun,
 		} = c.req.valid("json");
 
-		const { isSuperuser } = makeScopeChecker(ctx.scopes);
-		const maxConcurrency = isSuperuser
-			? SUPERUSER_MAX_CONCURRENCY
-			: MAX_CONCURRENCY;
-		if (concurrency !== undefined && concurrency > maxConcurrency) {
-			throw new RecaseError({
-				message: `Migration concurrency cannot exceed ${maxConcurrency}`,
-				code: ErrCode.InvalidRequest,
-				statusCode: 400,
-			});
-		}
-
 		const migration = await migrationRepo.find({ ctx, id });
 
 		if (!migration.operations)
@@ -119,7 +96,6 @@ export const handleRunMigration = createRoute({
 					controls: {
 						limit,
 						only,
-						concurrency,
 						retryItemStatuses,
 					},
 				};
@@ -129,11 +105,7 @@ export const handleRunMigration = createRoute({
 				}
 				const handle = await runMigrationTask.trigger(
 					payload,
-					getRunMigrationTriggerOptions({
-						orgId: ctx.org.id,
-						migrationId: id,
-						isDev,
-					}),
+					getMigrationTriggerOptions({ isDev }),
 				);
 				return { triggerRunId: handle.id };
 			},
@@ -176,7 +148,7 @@ export const handleRunMigration = createRoute({
 			migration_id: id,
 			dry_run: dryRun,
 			lazy_run: lazyRun,
-			concurrency,
+			concurrency: MIGRATION_RUN_CUSTOMER_CONCURRENCY,
 			run_id: migrationRunId,
 			trigger_run_id: triggerRunId,
 			public_access_token: publicAccessToken,

server/src/internal/migrations/v2/lazy/checkPendingMigrationsForCustomer.ts

@@ -2,6 +2,7 @@ import type { FullCustomer, MigrationItemRunData } from "@autumn/shared";
 import { customerFilterMatchesFullCustomer } from "@autumn/shared/api/customers/utils/match/index.js";
 import type { AutumnContext } from "@/honoUtils/HonoEnv.js";
 import { shouldRunMigrationInline } from "@/internal/migrations/v2/utils/shouldRunMigrationInline.js";
+import { MIGRATION_LAZY_TASK_PRIORITY_SECONDS } from "@/trigger/migrations/migrationTaskQueue.js";
 import {
 	executeRunMigrationCustomer,
 	runMigrationCustomerTask,
@@ -80,8 +81,6 @@ export const checkPendingMigrationsForCustomer = async ({
 		};
 
 		if (shouldRunMigrationInline()) {
-			// Inline loses trigger.dev's concurrencyKey serialization; the
-			// server-side item-run claim is the real authority either way.
 			const inlineCtx = { ...ctx, insideTriggerTask: true };
 			void executeRunMigrationCustomer({
 				ctx: inlineCtx,
@@ -99,7 +98,7 @@ export const checkPendingMigrationsForCustomer = async ({
 		}
 
 		await runMigrationCustomerTask.trigger(payload, {
-			concurrencyKey: `${migration.internal_id}:${fullCustomer.internal_id}`,
+			priority: MIGRATION_LAZY_TASK_PRIORITY_SECONDS,
 		});
 	}
 };

server/src/internal/migrations/v2/run/orchestrators/iterateScope.ts

@@ -1,44 +1,54 @@
+import type { MigrationRunScheduler } from "../types/migrationRunScheduler.js";
 import type { RunScopeItem } from "../types/runScope.js";
 
 export type IterateScopeItemResult<T> =
 	| { status: "ok"; item: RunScopeItem; value: T }
 	| { status: "failed"; item: RunScopeItem; error: Error };
 
+export type IterateScopeCompletion = "exhausted" | "slice_complete" | "stopped";
+
 export type IterateScopeSummary<T> = {
 	processed: number;
 	succeeded: number;
 	failed: number;
 	results: IterateScopeItemResult<T>[];
+	completion: IterateScopeCompletion;
+	cursor: string | null;
 };
 
-/**
- * Generic iteration over any kind-tagged scope iterator. Calls `perItem`
- * for every item, collecting per-item results into a summary.
- *
- * `concurrency` (default 1): max parallel `perItem` invocations. The
- * iterator's batch boundaries are preserved — work for batch N starts
- * only after the source yields it, but within and across batches up to
- * `concurrency` items run concurrently via a sliding worker pool.
- *
- * On error: keeps going with `onError: "continue"` (default), or rethrows
- * the first error with `onError: "throw"`. Either way every visited
- * item shows up in `results` so callers see the full audit trail.
- */
+/** Iterates scope items; a scheduler forces sequential execution and ends between items.
+ * Errors are collected by default or rethrown when `onError` is `throw`. */
 export const iterateScope = async <T>({
 	iterate,
 	perItem,
 	onError = "continue",
 	concurrency = 1,
+	scheduler,
+	shouldStop,
 }: {
 	iterate: () => AsyncGenerator<RunScopeItem[]>;
 	perItem: (item: RunScopeItem) => Promise<T>;
 	onError?: "throw" | "continue";
 	concurrency?: number;
+	scheduler?: MigrationRunScheduler;
+	shouldStop?: () => boolean;
 }): Promise<IterateScopeSummary<T>> => {
 	const results: IterateScopeItemResult<T>[] = [];
 	let succeeded = 0;
 	let failed = 0;
-	const maxParallel = Math.max(1, Math.floor(concurrency));
+	const maxParallel = scheduler ? 1 : Math.max(1, Math.floor(concurrency));
+	const sliceStartedAtMs = scheduler?.now();
+	let hasProcessedScheduledItem = false;
+	const summarize = (
+		completion: IterateScopeCompletion,
+	): IterateScopeSummary<T> => ({
+		processed: succeeded + failed,
+		succeeded,
+		failed,
+		results,
+		completion,
+		cursor: results[results.length - 1]?.item.internal_id ?? null,
+	});
 
 	const runItem = async (item: RunScopeItem) => {
 		try {
@@ -53,11 +63,23 @@ export const iterateScope = async <T>({
 		}
 	};
 
+	const scheduledSliceIsComplete = () =>
+		scheduler !== undefined &&
+		hasProcessedScheduledItem &&
+		sliceStartedAtMs !== undefined &&
+		scheduler.now() - sliceStartedAtMs >= scheduler.sliceDurationMs;
+
 	if (maxParallel === 1) {
 		for await (const batch of iterate()) {
-			for (const item of batch) await runItem(item);
+			for (const item of batch) {
+				if (shouldStop?.()) return summarize("stopped");
+				if (scheduledSliceIsComplete()) return summarize("slice_complete");
+				await runItem(item);
+				hasProcessedScheduledItem = true;
+				if (shouldStop?.()) return summarize("stopped");
+			}
 		}
-		return { processed: succeeded + failed, succeeded, failed, results };
+		return summarize("exhausted");
 	}
 
 	const inflight = new Set<Promise<void>>();
@@ -70,13 +92,21 @@ export const iterateScope = async <T>({
 
 	for await (const batch of iterate()) {
 		for (const item of batch) {
+			if (shouldStop?.()) {
+				await Promise.all(inflight);
+				return summarize("stopped");
+			}
 			schedule(item);
 			if (inflight.size >= maxParallel) {
 				await Promise.race(inflight);
+				if (shouldStop?.()) {
+					await Promise.all(inflight);
+					return summarize("stopped");
+				}
 			}
 		}
 	}
 	await Promise.all(inflight);
 
-	return { processed: succeeded + failed, succeeded, failed, results };
+	return summarize(shouldStop?.() ? "stopped" : "exhausted");
 };