From 6bb5eacf61e2dd4c3cf19d4930e279be17db83e8 Mon Sep 17 00:00:00 2001 From: Daniel Salazar Date: Mon, 10 Aug 2026 08:36:56 -0700 Subject: [PATCH] fix: metering kv spamc (#3532) --- .../metering/MeteringBufferStore.test.ts | 117 ++++++++++++++++++ .../stores/metering/MeteringBufferStore.ts | 80 +++++++++++- .../stores/systemKv/SystemKVStore.test.ts | 62 ++++++++++ src/backend/stores/systemKv/SystemKVStore.ts | 38 ++++-- 4 files changed, 283 insertions(+), 14 deletions(-) diff --git a/src/backend/stores/metering/MeteringBufferStore.test.ts b/src/backend/stores/metering/MeteringBufferStore.test.ts index 18369dbbd..b9dc0dc22 100644 --- a/src/backend/stores/metering/MeteringBufferStore.test.ts +++ b/src/backend/stores/metering/MeteringBufferStore.test.ts @@ -31,6 +31,7 @@ import { setupTestServer } from '../../testUtil.ts'; import type { SystemKVStore } from '../systemKv/SystemKVStore.ts'; import { bucketTag, + chunkAmounts, flattenAmounts, pairsToAmounts, parsePendingEntry, @@ -81,6 +82,21 @@ describe('MeteringBufferStore', () => { flattenAmounts({ total: 1, label: 'nope', nested: null }), ).toEqual({ total: 1 }); }); + + it('leaves a counter that already fits in one piece', () => { + const amounts = { total: 5, 'ai:chat.units': 2 }; + expect(chunkAmounts(amounts, 24)).toEqual([amounts]); + }); + + it('splits a counter too wide for one write, losing nothing', () => { + const amounts: Record = {}; + for (let i = 0; i < 7; i++) amounts[`ai${i}.units`] = i; + + const chunks = chunkAmounts(amounts, 3); + + expect(chunks.map((c) => Object.keys(c).length)).toEqual([3, 3, 1]); + expect(Object.assign({}, ...chunks)).toEqual(amounts); + }); }); describe('pending index entries', () => { @@ -180,6 +196,29 @@ describe('MeteringBufferStore', () => { incrSpy.mockRestore(); }); + it('splits a counter too wide for one expression across writes', async () => { + // Buffering is what makes this reachable: a single call meters a + // handful of paths, but a cycle's worth of calls for a busy counter + // adds up past what one update expression can hold. + const paths: Record = {}; + for (let i = 0; i < 60; i++) paths[`ai${i}.units`] = 1; + await target.incr({ key, pathAndAmountMap: paths }); + + const incrSpy = vi.spyOn(kv, 'incr'); + await target.flushCycle(); + + expect(incrSpy).toHaveBeenCalledTimes(3); + for (const [input] of incrSpy.mock.calls) { + expect( + Object.keys(input.pathAndAmountMap).length, + ).toBeLessThanOrEqual(24); + } + incrSpy.mockRestore(); + + const stored = flattenAmounts((await kv.get({ key })).res); + expect(Object.keys(stored)).toHaveLength(60); + }); + it('counts what is already stored when it first sees a counter', async () => { await kv.incr({ key, pathAndAmountMap: { total: 80 } }); @@ -551,6 +590,84 @@ describe('MeteringBufferStore', () => { ).toEqual({}); }); + it('gives up on a claim the store will never accept', async () => { + await target.incr({ key, pathAndAmountMap: { total: 5 } }); + // A rejection, not an outage: the next attempt would be rejected + // identically. Re-driving it every cycle for as long as the counter + // exists is what turned one bad counter into a write loop. + const rejected = Object.assign( + new Error( + 'Invalid UpdateExpression: Expression size has exceeded the maximum allowed size', + ), + { name: 'ValidationException' }, + ); + const boom = vi.spyOn(kv, 'incr').mockRejectedValue(rejected); + const logged = vi + .spyOn(console, 'error') + .mockImplementation(() => {}); + + await target.flushCycle(); + boom.mockRestore(); + + const tag = bucketTag(key); + expect( + await server.clients.redis.hgetall(`meter:pending:{${tag}}`), + ).toEqual({}); + expect(await server.clients.redis.keys('meter:p:*')).toEqual([]); + // The amounts are lost, so this must not be quiet. + expect(logged).toHaveBeenCalledWith( + expect.stringContaining('unwritable path'), + ); + logged.mockRestore(); + + // And a second cycle finds nothing left to re-drive. + const after = vi.spyOn(kv, 'incr'); + await target.flushCycle(); + expect(after).not.toHaveBeenCalled(); + after.mockRestore(); + }); + + it('does not write an applied chunk twice when a later one fails', async () => { + const paths: Record = {}; + for (let i = 0; i < 30; i++) paths[`ai${i}.units`] = 1; + await target.incr({ key, pathAndAmountMap: paths }); + + // Two chunks: let the first through and fail the second, the way a + // throttle landing mid-settle would. + const passThrough = kv.incr.bind(kv); + let call = 0; + const flaky = vi + .spyOn(kv, 'incr') + .mockImplementation((...args: Parameters) => { + if (++call === 2) + return Promise.reject(new Error('throttled')); + return passThrough(...args); + }); + + await target.flushCycle(); + flaky.mockRestore(); + + // Age the surviving claim so the sweep takes it, then let it finish. + const tag = bucketTag(key); + const pending = await server.clients.redis.hgetall( + `meter:pending:{${tag}}`, + ); + const nonce = Object.keys(pending)[0]!; + await server.clients.redis.hset( + `meter:pending:{${tag}}`, + nonce, + `${Date.now() - 60_000}:${key}`, + ); + + await target.flushCycle(); + + // Every path lands exactly once: the applied chunk came off the + // claim as it was written, so the re-drive carried only the rest. + const stored = flattenAmounts((await kv.get({ key })).res); + expect(Object.keys(stored)).toHaveLength(30); + expect([...new Set(Object.values(stored))]).toEqual([1]); + }); + it('leaves the claim in place when the write onward fails', async () => { await target.incr({ key, pathAndAmountMap: { total: 5 } }); const boom = vi diff --git a/src/backend/stores/metering/MeteringBufferStore.ts b/src/backend/stores/metering/MeteringBufferStore.ts index 32fdfcf63..30597adcb 100644 --- a/src/backend/stores/metering/MeteringBufferStore.ts +++ b/src/backend/stores/metering/MeteringBufferStore.ts @@ -71,6 +71,8 @@ const BUFFER_TTL_MS = 40 * 24 * 60 * 60 * 1000; /** How many counters are written onward at once, to keep the load even. */ const SETTLE_CONCURRENCY = 20; +const SETTLE_PATHS_PER_WRITE = 24; + /** Roughly a minute between compression reports, so they don't flood the log. */ const CYCLES_PER_COMPRESSION_REPORT = 12; @@ -152,6 +154,27 @@ const toScriptArgs = (amounts: Record): string[] => { return args; }; +/** Split counters into groups of at most `size` paths, preserving order. */ +export const chunkAmounts = ( + amounts: FlatAmounts, + size: number, +): FlatAmounts[] => { + const entries = Object.entries(amounts); + if (entries.length <= size) return [amounts]; + + const chunks: FlatAmounts[] = []; + for (let i = 0; i < entries.length; i += size) { + chunks.push(Object.fromEntries(entries.slice(i, i + size))); + } + return chunks; +}; + +const isPermanentSettleError = (err: Error): boolean => { + if (err.name === 'ValidationException') return true; + const status = (err as { statusCode?: unknown }).statusCode; + return typeof status === 'number' && status >= 400 && status < 500; +}; + // -- Scripts ---------------------------------------------------------- /** @@ -661,12 +684,49 @@ export class MeteringBufferStore extends PuterStore { nonce: string, amounts: FlatAmounts, ): Promise { - const { res } = await this.stores.kv.incr({ - key, - pathAndAmountMap: amounts, - }); + const total = Object.keys(amounts).length; + if (total === 0) { + // Nothing to write onward. Retire the claim rather than settling + // it, which would clear a base that is still good. + await this.#retireClaim(tag, nonce); + return; + } + const chunks = chunkAmounts(amounts, SETTLE_PATHS_PER_WRITE); - const flat = flattenAmounts(res); + // Each write returns the whole counter, so after the last chunk this + // holds every path the KV store now has — including the earlier chunks + // and anything another deployment contributed. + let settled: unknown; + let written = 0; + for (const chunk of chunks) { + try { + ({ res: settled } = await this.stores.kv.incr({ + key, + pathAndAmountMap: chunk, + })); + } catch (e) { + const err = e as Error; + if (!isPermanentSettleError(err)) throw e; + console.error( + `[metering] dropping ${total - written} unwritable path(s) of ${key}: ${err.message}`, + ); + await this.#retireClaim(tag, nonce); + return; + } + written += Object.keys(chunk).length; + + // This chunk is applied for good now, so take it off the claim: if a + // later one fails, the re-drive picks up only what is still + // outstanding instead of adding these amounts a second time. + if (chunks.length > 1) { + await this.clients.redis.hdel( + pendingKey(tag, nonce), + ...Object.keys(chunk), + ); + } + } + + const flat = flattenAmounts(settled); await this.#redis.meterSettle( baseKey(tag, key), pendingKey(tag, nonce), @@ -677,6 +737,16 @@ export class MeteringBufferStore extends PuterStore { ...toScriptArgs(flat), ); } + + /** + * Forget a claim and its index entry, leaving the base alone. Settling + * would also replace the base, which is only correct when something was + * actually written onward. + */ + async #retireClaim(tag: string, nonce: string): Promise { + await this.clients.redis.del(pendingKey(tag, nonce)); + await this.clients.redis.hdel(pendingIndexKey(tag), nonce); + } } // -- Pending index encoding ------------------------------------------- diff --git a/src/backend/stores/systemKv/SystemKVStore.test.ts b/src/backend/stores/systemKv/SystemKVStore.test.ts index a1eca3971..1c8623e7a 100644 --- a/src/backend/stores/systemKv/SystemKVStore.test.ts +++ b/src/backend/stores/systemKv/SystemKVStore.test.ts @@ -550,6 +550,68 @@ describe('SystemKVStore', () => { expect(after.res).toMatchObject({ a: { b: { c: 5 } } }); }); + it('does not try to build paths for an expression rejected on its size', async () => { + // Both failures arrive as a ValidationException, but this one is + // about the expression rather than the item: createPaths would + // write a layer per nested path — each against this same item, so + // each costing the whole item — and then re-send a byte-identical + // expression to be rejected again. In production that ran as a + // sweep every 5s and cost real money, so the guard is worth + // pinning: one attempt, then out. + const oversized = Object.assign( + new Error( + '1 validation error detected: Invalid UpdateExpression: Expression size has exceeded the maximum allowed size;', + ), + { name: 'ValidationException' }, + ); + const update = vi + .spyOn(server.clients.dynamo, 'update') + .mockRejectedValue(oversized); + + await expect( + target.incr( + { key: 'oversized', pathAndAmountMap: { 'a.b': 1 } }, + opts, + ), + ).rejects.toThrow(/Expression size/); + expect(update).toHaveBeenCalledTimes(1); + + update.mockRestore(); + }); + + it('still builds paths for a ValidationException about the item', async () => { + // The other side of the guard above: a genuinely missing nested + // parent must still be created and the update retried. + const missingPath = Object.assign( + new Error( + 'The document path provided in the update expression is invalid for update', + ), + { name: 'ValidationException' }, + ); + const real = server.clients.dynamo.update.bind( + server.clients.dynamo, + ); + let first = true; + const update = vi + .spyOn(server.clients.dynamo, 'update') + .mockImplementation((...args) => { + if (first) { + first = false; + return Promise.reject(missingPath); + } + return real(...args); + }); + + const result = await target.incr( + { key: 'guardedNest', pathAndAmountMap: { 'x.y.z': 4 } }, + opts, + ); + + expect(result.res).toMatchObject({ x: { y: { z: 4 } } }); + expect(update.mock.calls.length).toBeGreaterThan(1); + update.mockRestore(); + }); + it('decr subtracts via the same machinery', async () => { await target.incr( { key: 'counter3', pathAndAmountMap: { hits: 10 } }, diff --git a/src/backend/stores/systemKv/SystemKVStore.ts b/src/backend/stores/systemKv/SystemKVStore.ts index f3f986906..580add67f 100644 --- a/src/backend/stores/systemKv/SystemKVStore.ts +++ b/src/backend/stores/systemKv/SystemKVStore.ts @@ -201,6 +201,9 @@ const assertPaths = (paths: string[]): void => { for (const valPath of paths) assertPath(valPath); }; +const isOversizedExpression = (err: Error): boolean => + /expression size/i.test(err.message); + /** * Walk a value about to be stored and reject unsafe keys. Iterative so a deeply * nested value can't blow the stack; own keys only, matching what the document @@ -451,7 +454,8 @@ export class SystemKVStore extends PuterStore { probeUsage, writeUsage( response.ConsumedCapacity?.CapacityUnits as - number | undefined, + | number + | undefined, ), ), }; @@ -543,7 +547,8 @@ export class SystemKVStore extends PuterStore { probeUsage, writeUsage( (response.ConsumedCapacity?.CapacityUnits as - number | undefined) ?? 1, + | number + | undefined) ?? 1, ), ), }; @@ -575,7 +580,9 @@ export class SystemKVStore extends PuterStore { | { key: string; value: unknown }[] | { items: - string[] | unknown[] | { key: string; value: unknown }[]; + | string[] + | unknown[] + | { key: string; value: unknown }[]; cursor?: string; total?: number; } @@ -654,7 +661,8 @@ export class SystemKVStore extends PuterStore { usage, readUsage( (response.ConsumedCapacity?.CapacityUnits as - number | undefined) ?? 1, + | number + | undefined) ?? 1, ), ); return response; @@ -671,7 +679,8 @@ export class SystemKVStore extends PuterStore { const skip = await runQuery(remaining, startKey, 'COUNT'); remaining -= Number(skip.Count ?? 0); startKey = skip.LastEvaluatedKey as - Record | undefined; + | Record + | undefined; if (!startKey) { exhausted = remaining > 0; break; @@ -694,7 +703,8 @@ export class SystemKVStore extends PuterStore { >), ); nextKey = response.LastEvaluatedKey as - Record | undefined; + | Record + | undefined; pages++; if (normalizedLimit === undefined) { // Legacy full listing: follow continuation pages so the @@ -730,7 +740,8 @@ export class SystemKVStore extends PuterStore { const counted = await runQuery(0, countKey, 'COUNT'); total += Number(counted.Count ?? 0); countKey = counted.LastEvaluatedKey as - Record | undefined; + | Record + | undefined; } while (countKey); } @@ -904,7 +915,15 @@ export class SystemKVStore extends PuterStore { try { response = await runUpdate(); } catch (e) { - if ((e as Error)?.name !== 'ValidationException') throw e; + const err = e as Error; + if (err?.name !== 'ValidationException') throw e; + // An expression rejected for its own size is the one + // ValidationException createPaths cannot repair: it writes a layer + // per nested path — against this same item, so each of those writes + // costs the whole item — and then re-sends a byte-identical + // expression to be rejected again. Fail fast and let the caller + // send fewer paths at a time. + if (isOversizedExpression(err)) throw e; createPathsUsage = await this.createPaths( namespace, key, @@ -1068,7 +1087,8 @@ export class SystemKVStore extends PuterStore { probeUsage, writeUsage( (response.ConsumedCapacity?.CapacityUnits as - number | undefined) ?? 1, + | number + | undefined) ?? 1, ), ), };