From 881ab6e41d04b7f191d301a27cb843310680979e Mon Sep 17 00:00:00 2001 From: Daniel Salazar Date: Wed, 24 Jun 2026 11:29:22 -0700 Subject: [PATCH] feat: improve kv store incr/decr costs --- .../stores/systemKv/SystemKVStore.test.ts | 44 +++++++++++++ src/backend/stores/systemKv/SystemKVStore.ts | 63 ++++++++++++++----- 2 files changed, 93 insertions(+), 14 deletions(-) diff --git a/src/backend/stores/systemKv/SystemKVStore.test.ts b/src/backend/stores/systemKv/SystemKVStore.test.ts index 5c6cb69c4c..bfc5d77660 100644 --- a/src/backend/stores/systemKv/SystemKVStore.test.ts +++ b/src/backend/stores/systemKv/SystemKVStore.test.ts @@ -366,6 +366,50 @@ describe('SystemKVStore', () => { expect(result.res).toMatchObject({ page: { views: 4 } }); }); + it('folds expireAt into the incr and stamps it once', async () => { + const past = Math.floor(Date.now() / 1000) - 10; + // First bump creates the counter and stamps the (already-elapsed) + // ttl in the same write — no separate expireAt call. + await target.incr( + { key: 'ttlCounter', pathAndAmountMap: { hits: 1 }, expireAt: past }, + opts, + ); + const result = await target.get({ key: 'ttlCounter' }, opts); + expect(result.res).toBeNull(); + }); + + it('keeps the first expireAt stamp across later bumps (if_not_exists)', async () => { + const future = Math.floor(Date.now() / 1000) + 3600; + await target.incr( + { key: 'ttlKeep', pathAndAmountMap: { hits: 1 }, expireAt: future }, + opts, + ); + // A later bump passing an already-elapsed ttl must NOT override the + // first stamp, so the counter stays visible. + const past = Math.floor(Date.now() / 1000) - 10; + await target.incr( + { key: 'ttlKeep', pathAndAmountMap: { hits: 1 }, expireAt: past }, + opts, + ); + const result = await target.get({ key: 'ttlKeep' }, opts); + expect(result.res).toMatchObject({ hits: 2 }); + }); + + it('creates nested intermediate maps lazily on the first bump, then accumulates', async () => { + // First bump into a missing nested parent must still build the map + // (optimistic path: the direct update fails, createPaths runs, retry + // succeeds), and subsequent bumps keep accumulating. + await target.incr( + { key: 'lazyNest', pathAndAmountMap: { 'a.b.c': 2 } }, + opts, + ); + const after = await target.incr( + { key: 'lazyNest', pathAndAmountMap: { 'a.b.c': 3 } }, + opts, + ); + expect(after.res).toMatchObject({ a: { b: { c: 5 } } }); + }); + 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 44a5d1e250..47186c4eb0 100644 --- a/src/backend/stores/systemKv/SystemKVStore.ts +++ b/src/backend/stores/systemKv/SystemKVStore.ts @@ -506,7 +506,11 @@ export class SystemKVStore extends PuterStore { } async incr>( - { key, pathAndAmountMap }: { key: string; pathAndAmountMap: T }, + { + key, + pathAndAmountMap, + expireAt, + }: { key: string; pathAndAmountMap: T; expireAt?: number }, opts?: KVOpts, ): Promise< KVResult> @@ -529,12 +533,6 @@ export class SystemKVStore extends PuterStore { const actor = ensureActor(opts); const namespace = getNamespace(actor, opts?.appUuid); - const createPathsUsage = await this.createPaths( - namespace, - key, - Object.keys(pathAndAmountMap), - ); - const setStatements = Object.entries(pathAndAmountMap).map( ([valPath, _amt], idx) => { const attrName = ['value', ...valPath.split('.')] @@ -565,13 +563,50 @@ export class SystemKVStore extends PuterStore { {} as Record, ); - const response = await this.clients.dynamo.update( - this.tableName, - { key, namespace }, - `SET ${setStatements.join(', ')}`, - valueAttributeValues, - { ...valueAttributeNames, '#value': 'value' }, - ); + // Fold the TTL into the same UpdateItem so a counter bump is a single + // write instead of incr + a separate expireAt. if_not_exists keeps the + // first stamp for the key (re-stamping the same value was always a + // no-op) — but now we don't pay for that extra write on every bump. + if (expireAt !== undefined) { + const ttlSeconds = Number(expireAt); + if (Number.isNaN(ttlSeconds)) + throw new HttpError(400, 'kv: expireAt must be a number', { + legacyCode: 'bad_request', + }); + setStatements.push('#ttl = if_not_exists(#ttl, :ttl)'); + valueAttributeValues[':ttl'] = ttlSeconds; + valueAttributeNames['#ttl'] = 'ttl'; + } + + const updateExpression = `SET ${setStatements.join(', ')}`; + const expressionNames = { ...valueAttributeNames, '#value': 'value' }; + const runUpdate = () => + this.clients.dynamo.update( + this.tableName, + { key, namespace }, + updateExpression, + valueAttributeValues, + expressionNames, + ); + + // Most increments land on an item whose parent maps already exist (a + // day's counter is created once, then bumped on every event), so try + // the update directly and only pay the per-layer createPaths writes + // when a nested parent is genuinely missing — typically the first bump + // for a key. Mirrors the lazy ValidationException fallback in remove(). + let createPathsUsage = 0; + let response; + try { + response = await runUpdate(); + } catch (e) { + if ((e as Error)?.name !== 'ValidationException') throw e; + createPathsUsage = await this.createPaths( + namespace, + key, + Object.keys(pathAndAmountMap), + ); + response = await runUpdate(); + } const usage = writeUsage( Number(response.ConsumedCapacity?.CapacityUnits ?? 0) +