diff --git a/src/backend/stores/metering/MeteringBufferStore.test.ts b/src/backend/stores/metering/MeteringBufferStore.test.ts index 5095b88de..ca16048cf 100644 --- a/src/backend/stores/metering/MeteringBufferStore.test.ts +++ b/src/backend/stores/metering/MeteringBufferStore.test.ts @@ -524,6 +524,87 @@ describe('MeteringBufferStore', () => { expect((res as { total: number }).total).toBe(10); }); + it('misses nothing in an exact read when a flush lands between its reads', async () => { + await target.incr({ key, pathAndAmountMap: { total: 10 } }); + const write = holdCalls(kv, 'incr'); + const flushing = target.flushCycle(); + await vi.waitFor(() => expect(write.spy).toHaveBeenCalled()); + + // The store read sees the counter from before the flush's write; + // the flush then writes it and leaves its claim before any claim + // read that follows the store read. + let storeReadStarted = false; + let flushed!: Promise; + const realGet = kv.get.bind(kv); + const get = vi + .spyOn(kv, 'get') + .mockImplementation(async (args) => { + if (!args.consistentRead || args.key !== key) { + return realGet(args); + } + if (storeReadStarted) return realGet(args); + storeReadStarted = true; + const before = await realGet(args); + write.release(); + flushed = flushing.then(() => {}); + await flushed; + return before; + }); + const redis = server.clients.redis as unknown as { + meterRead: (...args: string[]) => Promise; + }; + const realRead = redis.meterRead.bind(redis); + const read = vi + .spyOn(redis, 'meterRead') + .mockImplementation(async (...args: string[]) => { + if (storeReadStarted) { + await vi.waitFor(() => expect(flushed).toBeDefined()); + await flushed; + } + return realRead(...args); + }); + try { + const { res } = await target.readExact({ key }); + expect((res as { total: number }).total).toBe(10); + } finally { + get.mockRestore(); + read.mockRestore(); + write.spy.mockRestore(); + await flushing; + } + expect(await storedTotal(key)).toBe(10); + }); + + it('counts nothing twice in an exact read when a flush lands during its store read', async () => { + await target.incr({ key, pathAndAmountMap: { total: 10 } }); + const write = holdCalls(kv, 'incr'); + const flushing = target.flushCycle(); + await vi.waitFor(() => expect(write.spy).toHaveBeenCalled()); + + // The claim was read with the amount still on it; the store read + // then sees it written. + let landed = false; + const realGet = kv.get.bind(kv); + const get = vi + .spyOn(kv, 'get') + .mockImplementation(async (args) => { + if (args.consistentRead && args.key === key && !landed) { + landed = true; + write.release(); + await flushing; + } + return realGet(args); + }); + try { + const { res } = await target.readExact({ key }); + expect((res as { total: number }).total).toBe(10); + } finally { + get.mockRestore(); + write.spy.mockRestore(); + await flushing; + } + }); + it('reads a written-through increment once the cache is back', async () => { await target.incr({ key, pathAndAmountMap: { total: 10 } }); await target.flushCycle(); diff --git a/src/backend/stores/metering/MeteringBufferStore.ts b/src/backend/stores/metering/MeteringBufferStore.ts index 04a3c790b..c79a27f0f 100644 --- a/src/backend/stores/metering/MeteringBufferStore.ts +++ b/src/backend/stores/metering/MeteringBufferStore.ts @@ -152,6 +152,12 @@ const UNWRITABLE_PATH_MEMO_LIMIT = 1000; */ const EXACT_READ_MEMO_LIMIT = 10_000; +/** + * Store reads an exact read makes while a claim keeps changing under it. The + * last one stands even so; it can only count an amount twice, never miss one. + */ +const EXACT_READ_ATTEMPTS = 3; + /** Write units past which settling a counter raises the large-item alarm. */ const LARGE_SETTLE_WRITE_UNITS = 100; @@ -339,6 +345,15 @@ const toScriptArgs = (amounts: Record): string[] => { return args; }; +/** Whether two sets of amounts hold the same paths and values. */ +const sameAmounts = (a: FlatAmounts, b: FlatAmounts): boolean => { + const paths = Object.keys(a); + return ( + paths.length === Object.keys(b).length && + paths.every((path) => a[path] === b[path]) + ); +}; + /** The subset of `amounts` at `paths`. */ const pickAmounts = (amounts: FlatAmounts, paths: string[]): FlatAmounts => Object.fromEntries(paths.map((path) => [path, amounts[path]!])); @@ -935,14 +950,26 @@ export class MeteringBufferStore extends PuterStore { `[metering] exact read could not flush ${key}: ${(e as Error).message}`, ); } - 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) }; + // the store yet. The claim is read before the store: an amount gone + // from it by then was written before the store read, so the read has + // it. One that leaves during the store read is counted twice, so the + // read is retried when the claim changes under it. + for (let attempt = 1; ; attempt++) { + const claim = await this.#outstandingClaim(key); + const stored = await this.stores.kv.get({ + key, + consistentRead: true, + }); + if ( + attempt < EXACT_READ_ATTEMPTS && + !sameAmounts(claim, await this.#outstandingClaim(key)) + ) { + continue; + } + 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. */