diff --git a/src/backend/stores/metering/MeteringBufferStore.test.ts b/src/backend/stores/metering/MeteringBufferStore.test.ts index 071204de9..7cf564c34 100644 --- a/src/backend/stores/metering/MeteringBufferStore.test.ts +++ b/src/backend/stores/metering/MeteringBufferStore.test.ts @@ -420,6 +420,75 @@ describe('MeteringBufferStore', () => { expect(await storedTotal(key)).toBe(0); }); + /** Holds `method` on `target` until the returned release is called. */ + const holdCalls = (on: T, method: keyof T) => { + let release: () => void = () => {}; + const gate = new Promise((r) => (release = r)); + const real = (on[method] as (...a: unknown[]) => unknown).bind(on); + const spy = vi + .spyOn(on, method as never) + .mockImplementation((async (...args: unknown[]) => { + await gate; + return real(...args); + }) as never); + return { spy, release }; + }; + const readTotal = async () => + ((await target.get({ key })).res as { total?: number } | null) + ?.total ?? 0; + + it('counts usage a flush has claimed but not yet written onward', 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()); + expect(await readTotal()).toBe(10); + write.release(); + await flushing; + } finally { + write.spy.mockRestore(); + } + expect(await readTotal()).toBe(10); + expect(await storedTotal(key)).toBe(10); + }); + + it('counts usage written onward until its settle lands', async () => { + await target.incr({ key, pathAndAmountMap: { total: 10 } }); + await target.flushCycle(); + await target.incr({ key, pathAndAmountMap: { total: 5 } }); + const redis = server.clients.redis as unknown as { + meterSettle: (...args: string[]) => Promise; + }; + const settle = holdCalls(redis, 'meterSettle'); + try { + const flushing = target.flushCycle(); + await vi.waitFor(() => expect(settle.spy).toHaveBeenCalled()); + expect(await storedTotal(key)).toBe(15); + expect(await readTotal()).toBe(15); + settle.release(); + await flushing; + } finally { + settle.spy.mockRestore(); + } + expect(await readTotal()).toBe(15); + }); + + it('counts a claim in array reads too', async () => { + await target.incr({ key, pathAndAmountMap: { total: 4 } }); + const write = holdCalls(kv, 'incr'); + try { + const flushing = target.flushCycle(); + await vi.waitFor(() => expect(write.spy).toHaveBeenCalled()); + const { res } = await target.get({ key: [key] }); + expect((res as [{ total: number }])[0].total).toBe(4); + write.release(); + await flushing; + } finally { + write.spy.mockRestore(); + } + }); + it('reads an unknown counter as absent', async () => { const { res } = await target.get({ key }); expect(res).toBeNull(); @@ -906,12 +975,16 @@ describe('MeteringBufferStore', () => { '1000000000', ); + // Another deployment's write, which only a replacing settle would + // bring into the base. + await kv.incr({ key, pathAndAmountMap: { total: 1000 } }); await target.incr({ key, pathAndAmountMap: { total: 5 } }); await target.flushCycle(); - expect(await storedTotal(key)).toBe(105); + expect(await storedTotal(key)).toBe(1105); + // The newer base keeps its view, plus what this flush wrote. expect(Number(await server.clients.redis.hget(base, 'total'))).toBe( - 100, + 105, ); }); @@ -1742,6 +1815,17 @@ describe('MeteringBufferStore', () => { order.push('dropFromClaim'); return realHdel(...args); }); + // A written batch leaves the claim through this script. + const redis = server.clients.redis as unknown as { + meterApply: (...args: string[]) => Promise; + }; + const realApply = redis.meterApply.bind(redis); + const applySpy = vi + .spyOn(redis, 'meterApply') + .mockImplementation((...args: string[]) => { + order.push('dropFromClaim'); + return realApply(...args); + }); const alarmSpy = vi .spyOn(server.clients.alarm, 'create') .mockImplementation(() => { @@ -1751,6 +1835,7 @@ describe('MeteringBufferStore', () => { await target.flushCycle(); incrSpy.mockRestore(); hdelSpy.mockRestore(); + applySpy.mockRestore(); alarmSpy.mockRestore(); expect(order).toEqual(['dropFromClaim', 'alarm']); diff --git a/src/backend/stores/metering/MeteringBufferStore.ts b/src/backend/stores/metering/MeteringBufferStore.ts index b46fe0486..2bfecf5ea 100644 --- a/src/backend/stores/metering/MeteringBufferStore.ts +++ b/src/backend/stores/metering/MeteringBufferStore.ts @@ -403,9 +403,39 @@ redis.call('PEXPIRE', KEYS[4], ARGV[1]) return { redis.call('HGETALL', KEYS[1]), redis.call('HGETALL', KEYS[2]) } `; -/** KEYS: delta, base. Reads both without marking the counter for flushing. */ +/** + * KEYS: delta, base, claimed. ARGV: counter, pending key prefix. + * + * Reads both without marking the counter for flushing, plus the counter's + * outstanding claim: amounts a flush has taken but not yet written onward are + * in neither the delta nor the base. The claim's key is built here, which is + * safe for the reason given at `RECONCILE_SCRIPT`. + */ const READ_SCRIPT = ` -return { redis.call('HGETALL', KEYS[1]), redis.call('HGETALL', KEYS[2]) } +local claim = {} +local nonce = redis.call('HGET', KEYS[3], ARGV[1]) +if nonce then claim = redis.call('HGETALL', ARGV[2] .. nonce) end +return { redis.call('HGETALL', KEYS[1]), redis.call('HGETALL', KEYS[2]), claim } +`; + +/** + * KEYS: pending, base. ARGV: path/amount pairs just written onward. + * + * Takes written paths off a claim and, when a base is cached, adds them to it + * in the same step, so a read counts them all the way to the settle that + * replaces the base with the store's own view. Without a base, reads go to the + * store, which already has them. + */ +const APPLY_SCRIPT = ` +redis.call('HDEL', KEYS[1], '${INTENT_FIELD}') +local based = redis.call('EXISTS', KEYS[2]) == 1 +for i = 1, #ARGV, 2 do + redis.call('HDEL', KEYS[1], ARGV[i]) + if based then + redis.call('HINCRBYFLOAT', KEYS[2], ARGV[i], ARGV[i + 1]) + end +end +return 1 `; /** @@ -653,7 +683,8 @@ return redis.call('HGETALL', KEYS[1]) type ScriptRunner = { meterIncr(...args: string[]): Promise<[string[], string[]]>; - meterRead(...args: string[]): Promise<[string[], string[]]>; + meterRead(...args: string[]): Promise<[string[], string[], string[]]>; + meterApply(...args: string[]): Promise; meterDrain(...args: string[]): Promise<[string[], number]>; meterClaim(...args: string[]): Promise; /** `[1, amounts]` when re-claimed, `[0, hadIndexEntry]` when not. */ @@ -1119,9 +1150,13 @@ export class MeteringBufferStore extends PuterStore { lua: INCR_SCRIPT, }); client.defineCommand('meterRead', { - numberOfKeys: 2, + numberOfKeys: 3, lua: READ_SCRIPT, }); + client.defineCommand('meterApply', { + numberOfKeys: 2, + lua: APPLY_SCRIPT, + }); client.defineCommand('meterDrain', { numberOfKeys: 2, lua: DRAIN_SCRIPT, @@ -1216,11 +1251,18 @@ export class MeteringBufferStore extends PuterStore { key: string, ): Promise<{ delta: FlatAmounts; base: FlatAmounts }> { const tag = bucketTag(key); - const [delta, base] = await this.#redis.meterRead( + const [delta, base, claim] = await this.#redis.meterRead( deltaKey(tag, key), baseKey(tag, key), + claimedKey(tag), + key, + pendingKey(tag, ''), ); - return { delta: pairsToAmounts(delta), base: pairsToAmounts(base) }; + // Claimed amounts count like buffered ones until their settle. + return { + delta: addAmounts(pairsToAmounts(delta), parseClaim(claim).amounts), + base: pairsToAmounts(base), + }; } async #readThrough(key: string): Promise { @@ -1618,7 +1660,11 @@ export class MeteringBufferStore extends PuterStore { // before anything else can fail: a re-drive then picks up only what // is still outstanding. The alarm below is the "anything else" — it // comes after. - await this.#dropFromClaim(tag, nonce, paths); + await this.#redis.meterApply( + pendingKey(tag, nonce), + baseKey(tag, key), + ...toScriptArgs(batch), + ); if (usage.write >= LARGE_SETTLE_WRITE_UNITS) { this.#recordLargeSettle(key, usage.write);