11import { db } from '@sim/db'
22import { createLogger } from '@sim/logger'
33import { getErrorMessage } from '@sim/utils/errors'
4+ import { generateShortId } from '@sim/utils/id'
45import { randomInt } from '@sim/utils/random'
56import type Redis from 'ioredis'
67import { 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'
2123import { withinDeadline } from '@/lib/core/utils/deadline'
2224import 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 */
9094const 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