diff --git a/apps/sim/lib/billing/constants.ts b/apps/sim/lib/billing/constants.ts index ebe5e7c17b5..c1fd884a293 100644 --- a/apps/sim/lib/billing/constants.ts +++ b/apps/sim/lib/billing/constants.ts @@ -39,6 +39,15 @@ export const DEFAULT_OVERAGE_THRESHOLD = 100 */ export const BILLING_LOCK_TIMEOUT_MS = 5_000 +/** + * Bound on one ledger sum. A large payer's period covers millions of rows, and from a cold + * cache or under heavy I/O the sum can run for tens of seconds; past this the database ends it + * and the read fails, so a caller that admits on the sum fails closed rather than waiting + * without limit. The usage gate derives its coalescing deadline from this bound, so the sum + * always ends at the database before the gate gives up on it. + */ +export const USAGE_LEDGER_STATEMENT_TIMEOUT_MS = 60_000 + /** * Available credit tiers. Each tier maps a credit amount to the underlying dollar * cost and carries that tier's fixed weekly refresh allowance. diff --git a/apps/sim/lib/billing/core/usage-gate-cache.test.ts b/apps/sim/lib/billing/core/usage-gate-cache.test.ts index 752ae97e762..3a073fc0640 100644 --- a/apps/sim/lib/billing/core/usage-gate-cache.test.ts +++ b/apps/sim/lib/billing/core/usage-gate-cache.test.ts @@ -16,6 +16,7 @@ import { checkIngestionUsageLimits, checkSearchUsageLimits, resetUsageGateCache, + USAGE_GATE_SETTLE_TIMEOUT_MS, USAGE_GATE_TTL_MS, } from '@/lib/billing/core/usage-gate-cache' @@ -178,4 +179,36 @@ describe('checkExecutionUsageLimits', () => { await checkExecutionUsageLimits(ATTRIBUTION) expect(mockCheck).toHaveBeenCalledTimes(2) }) + + it('waits for a slow ledger read past the singleflight default instead of blocking', async () => { + vi.useFakeTimers() + try { + mockCheck.mockImplementationOnce(() => sleep(45_000).then(() => ({ isExceeded: false }))) + const pending = checkExecutionUsageLimits(ATTRIBUTION) + await vi.advanceTimersByTimeAsync(45_000) + await expect(pending).resolves.toEqual({ isExceeded: false }) + expect(mockCheck).toHaveBeenCalledTimes(1) + } finally { + vi.useRealTimers() + } + }) + + it('gives up on a read that never answers at the gate deadline, then reads fresh', async () => { + vi.useFakeTimers() + try { + mockCheck.mockReturnValueOnce(new Promise(() => {})) + const hung = checkExecutionUsageLimits(ATTRIBUTION) + const rejection = expect(hung).rejects.toThrow( + `did not settle within ${USAGE_GATE_SETTLE_TIMEOUT_MS}ms` + ) + await vi.advanceTimersByTimeAsync(USAGE_GATE_SETTLE_TIMEOUT_MS) + await rejection + await expect(checkExecutionUsageLimits(ATTRIBUTION)).resolves.toEqual({ + isExceeded: false, + }) + expect(mockCheck).toHaveBeenCalledTimes(2) + } finally { + vi.useRealTimers() + } + }) }) diff --git a/apps/sim/lib/billing/core/usage-gate-cache.ts b/apps/sim/lib/billing/core/usage-gate-cache.ts index 4e3a1284226..dbcf326f3d4 100644 --- a/apps/sim/lib/billing/core/usage-gate-cache.ts +++ b/apps/sim/lib/billing/core/usage-gate-cache.ts @@ -1,4 +1,5 @@ import { LRUCache } from 'lru-cache' +import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants' import { type AttributedUsageLimitsResult, type BillingAttributionSnapshot, @@ -20,6 +21,17 @@ import { coalesceLocally } from '@/lib/concurrency/singleflight' */ export const USAGE_GATE_TTL_MS = 5 * 60 * 1000 +/** + * How long a coalesced usage read may take before its callers give up on it. The read's cost is + * the ledger sum, which the database ends at {@link USAGE_LEDGER_STATEMENT_TIMEOUT_MS}; the + * remainder is a few indexed lookups and the connection waits around them. The singleflight + * default of 30 s exists to bound a hung producer, and a slow sum is not a hung one: given up on + * early, it keeps running detached while every joined caller fails and the next caller starts a + * second sum alongside it. Derived from the statement bound so the database always ends the sum + * first, and the gate only gives up on a connection that never answers. + */ +export const USAGE_GATE_SETTLE_TIMEOUT_MS = USAGE_LEDGER_STATEMENT_TIMEOUT_MS + 15_000 + /** * Recent gate answers, admitted and refused, with `LRUCache` supplying the TTL * and the size bound. Each entry point decides which of them it may serve. @@ -62,9 +74,9 @@ function gateKey(attribution: BillingAttributionSnapshot): string { * the cache. A read that throws writes nothing. * * `coalesceLocally` collapses concurrent misses onto one ledger read and bounds - * a hung read at its settle deadline. The write stays on the value this caller - * received, so a producer that timed out and later resolved cannot overwrite a - * fresher answer. + * a hung read at {@link USAGE_GATE_SETTLE_TIMEOUT_MS}. The write stays on the + * value this caller received, so a producer that timed out and later resolved + * cannot overwrite a fresher answer. * * There is deliberately no invalidator: usage and limit changes land in other * processes (execution workers, Stripe webhooks), so the TTL is the real bound. @@ -77,8 +89,10 @@ async function checkUsageLimitsThroughCache( const cached = gateCache.get(key) if (cached !== undefined && (cacheRefusals || !cached.isExceeded)) return cached - const result = await coalesceLocally(`usage-gate:${key}`, () => - checkAttributedUsageLimits(attribution) + const result = await coalesceLocally( + `usage-gate:${key}`, + () => checkAttributedUsageLimits(attribution), + USAGE_GATE_SETTLE_TIMEOUT_MS ) if (cacheRefusals || !result.isExceeded) gateCache.set(key, result) return result diff --git a/apps/sim/lib/billing/core/usage-log.test.ts b/apps/sim/lib/billing/core/usage-log.test.ts index cd612889af2..ccb88ae4c0b 100644 --- a/apps/sim/lib/billing/core/usage-log.test.ts +++ b/apps/sim/lib/billing/core/usage-log.test.ts @@ -33,9 +33,11 @@ vi.mock('@/lib/billing/subscriptions/utils', () => ({ isOrgScopedSubscription: mockIsOrgScopedSubscription, })) +import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants' import { CUMULATIVE_COST_EPSILON, CumulativeUsageContextMismatchError, + getBillingPeriodUsageCost, getUserUsageLogs, getWorkspaceUsageLogs, recordCumulativeUsage, @@ -554,3 +556,35 @@ describe('usage-log query scopes', () => { }) }) }) + +describe('getBillingPeriodUsageCost', () => { + beforeEach(() => { + vi.clearAllMocks() + installSharedDbMocks() + }) + + it('bounds the ledger sum with its own statement timeout inside one transaction', async () => { + const execute = vi.fn().mockResolvedValue([]) + const where = vi.fn().mockResolvedValue([{ cost: '12.5' }]) + const tx = { execute, select: vi.fn(() => ({ from: vi.fn(() => ({ where })) })) } + mockTransaction.mockImplementation((callback: (client: typeof tx) => Promise) => + callback(tx) + ) + + const cost = await getBillingPeriodUsageCost( + { type: 'organization', id: 'org-1' }, + { start: new Date('2026-05-01T00:00:00Z'), end: new Date('2027-05-01T00:00:00Z') } + ) + + expect(cost).toBe(12.5) + expect(mockTransaction).toHaveBeenCalledTimes(1) + const executed = execute.mock.calls.map( + ([statement]) => (statement as { toSQL: () => { sql: string } }).toSQL().sql + ) + expect(executed).toContain( + `SET LOCAL statement_timeout = '${USAGE_LEDGER_STATEMENT_TIMEOUT_MS}ms'` + ) + /** The bound is set before the sum runs, not after. */ + expect(execute.mock.invocationCallOrder[0]).toBeLessThan(where.mock.invocationCallOrder[0]) + }) +}) diff --git a/apps/sim/lib/billing/core/usage-log.ts b/apps/sim/lib/billing/core/usage-log.ts index b79653f0e49..0849352ef35 100644 --- a/apps/sim/lib/billing/core/usage-log.ts +++ b/apps/sim/lib/billing/core/usage-log.ts @@ -15,6 +15,7 @@ import { textKey, timestampKey, } from '@/lib/api/list-query' +import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants' import { defaultBillingPeriod } from '@/lib/billing/core/billing-period' import { getHighestPrioritySubscription } from '@/lib/billing/core/plan' import { @@ -215,6 +216,10 @@ async function resolveBillingContext( /** * Returns attributed ledger usage for a billing entity/period. The ledger is * the sole source of truth for usage — there is no userStats baseline. + * + * The sum runs in a transaction of its own on the given client so that it can + * be bounded by {@link USAGE_LEDGER_STATEMENT_TIMEOUT_MS} for that statement + * alone: `SET LOCAL` ends with the transaction and never reaches the pool. */ export async function getBillingPeriodUsageCost( billingEntity: BillingEntity, @@ -238,12 +243,17 @@ export async function getBillingPeriodUsageCost( ) } - const [row] = await executor - .select({ - cost: sql`COALESCE(SUM(${usageLog.cost}), 0)`, - }) - .from(usageLog) - .where(and(...conditions)) + const [row] = await executor.transaction(async (tx) => { + await tx.execute( + sql.raw(`SET LOCAL statement_timeout = '${USAGE_LEDGER_STATEMENT_TIMEOUT_MS}ms'`) + ) + return tx + .select({ + cost: sql`COALESCE(SUM(${usageLog.cost}), 0)`, + }) + .from(usageLog) + .where(and(...conditions)) + }) return Number.parseFloat(row?.cost ?? '0') } diff --git a/apps/sim/lib/billing/enterprise-provisioning.test.ts b/apps/sim/lib/billing/enterprise-provisioning.test.ts index 1368760c0ee..b69eb60f452 100644 --- a/apps/sim/lib/billing/enterprise-provisioning.test.ts +++ b/apps/sim/lib/billing/enterprise-provisioning.test.ts @@ -199,6 +199,8 @@ describe('Enterprise issuance preflight', () => { queueTableRows(schemaMock.workspace, []) queueTableRows(schemaMock.workspace, []) queueTableRows(schemaMock.subscription, []) + /** The run count resolves first; the ledger sum opens its bounded transaction before it reads. */ + queueTableRows(schemaMock.usageLog, [{ workflowRuns: 0 }]) queueTableRows(schemaMock.usageLog, [{ cost: '150' }]) const result = await getEnterpriseIssuancePreflight({