From a8f7c94f2e8df1838102c2bdddaeb6329a9582fc Mon Sep 17 00:00:00 2001 From: 404oops Date: Sat, 31 Oct 2026 16:56:18 +0100 Subject: [PATCH] fix(metering): count claims on the write path and read missing bases exactly The running total `incr` returns and `readExact` also missed usage a flush had claimed but not yet written. Both feed how a charge is split between the allowance and purchased credit, so an account could be charged to its allowance for usage that had already used it up. Both now add the counter's outstanding claim. A read or seed with no cached base asked the store with an ordinary read, which could miss a write a settle had just made; it now reads consistently. A write-through after a cache failure drops the cached base, which doesn't have the write, so the next read reseeds it. The dirty-set reconcile test no longer depends on how many flush cycles earlier tests ran. Co-Authored-By: Claude Opus 5.5 --- .../metering/MeteringBufferStore.test.ts | 73 +++++++++++++++++++ .../stores/metering/MeteringBufferStore.ts | 72 ++++++++++++++---- 2 files changed, 132 insertions(+), 13 deletions(-) diff --git a/src/backend/stores/metering/MeteringBufferStore.test.ts b/src/backend/stores/metering/MeteringBufferStore.test.ts index 7cf564c34..5095b88de 100644 --- a/src/backend/stores/metering/MeteringBufferStore.test.ts +++ b/src/backend/stores/metering/MeteringBufferStore.test.ts @@ -489,6 +489,73 @@ describe('MeteringBufferStore', () => { } }); + it('counts a claim in the running total incr returns', async () => { + await target.incr({ key, pathAndAmountMap: { total: 10 } }); + const write = holdCalls(kv, 'incr'); + try { + const flushing = target.flushCycle(); + await vi.waitFor(() => expect(write.spy).toHaveBeenCalled()); + const { res } = await target.incr({ + key, + pathAndAmountMap: { total: 5 }, + }); + expect((res as { total: number }).total).toBe(15); + write.release(); + await flushing; + } finally { + write.spy.mockRestore(); + } + }); + + it('counts another flush’s claim in an exact read', async () => { + await target.incr({ key, pathAndAmountMap: { total: 10 } }); + const write = holdCalls(kv, 'incr'); + try { + const flushing = target.flushCycle(); + await vi.waitFor(() => expect(write.spy).toHaveBeenCalled()); + const { res } = await target.readExact({ key }); + expect((res as { total: number }).total).toBe(10); + write.release(); + await flushing; + } finally { + write.spy.mockRestore(); + } + const { res } = await target.readExact({ key }); + expect((res as { total: number }).total).toBe(10); + }); + + it('reads a written-through increment once the cache is back', async () => { + await target.incr({ key, pathAndAmountMap: { total: 10 } }); + await target.flushCycle(); + const redis = server.clients.redis as unknown as { + meterIncr: (...args: string[]) => Promise; + }; + const failing = vi + .spyOn(redis, 'meterIncr') + .mockRejectedValueOnce(new Error('cache down')); + const { res, exact } = await target.incr({ + key, + pathAndAmountMap: { total: 5 }, + }); + failing.mockRestore(); + expect(exact).toBe(true); + expect((res as { total: number }).total).toBe(15); + expect(await readTotal()).toBe(15); + }); + + it('reads the store consistently when no base is cached', async () => { + await kv.incr({ key, pathAndAmountMap: { total: 3 } }); + const get = vi.spyOn(kv, 'get'); + try { + expect(await readTotal()).toBe(3); + expect(get).toHaveBeenCalledWith( + expect.objectContaining({ key, consistentRead: true }), + ); + } finally { + get.mockRestore(); + } + }); + it('reads an unknown counter as absent', async () => { const { res } = await target.get({ key }); expect(res).toBeNull(); @@ -1868,7 +1935,13 @@ describe('MeteringBufferStore', () => { // right up until the TTL takes the amounts with it. const tag = bucketTag(key); await server.clients.redis.del(`meter:dirty:{${tag}}`); + // Every few cycles a flush reconciles on its own; keep this one + // from doing so, whatever the earlier tests left the count at. + const scheduled = vi + .spyOn(target, 'reconcile') + .mockResolvedValue(undefined); expect(await target.flushCycle()).toBe(0); + scheduled.mockRestore(); expect(await storedTotal(key)).toBe(0); const logged = vi diff --git a/src/backend/stores/metering/MeteringBufferStore.ts b/src/backend/stores/metering/MeteringBufferStore.ts index 2bfecf5ea..04a3c790b 100644 --- a/src/backend/stores/metering/MeteringBufferStore.ts +++ b/src/backend/stores/metering/MeteringBufferStore.ts @@ -380,11 +380,12 @@ export const isBilledCounter = (key: string): boolean => // -- Scripts ---------------------------------------------------------- /** - * KEYS: delta, base, dirty set, tracked set. ARGV: ttl, member, then - * path/amount pairs. + * KEYS: delta, base, dirty set, tracked set, claimed. ARGV: ttl, member, + * pending key prefix, then path/amount pairs. * - * Returns the delta and the base together so a running total can be served - * without a second round trip. Amounts are floats — costs are fractional. + * Returns the delta, the base and any outstanding claim together, so a running + * total can be served without a second round trip (see `READ_SCRIPT` for the + * claim). Amounts are floats — costs are fractional. * * The counter joins both sets here. The dirty set is the working list a flush * takes from and hands back; the tracked set is only ever added to here and @@ -392,7 +393,7 @@ export const isBilledCounter = (key: string): boolean => * working list does not. */ const INCR_SCRIPT = ` -for i = 3, #ARGV, 2 do +for i = 4, #ARGV, 2 do redis.call('HINCRBYFLOAT', KEYS[1], ARGV[i], ARGV[i + 1]) end redis.call('PEXPIRE', KEYS[1], ARGV[1]) @@ -400,7 +401,10 @@ redis.call('SADD', KEYS[3], ARGV[2]) redis.call('PEXPIRE', KEYS[3], ARGV[1]) redis.call('SADD', KEYS[4], ARGV[2]) redis.call('PEXPIRE', KEYS[4], ARGV[1]) -return { redis.call('HGETALL', KEYS[1]), redis.call('HGETALL', KEYS[2]) } +local claim = {} +local nonce = redis.call('HGET', KEYS[5], ARGV[2]) +if nonce then claim = redis.call('HGETALL', ARGV[3] .. nonce) end +return { redis.call('HGETALL', KEYS[1]), redis.call('HGETALL', KEYS[2]), claim } `; /** @@ -682,7 +686,7 @@ return redis.call('HGETALL', KEYS[1]) `; type ScriptRunner = { - meterIncr(...args: string[]): Promise<[string[], string[]]>; + meterIncr(...args: string[]): Promise<[string[], string[], string[]]>; meterRead(...args: string[]): Promise<[string[], string[], string[]]>; meterApply(...args: string[]): Promise; meterDrain(...args: string[]): Promise<[string[], number]>; @@ -801,6 +805,8 @@ export class MeteringBufferStore extends PuterStore { `[metering] buffered incr failed, writing through: ${(e as Error).message}`, ); const { res } = await this.#writeThrough(input); + // The cached base doesn't have this write; the next read reseeds it. + await this.forgetBase(input.key); return { res, exact: true }; } } @@ -880,6 +886,7 @@ export class MeteringBufferStore extends PuterStore { // is incomplete rather than silently read a buffered-only view. const { res: fetched } = await this.stores.kv.get({ key: misses.map((i) => keys[i]!), + consistentRead: true, }); const fetchedValues = fetched as (unknown | null)[]; misses.forEach((i, idx) => { @@ -928,7 +935,34 @@ export class MeteringBufferStore extends PuterStore { `[metering] exact read could not flush ${key}: ${(e as Error).message}`, ); } - return this.stores.kv.get({ key, consistentRead: true }); + const [stored, claim] = await Promise.all([ + this.stores.kv.get({ key, consistentRead: true }), + this.#outstandingClaim(key), + ]); + // Another deployment's claim (or ours, if the flush failed) isn't in + // the store yet. + if (Object.keys(claim).length === 0) return stored; + return { res: mergedOrNull(flattenAmounts(stored.res), claim) }; + } + + /** Amounts a flush has claimed for `key` and not yet written onward. */ + async #outstandingClaim(key: string): Promise { + try { + const tag = bucketTag(key); + const [, , claim] = await this.#redis.meterRead( + deltaKey(tag, key), + baseKey(tag, key), + claimedKey(tag), + key, + pendingKey(tag, ''), + ); + return parseClaim(claim).amounts; + } catch (e) { + console.warn( + `[metering] claim unreadable for ${key}: ${(e as Error).message}`, + ); + return {}; + } } #rememberExactRead(key: string, now: number): void { @@ -1146,7 +1180,7 @@ export class MeteringBufferStore extends PuterStore { this.#definedScripts = true; const client = this.clients.redis; client.defineCommand('meterIncr', { - numberOfKeys: 4, + numberOfKeys: 5, lua: INCR_SCRIPT, }); client.defineCommand('meterRead', { @@ -1207,18 +1241,23 @@ export class MeteringBufferStore extends PuterStore { pathAndAmountMap, }: IncrInput): Promise<{ delta: FlatAmounts; base: FlatAmounts }> { const tag = bucketTag(key); - const [delta, base] = await this.#redis.meterIncr( + const [delta, base, claim] = await this.#redis.meterIncr( deltaKey(tag, key), baseKey(tag, key), dirtyKey(tag), trackedKey(tag), + claimedKey(tag), String(BUFFER_TTL_MS), key, + pendingKey(tag, ''), ...toScriptArgs(pathAndAmountMap), ); this.#absorbedCount++; absorbedCounter.add(1); - return { delta: pairsToAmounts(delta), base: pairsToAmounts(base) }; + return { + delta: addAmounts(pairsToAmounts(delta), parseClaim(claim).amounts), + base: pairsToAmounts(base), + }; } /** @@ -1231,7 +1270,9 @@ export class MeteringBufferStore extends PuterStore { * is over within seconds. */ async #seedBase(key: string): Promise { - const { res } = await this.stores.kv.get({ key }); + // Consistent, or a write that just landed could be missing from the + // base until the next settle replaces it. + const { res } = await this.stores.kv.get({ key, consistentRead: true }); const persisted = flattenAmounts(res); const seeded = await this.#redis.meterSeed( baseKey(bucketTag(key), key), @@ -1267,10 +1308,15 @@ export class MeteringBufferStore extends PuterStore { async #readThrough(key: string): Promise { const { delta, base } = await this.#readBuffered(key); + // Without a cached base, a write a settle just made is in the store + // and nowhere else, so the read has to be consistent. const resolved = Object.keys(base).length > 0 ? base - : flattenAmounts((await this.stores.kv.get({ key })).res); + : flattenAmounts( + (await this.stores.kv.get({ key, consistentRead: true })) + .res, + ); return mergedOrNull(resolved, delta); }