From efc48263fcf76efeeea1ec849e9ab30d3c979b7f Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 15 Oct 2026 21:43:31 -0700 Subject: [PATCH] fix(billing): release the reporting prefix lease once its prefix is stored --- .../core/reporting-usage-cache.integration.ts | 27 ++++++++++++++++ .../lib/billing/core/reporting-usage-cache.ts | 32 ++++++++++++------- 2 files changed, 47 insertions(+), 12 deletions(-) diff --git a/apps/sim/lib/billing/core/reporting-usage-cache.integration.ts b/apps/sim/lib/billing/core/reporting-usage-cache.integration.ts index e83bed744cf..062c9de0f98 100644 --- a/apps/sim/lib/billing/core/reporting-usage-cache.integration.ts +++ b/apps/sim/lib/billing/core/reporting-usage-cache.integration.ts @@ -186,6 +186,10 @@ describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => { transaction.mockClear() summedPeriods.length = 0 redisConfigMockFns.mockGetRedisClient.mockReturnValue(redis) + redisConfigMockFns.mockReleaseLock.mockImplementation(async (key: string, value: string) => { + if ((await redis.get(key)) !== value) return false + return (await redis.del(key)) === 1 + }) payer = { type: 'organization', id: generateId() } await connection`INSERT INTO usage_log (id, cost, billing_entity_type, billing_entity_id, created_at) VALUES (${generateId()}, 4.25, 'organization', ${payer.id}, '2026-03-01'), @@ -390,6 +394,29 @@ describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => { expect(await exactSum(payer, period)).toBe(9) }) + it('refreshes a prefix discarded by a top-up once, instead of summing the window on every read', async () => { + const period = openReportingWindow() + await insertUsage(payer, 3, new Date(Date.now() - MINUTE_MS)) + const eventKey = await startCumulativeCharge( + payer, + 2, + new Date(period.start.getTime() + DAY_MS) + ) + await expect(readInNewProcess(payer, period)).resolves.toBe(5) + await prefixStored(payer, period) + await vi.waitFor(async () => expect(await redis.exists(prefixLeaseKey(payer, period))).toBe(0)) + + await recordCumulativeUsage(cumulativeCharge(payer, eventKey, 6)) + await vi.waitFor(async () => expect(await redis.exists(payerEpochKey(payer))).toBe(1)) + await redis.del(sharedKey(payer, period)) + await expect(readInNewProcess(payer, period)).resolves.toBe(9) + await redis.del(sharedKey(payer, period)) + await expect(readInNewProcess(payer, period)).resolves.toBe(9) + + expect(await exactSum(payer, period)).toBe(9) + expect(windowStartSums(period)).toHaveLength(2) + }) + it('leaves the prefix in place for a top-up to a row still in the tail, and reads it exactly', async () => { const period = openReportingWindow() await insertUsage(payer, 2, new Date(period.start.getTime() + DAY_MS)) diff --git a/apps/sim/lib/billing/core/reporting-usage-cache.ts b/apps/sim/lib/billing/core/reporting-usage-cache.ts index ba0ecf7c30a..60dc97d23fe 100644 --- a/apps/sim/lib/billing/core/reporting-usage-cache.ts +++ b/apps/sim/lib/billing/core/reporting-usage-cache.ts @@ -1,6 +1,7 @@ import { db } from '@sim/db' import { createLogger } from '@sim/logger' import { getErrorMessage } from '@sim/utils/errors' +import { generateShortId } from '@sim/utils/id' import { randomInt } from '@sim/utils/random' import type Redis from 'ioredis' import { LRUCache } from 'lru-cache' @@ -18,6 +19,7 @@ import { getBillingPeriodUsageCost, type UsageQueryPeriod, } from '@/lib/billing/core/usage-log' +import { releaseLock } from '@/lib/core/config/redis' import { withinDeadline } from '@/lib/core/utils/deadline' import type { DbClient } from '@/lib/db/types' @@ -85,7 +87,9 @@ const PREFIX_REFRESH_AFTER_MS = 3 * 60_000 * The refresh lease is held for as long as one ledger sum may run, so a second process starts * the same sum only once the first has certainly ended, and a holder that died frees it then. * Nothing tells a stranded lease from a live one, and nothing needs to: while it is held, every - * other reader serves the stored prefix or, with none, sums the whole window. + * other reader serves the stored prefix or, with none, sums the whole window. Its holder + * releases it once the new prefix is stored, so a prefix discarded by an epoch bump right after + * is refreshed by the next read rather than leaving every reader summing the whole window. */ const PREFIX_LEASE_MS = USAGE_LEDGER_STATEMENT_TIMEOUT_MS @@ -259,23 +263,24 @@ async function readStoredPrefix( } /** - * Whether this process took the refresh lease. It is never released: the stored prefix ends the - * need for it, and expiry frees it for the next refresh. A lease whose `SET` failed or timed out - * may still have landed, held by nobody, for {@link PREFIX_LEASE_MS}. Not `acquireLock`, which + * This process's token for the refresh lease, or `null` when it did not take it. A refresh that + * fails keeps the lease until expiry, as a backoff. A lease whose `SET` failed or timed out may + * still have landed, held by nobody, for {@link PREFIX_LEASE_MS}. Not `acquireLock`, which * issues on a client that is not ready and has no read deadline. */ -async function acquirePrefixLease(redis: Redis, key: string): Promise { +async function acquirePrefixLease(redis: Redis, key: string): Promise { + const token = generateShortId() try { const result = await withinDeadline( - () => redis.set(sharedPrefixLeaseKey(key), '1', 'PX', PREFIX_LEASE_MS, 'NX'), + () => redis.set(sharedPrefixLeaseKey(key), token, 'PX', PREFIX_LEASE_MS, 'NX'), Date.now() + SHARED_READ_TIMEOUT_MS ) - return result === 'OK' + return result === 'OK' ? token : null } catch (error) { logger.warn('Shared reporting usage prefix lease failed', { error: getErrorMessage(error), }) - return false + return null } } @@ -283,9 +288,10 @@ async function acquirePrefixLease(redis: Redis, key: string): Promise { * Fire-and-forget, and only on a ready client, like {@link writeSharedReportingUsageCost}. A * plain `SET` replaces the stored prefix: every prefix is a lower bound for its own cutoff, so * two refreshes racing can leave either one, and the tail always resumes from the cutoff stored - * with it. + * with it. Once it lands, the lease is released only while it still holds `lease`, so one that + * expired and was taken by another process is left to that process. */ -function writeStoredPrefix(key: string, prefix: LedgerPrefix): void { +function writeStoredPrefix(key: string, prefix: LedgerPrefix, lease: string): void { try { const redis = readyRedisClient() if (!redis) return @@ -296,6 +302,7 @@ function writeStoredPrefix(key: string, prefix: LedgerPrefix): void { 'PX', prefix.cutoff + PREFIX_CUTOFF_LAG_MS + PREFIX_RETAIN_MS - Date.now() ) + .then(() => releaseLock(sharedPrefixLeaseKey(key), lease)) .catch(warnSharedWriteFailed) } catch (error) { warnSharedWriteFailed(error) @@ -335,10 +342,11 @@ async function resolveLedgerPrefix( const { epochs, prefix: stored } = state if (stored && cutoff - stored.cutoff < PREFIX_REFRESH_AFTER_MS) return stored - if (await acquirePrefixLease(redis, key)) { + const lease = await acquirePrefixLease(redis, key) + if (lease) { const cost = getBillingPeriodUsageCost(entity, { ...period, end: new Date(cutoff) }).then( (prefixCost) => { - writeStoredPrefix(key, { ...epochs, cutoff, cost: prefixCost }) + writeStoredPrefix(key, { ...epochs, cutoff, cost: prefixCost }, lease) return prefixCost } )