Skip to content

Commit e25a6cf

Browse files
authored
fix(billing): release the reporting prefix lease once its prefix is stored (#9108)
1 parent e7076a9 commit e25a6cf

2 files changed

Lines changed: 47 additions & 12 deletions

File tree

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -186,6 +186,10 @@ describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => {
186186
transaction.mockClear()
187187
summedPeriods.length = 0
188188
redisConfigMockFns.mockGetRedisClient.mockReturnValue(redis)
189+
redisConfigMockFns.mockReleaseLock.mockImplementation(async (key: string, value: string) => {
190+
if ((await redis.get(key)) !== value) return false
191+
return (await redis.del(key)) === 1
192+
})
189193
payer = { type: 'organization', id: generateId() }
190194
await connection`INSERT INTO usage_log (id, cost, billing_entity_type, billing_entity_id, created_at)
191195
VALUES (${generateId()}, 4.25, 'organization', ${payer.id}, '2026-03-01'),
@@ -390,6 +394,29 @@ describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => {
390394
expect(await exactSum(payer, period)).toBe(9)
391395
})
392396

397+
it('refreshes a prefix discarded by a top-up once, instead of summing the window on every read', async () => {
398+
const period = openReportingWindow()
399+
await insertUsage(payer, 3, new Date(Date.now() - MINUTE_MS))
400+
const eventKey = await startCumulativeCharge(
401+
payer,
402+
2,
403+
new Date(period.start.getTime() + DAY_MS)
404+
)
405+
await expect(readInNewProcess(payer, period)).resolves.toBe(5)
406+
await prefixStored(payer, period)
407+
await vi.waitFor(async () => expect(await redis.exists(prefixLeaseKey(payer, period))).toBe(0))
408+
409+
await recordCumulativeUsage(cumulativeCharge(payer, eventKey, 6))
410+
await vi.waitFor(async () => expect(await redis.exists(payerEpochKey(payer))).toBe(1))
411+
await redis.del(sharedKey(payer, period))
412+
await expect(readInNewProcess(payer, period)).resolves.toBe(9)
413+
await redis.del(sharedKey(payer, period))
414+
await expect(readInNewProcess(payer, period)).resolves.toBe(9)
415+
416+
expect(await exactSum(payer, period)).toBe(9)
417+
expect(windowStartSums(period)).toHaveLength(2)
418+
})
419+
393420
it('leaves the prefix in place for a top-up to a row still in the tail, and reads it exactly', async () => {
394421
const period = openReportingWindow()
395422
await insertUsage(payer, 2, new Date(period.start.getTime() + DAY_MS))

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

Lines changed: 20 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { db } from '@sim/db'
22
import { createLogger } from '@sim/logger'
33
import { getErrorMessage } from '@sim/utils/errors'
4+
import { generateShortId } from '@sim/utils/id'
45
import { randomInt } from '@sim/utils/random'
56
import type Redis from 'ioredis'
67
import { LRUCache } from 'lru-cache'
@@ -18,6 +19,7 @@ import {
1819
getBillingPeriodUsageCost,
1920
type UsageQueryPeriod,
2021
} from '@/lib/billing/core/usage-log'
22+
import { releaseLock } from '@/lib/core/config/redis'
2123
import { withinDeadline } from '@/lib/core/utils/deadline'
2224
import type { DbClient } from '@/lib/db/types'
2325

@@ -85,7 +87,9 @@ const PREFIX_REFRESH_AFTER_MS = 3 * 60_000
8587
* The refresh lease is held for as long as one ledger sum may run, so a second process starts
8688
* the same sum only once the first has certainly ended, and a holder that died frees it then.
8789
* Nothing tells a stranded lease from a live one, and nothing needs to: while it is held, every
88-
* other reader serves the stored prefix or, with none, sums the whole window.
90+
* other reader serves the stored prefix or, with none, sums the whole window. Its holder
91+
* releases it once the new prefix is stored, so a prefix discarded by an epoch bump right after
92+
* is refreshed by the next read rather than leaving every reader summing the whole window.
8993
*/
9094
const PREFIX_LEASE_MS = USAGE_LEDGER_STATEMENT_TIMEOUT_MS
9195

@@ -259,33 +263,35 @@ async function readStoredPrefix(
259263
}
260264

261265
/**
262-
* Whether this process took the refresh lease. It is never released: the stored prefix ends the
263-
* need for it, and expiry frees it for the next refresh. A lease whose `SET` failed or timed out
264-
* may still have landed, held by nobody, for {@link PREFIX_LEASE_MS}. Not `acquireLock`, which
266+
* This process's token for the refresh lease, or `null` when it did not take it. A refresh that
267+
* fails keeps the lease until expiry, as a backoff. A lease whose `SET` failed or timed out may
268+
* still have landed, held by nobody, for {@link PREFIX_LEASE_MS}. Not `acquireLock`, which
265269
* issues on a client that is not ready and has no read deadline.
266270
*/
267-
async function acquirePrefixLease(redis: Redis, key: string): Promise<boolean> {
271+
async function acquirePrefixLease(redis: Redis, key: string): Promise<string | null> {
272+
const token = generateShortId()
268273
try {
269274
const result = await withinDeadline(
270-
() => redis.set(sharedPrefixLeaseKey(key), '1', 'PX', PREFIX_LEASE_MS, 'NX'),
275+
() => redis.set(sharedPrefixLeaseKey(key), token, 'PX', PREFIX_LEASE_MS, 'NX'),
271276
Date.now() + SHARED_READ_TIMEOUT_MS
272277
)
273-
return result === 'OK'
278+
return result === 'OK' ? token : null
274279
} catch (error) {
275280
logger.warn('Shared reporting usage prefix lease failed', {
276281
error: getErrorMessage(error),
277282
})
278-
return false
283+
return null
279284
}
280285
}
281286

282287
/**
283288
* Fire-and-forget, and only on a ready client, like {@link writeSharedReportingUsageCost}. A
284289
* plain `SET` replaces the stored prefix: every prefix is a lower bound for its own cutoff, so
285290
* two refreshes racing can leave either one, and the tail always resumes from the cutoff stored
286-
* with it.
291+
* with it. Once it lands, the lease is released only while it still holds `lease`, so one that
292+
* expired and was taken by another process is left to that process.
287293
*/
288-
function writeStoredPrefix(key: string, prefix: LedgerPrefix): void {
294+
function writeStoredPrefix(key: string, prefix: LedgerPrefix, lease: string): void {
289295
try {
290296
const redis = readyRedisClient()
291297
if (!redis) return
@@ -296,6 +302,7 @@ function writeStoredPrefix(key: string, prefix: LedgerPrefix): void {
296302
'PX',
297303
prefix.cutoff + PREFIX_CUTOFF_LAG_MS + PREFIX_RETAIN_MS - Date.now()
298304
)
305+
.then(() => releaseLock(sharedPrefixLeaseKey(key), lease))
299306
.catch(warnSharedWriteFailed)
300307
} catch (error) {
301308
warnSharedWriteFailed(error)
@@ -335,10 +342,11 @@ async function resolveLedgerPrefix(
335342
const { epochs, prefix: stored } = state
336343
if (stored && cutoff - stored.cutoff < PREFIX_REFRESH_AFTER_MS) return stored
337344

338-
if (await acquirePrefixLease(redis, key)) {
345+
const lease = await acquirePrefixLease(redis, key)
346+
if (lease) {
339347
const cost = getBillingPeriodUsageCost(entity, { ...period, end: new Date(cutoff) }).then(
340348
(prefixCost) => {
341-
writeStoredPrefix(key, { ...epochs, cutoff, cost: prefixCost })
349+
writeStoredPrefix(key, { ...epochs, cutoff, cost: prefixCost }, lease)
342350
return prefixCost
343351
}
344352
)

0 commit comments

Comments
 (0)