fix: metering kv spamc (#3532)

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<string, number> = {};
+            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<string, number> = {};
+            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<string, number> = {};
+            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<typeof kv.incr>) => {
+                    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

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, number>): 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<void> {
-        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<void> {
+        await this.clients.redis.del(pendingKey(tag, nonce));
+        await this.clients.redis.hdel(pendingIndexKey(tag), nonce);
+    }
 }
 
 // -- Pending index encoding -------------------------------------------

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 } },

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<string, unknown> | undefined;
+                    | Record<string, unknown>
+                    | undefined;
                 if (!startKey) {
                     exhausted = remaining > 0;
                     break;
@@ -694,7 +703,8 @@ export class SystemKVStore extends PuterStore {
                     >),
                 );
                 nextKey = response.LastEvaluatedKey as
-                    Record<string, unknown> | undefined;
+                    | Record<string, unknown>
+                    | 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<string, unknown> | undefined;
+                    | Record<string, unknown>
+                    | 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,
                     ),
                 ),
             };