Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions apps/sim/lib/billing/constants.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
33 changes: 33 additions & 0 deletions apps/sim/lib/billing/core/usage-gate-cache.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import {
checkIngestionUsageLimits,
checkSearchUsageLimits,
resetUsageGateCache,
USAGE_GATE_SETTLE_TIMEOUT_MS,
USAGE_GATE_TTL_MS,
} from '@/lib/billing/core/usage-gate-cache'

Expand Down Expand Up @@ -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()
}
})
})
24 changes: 19 additions & 5 deletions apps/sim/lib/billing/core/usage-gate-cache.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { LRUCache } from 'lru-cache'
import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants'
import {
type AttributedUsageLimitsResult,
type BillingAttributionSnapshot,
Expand All @@ -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.
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand Down
34 changes: 34 additions & 0 deletions apps/sim/lib/billing/core/usage-log.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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])
})
})
22 changes: 16 additions & 6 deletions apps/sim/lib/billing/core/usage-log.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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,
Expand All @@ -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')
}
Expand Down
2 changes: 2 additions & 0 deletions apps/sim/lib/billing/enterprise-provisioning.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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({
Expand Down
Loading