fix(billing): let a slow ledger read finish instead of blocking execution at the singleflight default (#8146)

* fix(billing): let a slow ledger read finish instead of blocking execution at the singleflight default

* fix(billing): bound the ledger sum at the database and derive the gate deadline from it

* fix(billing): keep the ledger statement bound in the billing constants module

Author: waleedlatif1

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.

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()
+    }
+  })
 })

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

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<unknown>) =>
+      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])
+  })
+})

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<string>`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<string>`COALESCE(SUM(${usageLog.cost}), 0)`,
+      })
+      .from(usageLog)
+      .where(and(...conditions))
+  })
 
   return Number.parseFloat(row?.cost ?? '0')
 }

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({