improvement(billing): paginate and parallelize the cycle-close sweep (#7103)

* improvement(billing): paginate and parallelize the cycle-close sweep

The sweep previously loaded every lagging candidate in one select and
closed them serially — correct at any scale (per-subscription guarded
marker claims make any interleaving safe) but O(fleet) in memory and
wall clock. It now iterates candidates in keyset pages and closes each
page with bounded concurrency; a total mapper keeps one failure from
rejecting its page, and an interrupted sweep resumes from the durable
markers on the next run.

The member-lock select gains a deterministic ORDER BY: today
member_user_id_unique keeps org rosters disjoint and every other locker
is single-row, so no deadlock is currently constructible — the ordering
is hygiene that also holds if membership exclusivity is ever relaxed.

A partial index on exactly the sweep's candidate predicate (built
CONCURRENTLY behind the runner's COMMIT breakpoint) keeps the scan
O(lagging) rather than O(fleet); closes shrink it as they land.

Validated end-to-end against Stripe test mode: one sweep over 304
staged lagging subscriptions (300 synthetic + 2 multi-member orgs
closing concurrently + 2 real test-clock subscriptions with genuine
overage) closed 304/304 with failed=0 in 667ms — the same fleet took
~302s serially — while both real subscriptions produced exactly one
paid overage invoice each at the expected amounts, with zero invoices
for the no-overage fleet.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* fix(billing): restore the runner's lock_timeout and inline entitled statuses for the sweep index

The concurrent-index migration now restores lock_timeout='5s' after its
post-COMMIT section, matching the runner convention every prior
CONCURRENTLY migration follows — without it, later pending files' DDL in
the same session would wait indefinitely instead of hitting the runner's
55P03 retry path.

The sweep's status filter becomes inlined literals derived from
ENTITLED_SUBSCRIPTION_STATUSES: the planner can only prove a query
implies a partial index's predicate from literals, and the parameterized
generic plan demonstrably seq-scans where the literal form uses
subscription_cycle_close_lagging_idx.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

apps/sim/lib/billing/cycle-close.test.ts

@@ -1,7 +1,13 @@
 /**
  * @vitest-environment node
  */
-import { dbChainMockFns, queueTableRows, resetDbChainMock, schemaMock } from '@sim/testing'
+import {
+  dbChainMockFns,
+  drizzleOrmMock,
+  queueTableRows,
+  resetDbChainMock,
+  schemaMock,
+} from '@sim/testing'
 import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
 const {
@@ -223,6 +229,10 @@ describe('closeElapsedBillingPeriod', () => {
       metadata: expect.objectContaining({ type: 'overage_billing', organizationId: 'org-1' }),
     })
 
+    // Member row locks are acquired in sorted order so parallel closes of
+    // organizations sharing a member cannot deadlock.
+    expect(drizzleOrmMock.asc).toHaveBeenCalledWith(schemaMock.userStats.userId)
+
     // Bookkeeping: last-period CASE write + billedOverage reset on member rows.
     const bookkeepingSet = dbChainMockFns.set.mock.calls.find(
       (call) => (call[0] as Record<string, unknown>).billedOverageThisPeriod === '0'
@@ -662,4 +672,39 @@ describe('sweepBillingCycleCloses', () => {
     expect(summary.initialized).toBe(2)
     expect(summary.failed).toBe(0)
   })
+
+  it('iterates candidates in keyset pages, fetching until a short page', async () => {
+    queueTableRows(
+      schemaMock.subscription,
+      Array.from({ length: 250 }, (_, i) =>
+        subRow({ id: `sub-p1-${String(i).padStart(3, '0')}`, lastClosedPeriodStart: null })
+      )
+    )
+    queueTableRows(schemaMock.subscription, [
+      subRow({ id: 'sub-p2-0', lastClosedPeriodStart: null }),
+      subRow({ id: 'sub-p2-1', lastClosedPeriodStart: null }),
+      subRow({ id: 'sub-p2-2', lastClosedPeriodStart: null }),
+    ])
+
+    const summary = await sweepBillingCycleCloses()
+
+    expect(summary).toEqual({ candidates: 253, closed: 0, initialized: 253, failed: 0 })
+    // A full page signals another fetch; the short second page ends the loop.
+    expect(dbChainMockFns.limit).toHaveBeenCalledTimes(2)
+  })
+
+  it('isolates a failing close inside a page', async () => {
+    // 'sub-bad' has a lagging marker past the grace, so its close proceeds
+    // into the org-scope lookup and blows up; 'sub-ok' initializes. The page
+    // completes and the failure is counted, never rethrown.
+    queueTableRows(schemaMock.subscription, [
+      subRow({ id: 'sub-bad' }),
+      subRow({ id: 'sub-ok', lastClosedPeriodStart: null }),
+    ])
+    mockIsSubscriptionOrgScoped.mockRejectedValueOnce(new Error('boom'))
+
+    const summary = await sweepBillingCycleCloses()
+
+    expect(summary).toEqual({ candidates: 2, closed: 0, initialized: 1, failed: 1 })
+  })
 })

apps/sim/lib/billing/cycle-close.ts

@@ -9,7 +9,7 @@ import {
 } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { getErrorMessage } from '@sim/utils/errors'
-import { and, eq, inArray, isNull, lt, or, sql } from 'drizzle-orm'
+import { and, asc, eq, gt, inArray, isNull, lt, or, sql } from 'drizzle-orm'
 import { BILLING_LOCK_TIMEOUT_MS } from '@/lib/billing/constants'
 import { computeOrgOverageAmount, isSubscriptionOrgScoped } from '@/lib/billing/core/billing'
 import { resolveSubscriptionUsagePeriod } from '@/lib/billing/core/reporting-period'
@@ -23,6 +23,7 @@ import { ENTITLED_SUBSCRIPTION_STATUSES, getPlanPricing } from '@/lib/billing/su
 import { toDecimal, toNumber } from '@/lib/billing/utils/decimal'
 import { OUTBOX_EVENT_TYPES } from '@/lib/billing/webhooks/outbox-handlers'
 import { enqueueOutboxEvent } from '@/lib/core/outbox/service'
+import { mapWithConcurrency } from '@/lib/core/utils/concurrency'
 import type { DbOrTx } from '@/lib/db/types'
 import { captureServerEvent } from '@/lib/posthog/server'
 
@@ -490,12 +491,18 @@ export async function closeElapsedBillingPeriod(
     }> => {
       await tx.execute(sql.raw(`SET LOCAL lock_timeout = '${BILLING_LOCK_TIMEOUT_MS}ms'`))
 
-      // Canonical lock order: member userStats rows, then the organization row.
+      // Canonical lock order: member userStats rows (sorted, so any closers
+      // with overlapping row sets acquire locks in one global order — today
+      // `member_user_id_unique` keeps org rosters disjoint and other lockers
+      // are single-row, so this is deterministic-order hygiene that also
+      // holds if membership exclusivity is ever relaxed), then the
+      // organization row.
       if (memberIds.length > 0) {
         await tx
           .select({ userId: userStats.userId })
           .from(userStats)
           .where(inArray(userStats.userId, memberIds))
+          .orderBy(asc(userStats.userId))
           .for('update')
       }
       let orgCreditBalance = 0
@@ -804,49 +811,101 @@ export interface CycleCloseSweepSummary {
   failed: number
 }
 
+/**
+ * Candidates fetched per page of the sweep's keyset iteration. Bounds sweep
+ * memory to one page of subscription rows regardless of fleet size.
+ */
+const SWEEP_PAGE_SIZE = 250
+
+/**
+ * Concurrent closes per page. Closes of different subscriptions are
+ * independent transactions; the per-subscription guarded marker claim and the
+ * deterministic member-lock ordering make any interleaving safe, so this
+ * bound exists only to cap database pressure from one sweep.
+ */
+const SWEEP_CLOSE_CONCURRENCY = 10
+
+/**
+ * Entitled statuses inlined as SQL literals rather than bind parameters: the
+ * candidate scan targets the partial index on exactly this predicate, and the
+ * planner can only prove a query implies a partial index's predicate from
+ * literals — a parameterized generic plan would fall back to scanning the
+ * whole table.
+ */
+const ENTITLED_STATUS_LITERALS = sql.raw(
+  ENTITLED_SUBSCRIPTION_STATUSES.map((status) => `'${status}'`).join(', ')
+)
+
 /**
  * Daily catch-all that closes every elapsed billing period. Candidates are
  * entitled subscriptions whose close marker lags their current `periodStart`
- * — i.e. the period advanced (via Stripe sync) since the last close. Each
- * close is independently atomic, so one failure never blocks the rest.
+ * — i.e. the period advanced (via Stripe sync) since the last close.
+ *
+ * Iterates candidates in keyset pages (matching the partial index on exactly
+ * this predicate) and closes each page with bounded concurrency. Each close
+ * is independently atomic and error-isolated, so one failure never blocks
+ * the rest; a page that closes successfully leaves the candidate set, which
+ * also makes an interrupted sweep (deploy, crash) resume where it left off
+ * on the next run.
  */
 export async function sweepBillingCycleCloses(): Promise<CycleCloseSweepSummary> {
-  const candidates = await db
-    .select()
-    .from(subscriptionTable)
-    .where(
-      and(
-        inArray(subscriptionTable.status, ENTITLED_SUBSCRIPTION_STATUSES),
-        sql`${subscriptionTable.periodStart} IS NOT NULL`,
-        or(
-          isNull(subscriptionTable.lastClosedPeriodStart),
-          lt(subscriptionTable.lastClosedPeriodStart, subscriptionTable.periodStart)
-        )
-      )
-    )
-
   const summary: CycleCloseSweepSummary = {
-    candidates: candidates.length,
+    candidates: 0,
     closed: 0,
     initialized: 0,
     failed: 0,
   }
+  const startedAt = Date.now()
+  let cursor = ''
 
-  for (const sub of candidates) {
-    try {
-      const result = await closeElapsedBillingPeriod(sub)
-      if (result.status === 'closed') summary.closed++
-      if (result.status === 'initialized') summary.initialized++
-    } catch (error) {
-      summary.failed++
-      logger.error('Cycle close failed for subscription', {
-        subscriptionId: sub.id,
-        plan: sub.plan,
-        error: getErrorMessage(error),
-      })
+  while (true) {
+    const page = await db
+      .select()
+      .from(subscriptionTable)
+      .where(
+        and(
+          sql`${subscriptionTable.status} in (${ENTITLED_STATUS_LITERALS})`,
+          sql`${subscriptionTable.periodStart} IS NOT NULL`,
+          or(
+            isNull(subscriptionTable.lastClosedPeriodStart),
+            lt(subscriptionTable.lastClosedPeriodStart, subscriptionTable.periodStart)
+          ),
+          gt(subscriptionTable.id, cursor)
+        )
+      )
+      .orderBy(asc(subscriptionTable.id))
+      .limit(SWEEP_PAGE_SIZE)
+
+    if (page.length === 0) break
+    summary.candidates += page.length
+    cursor = page[page.length - 1].id
+
+    // Total mapper: every close resolves to a status so one failure never
+    // rejects the page (mapWithConcurrency is all-or-nothing on rejection).
+    const results = await mapWithConcurrency(page, SWEEP_CLOSE_CONCURRENCY, async (sub) => {
+      try {
+        return (await closeElapsedBillingPeriod(sub)).status
+      } catch (error) {
+        logger.error('Cycle close failed for subscription', {
+          subscriptionId: sub.id,
+          plan: sub.plan,
+          error: getErrorMessage(error),
+        })
+        return 'failed' as const
+      }
+    })
+    for (const status of results) {
+      if (status === 'closed') summary.closed++
+      if (status === 'initialized') summary.initialized++
+      if (status === 'failed') summary.failed++
     }
+
+    if (page.length < SWEEP_PAGE_SIZE) break
   }
 
-  logger.info('Billing cycle-close sweep finished', { ...summary })
+  logger.info('Billing cycle-close sweep finished', {
+    ...summary,
+    durationMs: Date.now() - startedAt,
+  })
   return summary
 }

packages/db/migrations/0307_add_subscription_cycle_close_lagging_idx.sql

@@ -0,0 +1,7 @@
+-- Concurrent index operations cannot run inside the migration runner's transaction.
+COMMIT;--> statement-breakpoint
+SET lock_timeout = 0;--> statement-breakpoint
+-- migration-safe: replay removes an invalid build created by this migration; concurrent operations preserve writes.
+DROP INDEX CONCURRENTLY IF EXISTS "subscription_cycle_close_lagging_idx";--> statement-breakpoint
+CREATE INDEX CONCURRENTLY IF NOT EXISTS "subscription_cycle_close_lagging_idx" ON "subscription" USING btree ("id") WHERE "subscription"."status" in ('active', 'past_due') and "subscription"."period_start" is not null and ("subscription"."last_closed_period_start" is null or "subscription"."last_closed_period_start" < "subscription"."period_start");--> statement-breakpoint
+SET lock_timeout = '5s';

packages/db/migrations/meta/_journal.json

@@ -2143,6 +2143,13 @@
       "when": 1787702824793,
       "tag": "0306_knowledge_pipeline_hardening",
       "breakpoints": true
+    },
+    {
+      "idx": 307,
+      "version": "7",
+      "when": 1787722270619,
+      "tag": "0307_add_subscription_cycle_close_lagging_idx",
+      "breakpoints": true
     }
   ]
 }

packages/db/schema.ts

@@ -1377,6 +1377,20 @@ export const subscription = pgTable(
       table.referenceId,
       table.status
     ),
+    /**
+     * Partial index for the cycle-close sweep's keyset iteration: exactly the
+     * entitled subscriptions whose close marker lags the current period. The
+     * predicate must mirror the sweep query in `lib/billing/cycle-close.ts`
+     * (status list = ENTITLED_SUBSCRIPTION_STATUSES, hardcoded here because
+     * packages cannot import from apps); a drifted predicate degrades to a
+     * seq scan, never a wrong result. The index stays tiny — closes remove
+     * rows from it — so the candidate scan is O(lagging), not O(fleet).
+     */
+    cycleCloseLaggingIdx: index('subscription_cycle_close_lagging_idx')
+      .on(table.id)
+      .where(
+        sql`${table.status} in ('active', 'past_due') and ${table.periodStart} is not null and (${table.lastClosedPeriodStart} is null or ${table.lastClosedPeriodStart} < ${table.periodStart})`
+      ),
     enterpriseMetadataCheck: check(
       'check_enterprise_metadata',
       sql`plan != 'enterprise' OR metadata IS NOT NULL`