Skip to content

Commit 25da162

Browse files
authored
perf(billing): sum reporting-window usage from a shared prefix plus a short tail (#8974)
* perf(billing): sum reporting-window usage from a shared prefix plus a short tail The soft-gate reporting-usage cache re-ran a payer's year-long ledger sum on every 30 s refill, in every process that missed. Hold the window's sum before a cutoff (now - 5 min) in Redis, refreshed by one process under a SET NX PX lease about every 10 minutes, and sum only the rows since the cutoff on each refill. Readers serve a stale prefix while the lease holder refreshes; a reader with no prefix waits up to 3 s for the holder's before summing the window. Without a ready Redis, or for a window that does not span the cutoff, the read is the exact whole-window sum as before. Exact billing callers are unchanged. * perf(billing): refresh the usage prefix every 3 minutes, discard it on account deletion - Refresh the shared prefix about every 3 minutes instead of 10, so a top-up to a row older than the cutoff lags at most that long, and keep a stale prefix only for one refresh interval plus one lease. - With no prefix stored and the lease held elsewhere, sum the whole window at once instead of waiting, so no read is slower than before. - The lease holder sums its tail alongside the prefix. - Tag each prefix with a shared epoch read in the same MGET, and replace the epoch once an account deletion commits, so a deleted member's cascaded rows are never counted from a stored prefix. * fix(billing): discard a payer's usage prefix when a top-up raises a row inside it - a committed cumulative top-up to a row older than the prefix cutoff lag bumps a per-payer prefix epoch, so the next soft-gate read sums the window again and sees the top-up as fast as an uncached read - prefixes are stored under both the global and the payer epoch; a mismatch on either counts as no prefix - epoch keys, readiness, and invalidation move to reporting-usage-epoch so usage-log can bump without an import cycle - simplify the prefix-plus-tail sum and drop history from comments
1 parent 2abff53 commit 25da162

5 files changed

Lines changed: 645 additions & 38 deletions

File tree

‎apps/sim/lib/billing/core/reporting-usage-cache.integration.ts‎

Lines changed: 301 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,27 +10,47 @@ import * as schema from '@sim/db/schema'
1010
import { readTestDatabaseUrl, readTestRedisUrl } from '@sim/db/testing/test-infrastructure'
1111
import { redisConfigMock, redisConfigMockFns } from '@sim/testing/mocks/redis-config.mock'
1212
import { generateId } from '@sim/utils/id'
13+
import { randomInt, randomItem } from '@sim/utils/random'
1314
import { drizzle } from 'drizzle-orm/postgres-js'
1415
import Redis from 'ioredis'
1516
import postgres from 'postgres'
1617
import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'
1718

18-
const { transaction } = vi.hoisted(() => ({ transaction: vi.fn() }))
19+
const { transaction, summedPeriods } = vi.hoisted(() => ({
20+
transaction: vi.fn(),
21+
summedPeriods: [] as UsageQueryPeriod[],
22+
}))
1923
const databaseUrl = readTestDatabaseUrl()
2024
const redisUrl = readTestRedisUrl()
2125

2226
vi.mock('@sim/db', () => ({ db: { transaction }, dbReplica: {} }))
2327
vi.mock('@/lib/core/config/redis', () => redisConfigMock)
28+
vi.mock('@/lib/billing/core/usage-log', async (importOriginal) => {
29+
const actual = await importOriginal<typeof import('@/lib/billing/core/usage-log')>()
30+
return {
31+
...actual,
32+
getBillingPeriodUsageCost: (...args: Parameters<typeof actual.getBillingPeriodUsageCost>) => {
33+
summedPeriods.push(args[1])
34+
return actual.getBillingPeriodUsageCost(...args)
35+
},
36+
}
37+
})
2438

2539
import {
2640
REPORTING_USAGE_CACHE_TTL_MS,
2741
readSoftGateUsageCost,
2842
} from '@/lib/billing/core/reporting-usage-cache'
29-
import type { BillingEntity, UsageQueryPeriod } from '@/lib/billing/core/usage-log'
43+
import { invalidateReportingUsagePrefixes } from '@/lib/billing/core/reporting-usage-epoch'
44+
import {
45+
type BillingEntity,
46+
type RecordCumulativeUsageParams,
47+
recordCumulativeUsage,
48+
type UsageQueryPeriod,
49+
} from '@/lib/billing/core/usage-log'
3050

3151
const schemaName = `reporting_usage_${generateId().replaceAll('-', '')}`
3252
const connection = postgres(databaseUrl, {
33-
max: 2,
53+
max: 4,
3454
prepare: false,
3555
connection: { search_path: schemaName },
3656
onnotice: () => undefined,
@@ -43,28 +63,128 @@ const REPORTING: UsageQueryPeriod = {
4363
source: 'reporting',
4464
}
4565

46-
function sharedKey(payer: BillingEntity): string {
47-
return `usage:reporting:v1:${payer.type}:${payer.id}:reporting:${REPORTING.start.toISOString()}:${REPORTING.end.toISOString()}`
66+
const DAY_MS = 86_400_000
67+
const MINUTE_MS = 60_000
68+
69+
/** A year-long window that started 100 days ago, so the clock is always inside it. */
70+
function openReportingWindow(): UsageQueryPeriod {
71+
const start = Math.floor((Date.now() - 100 * DAY_MS) / DAY_MS) * DAY_MS
72+
return { start: new Date(start), end: new Date(start + 365 * DAY_MS), source: 'reporting' }
73+
}
74+
75+
function windowKey(payer: BillingEntity, period: UsageQueryPeriod): string {
76+
return `${payer.type}:${payer.id}:reporting:${period.start.toISOString()}:${period.end.toISOString()}`
77+
}
78+
79+
function sharedKey(payer: BillingEntity, period: UsageQueryPeriod = REPORTING): string {
80+
return `usage:reporting:v1:${windowKey(payer, period)}`
81+
}
82+
83+
function prefixKey(payer: BillingEntity, period: UsageQueryPeriod): string {
84+
return `usage:reporting:v1:prefix:${windowKey(payer, period)}`
85+
}
86+
87+
function prefixLeaseKey(payer: BillingEntity, period: UsageQueryPeriod): string {
88+
return `usage:reporting:v1:prefix-lease:${windowKey(payer, period)}`
89+
}
90+
91+
function payerEpochKey(payer: BillingEntity): string {
92+
return `usage:reporting:v1:prefix-epoch:${payer.type}:${payer.id}`
93+
}
94+
95+
type ReportingUsageCacheModule = typeof import('@/lib/billing/core/reporting-usage-cache')
96+
97+
/** Fresh module instances, so each in-process cache is cold, as in a new process. */
98+
async function newProcesses(count: number): Promise<ReportingUsageCacheModule[]> {
99+
const processes: ReportingUsageCacheModule[] = []
100+
for (let index = 0; index < count; index++) {
101+
vi.resetModules()
102+
processes.push(await import('@/lib/billing/core/reporting-usage-cache'))
103+
}
104+
return processes
105+
}
106+
107+
async function readInNewProcess(payer: BillingEntity, period: UsageQueryPeriod): Promise<number> {
108+
const [cache] = await newProcesses(1)
109+
return cache.readSoftGateUsageCost(payer, period)
110+
}
111+
112+
/** Sums that start at the window's start: a whole-window sum or a prefix, never a tail. */
113+
function windowStartSums(period: UsageQueryPeriod): UsageQueryPeriod[] {
114+
return summedPeriods.filter((summed) => summed.start.getTime() === period.start.getTime())
115+
}
116+
117+
async function insertUsage(payer: BillingEntity, cost: number, createdAt: Date): Promise<string> {
118+
const id = generateId()
119+
await connection`INSERT INTO usage_log (id, cost, billing_entity_type, billing_entity_id, created_at)
120+
VALUES (${id}, ${cost}, ${payer.type}, ${payer.id}, ${createdAt.toISOString()})`
121+
return id
122+
}
123+
124+
function cumulativeCharge(
125+
payer: BillingEntity,
126+
eventKey: string,
127+
cost: number
128+
): RecordCumulativeUsageParams {
129+
return {
130+
userId: 'actor',
131+
workspaceId: 'workspace',
132+
billingEntity: payer,
133+
billingPeriod: { start: REPORTING.start, end: REPORTING.end },
134+
source: 'workspace-chat',
135+
model: 'test-model',
136+
eventKey,
137+
cost,
138+
}
139+
}
140+
141+
/** A cumulative charge whose row was first written at `createdAt`. */
142+
async function startCumulativeCharge(
143+
payer: BillingEntity,
144+
cost: number,
145+
createdAt: Date
146+
): Promise<string> {
147+
const eventKey = `update-cost:${generateId()}`
148+
await recordCumulativeUsage(cumulativeCharge(payer, eventKey, cost))
149+
await connection`UPDATE usage_log SET created_at = ${createdAt.toISOString()}
150+
WHERE event_key = ${eventKey}`
151+
return eventKey
152+
}
153+
154+
async function exactSum(payer: BillingEntity, period: UsageQueryPeriod): Promise<number> {
155+
const [row] = await connection<{ cost: string }[]>`SELECT COALESCE(SUM(cost), 0) AS cost
156+
FROM usage_log WHERE billing_entity_type = ${payer.type} AND billing_entity_id = ${payer.id}
157+
AND created_at >= ${period.start.toISOString()} AND created_at < ${period.end.toISOString()}`
158+
return Number(row?.cost ?? 0)
48159
}
49160

50161
describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => {
51162
let redis: Redis
52163
let payer: BillingEntity
53164

165+
async function prefixStored(payer: BillingEntity, period: UsageQueryPeriod): Promise<void> {
166+
await vi.waitFor(async () => expect(await redis.exists(prefixKey(payer, period))).toBe(1))
167+
}
168+
54169
beforeAll(async () => {
55170
redis = new Redis(redisUrl!, { lazyConnect: true, maxRetriesPerRequest: 0 })
56171
await redis.connect()
57172
await connection.unsafe(`CREATE SCHEMA "${schemaName}"`)
58173
await connection.unsafe(`CREATE TABLE usage_log (
59-
id text PRIMARY KEY, cost numeric NOT NULL, billing_entity_type text,
174+
id text PRIMARY KEY, user_id text, category text, source text, description text,
175+
metadata jsonb, cost numeric NOT NULL, event_key text, billing_entity_type text,
60176
billing_entity_id text, billing_period_start timestamp, billing_period_end timestamp,
61-
created_at timestamp NOT NULL
62-
)`)
177+
workspace_id text, workflow_id text, execution_id text,
178+
created_at timestamp NOT NULL DEFAULT now()
179+
);
180+
CREATE UNIQUE INDEX usage_log_event_key_unique ON usage_log(event_key)
181+
WHERE event_key IS NOT NULL`)
63182
transaction.mockImplementation((callback) => database.transaction(callback))
64183
})
65184

66185
beforeEach(async () => {
67186
transaction.mockClear()
187+
summedPeriods.length = 0
68188
redisConfigMockFns.mockGetRedisClient.mockReturnValue(redis)
69189
payer = { type: 'organization', id: generateId() }
70190
await connection`INSERT INTO usage_log (id, cost, billing_entity_type, billing_entity_id, created_at)
@@ -75,7 +195,15 @@ describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => {
75195

76196
afterEach(async () => {
77197
vi.useRealTimers()
78-
await redis.del(sharedKey(payer))
198+
await redis.del(
199+
...[REPORTING, openReportingWindow()].flatMap((period) => [
200+
sharedKey(payer, period),
201+
prefixKey(payer, period),
202+
prefixLeaseKey(payer, period),
203+
]),
204+
payerEpochKey(payer),
205+
'usage:reporting:v1:prefix-epoch'
206+
)
79207
})
80208

81209
afterAll(async () => {
@@ -149,6 +277,170 @@ describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => {
149277
expect(await redis.get(sharedKey(payer))).toBe('9.99')
150278
})
151279

280+
it('sums the whole window in one query when Redis is absent', async () => {
281+
redisConfigMockFns.mockGetRedisClient.mockReturnValue(null)
282+
const period = openReportingWindow()
283+
await insertUsage(payer, 2, new Date(period.start.getTime() + DAY_MS))
284+
await insertUsage(payer, 3, new Date(Date.now() - MINUTE_MS))
285+
286+
await expect(readInNewProcess(payer, period)).resolves.toBe(5)
287+
expect(summedPeriods).toEqual([period])
288+
})
289+
290+
it('leases a due refresh to one of the processes missing together', async () => {
291+
const period = openReportingWindow()
292+
await insertUsage(payer, 2, new Date(period.start.getTime() + DAY_MS))
293+
await insertUsage(payer, 3, new Date(Date.now() - MINUTE_MS))
294+
await expect(readInNewProcess(payer, period)).resolves.toBe(5)
295+
await insertUsage(payer, 4, new Date(period.start.getTime() + 2 * DAY_MS))
296+
297+
const now = Date.now()
298+
vi.useFakeTimers({ toFake: ['Date'] })
299+
vi.setSystemTime(now + 4 * MINUTE_MS)
300+
await redis.del(sharedKey(payer, period), prefixLeaseKey(payer, period))
301+
const processes = await newProcesses(5)
302+
const costs = await Promise.all(
303+
processes.map((cache) => cache.readSoftGateUsageCost(payer, period))
304+
)
305+
306+
expect(windowStartSums(period)).toHaveLength(2)
307+
expect(costs).toContain(9)
308+
for (const cost of costs) expect([5, 9]).toContain(cost)
309+
})
310+
311+
it('sums the whole window at once when no prefix is stored and another process is summing one', async () => {
312+
const period = openReportingWindow()
313+
await insertUsage(payer, 2, new Date(period.start.getTime() + DAY_MS))
314+
await insertUsage(payer, 3, new Date(Date.now() - MINUTE_MS))
315+
await redis.set(prefixLeaseKey(payer, period), '1', 'PX', MINUTE_MS)
316+
const [cache] = await newProcesses(1)
317+
318+
const startedAt = performance.now()
319+
await expect(cache.readSoftGateUsageCost(payer, period)).resolves.toBe(5)
320+
expect(performance.now() - startedAt).toBeLessThan(1_000)
321+
expect(summedPeriods).toEqual([period])
322+
})
323+
324+
it('never reads above the ledger as it grows, and matches it once the prefix is refreshed', async () => {
325+
const period = openReportingWindow()
326+
const oldRows = [
327+
await insertUsage(payer, 1.25, new Date(period.start.getTime() + DAY_MS)),
328+
await insertUsage(payer, 0.5, new Date(Date.now() - 10 * DAY_MS)),
329+
]
330+
await expect(readInNewProcess(payer, period)).resolves.toBe(1.75)
331+
332+
for (let round = 0; round < 8; round++) {
333+
const recent = new Date(Date.now() - randomInt(0, 4 * MINUTE_MS))
334+
await insertUsage(payer, randomInt(1, 500) / 100, recent)
335+
if (round % 2 === 0) {
336+
const late = new Date(period.start.getTime() + randomInt(1, 90) * DAY_MS)
337+
await insertUsage(payer, randomInt(1, 500) / 100, late)
338+
} else {
339+
await connection`UPDATE usage_log SET cost = cost + ${randomInt(1, 500) / 100}
340+
WHERE id = ${randomItem(oldRows)}`
341+
}
342+
await redis.del(sharedKey(payer, period))
343+
344+
const served = await readInNewProcess(payer, period)
345+
expect(served).toBeLessThanOrEqual((await exactSum(payer, period)) + 1e-9)
346+
}
347+
expect(windowStartSums(period)).toHaveLength(1)
348+
349+
const now = Date.now()
350+
vi.useFakeTimers({ toFake: ['Date'] })
351+
vi.setSystemTime(now + 4 * MINUTE_MS)
352+
await redis.del(sharedKey(payer, period), prefixLeaseKey(payer, period))
353+
354+
const refreshed = await readInNewProcess(payer, period)
355+
expect(refreshed).toBeCloseTo(await exactSum(payer, period), 9)
356+
expect(windowStartSums(period)).toHaveLength(2)
357+
})
358+
359+
it('never reads above the ledger after ledger rows are deleted and prefixes are invalidated', async () => {
360+
const period = openReportingWindow()
361+
await insertUsage(payer, 2, new Date(period.start.getTime() + DAY_MS))
362+
const departedRow = await insertUsage(payer, 6, new Date(period.start.getTime() + 2 * DAY_MS))
363+
await insertUsage(payer, 3, new Date(Date.now() - MINUTE_MS))
364+
await expect(readInNewProcess(payer, period)).resolves.toBe(11)
365+
await prefixStored(payer, period)
366+
367+
await connection`DELETE FROM usage_log WHERE id = ${departedRow}`
368+
await invalidateReportingUsagePrefixes()
369+
await redis.del(sharedKey(payer, period))
370+
371+
await expect(readInNewProcess(payer, period)).resolves.toBe(await exactSum(payer, period))
372+
})
373+
374+
it('sees a top-up to a row inside the stored prefix on the next read', async () => {
375+
const period = openReportingWindow()
376+
await insertUsage(payer, 3, new Date(Date.now() - MINUTE_MS))
377+
const eventKey = await startCumulativeCharge(
378+
payer,
379+
2,
380+
new Date(period.start.getTime() + DAY_MS)
381+
)
382+
await expect(readInNewProcess(payer, period)).resolves.toBe(5)
383+
await prefixStored(payer, period)
384+
385+
await recordCumulativeUsage(cumulativeCharge(payer, eventKey, 6))
386+
await vi.waitFor(async () => expect(await redis.exists(payerEpochKey(payer))).toBe(1))
387+
await redis.del(sharedKey(payer, period))
388+
389+
await expect(readInNewProcess(payer, period)).resolves.toBe(9)
390+
expect(await exactSum(payer, period)).toBe(9)
391+
})
392+
393+
it('leaves the prefix in place for a top-up to a row still in the tail, and reads it exactly', async () => {
394+
const period = openReportingWindow()
395+
await insertUsage(payer, 2, new Date(period.start.getTime() + DAY_MS))
396+
const eventKey = await startCumulativeCharge(payer, 3, new Date(Date.now() - MINUTE_MS))
397+
await expect(readInNewProcess(payer, period)).resolves.toBe(5)
398+
await prefixStored(payer, period)
399+
400+
await recordCumulativeUsage(cumulativeCharge(payer, eventKey, 7))
401+
await redis.del(sharedKey(payer, period))
402+
403+
await expect(readInNewProcess(payer, period)).resolves.toBe(9)
404+
expect(await redis.exists(payerEpochKey(payer))).toBe(0)
405+
expect(windowStartSums(period)).toHaveLength(1)
406+
})
407+
408+
it('never over-counts a top-up that commits while a prefix refresh sums', async () => {
409+
const period = openReportingWindow()
410+
await insertUsage(payer, 3, new Date(Date.now() - MINUTE_MS))
411+
const eventKey = await startCumulativeCharge(
412+
payer,
413+
2,
414+
new Date(period.start.getTime() + DAY_MS)
415+
)
416+
const [refreshing] = await newProcesses(1)
417+
transaction.mockImplementationOnce(async (callback) => {
418+
const result = await database.transaction(callback)
419+
await recordCumulativeUsage(cumulativeCharge(payer, eventKey, 6))
420+
return result
421+
})
422+
423+
const served = await refreshing.readSoftGateUsageCost(payer, period)
424+
expect(served).toBeLessThanOrEqual(await exactSum(payer, period))
425+
await prefixStored(payer, period)
426+
await vi.waitFor(async () => expect(await redis.exists(payerEpochKey(payer))).toBe(1))
427+
await redis.del(sharedKey(payer, period))
428+
429+
await expect(readInNewProcess(payer, period)).resolves.toBe(await exactSum(payer, period))
430+
})
431+
432+
it('serves a stale prefix while another process holds the refresh lease', async () => {
433+
const period = openReportingWindow()
434+
await insertUsage(payer, 2, new Date(period.start.getTime() + DAY_MS))
435+
await insertUsage(payer, 3, new Date(Date.now() - MINUTE_MS))
436+
const staleCutoff = Date.now() - 10 * MINUTE_MS
437+
await redis.set(prefixKey(payer, period), `::${staleCutoff}:7`, 'PX', MINUTE_MS)
438+
await redis.set(prefixLeaseKey(payer, period), '1', 'PX', MINUTE_MS)
439+
440+
await expect(readInNewProcess(payer, period)).resolves.toBe(10)
441+
expect(windowStartSums(period)).toHaveLength(0)
442+
})
443+
152444
it('sums the ledger when the Redis client cannot be built', async () => {
153445
redisConfigMockFns.mockGetRedisClient.mockImplementation(() => {
154446
throw new Error('Invalid Redis configuration')

0 commit comments

Comments
 (0)