From c7edfc0c74e48e0d71849181ea44b1c3ac966d7b Mon Sep 17 00:00:00 2001 From: Daniel Salazar Date: Thu, 6 Aug 2026 10:26:47 -0700 Subject: [PATCH] perf: improve metering perf + add extra compat for monthly charges (#3513) closes PUT-1445 and PUT-1446 --- src/backend/clients/event/EventClient.test.ts | 32 + src/backend/clients/event/EventClient.ts | 21 + src/backend/clients/event/types.ts | 36 +- .../services/metering/MeteringService.test.ts | 506 ++++++++++++- .../services/metering/MeteringService.ts | 307 ++++++-- src/backend/services/metering/consts.ts | 7 + src/backend/services/metering/types.ts | 19 +- src/backend/stores/index.ts | 5 + .../metering/MeteringBufferStore.test.ts | 593 ++++++++++++++++ .../stores/metering/MeteringBufferStore.ts | 667 ++++++++++++++++++ src/backend/stores/systemKv/SystemKVStore.ts | 14 +- 11 files changed, 2138 insertions(+), 69 deletions(-) create mode 100644 src/backend/stores/metering/MeteringBufferStore.test.ts create mode 100644 src/backend/stores/metering/MeteringBufferStore.ts diff --git a/src/backend/clients/event/EventClient.test.ts b/src/backend/clients/event/EventClient.test.ts index 9b5a24d79..4e09043ab 100644 --- a/src/backend/clients/event/EventClient.test.ts +++ b/src/backend/clients/event/EventClient.test.ts @@ -214,6 +214,38 @@ describe('EventClient', () => { }); }); + describe('hasListeners', () => { + it('is false for a key nobody subscribed to', () => { + expect(target.hasListeners(key)).toBe(false); + }); + + it('is true for an exact subscriber', () => { + target.on(key, vi.fn()); + expect(target.hasListeners(key)).toBe(true); + }); + + it('is true when only a wildcard prefix matches', () => { + target.on(`${key}.*`, vi.fn()); + expect(target.hasListeners(`${key}.child.deep`)).toBe(true); + }); + + it('is true for an extension-registered listener', () => { + extensionStore.events[key] = [vi.fn()]; + try { + expect(target.hasListeners(key)).toBe(true); + } finally { + delete extensionStore.events[key]; + } + }); + + it('goes back to false once the last listener is removed', () => { + const listener = vi.fn(); + target.on(key, listener); + target.off(key, listener); + expect(target.hasListeners(key)).toBe(false); + }); + }); + describe('extension-registered listeners', () => { afterEach(() => { delete extensionStore.events[key]; diff --git a/src/backend/clients/event/EventClient.ts b/src/backend/clients/event/EventClient.ts index 260b4feff..01a73c0ea 100644 --- a/src/backend/clients/event/EventClient.ts +++ b/src/backend/clients/event/EventClient.ts @@ -129,6 +129,27 @@ export class EventClient extends PuterClient { }); } + /** + * Whether anything would run if `key` were emitted, wildcards included. + * + * For emitters whose work is only worth doing when someone is listening — + * gathering the payload costs a round trip, say. Emitting into the void is + * otherwise perfectly cheap and doesn't need this. + */ + hasListeners(key: T): boolean { + const parts = key.split('.'); + for (let i = 0; i < parts.length; i++) { + const matchKey = ( + i === parts.length - 1 + ? key + : `${parts.slice(0, i + 1).join('.')}.*` + ) as ListenKey; + if (this.#eventListeners[matchKey]?.length) return true; + if (extensionStore.events[matchKey]?.length) return true; + } + return false; + } + /** * Subscribe to an event by exact key OR a wildcard prefix. * diff --git a/src/backend/clients/event/types.ts b/src/backend/clients/event/types.ts index dc0bfa1dc..a65e61f22 100644 --- a/src/backend/clients/event/types.ts +++ b/src/backend/clients/event/types.ts @@ -20,6 +20,7 @@ import type { Response as ExpressResponse, } from 'express'; import { Actor } from '../../core'; +import type { UsageInput } from '../../services/metering/types'; import { FSEntry } from '../../stores/fs/FSEntry'; // GUI write events spread an entry plus per-event metadata into `response`. @@ -305,6 +306,32 @@ export type EventMap = { }; 'storage.quota.bonus': { userId: number; extra: number }; + // ---- Metering ---- + // Recurring charges are pure mechanism here: the metering service knows + // when to ask (once per user per month, the first time that month's usage + // record is touched) and how to record the answer, and an extension + // decides what the account owes. Listeners push onto `charges`; the + // service applies them as ordinary usage once every listener has run. + // Nothing listening → no charges and no claim is ever taken. + // + // Emitted after the claim is settled, so a listener that calls back into + // metering can't re-trigger it. + 'metering.monthly.charges': { + /** + * User-scoped, with no app on it: the account is what recurs, and the + * app that happened to trigger the month's first call has nothing to do + * with what is owed. Price against what the user owns. + */ + actor: Actor; + /** The month being charged for, `YYYY-MM` in UTC. */ + month: string; + /** + * Push what the user owes here. Every listener's charges are merged + * into one amount map and recorded as a single increment. + */ + charges: UsageInput[]; + }; + // ---- Outer / GUI broadcast ---- 'outer.cacheUpdate': { cacheKey: string[]; @@ -482,10 +509,11 @@ export type EventKey = keyof EventMap & string; // Generates a wildcard for every non-final dot-separated prefix of K. export type WildcardPrefixes = K extends `${infer Head}.${infer Tail}` - ? | `${Head}.*` - | (Tail extends `${string}.${string}` - ? `${Head}.${WildcardPrefixes}` - : never) + ? + | `${Head}.*` + | (Tail extends `${string}.${string}` + ? `${Head}.${WildcardPrefixes}` + : never) : never; export type ListenKey = EventKey | WildcardPrefixes; diff --git a/src/backend/services/metering/MeteringService.test.ts b/src/backend/services/metering/MeteringService.test.ts index 9ba3804a8..33c20fb7b 100644 --- a/src/backend/services/metering/MeteringService.test.ts +++ b/src/backend/services/metering/MeteringService.test.ts @@ -21,6 +21,7 @@ import { POLICY_PREFIX, } from './consts.ts'; import type { MeteringService } from './MeteringService.ts'; +import type { UsageInput } from './types.ts'; import { toMicroCents } from './utils.ts'; const escape = (usageType: string) => usageType.replace(/\./g, PERIOD_ESCAPE); @@ -61,6 +62,11 @@ describe('MeteringService', () => { defs: [...internals.defaultSubscriptionResolvers], pols: [...internals.extraPolicies], }; + + // Usage counters accumulate in a buffer that a background loop writes + // onward. Stop that loop so tests settle it explicitly and never race + // a cycle firing mid-assertion. + await server.stores.meteringBuffer.onServerShutdown(); }); afterEach(() => { @@ -574,14 +580,17 @@ describe('MeteringService', () => { }); it('returns zero and writes nothing when every item is skipped', async () => { - const incrSpy = vi.spyOn(server.stores.kv, 'incr'); + const incrSpy = vi.spyOn(server.stores.meteringBuffer, 'incr'); + const auxSpy = vi.spyOn(server.stores.meteringBuffer, 'incrAux'); const result = await target.batchIncrementUsages(actor, [ { usageType: '', usageAmount: 1, costOverride: 10 }, { usageType: 'kv:write', usageAmount: 0, costOverride: 20 }, ]); expect(result).toEqual({ total: 0 }); expect(incrSpy).not.toHaveBeenCalled(); + expect(auxSpy).not.toHaveBeenCalled(); incrSpy.mockRestore(); + auxSpy.mockRestore(); }); it('raises an alarm for any negative costOverride in the batch', async () => { @@ -780,7 +789,8 @@ describe('MeteringService', () => { expect(result.total).toBe(500); const adj = (result as Record) .manual_adjustment as - { cost: number; units: number; count: number } | undefined; + | { cost: number; units: number; count: number } + | undefined; expect(adj).toMatchObject({ cost: 500, units: 500, count: 1 }); }); @@ -946,20 +956,33 @@ describe('MeteringService', () => { // ── getGlobalUsage ─────────────────────────────────────────────── describe('getGlobalUsage', () => { + // The global view is read straight from the store, and aggregate + // counters are written onward a cycle at a time. Flush until the view + // stops moving so a baseline isn't polluted by usage other tests left + // buffered. + const settledGlobalUsage = async () => { + let previous = Number.NaN; + for (let attempt = 0; attempt < 20; attempt++) { + await server.stores.meteringBuffer.flushCycle(); + const usage = await target.getGlobalUsage(); + if (usage.total === previous) return usage; + previous = usage.total; + } + throw new Error('global usage never settled'); + }; + it('aggregates increments across actors into the same global view', async () => { - const before = await target.getGlobalUsage(); + const before = await settledGlobalUsage(); const user1: Actor = { user: makeUser() }; const user2: Actor = { user: makeUser() }; await target.incrementUsage(user1, 'kv:read', 1, 100); await target.incrementUsage(user2, 'kv:read', 1, 200); - await waitFor(async () => { - const now = await target.getGlobalUsage(); - expect(now.total - before.total).toBe(300); - const beforeRead = (before['kv:read']?.cost ?? 0) as number; - const nowRead = (now['kv:read']?.cost ?? 0) as number; - expect(nowRead - beforeRead).toBe(300); - }); + const now = await settledGlobalUsage(); + expect(now.total - before.total).toBe(300); + const beforeRead = (before['kv:read']?.cost ?? 0) as number; + const nowRead = (now['kv:read']?.cost ?? 0) as number; + expect(nowRead - beforeRead).toBe(300); }); }); @@ -968,6 +991,9 @@ describe('MeteringService', () => { describe('KV layout', () => { it('writes the actor monthly record at the expected key shape', async () => { await target.incrementUsage(actor, 'kv:read', 1, 100); + // Counters are written onward a cycle at a time, so settle first + // and then assert where the data actually landed. + await server.stores.meteringBuffer.flushCycle(); const month = `${new Date().getUTCFullYear()}-${String( new Date().getUTCMonth() + 1, ).padStart(2, '0')}`; @@ -984,6 +1010,466 @@ describe('MeteringService', () => { }); }); + // ── Buffered counters ──────────────────────────────────────────── + + describe('buffered usage counters', () => { + const actorKey = (usageActor: Actor) => { + const now = new Date(); + const month = `${now.getUTCFullYear()}-${String( + now.getUTCMonth() + 1, + ).padStart(2, '0')}`; + return `${METRICS_PREFIX}:actor:${usageActor.user!.uuid}:${month}`; + }; + + it('accumulates a running total without a write per call', async () => { + const bufActor: Actor = { user: makeUser() }; + const key = actorKey(bufActor); + + const first = await target.incrementUsage( + bufActor, + 'ai:chat', + 1, + 100, + ); + const second = await target.incrementUsage( + bufActor, + 'ai:chat', + 1, + 150, + ); + + expect(first.total).toBe(100); + expect(second.total).toBe(250); + // Nothing recorded yet — the flush loop is the only writer. + const { res: beforeFlush } = await server.stores.kv.get({ key }); + expect(beforeFlush).toBeNull(); + + await server.stores.meteringBuffer.flushCycle(); + const { res: afterFlush } = await server.stores.kv.get({ key }); + expect(afterFlush).toMatchObject({ total: 250 }); + }); + + it('takes an exact reading once usage approaches the allowance', async () => { + const bufActor: Actor = { user: makeUser() }; + const key = actorKey(bufActor); + const allowance = (await target.getActorSubscription(bufActor)) + .monthUsageAllowance; + + await target.incrementUsage( + bufActor, + 'ai:chat', + 1, + Math.round(allowance * 0.85), + ); + await server.stores.meteringBuffer.flushCycle(); + + // Usage recorded elsewhere for the same account, which this + // deployment's buffered view has no way to know about. + const elsewhere = Math.round(allowance * 0.45); + await server.stores.kv.incr({ + key, + pathAndAmountMap: { total: elsewhere }, + }); + + const step = Math.round(allowance * 0.06); + const usage = await target.incrementUsage( + bufActor, + 'ai:chat', + 1, + step, + ); + + expect(usage.total).toBe( + Math.round(allowance * 0.85) + elsewhere + step, + ); + }); + + it('stays with the buffered total while far from the allowance', async () => { + const bufActor: Actor = { user: makeUser() }; + const key = actorKey(bufActor); + const allowance = (await target.getActorSubscription(bufActor)) + .monthUsageAllowance; + const started = Math.round(allowance * 0.1); + + await target.incrementUsage(bufActor, 'ai:chat', 1, started); + await server.stores.meteringBuffer.flushCycle(); + + await server.stores.kv.incr({ + key, + pathAndAmountMap: { total: Math.round(allowance * 0.45) }, + }); + + const usage = await target.incrementUsage( + bufActor, + 'ai:chat', + 1, + 5, + ); + + // Well inside the allowance the decision is the same either way, + // so this deliberately does not pay for an exact reading. + expect(usage.total).toBe(started + 5); + }); + }); + + // ── Monthly recurring charges ──────────────────────────────────── + + describe('monthly recurring charges', () => { + type ChargeEvent = { charges: UsageInput[]; month: string }; + type ChargeListener = ( + key: unknown, + data: ChargeEvent, + ) => void | Promise; + + const monthKey = (chargeActor: Actor) => { + const now = new Date(); + const month = `${now.getUTCFullYear()}-${String( + now.getUTCMonth() + 1, + ).padStart(2, '0')}`; + return `${METRICS_PREFIX}:actor:${chargeActor.user!.uuid}:${month}`; + }; + + const claimOf = async (chargeActor: Actor) => { + const { res } = await server.stores.kv.get({ + key: monthKey(chargeActor), + }); + return (res as { monthlyChargesApplied?: number } | null) + ?.monthlyChargesApplied; + }; + + // Every deployment settles a month once and then remembers it; a + // second deployment (or this one after a restart) starts with an empty + // memory and has to ask the KV store. + const forgetSettled = () => + ( + target as unknown as { settledActors: Set } + ).settledActors.clear(); + + const registered: ChargeListener[] = []; + const listen = (fn: ChargeListener) => { + server.clients.event.on( + 'metering.monthly.charges', + fn as Parameters[1], + ); + registered.push(fn); + return fn; + }; + const chargeOnce = (cost: number) => + listen( + vi.fn((_key, data: ChargeEvent) => { + data.charges.push({ + usageType: 'workers:monthly', + usageAmount: 1, + costOverride: cost, + }); + }), + ); + + afterEach(() => { + for (const fn of registered) { + server.clients.event.off( + 'metering.monthly.charges', + fn as Parameters[1], + ); + } + registered.length = 0; + }); + + it('applies a listener charge on the first write and returns it in the total', async () => { + const listener = chargeOnce(700); + + const usage = await target.incrementUsage(actor, 'kv:read', 1, 100); + + expect(listener).toHaveBeenCalledTimes(1); + expect(usage.total).toBe(800); + expect(usage['workers:monthly']).toMatchObject({ + cost: 700, + units: 1, + count: 1, + }); + }); + + it('applies the charge on a read when the read comes first', async () => { + chargeOnce(500); + + const { usage } = + await target.getActorCurrentMonthUsageDetails(actor); + + expect(usage.total).toBe(500); + }); + + it('charges once per month however many calls follow', async () => { + const listener = chargeOnce(400); + + await target.incrementUsage(actor, 'kv:read', 1, 10); + await target.incrementUsage(actor, 'kv:read', 1, 10); + const usage = await target.getActorCurrentMonthUsageDetails(actor); + + expect(listener).toHaveBeenCalledTimes(1); + expect(usage.usage.total).toBe(420); + }); + + it('charges once when several calls race for the same actor', async () => { + const listener = chargeOnce(300); + + await Promise.all( + Array.from({ length: 8 }, () => + target.incrementUsage(actor, 'kv:read', 1, 10), + ), + ); + + expect(listener).toHaveBeenCalledTimes(1); + expect(await claimOf(actor)).toBe(1); + }); + + it('does not charge again for a month another deployment already claimed', async () => { + // Settle a buffered view first, so the claim the other deployment + // takes next is one this one genuinely cannot see. + await target.incrementUsage(actor, 'kv:read', 1, 10); + await server.stores.meteringBuffer.flushCycle(); + await server.stores.kv.incr({ + key: monthKey(actor), + pathAndAmountMap: { monthlyChargesApplied: 1 }, + }); + + const listener = chargeOnce(900); + const usage = await target.incrementUsage(actor, 'kv:read', 1, 50); + + expect(listener).not.toHaveBeenCalled(); + expect(usage.total).toBe(60); + // The claim counts every attempt, so the loser is visible as 2. + expect(await claimOf(actor)).toBe(2); + }); + + it('skips the claim entirely when nothing is listening', async () => { + await target.incrementUsage(actor, 'kv:read', 1, 100); + await server.stores.meteringBuffer.flushCycle(); + + expect(await claimOf(actor)).toBeUndefined(); + }); + + it('leaves the month settled when a listener throws, and the call still succeeds', async () => { + const listener = listen( + vi.fn(() => { + throw new Error('pricing lookup failed'); + }), + ); + + const usage = await target.incrementUsage(actor, 'kv:read', 1, 100); + forgetSettled(); + await target.incrementUsage(actor, 'kv:read', 1, 100); + + expect(usage.total).toBe(100); + expect(listener).toHaveBeenCalledTimes(1); + }); + + it('retries on the next call when the claim write fails', async () => { + const listener = chargeOnce(600); + const incr = vi + .spyOn(server.stores.kv, 'incr') + .mockRejectedValueOnce(new Error('kv unavailable')); + + const first = await target.incrementUsage(actor, 'kv:read', 1, 100); + expect(listener).not.toHaveBeenCalled(); + expect(first.total).toBe(100); + + incr.mockRestore(); + const second = await target.incrementUsage( + actor, + 'kv:read', + 1, + 100, + ); + + expect(listener).toHaveBeenCalledTimes(1); + expect(second.total).toBe(800); + }); + + // An actor with usage earlier in the month has a buffered view already + // built, and a claim written straight to the KV store does not show up + // in it until the next flush. That is the window where re-entry has + // nothing but the in-flight guard to stop it, so these start there. + const warmBufferedView = async () => { + await target.incrementUsage(actor, 'kv:read', 1, 10); + await server.stores.meteringBuffer.flushCycle(); + }; + + it('charges once when the listener meters through the service itself', async () => { + await warmBufferedView(); + const listener = listen( + vi.fn(async () => { + await target.incrementUsage( + actor, + 'workers:monthly', + 1, + 20, + ); + }), + ); + + const usage = await target.incrementUsage(actor, 'kv:read', 1, 100); + + expect(listener).toHaveBeenCalledTimes(1); + expect(await claimOf(actor)).toBe(1); + // Metering itself rather than pushing onto `charges` means the + // cost lands on the record but misses the total this call already + // computed — visible from the next read on. + expect(usage.total).toBe(110); + const after = await target.getActorCurrentMonthUsageDetails(actor); + expect(after.usage.total).toBe(130); + }); + + it('charges once even if the settled memory is dropped mid-claim', async () => { + // The memo is capped and cleared wholesale when it fills, which can + // land in the window where a listener is still running. + await warmBufferedView(); + const listener = listen( + vi.fn(async (_key, data: ChargeEvent) => { + forgetSettled(); + await target.incrementUsage(actor, 'kv:read', 1, 5); + data.charges.push({ + usageType: 'workers:monthly', + usageAmount: 1, + costOverride: 200, + }); + }), + ); + + const usage = await target.incrementUsage(actor, 'kv:read', 1, 100); + + expect(listener).toHaveBeenCalledTimes(1); + expect(usage.total).toBe(315); + expect(await claimOf(actor)).toBe(1); + }); + + it('merges every listener into one amount map and one increment', async () => { + listen( + vi.fn((_key, data: ChargeEvent) => { + data.charges.push( + { + usageType: 'workers:monthly', + usageAmount: 3, + costOverride: 300, + }, + { + usageType: 'domains:monthly', + usageAmount: 1, + costOverride: 100, + }, + ); + }), + ); + listen( + vi.fn((_key, data: ChargeEvent) => { + data.charges.push({ + usageType: 'workers:monthly', + usageAmount: 2, + costOverride: 200, + }); + }), + ); + + const incr = vi.spyOn(server.stores.meteringBuffer, 'incr'); + const usage = await target.getActorCurrentMonthUsageDetails(actor); + + // Four charges across two listeners, settling as a single write. + expect(incr).toHaveBeenCalledTimes(1); + expect(incr.mock.calls[0]![0].pathAndAmountMap).toEqual({ + total: 600, + 'workers:monthly.units': 5, + 'workers:monthly.cost': 500, + 'workers:monthly.count': 2, + 'domains:monthly.units': 1, + 'domains:monthly.cost': 100, + 'domains:monthly.count': 1, + }); + expect(usage.usage.total).toBe(600); + incr.mockRestore(); + }); + + it('bills the user, not the app that happened to trigger it', async () => { + chargeOnce(700); + const appActor: Actor = { ...actor, app: { uid: 'app-abc' } }; + + await target.incrementUsage(appActor, 'kv:read', 1, 50); + await server.stores.meteringBuffer.flushCycle(); + + // The app wears only what it actually spent... + const appUsage = await target.getActorCurrentMonthAppUsageDetails( + appActor, + 'app-abc', + ); + expect(appUsage.total).toBe(50); + // ...while the recurring charge sits in the user's own bucket. + const global = await target.getActorCurrentMonthAppUsageDetails( + actor, + GLOBAL_APP_KEY, + ); + expect(global.total).toBe(700); + + const { usage } = await target.getActorCurrentMonthUsageDetails( + actor, + ); + expect(usage.total).toBe(750); + }); + + it('charges the user once across several of their apps', async () => { + const listener = chargeOnce(800); + + await target.incrementUsage( + { ...actor, app: { uid: 'app-one' } }, + 'kv:read', + 1, + 10, + ); + await target.incrementUsage( + { ...actor, app: { uid: 'app-two' } }, + 'kv:read', + 1, + 10, + ); + + expect(listener).toHaveBeenCalledTimes(1); + expect(await claimOf(actor)).toBe(1); + }); + + it('hands listeners a user-scoped actor', async () => { + let seen: Actor | undefined; + listen( + vi.fn((_key, data: ChargeEvent & { actor: Actor }) => { + seen = data.actor; + }), + ); + + await target.incrementUsage( + { ...actor, app: { uid: 'app-abc' } }, + 'kv:read', + 1, + 10, + ); + + expect(seen?.user.uuid).toBe(actor.user.uuid); + expect(seen?.app).toBeUndefined(); + }); + + it('ignores charges a listener pushed with no usage type', async () => { + listen( + vi.fn((_key, data: ChargeEvent) => { + data.charges.push({ + usageType: '', + usageAmount: 1, + costOverride: 100, + }); + }), + ); + + const usage = await target.incrementUsage(actor, 'kv:read', 1, 50); + + expect(usage.total).toBe(50); + expect(await claimOf(actor)).toBe(1); + }); + }); + // ── Resolver registration ──────────────────────────────────────── describe('resolver registration', () => { diff --git a/src/backend/services/metering/MeteringService.ts b/src/backend/services/metering/MeteringService.ts index e601c41a7..546ad6464 100644 --- a/src/backend/services/metering/MeteringService.ts +++ b/src/backend/services/metering/MeteringService.ts @@ -27,11 +27,18 @@ import { DEFAULT_TEMP_SUBSCRIPTION, GLOBAL_APP_KEY, METRICS_PREFIX, + MONTHLY_CHARGE_CLAIM, PERIOD_ESCAPE, POLICY_PREFIX, UNLIMITED_SUBSCRIPTION, } from './consts'; -import type { AppTotals, UsageAddons, UsageByType, UsageRecord } from './types'; +import type { + AppTotals, + UsageAddons, + UsageByType, + UsageInput, + UsageRecord, +} from './types'; import { toMicroCents } from './utils'; import { SUB_POLICIES } from '../../data/subPolicies/index.js'; @@ -44,12 +51,6 @@ export type SubscriptionResolver = ( actor: Actor, ) => Promise | string | null | undefined; -interface UsageInput { - usageType: string; - usageAmount: number; - costOverride?: number; -} - // -- Helpers ---------------------------------------------------------- /** @@ -77,15 +78,48 @@ function actorLabel(actor: Actor): string { * fan that out into several aggregated KV records. */ export class MeteringService extends PuterService { + /** + * How wide the global and per-app aggregates are spread. These are counters + * many actors increment at once, so spreading them keeps any single record + * from being written by everyone — including from several deployments + * concurrently, where writes to one record can otherwise lose an increment. + * The width is why reading an aggregate has to sum every shard. + */ static GLOBAL_SHARD_COUNT = 10000; static APP_SHARD_COUNT = 10000; + static MAX_GLOBAL_USAGE_PER_MINUTE = toMicroCents(0.2); + /** + * Share of the allowance past which an approximate running total is no + * longer good enough to decide on. + */ + static PRECISION_THRESHOLD = 0.9; + + /** + * How many actors this deployment remembers as settled for the month. The + * claim in the KV store is what makes monthly charges once-only; this + * memory only saves the round trip that would discover that, so forgetting + * it costs a claim write and nothing else. + */ + static MONTHLY_CHARGE_MEMO_LIMIT = 100_000; + private rateCheckTimer: ReturnType | null = null; private extraPolicies: SubscriptionPolicy[] = []; private subscriptionResolvers: SubscriptionResolver[] = []; private defaultSubscriptionResolvers: SubscriptionResolver[] = []; + /** Actors settled for `settledMonth`; see MONTHLY_CHARGE_MEMO_LIMIT. */ + private settledMonth: string | null = null; + private settledActors = new Set(); + /** + * Actors with a claim in flight. Unlike `settledActors` this is never + * dropped early, because it is what stops a second claim inside the first: + * applying the charges goes back through `batchIncrementUsages`, which + * arrives here again for the same actor and month. + */ + private claimsInFlight = new Set(); + // -- Lifecycle ---------------------------------------------------- override onServerStart(): void { @@ -210,17 +244,15 @@ export class MeteringService extends PuterService { }; const actorUsageKey = `${METRICS_PREFIX}:actor:${userId}:${currentMonth}`; - const actorUsagesPromise = this.stores.kv - .incr({ - key: actorUsageKey, - pathAndAmountMap, - }) - .then((r) => r.res as unknown as UsageByType); + const actorUsagesPromise = this.stores.meteringBuffer.incr({ + key: actorUsageKey, + pathAndAmountMap, + }); // Aux writes — fire and forget this.handleAuxPromise( `puterConsumption ${userId}/${appId}`, - this.stores.kv.incr({ + this.stores.meteringBuffer.incrAux({ key: this.globalUsageKey(userId, appId, currentMonth), pathAndAmountMap, }), @@ -228,7 +260,7 @@ export class MeteringService extends PuterService { this.handleAuxPromise( `actorAppUsage ${userId}/${appId}`, - this.stores.kv.incr({ + this.stores.meteringBuffer.incrAux({ key: `${METRICS_PREFIX}:actor:${userId}:app:${appId}:${currentMonth}`, pathAndAmountMap, }), @@ -237,7 +269,7 @@ export class MeteringService extends PuterService { if (appId !== GLOBAL_APP_KEY) { this.handleAuxPromise( `appUsage ${appId}/${userId}`, - this.stores.kv.incr({ + this.stores.meteringBuffer.incrAux({ key: this.appUsageKey(appId, userId, currentMonth), pathAndAmountMap, }), @@ -246,7 +278,7 @@ export class MeteringService extends PuterService { this.handleAuxPromise( `actorAppTotals ${userId}`, - this.stores.kv.incr({ + this.stores.meteringBuffer.incrAux({ key: `${METRICS_PREFIX}:actor:${userId}:apps:${currentMonth}`, pathAndAmountMap: { [`${appId}.total`]: totalCost, @@ -255,13 +287,19 @@ export class MeteringService extends PuterService { }), ); - const [actorUsages, actorSubscription, actorAddons] = + const [usageResult, actorSubscription, actorAddons] = await Promise.all([ actorUsagesPromise, this.getActorSubscription(actor), this.getActorAddons(actor), ]); + const actorUsages = await this.exactUsageNearAllowance( + actorUsageKey, + usageResult, + actorSubscription.monthUsageAllowance, + ); + await this.maybeConsumeAddonCredits( userId, actorUsages.total, @@ -282,7 +320,13 @@ export class MeteringService extends PuterService { costOverride, }); - return actorUsages; + return ( + (await this.applyMonthlyCharges( + actor, + currentMonth, + actorUsages, + )) ?? actorUsages + ); } catch (e) { console.error('[metering] incrementUsage failed', { actor, @@ -379,37 +423,35 @@ export class MeteringService extends PuterService { const userId = actor.user.uuid!; const actorUsageKey = `${METRICS_PREFIX}:actor:${userId}:${currentMonth}`; - const actorUsagesPromise = this.stores.kv - .incr({ - key: actorUsageKey, - pathAndAmountMap: aggregated, - }) - .then((r) => r.res as unknown as UsageByType); + const actorUsagesPromise = this.stores.meteringBuffer.incr({ + key: actorUsageKey, + pathAndAmountMap: aggregated, + }); this.handleAuxPromise( `puterConsumption ${userId}/${appId}`, - this.stores.kv.incr({ + this.stores.meteringBuffer.incrAux({ key: this.globalUsageKey(userId, appId, currentMonth), pathAndAmountMap: aggregated, }), ); this.handleAuxPromise( `actorAppUsage ${userId}/${appId}`, - this.stores.kv.incr({ + this.stores.meteringBuffer.incrAux({ key: `${METRICS_PREFIX}:actor:${userId}:app:${appId}:${currentMonth}`, pathAndAmountMap: aggregated, }), ); this.handleAuxPromise( `appUsage ${appId}/${userId}`, - this.stores.kv.incr({ + this.stores.meteringBuffer.incrAux({ key: this.appUsageKey(appId, userId, currentMonth), pathAndAmountMap: aggregated, }), ); this.handleAuxPromise( `actorAppTotals ${userId}`, - this.stores.kv.incr({ + this.stores.meteringBuffer.incrAux({ key: `${METRICS_PREFIX}:actor:${userId}:apps:${currentMonth}`, pathAndAmountMap: { [`${appId}.total`]: totalBatchCost, @@ -418,13 +460,19 @@ export class MeteringService extends PuterService { }), ); - const [actorUsages, actorSubscription, actorAddons] = + const [usageResult, actorSubscription, actorAddons] = await Promise.all([ actorUsagesPromise, this.getActorSubscription(actor), this.getActorAddons(actor), ]); + const actorUsages = await this.exactUsageNearAllowance( + actorUsageKey, + usageResult, + actorSubscription.monthUsageAllowance, + ); + await this.maybeConsumeAddonCredits( userId, actorUsages.total, @@ -443,7 +491,13 @@ export class MeteringService extends PuterService { batchUsages: usages, }); - return actorUsages; + return ( + (await this.applyMonthlyCharges( + actor, + currentMonth, + actorUsages, + )) ?? actorUsages + ); } catch (e) { console.error('[metering] batchIncrementUsages failed', { actor, @@ -489,12 +543,23 @@ export class MeteringService extends PuterService { `${METRICS_PREFIX}:actor:${actor.user.uuid}:apps:${currentMonth}`, ]; - const { res } = await this.stores.kv.get({ key: keys }); + const { res } = await this.stores.meteringBuffer.get({ key: keys }); const [usage, appTotals] = (res ?? []) as [ UsageByType | null, Record | null, ]; + // Reading the month is one of the two things that settles its + // recurring charges. The per-app breakdown is written by the same + // increment but read above it, so it picks them up a read later than + // the total does. + const charged = await this.applyMonthlyCharges( + actor, + currentMonth, + usage, + ); + const resolvedUsage = charged ?? usage ?? ({ total: 0 } as UsageByType); + const appId = actor.app?.uid; if (appTotals && appId) { const filtered: Record = {}; @@ -511,16 +576,10 @@ export class MeteringService extends PuterService { } }); if (others) filtered['others'] = others; - return { - usage: usage || ({ total: 0 } as UsageByType), - appTotals: filtered, - }; + return { usage: resolvedUsage, appTotals: filtered }; } - return { - usage: usage || ({ total: 0 } as UsageByType), - appTotals: appTotals || {}, - }; + return { usage: resolvedUsage, appTotals: appTotals || {} }; } async setActorCurrentMonthUsageTotal( @@ -551,7 +610,9 @@ export class MeteringService extends PuterService { const appId = actor.app?.uid || GLOBAL_APP_KEY; const actorUsageKey = `${METRICS_PREFIX}:actor:${userId}:${currentMonth}`; - const { res: current } = await this.stores.kv.get({ + // Setting an absolute total is only meaningful against an exact + // starting point, so this one reads through everything pending. + const { res: current } = await this.stores.meteringBuffer.readExact({ key: actorUsageKey, }); const currentTotal = (current as UsageByType | null)?.total ?? 0; @@ -569,26 +630,29 @@ export class MeteringService extends PuterService { }; const updated = ( - await this.stores.kv.incr({ key: actorUsageKey, pathAndAmountMap }) + await this.stores.meteringBuffer.incr({ + key: actorUsageKey, + pathAndAmountMap, + }) ).res as unknown as UsageByType; this.handleAuxPromise( `puterConsumption ${userId}/${appId}`, - this.stores.kv.incr({ + this.stores.meteringBuffer.incrAux({ key: this.globalUsageKey(userId, appId, currentMonth), pathAndAmountMap, }), ); this.handleAuxPromise( `actorAppUsage ${userId}/${appId}`, - this.stores.kv.incr({ + this.stores.meteringBuffer.incrAux({ key: `${METRICS_PREFIX}:actor:${userId}:app:${appId}:${currentMonth}`, pathAndAmountMap, }), ); this.handleAuxPromise( `actorAppTotals ${userId}`, - this.stores.kv.incr({ + this.stores.meteringBuffer.incrAux({ key: `${METRICS_PREFIX}:actor:${userId}:apps:${currentMonth}`, pathAndAmountMap: { [`${appId}.total`]: delta, @@ -630,7 +694,7 @@ export class MeteringService extends PuterService { const currentMonth = this.monthYearString(); const key = `${METRICS_PREFIX}:actor:${actor.user.uuid}:app:${resolvedAppId}:${currentMonth}`; - const { res } = await this.stores.kv.get({ key }); + const { res } = await this.stores.meteringBuffer.get({ key }); return (res as UsageByType) || ({ total: 0 } as UsageByType); } @@ -744,7 +808,7 @@ export class MeteringService extends PuterService { const currentMonth = this.monthYearString(); const key = `${METRICS_PREFIX}:actor:${actor.user.uuid}:app:${appId}:${currentMonth}`; - const { res } = await this.stores.kv.get({ key }); + const { res } = await this.stores.meteringBuffer.get({ key }); return (res ?? { total: 0 }) as UsageByType; } @@ -832,6 +896,29 @@ export class MeteringService extends PuterService { return `${METRICS_PREFIX}:app:${appId}:${hash}:${currentMonth}`; } + /** + * Well under the allowance an approximate running total leads to the same + * decisions as an exact one, so it isn't worth paying for precision. Close + * to the limit it is — that's where the decisions below actually turn on + * the number. + */ + private async exactUsageNearAllowance( + key: string, + usage: { res: unknown; exact: boolean }, + monthUsageAllowance: number, + ): Promise { + const approximate = usage.res as UsageByType; + if (usage.exact || !(monthUsageAllowance > 0)) return approximate; + if ( + (approximate.total || 0) < + monthUsageAllowance * MeteringService.PRECISION_THRESHOLD + ) + return approximate; + + const { res } = await this.stores.meteringBuffer.readExact({ key }); + return (res as UsageByType) ?? approximate; + } + private handleAuxPromise(label: string, promise: Promise): void { promise.catch((e: Error) => { console.warn( @@ -855,6 +942,132 @@ export class MeteringService extends PuterService { return null; } + // -- Internals: monthly charges ----------------------------------- + + /** + * Charges that recur monthly are applied the first time an actor touches + * the month rather than swept for on a schedule: an actor who never comes + * back is never looked at, and the work lands on the one request that was + * already reading or writing that month's record anyway. + * + * `usage` is the record the caller has in hand. Once it carries the claim + * this costs nothing at all, which is the case for every request but the + * first. Returns the usage including the charges when this call is the one + * that applied them, and null otherwise — including on failure, since a + * charge that couldn't be applied shouldn't take the request down with it. + */ + private async applyMonthlyCharges( + actor: Actor, + currentMonth: string, + usage: UsageByType | null, + ): Promise { + const userId = actor?.user?.uuid; + if (!userId || isSystemActor(actor)) return null; + if (!this.clients.event.hasListeners('metering.monthly.charges')) + return null; + + // Scoped to the month as well as the actor: a claim in flight across + // midnight says nothing about the month that just started. + const claimId = `${userId}:${currentMonth}`; + // Checked before anything that can be forgotten, and answered with + // null rather than the running claim — a caller that awaited it could + // be the claim itself, one frame down. + if (this.claimsInFlight.has(claimId)) return null; + + if (usage?.[MONTHLY_CHARGE_CLAIM]) { + this.rememberSettled(claimId, currentMonth); + return null; + } + if (this.isSettled(claimId, currentMonth)) return null; + + // Nothing awaits between the check and the add, so two callers can't + // both get past it. + this.claimsInFlight.add(claimId); + try { + return await this.claimAndCharge(actor, userId, currentMonth); + } finally { + this.claimsInFlight.delete(claimId); + } + } + + /** + * Take the month's claim, and if it was ours, ask what the user owes and + * record it. + * + * The claim goes straight to the KV store rather than through the metering + * buffer: the buffer answers from this deployment's own view, and the point + * of this counter is to be the one value every deployment agrees on. + * Exactly one caller anywhere sees it come back as 1. + */ + private async claimAndCharge( + actor: Actor, + userId: string, + currentMonth: string, + ): Promise { + let claim: number; + try { + const { res } = await this.stores.kv.incr({ + key: `${METRICS_PREFIX}:actor:${userId}:${currentMonth}`, + pathAndAmountMap: { [MONTHLY_CHARGE_CLAIM]: 1 }, + }); + claim = Number( + (res as Record)?.[MONTHLY_CHARGE_CLAIM] ?? 0, + ); + } catch (e) { + // Unclaimed, so the next request retries. Charging late beats + // charging never, and beats failing the request outright. + console.warn( + `[metering] monthly charge claim failed for ${userId}: ${(e as Error).message}`, + ); + return null; + } + + this.rememberSettled(`${userId}:${currentMonth}`, currentMonth); + // Every attempt bumps the counter, so exactly one caller anywhere ever + // reads 1 back. Everyone else lost the race and must not charge. + if (claim !== 1) return null; + + // The account owes this, not whichever app happened to make the first + // call of the month. Dropping the app bills it to the user's own + // bucket instead of landing it in that app's usage — which its + // developer reads — and hands listeners a subject they can price + // against what the user owns. + const userActor: Actor = { user: actor.user }; + + const charges: UsageInput[] = []; + await this.clients.event.emitAndWait( + 'metering.monthly.charges', + { actor: userActor, month: currentMonth, charges }, + {}, + ); + + const valid = charges.filter( + (charge) => + charge?.usageType && Number.isFinite(charge.usageAmount), + ); + if (valid.length === 0) return null; + // One call, so every charge is folded into a single amount map and + // settles as one write however many listeners contributed. + return this.batchIncrementUsages(userActor, valid); + } + + private rememberSettled(claimId: string, month: string): void { + if (this.settledMonth !== month) { + this.settledMonth = month; + this.settledActors.clear(); + } + if ( + this.settledActors.size >= MeteringService.MONTHLY_CHARGE_MEMO_LIMIT + ) { + this.settledActors.clear(); + } + this.settledActors.add(claimId); + } + + private isSettled(claimId: string, month: string): boolean { + return this.settledMonth === month && this.settledActors.has(claimId); + } + private async maybeConsumeAddonCredits( userId: string, totalUsage: number, diff --git a/src/backend/services/metering/consts.ts b/src/backend/services/metering/consts.ts index 3f4181ba8..fe97300af 100644 --- a/src/backend/services/metering/consts.ts +++ b/src/backend/services/metering/consts.ts @@ -23,6 +23,13 @@ export const METRICS_PREFIX = 'metering'; export const POLICY_PREFIX = 'policy'; /** Dots in usage types are escaped so they don't collide with kv nested paths */ export const PERIOD_ESCAPE = '_dot_'; +/** + * Field on an actor's monthly usage record holding the claim for that month's + * recurring charges. Lives on the record itself so every read that already + * fetches usage can tell whether the charges are settled without a second + * lookup. Must match the `monthlyChargesApplied` member of `UsageByType`. + */ +export const MONTHLY_CHARGE_CLAIM = 'monthlyChargesApplied'; export const DEFAULT_FREE_SUBSCRIPTION = 'user_free'; export const DEFAULT_TEMP_SUBSCRIPTION = 'temp_free'; diff --git a/src/backend/services/metering/types.ts b/src/backend/services/metering/types.ts index a0f9d2ff2..686c58724 100644 --- a/src/backend/services/metering/types.ts +++ b/src/backend/services/metering/types.ts @@ -32,9 +32,22 @@ export interface UsageRecord { units: number; } -export type UsageByType = { total: number } & Partial< - Record, UsageRecord> ->; +/** One metered event: what was used, how much of it, and what it cost. */ +export interface UsageInput { + usageType: string; + usageAmount: number; + costOverride?: number; +} + +export type UsageByType = { + total: number; + /** + * Claim counter for the month's recurring charges — see + * `MONTHLY_CHARGE_CLAIM`. Absent until the first read or write of the + * month; 1 for whoever claimed it, higher for anyone who raced and lost. + */ + monthlyChargesApplied?: number; +} & Partial, UsageRecord>>; export interface AppTotals { total: number; diff --git a/src/backend/stores/index.ts b/src/backend/stores/index.ts index 39b3eda25..d5052c816 100644 --- a/src/backend/stores/index.ts +++ b/src/backend/stores/index.ts @@ -20,6 +20,7 @@ import { AppStore } from './app/AppStore.js'; import { FSEntryStore } from './fs/FSEntryStore.js'; import { GroupStore } from './group/GroupStore.js'; +import { MeteringBufferStore } from './metering/MeteringBufferStore.js'; import { NotificationStore } from './notification/NotificationStore.js'; import { OIDCStore } from './oidc/OIDCStore.js'; import { PermissionStore } from './permission/PermissionStore.js'; @@ -41,6 +42,7 @@ import type { IPuterStoreRegistry } from './types.js'; declare module './types.js' { interface IPuterStoreInstances { kv: SystemKVStore; + meteringBuffer: MeteringBufferStore; user: UserStore; app: AppStore; fsEntry: FSEntryStore; @@ -57,6 +59,8 @@ declare module './types.js' { // Ordering matters: stores declared later see earlier ones as peers. // PermissionStore depends on `kv`, so `kv` must come first. +// MeteringBufferStore sits in front of `kv` for metering counters, so it too +// has to come after it. // UserStore / AppStore are leaves (db + redis only); sit early so other // stores/services can lean on them for cached lookups. // FSEntryStore depends on `kv` (pending-upload sessions live there). @@ -64,6 +68,7 @@ declare module './types.js' { // SessionStore / ShareStore are leaves — only use clients.db. export const puterStores = { kv: SystemKVStore, + meteringBuffer: MeteringBufferStore, user: UserStore, app: AppStore, fsEntry: FSEntryStore, diff --git a/src/backend/stores/metering/MeteringBufferStore.test.ts b/src/backend/stores/metering/MeteringBufferStore.test.ts new file mode 100644 index 000000000..9f63aa8ad --- /dev/null +++ b/src/backend/stores/metering/MeteringBufferStore.test.ts @@ -0,0 +1,593 @@ +/* + * Copyright (C) 2024-present Puter Technologies Inc. + * + * This file is part of Puter. + * + * Puter is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published + * by the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +import { + afterAll, + beforeAll, + beforeEach, + describe, + expect, + it, + vi, +} from 'vitest'; +import { PuterServer } from '../../server.ts'; +import { setupTestServer } from '../../testUtil.ts'; +import type { SystemKVStore } from '../systemKv/SystemKVStore.ts'; +import { + bucketTag, + flattenAmounts, + pairsToAmounts, + parsePendingEntry, + unflattenAmounts, + type MeteringBufferStore, +} from './MeteringBufferStore.ts'; + +describe('MeteringBufferStore', () => { + // -- Pure shape helpers ------------------------------------------- + + describe('amount shapes', () => { + it('reads a hash reply into amounts', () => { + expect(pairsToAmounts(['total', '5', 'a.units', '2.5'])).toEqual({ + total: 5, + 'a.units': 2.5, + }); + }); + + it('nests flat paths the way stored counters are nested', () => { + expect( + unflattenAmounts({ + total: 7, + 'openai:gpt-4o.units': 12, + 'openai:gpt-4o.cost': 3, + }), + ).toEqual({ + total: 7, + 'openai:gpt-4o': { units: 12, cost: 3 }, + }); + }); + + it('round-trips a stored counter through flatten and back', () => { + const stored = { + total: 42, + 'ai:chat': { units: 100, cost: 42, count: 3 }, + }; + expect(unflattenAmounts(flattenAmounts(stored))).toEqual(stored); + }); + + it('keeps fractional amounts intact', () => { + expect(unflattenAmounts(flattenAmounts({ total: 0.25 }))).toEqual({ + total: 0.25, + }); + }); + + it('ignores non-numeric leaves when flattening', () => { + expect( + flattenAmounts({ total: 1, label: 'nope', nested: null }), + ).toEqual({ total: 1 }); + }); + }); + + describe('pending index entries', () => { + it('round-trips a key that contains separators', () => { + const key = 'metering:actor:abc-123:app:os-global:2026-08'; + const parsed = parsePendingEntry(`1750000000000:${key}`); + expect(parsed).toEqual({ claimedAt: 1750000000000, key }); + }); + + it('rejects entries it cannot trust', () => { + expect(parsePendingEntry('garbage')).toBeNull(); + expect(parsePendingEntry('notanumber:key')).toBeNull(); + expect(parsePendingEntry('123:')).toBeNull(); + }); + }); + + // -- Behaviour against a booted backend --------------------------- + + let server: PuterServer; + let target: MeteringBufferStore; + let kv: SystemKVStore; + + /** What the KV store itself holds, bypassing the buffer entirely. */ + const storedTotal = async (key: string): Promise => { + const { res } = await kv.get({ key }); + return (res as { total?: number } | null)?.total ?? 0; + }; + + const cacheKeys = async (): Promise => + (await server.clients.redis.keys('meter:*')).sort(); + + beforeAll(async () => { + server = await setupTestServer(); + target = server.stores.meteringBuffer; + kv = server.stores.kv; + // Stop the background flush loop so every test below drives flushing + // explicitly. Otherwise a cycle firing mid-test would settle counters + // the test is asserting are still buffered. The loop gets its own + // server further down. + await target.onServerShutdown(); + }); + + afterAll(async () => { + await server?.shutdown(); + }); + + // A fresh counter per test, and a clean cache, so nothing leaks between + // tests through either layer. + let key: string; + beforeEach(async () => { + key = `metering:actor:buf-${Math.random().toString(36).slice(2)}:2026-08`; + const stale = await server.clients.redis.keys('meter:*'); + if (stale.length) await server.clients.redis.del(...stale); + }); + + describe('incr', () => { + it('returns a running total without writing on every call', async () => { + const first = await target.incr({ + key, + pathAndAmountMap: { total: 10, 'ai:chat.count': 1 }, + }); + const second = await target.incr({ + key, + pathAndAmountMap: { total: 5, 'ai:chat.count': 1 }, + }); + + expect(first.res).toEqual({ total: 10, 'ai:chat': { count: 1 } }); + expect(second.res).toEqual({ total: 15, 'ai:chat': { count: 2 } }); + expect(first.exact).toBe(false); + expect(await storedTotal(key)).toBe(0); + }); + + it('writes the accumulated total onward on flush', async () => { + await target.incr({ key, pathAndAmountMap: { total: 10 } }); + await target.incr({ key, pathAndAmountMap: { total: 5 } }); + + await target.flushCycle(); + + expect(await storedTotal(key)).toBe(15); + }); + + it('collapses many increments into a single write', async () => { + const incrSpy = vi.spyOn(kv, 'incr'); + for (let i = 0; i < 5; i++) { + await target.incr({ + key, + pathAndAmountMap: { total: 2, 'ai:chat.count': 1 }, + }); + } + expect(incrSpy).not.toHaveBeenCalled(); + + await target.flushCycle(); + + expect(incrSpy).toHaveBeenCalledTimes(1); + const { res } = await kv.get({ key }); + expect(res).toEqual({ total: 10, 'ai:chat': { count: 5 } }); + incrSpy.mockRestore(); + }); + + it('counts what is already stored when it first sees a counter', async () => { + await kv.incr({ key, pathAndAmountMap: { total: 80 } }); + + const { res } = await target.incr({ + key, + pathAndAmountMap: { total: 5 }, + }); + + expect((res as { total: number }).total).toBe(85); + }); + + it('keeps fractional costs exact through a flush', async () => { + await target.incr({ key, pathAndAmountMap: { total: 0.25 } }); + const { res } = await target.incr({ + key, + pathAndAmountMap: { total: 0.5 }, + }); + + expect((res as { total: number }).total).toBe(0.75); + await target.flushCycle(); + expect(await storedTotal(key)).toBe(0.75); + }); + + it('records the increment even when the cache is unreachable', async () => { + const boom = vi + .spyOn( + server.clients.redis as unknown as { + meterIncr: () => Promise; + }, + 'meterIncr', + ) + .mockRejectedValue(new Error('cache down')); + + const { res, exact } = await target.incr({ + key, + pathAndAmountMap: { total: 9 }, + }); + + // Written straight through, so it is authoritative and exact. + expect(exact).toBe(true); + expect((res as { total: number }).total).toBe(9); + expect(await storedTotal(key)).toBe(9); + boom.mockRestore(); + }); + }); + + describe('incrAux', () => { + it('defers aggregate writes to the flush loop', async () => { + await target.incrAux({ key, pathAndAmountMap: { total: 7 } }); + expect(await storedTotal(key)).toBe(0); + + await target.flushCycle(); + expect(await storedTotal(key)).toBe(7); + }); + + it('records the aggregate even when the cache is unreachable', async () => { + const boom = vi + .spyOn( + server.clients.redis as unknown as { + meterIncr: () => Promise; + }, + 'meterIncr', + ) + .mockRejectedValue(new Error('cache down')); + + await target.incrAux({ key, pathAndAmountMap: { total: 4 } }); + + expect(await storedTotal(key)).toBe(4); + boom.mockRestore(); + }); + }); + + describe('get', () => { + it('reads back increments that have not been written onward', async () => { + await target.incr({ key, pathAndAmountMap: { total: 12 } }); + + const { res } = await target.get({ key }); + expect((res as { total: number }).total).toBe(12); + expect(await storedTotal(key)).toBe(0); + }); + + it('reads an unknown counter as absent', async () => { + const { res } = await target.get({ key }); + expect(res).toBeNull(); + }); + + it('sees a counter written straight to the store', async () => { + await kv.incr({ key, pathAndAmountMap: { total: 3 } }); + const { res } = await target.get({ key }); + expect((res as { total: number }).total).toBe(3); + }); + + it('keeps array reads aligned with the keys asked for', async () => { + const other = `${key}:other`; + await target.incr({ key, pathAndAmountMap: { total: 1 } }); + + const { res } = await target.get({ key: [key, other] }); + expect(Array.isArray(res)).toBe(true); + expect((res as [{ total: number }, null])[0].total).toBe(1); + expect((res as [unknown, null])[1]).toBeNull(); + }); + + it('does not mark a counter for flushing just by reading it', async () => { + await target.get({ key }); + expect(await cacheKeys()).toEqual([]); + }); + + it('is consistent with the total incr returned', async () => { + const { res: fromIncr } = await target.incr({ + key, + pathAndAmountMap: { total: 6, 'ai:chat.units': 2 }, + }); + const { res: fromGet } = await target.get({ key }); + expect(fromGet).toEqual(fromIncr); + }); + }); + + describe('readExact', () => { + it('writes buffered increments onward before reading back', async () => { + await target.incr({ key, pathAndAmountMap: { total: 33 } }); + expect(await storedTotal(key)).toBe(0); + + const { res } = await target.readExact({ key }); + + expect((res as { total: number }).total).toBe(33); + expect(await storedTotal(key)).toBe(33); + }); + + it('surfaces usage another deployment recorded', async () => { + await target.incr({ key, pathAndAmountMap: { total: 10 } }); + await target.flushCycle(); + + // Stand in for another deployment flushing the same counter; the + // buffered view has no way to know about it. + await kv.incr({ key, pathAndAmountMap: { total: 40 } }); + + const { res: buffered } = await target.get({ key }); + expect((buffered as { total: number }).total).toBe(10); + + const { res: exact } = await target.readExact({ key }); + expect((exact as { total: number }).total).toBe(50); + }); + + it('reads an untouched counter without inventing one', async () => { + const { res } = await target.readExact({ key }); + expect(res).toBeNull(); + }); + }); + + describe('flush claim', () => { + it('hands a delta to exactly one of two concurrent flushes', async () => { + await target.incrAux({ key, pathAndAmountMap: { total: 10 } }); + + await Promise.all([target.flushCycle(), target.flushCycle()]); + + expect(await storedTotal(key)).toBe(10); + }); + + it('keeps increments that arrive after a flush', async () => { + await target.incrAux({ key, pathAndAmountMap: { total: 10 } }); + await target.flushCycle(); + await target.incrAux({ key, pathAndAmountMap: { total: 4 } }); + await target.flushCycle(); + + expect(await storedTotal(key)).toBe(14); + }); + + it('leaves nothing buffered once a cycle has drained', async () => { + await target.incrAux({ key, pathAndAmountMap: { total: 1 } }); + await target.flushCycle(); + + const tag = bucketTag(key); + expect( + await server.clients.redis.exists(`meter:d:{${tag}}:${key}`), + ).toBe(0); + expect( + await server.clients.redis.scard(`meter:dirty:{${tag}}`), + ).toBe(0); + expect( + await server.clients.redis.hgetall(`meter:pending:{${tag}}`), + ).toEqual({}); + }); + + it('maintains the base from what the store returned', async () => { + await target.incr({ key, pathAndAmountMap: { total: 10 } }); + await target.flushCycle(); + + const tag = bucketTag(key); + expect( + await server.clients.redis.hget( + `meter:b:{${tag}}:${key}`, + 'total', + ), + ).toBe('10'); + }); + + it('does not let the base total move backwards', async () => { + await target.incr({ key, pathAndAmountMap: { total: 100 } }); + await target.flushCycle(); + + const tag = bucketTag(key); + const base = `meter:b:{${tag}}:${key}`; + // A settle carrying an older, smaller view must not win. + await server.clients.redis.hset( + `meter:p:{${tag}}:stalenonce`, + 'total', + '1', + ); + await server.clients.redis.hset( + `meter:pending:{${tag}}`, + 'stalenonce', + `${Date.now() - 60_000}:${key}`, + ); + await target.flushCycle(); + + expect(Number(await server.clients.redis.hget(base, 'total'))).toBe( + 101, + ); + }); + + it('flushes many counters in one cycle', async () => { + const keys = Array.from( + { length: 25 }, + (_, i) => `${key}-many-${i}`, + ); + for (const k of keys) { + await target.incrAux({ + key: k, + pathAndAmountMap: { total: 2 }, + }); + } + + await target.flushCycle(); + + for (const k of keys) expect(await storedTotal(k)).toBe(2); + }); + + it('groups every key it touches into one hash slot', async () => { + await target.incrAux({ key, pathAndAmountMap: { total: 1 } }); + const expected = `{${bucketTag(key)}}`; + + const keys = await cacheKeys(); + expect(keys.length).toBeGreaterThan(0); + for (const cacheKey of keys) { + expect(cacheKey).toContain(expected); + } + }); + + it('reports nothing to do on an idle cycle', async () => { + expect(await target.flushCycle()).toBe(0); + }); + }); + + describe('flush timer', () => { + // Its own backend, because this is the one test that needs the flush + // loop the server installs at boot to be left running. + it('writes buffered counters onward on its own', async () => { + const own = await setupTestServer(); + try { + const timerKey = `metering:actor:timer-${Math.random() + .toString(36) + .slice(2)}:2026-08`; + await own.stores.meteringBuffer.incrAux({ + key: timerKey, + pathAndAmountMap: { total: 6 }, + }); + + // Nothing here calls flushCycle — the loop has to do it. + await vi.waitFor( + async () => { + const { res } = await own.stores.kv.get({ + key: timerKey, + }); + expect((res as { total?: number } | null)?.total).toBe( + 6, + ); + }, + { timeout: 15_000, interval: 250 }, + ); + } finally { + await own.shutdown(); + } + }, 25_000); + }); + + describe('orphan recovery', () => { + it('re-drives a claim whose flush never finished', async () => { + const tag = bucketTag(key); + const nonce = 'orphanednonce'; + + await server.clients.redis.hset( + `meter:p:{${tag}}:${nonce}`, + 'total', + '25', + ); + await server.clients.redis.hset( + `meter:pending:{${tag}}`, + nonce, + `${Date.now() - 60_000}:${key}`, + ); + + await target.flushCycle(); + + expect(await storedTotal(key)).toBe(25); + expect( + await server.clients.redis.exists(`meter:p:{${tag}}:${nonce}`), + ).toBe(0); + expect( + await server.clients.redis.hgetall(`meter:pending:{${tag}}`), + ).toEqual({}); + }); + + it('leaves a claim that is still fresh alone', async () => { + const tag = bucketTag(key); + await server.clients.redis.hset( + `meter:p:{${tag}}:freshnonce`, + 'total', + '25', + ); + await server.clients.redis.hset( + `meter:pending:{${tag}}`, + 'freshnonce', + `${Date.now()}:${key}`, + ); + + await target.flushCycle(); + + expect(await storedTotal(key)).toBe(0); + }); + + it('drops an index entry whose data is gone', async () => { + const tag = bucketTag(key); + await server.clients.redis.hset( + `meter:pending:{${tag}}`, + 'lostnonce', + `${Date.now() - 60_000}:${key}`, + ); + + await target.flushCycle(); + + expect( + await server.clients.redis.hgetall(`meter:pending:{${tag}}`), + ).toEqual({}); + }); + + it('leaves the claim in place when the write onward fails', async () => { + await target.incr({ key, pathAndAmountMap: { total: 5 } }); + const boom = vi + .spyOn(kv, 'incr') + .mockRejectedValue(new Error('store unavailable')); + + await target.flushCycle(); + boom.mockRestore(); + + // The delta was claimed but never written, so it must still be + // recoverable rather than silently dropped. + const tag = bucketTag(key); + const pending = await server.clients.redis.hgetall( + `meter:pending:{${tag}}`, + ); + expect(Object.keys(pending)).toHaveLength(1); + expect(await storedTotal(key)).toBe(0); + }); + }); + + describe('month boundaries', () => { + it('keeps the counter for each month independent', async () => { + const august = `${key}-2026-08`; + const september = `${key}-2026-09`; + + await target.incr({ key: august, pathAndAmountMap: { total: 10 } }); + const carried = await target.incr({ + key: september, + pathAndAmountMap: { total: 3 }, + }); + + // A new month is a new counter, so nothing carries over. + expect((carried.res as { total: number }).total).toBe(3); + + await target.flushCycle(); + expect(await storedTotal(august)).toBe(10); + expect(await storedTotal(september)).toBe(3); + }); + + it('starts a fresh base when the month rolls over mid-flight', async () => { + const august = `${key}-2026-08`; + const september = `${key}-2026-09`; + + await target.incr({ key: august, pathAndAmountMap: { total: 10 } }); + await target.flushCycle(); + await target.incr({ + key: september, + pathAndAmountMap: { total: 1 }, + }); + await target.flushCycle(); + + expect( + await server.clients.redis.hget( + `meter:b:{${bucketTag(august)}}:${august}`, + 'total', + ), + ).toBe('10'); + expect( + await server.clients.redis.hget( + `meter:b:{${bucketTag(september)}}:${september}`, + 'total', + ), + ).toBe('1'); + }); + }); +}); diff --git a/src/backend/stores/metering/MeteringBufferStore.ts b/src/backend/stores/metering/MeteringBufferStore.ts new file mode 100644 index 000000000..b87896e72 --- /dev/null +++ b/src/backend/stores/metering/MeteringBufferStore.ts @@ -0,0 +1,667 @@ +/* + * Copyright (C) 2024-present Puter Technologies Inc. + * + * This file is part of Puter. + * + * Puter is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published + * by the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +import { randomUUID } from 'node:crypto'; +import murmurhash from 'murmurhash'; +import type { RecursiveRecord } from '../systemKv/SystemKVStore'; +import { PuterStore } from '../types'; + +// -- Types ------------------------------------------------------------ + +/** Matches the operation shape of `SystemKVStore.incr`. */ +interface IncrInput { + key: string; + pathAndAmountMap: Record; +} + +type FlatAmounts = Record; + +// -- Constants -------------------------------------------------------- + +/** + * Buffered keys are grouped into buckets, and the bucket is the hash tag, so a + * counter's delta, its base, and its bucket's bookkeeping all live in one slot + * and a single script can touch them together. Tagging on the counter's own key + * instead would put the per-bucket sets in a different slot from the deltas + * they track, which a clustered cache refuses. + * + * A counter always maps to exactly one bucket — buckets group different + * counters together, they never split one. + */ +const BUCKET_COUNT = 64; + +/** + * How long an increment can sit in the cache before it is written onward. This + * is the durability bound: a host lost without warning takes at most this much + * usage with it. + */ +const FLUSH_INTERVAL_MS = 5000; + +/** + * Deployments start at fixed offsets, so without jitter two that boot near each + * other would contend on every cycle rather than occasionally. + */ +const FLUSH_JITTER_MS = 500; + +/** Per bucket, per cycle. Anything above this flushes on the next cycle. */ +const CLAIMS_PER_BUCKET = 100; + +/** A claim left unfinished this long is assumed abandoned and re-driven. */ +const ORPHAN_AGE_MS = 30_000; + +/** Long enough that a counter for the current month can never expire. */ +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; + +/** Roughly a minute between compression reports, so they don't flood the log. */ +const CYCLES_PER_COMPRESSION_REPORT = 12; + +// -- Keys ------------------------------------------------------------- + +export const bucketTag = (key: string): string => + `m${murmurhash.v3(key) % BUCKET_COUNT}`; + +const deltaKey = (tag: string, key: string): string => + `meter:d:{${tag}}:${key}`; +const baseKey = (tag: string, key: string): string => `meter:b:{${tag}}:${key}`; +const dirtyKey = (tag: string): string => `meter:dirty:{${tag}}`; +const pendingKey = (tag: string, nonce: string): string => + `meter:p:{${tag}}:${nonce}`; +const pendingIndexKey = (tag: string): string => `meter:pending:{${tag}}`; + +// -- Shape helpers ---------------------------------------------------- + +/** `['a', '1', 'b', '2']` (a cache hash reply) to `{ a: 1, b: 2 }`. */ +export const pairsToAmounts = (pairs: string[]): FlatAmounts => { + const out: FlatAmounts = {}; + for (let i = 0; i < pairs.length - 1; i += 2) { + out[pairs[i]!] = Number(pairs[i + 1]); + } + return out; +}; + +/** + * Cache hashes are flat, but stored counters nest on `.` the same way `incr` + * paths do — so `{'a.units': 1}` becomes `{a: {units: 1}}`. + */ +export const unflattenAmounts = ( + flat: FlatAmounts, +): RecursiveRecord => { + const out: Record = {}; + for (const [path, amount] of Object.entries(flat)) { + const parts = path.split('.'); + let node = out; + for (const part of parts.slice(0, -1)) { + if (typeof node[part] !== 'object' || node[part] === null) + node[part] = {}; + node = node[part] as Record; + } + node[parts[parts.length - 1]!] = amount; + } + return out as RecursiveRecord; +}; + +/** Inverse of `unflattenAmounts`, for values read back from the KV store. */ +export const flattenAmounts = (value: unknown): FlatAmounts => { + const out: FlatAmounts = {}; + const walk = (node: unknown, prefix: string): void => { + if (typeof node === 'number') { + if (prefix) out[prefix] = node; + return; + } + if (!node || typeof node !== 'object' || Array.isArray(node)) return; + for (const [k, v] of Object.entries(node)) { + walk(v, prefix ? `${prefix}.${k}` : k); + } + }; + walk(value, ''); + return out; +}; + +const addAmounts = (a: FlatAmounts, b: FlatAmounts): FlatAmounts => { + const out: FlatAmounts = { ...a }; + for (const [path, amount] of Object.entries(b)) { + out[path] = (out[path] ?? 0) + amount; + } + return out; +}; + +const toScriptArgs = (amounts: Record): string[] => { + const args: string[] = []; + for (const [path, amount] of Object.entries(amounts)) { + args.push(path, String(amount)); + } + return args; +}; + +// -- Scripts ---------------------------------------------------------- + +/** + * KEYS: delta, base, dirty set. ARGV: ttl, member, then path/amount pairs. + * + * Returns the delta and the base together so a running total can be served + * without a second round trip. Amounts are floats — costs are fractional. + */ +const INCR_SCRIPT = ` +for i = 3, #ARGV, 2 do + redis.call('HINCRBYFLOAT', KEYS[1], ARGV[i], ARGV[i + 1]) +end +redis.call('PEXPIRE', KEYS[1], ARGV[1]) +redis.call('SADD', KEYS[3], ARGV[2]) +redis.call('PEXPIRE', KEYS[3], ARGV[1]) +return { redis.call('HGETALL', KEYS[1]), redis.call('HGETALL', KEYS[2]) } +`; + +/** KEYS: delta, base. Reads both without marking the counter for flushing. */ +const READ_SCRIPT = ` +return { redis.call('HGETALL', KEYS[1]), redis.call('HGETALL', KEYS[2]) } +`; + +/** + * KEYS: delta, pending, pending index. ARGV: nonce, index value, ttl. + * + * The rename is the claim, and it is atomic: if two flushes race for one delta, + * one takes all of it and the other sees nothing. Increments arriving mid-flush + * start a fresh delta and go out on the next cycle. + */ +const CLAIM_SCRIPT = ` +if redis.call('EXISTS', KEYS[1]) == 0 then return nil end +redis.call('RENAME', KEYS[1], KEYS[2]) +redis.call('PEXPIRE', KEYS[2], ARGV[3]) +redis.call('HSET', KEYS[3], ARGV[1], ARGV[2]) +redis.call('PEXPIRE', KEYS[3], ARGV[3]) +return redis.call('HGETALL', KEYS[2]) +`; + +/** + * KEYS: base, pending, pending index. ARGV: nonce, ttl, new total (or ''), then + * the authoritative path/value pairs. + * + * Replaces the base with what the KV store now holds, which is how the base + * picks up other deployments' contributions without a separate read. The total + * may not move backwards: two flushes settling out of order would otherwise + * briefly under-report, and this total decides whether someone may spend. + */ +const SETTLE_SCRIPT = ` +local current = redis.call('HGET', KEYS[1], 'total') +local replace = true +if ARGV[3] ~= '' and current then + if tonumber(current) > tonumber(ARGV[3]) then replace = false end +end +if replace then + redis.call('DEL', KEYS[1]) + for i = 4, #ARGV, 2 do + redis.call('HSET', KEYS[1], ARGV[i], ARGV[i + 1]) + end + redis.call('PEXPIRE', KEYS[1], ARGV[2]) +end +redis.call('DEL', KEYS[2]) +redis.call('HDEL', KEYS[3], ARGV[1]) +return 1 +`; + +/** KEYS: base. ARGV: ttl, then path/value pairs. Returns the base. */ +const SEED_SCRIPT = ` +if redis.call('EXISTS', KEYS[1]) == 0 then + for i = 2, #ARGV, 2 do + redis.call('HSET', KEYS[1], ARGV[i], ARGV[i + 1]) + end + redis.call('PEXPIRE', KEYS[1], ARGV[1]) +end +return redis.call('HGETALL', KEYS[1]) +`; + +type ScriptRunner = { + meterIncr(...args: string[]): Promise<[string[], string[]]>; + meterRead(...args: string[]): Promise<[string[], string[]]>; + meterClaim(...args: string[]): Promise; + meterSettle(...args: string[]): Promise; + meterSeed(...args: string[]): Promise; +}; + +// -- MeteringBufferStore ---------------------------------------------- + +/** + * Absorbs metering counter increments in the cache and writes each counter + * onward once per cycle instead of once per call. + * + * This stands in for the `incr` and `get` operations of `stores.kv` for + * metering counters, and it is the only place that knows any buffering happens: + * callers pass the same arguments and read the same running totals either way. + * + * The running total a caller gets back is this deployment's view — its own + * buffered increments plus everything the KV store held as of the last cycle. + * For an actor served by a single deployment that is exact at all times. When a + * decision genuinely turns on the number, `readExact` is the way to be sure. + */ +export class MeteringBufferStore extends PuterStore { + #flushTimer: ReturnType | null = null; + #stopped = false; + #definedScripts = false; + #absorbedCount = 0; + #flushedCount = 0; + #cyclesSinceReport = 0; + + // -- Lifecycle ---------------------------------------------------- + + override onServerStart(): void { + this.#defineScripts(); + this.#scheduleFlush(); + } + + override async onServerShutdown(): Promise { + this.#stopped = true; + if (this.#flushTimer) { + clearTimeout(this.#flushTimer); + this.#flushTimer = null; + } + + // Best-effort drain. The flush interval, not this, is what bounds how + // much can be lost — a host that dies gives no signal at all. + try { + for (let pass = 0; pass < 10; pass++) { + if ((await this.flushCycle()) === 0) break; + } + } catch (e) { + console.warn('[metering] shutdown flush failed', e); + } + } + + // -- Public API --------------------------------------------------- + + /** + * Increment a counter whose running total the caller acts on. Returns the + * counter's value after the increment, as `stores.kv.incr` does, plus + * whether that value is known to account for every deployment. + */ + async incr( + input: IncrInput, + ): Promise<{ res: RecursiveRecord; exact: boolean }> { + try { + const { delta, base } = await this.#buffer(input); + const resolved = + Object.keys(base).length > 0 + ? base + : await this.#seedBase(input.key); + return { + res: unflattenAmounts(addAmounts(resolved, delta)), + exact: false, + }; + } catch (e) { + // A counter that decides whether someone may spend is never worth + // dropping, so fall back to writing it directly. That value is + // authoritative, hence exact. + console.warn( + `[metering] buffered incr failed, writing through: ${(e as Error).message}`, + ); + const { res } = await this.#writeThrough(input); + return { res, exact: true }; + } + } + + /** + * Increment an aggregate counter that only feeds reporting. Nothing reads + * these to make a decision, so there is no running total to return. + */ + async incrAux(input: IncrInput): Promise { + try { + await this.#buffer(input); + } catch (e) { + console.warn( + `[metering] buffered aux incr failed, writing through: ${(e as Error).message}`, + ); + await this.#writeThrough(input); + } + } + + /** + * Read counters, including increments buffered but not yet written onward. + * Mirrors `stores.kv.get`: one key in, one value out; an array in, an array + * out. + */ + async get({ + key, + }: { + key: string | string[]; + }): Promise<{ res: unknown | null | (unknown | null)[] }> { + const keys = Array.isArray(key) ? key : [key]; + const values = await Promise.all( + keys.map((k) => + this.#readThrough(k).catch((e: Error) => { + console.warn( + `[metering] buffered read failed, reading through: ${e.message}`, + ); + return this.stores.kv + .get({ key: k }) + .then(({ res }) => res ?? null); + }), + ), + ); + return { res: Array.isArray(key) ? values : values[0]! }; + } + + /** + * Read a counter with everything this deployment has buffered for it + * written onward first, then read back strongly consistently. This is the + * only read that costs extra; it exists for decisions taken close enough to + * a limit that an approximate total would be the wrong answer. + */ + async readExact({ key }: { key: string }): Promise<{ res: unknown }> { + try { + await this.#flushOne(bucketTag(key), key); + } catch (e) { + console.warn( + `[metering] exact read could not flush ${key}: ${(e as Error).message}`, + ); + } + return this.stores.kv.get({ key, consistentRead: true }); + } + + // -- Internals: client & scripts ---------------------------------- + + get #redis(): ScriptRunner { + this.#defineScripts(); + return this.clients.redis as unknown as ScriptRunner; + } + + #defineScripts(): void { + if (this.#definedScripts) return; + this.#definedScripts = true; + const client = this.clients.redis; + client.defineCommand('meterIncr', { + numberOfKeys: 3, + lua: INCR_SCRIPT, + }); + client.defineCommand('meterRead', { + numberOfKeys: 2, + lua: READ_SCRIPT, + }); + client.defineCommand('meterClaim', { + numberOfKeys: 3, + lua: CLAIM_SCRIPT, + }); + client.defineCommand('meterSettle', { + numberOfKeys: 3, + lua: SETTLE_SCRIPT, + }); + client.defineCommand('meterSeed', { + numberOfKeys: 1, + lua: SEED_SCRIPT, + }); + } + + // -- Internals: writes -------------------------------------------- + + #writeThrough(input: IncrInput): Promise<{ res: RecursiveRecord }> { + return this.stores.kv.incr(input) as Promise<{ + res: RecursiveRecord; + }>; + } + + async #buffer({ + key, + pathAndAmountMap, + }: IncrInput): Promise<{ delta: FlatAmounts; base: FlatAmounts }> { + const tag = bucketTag(key); + const [delta, base] = await this.#redis.meterIncr( + deltaKey(tag, key), + baseKey(tag, key), + dirtyKey(tag), + String(BUFFER_TTL_MS), + key, + ...toScriptArgs(pathAndAmountMap), + ); + this.#absorbedCount++; + return { delta: pairsToAmounts(delta), base: pairsToAmounts(base) }; + } + + /** + * The base mirrors what the KV store holds, so a counter the cache has not + * seen yet — a new month, or one whose entry has since gone — has to be + * fetched once. From the first flush onward the base is maintained from + * what the KV store returns, and increments are served from the cache + * alone. A counter with nothing stored against it yet seeds nothing, so it + * keeps reading for that first cycle; reads are the cheap direction and it + * is over within seconds. + */ + async #seedBase(key: string): Promise { + const { res } = await this.stores.kv.get({ key }); + const persisted = flattenAmounts(res); + const seeded = await this.#redis.meterSeed( + baseKey(bucketTag(key), key), + String(BUFFER_TTL_MS), + ...toScriptArgs(persisted), + ); + return pairsToAmounts(seeded); + } + + // -- Internals: reads --------------------------------------------- + + async #readThrough(key: string): Promise { + const tag = bucketTag(key); + const [delta, base] = await this.#redis.meterRead( + deltaKey(tag, key), + baseKey(tag, key), + ); + const deltaAmounts = pairsToAmounts(delta); + const baseAmounts = pairsToAmounts(base); + const resolved = + Object.keys(baseAmounts).length > 0 + ? baseAmounts + : flattenAmounts((await this.stores.kv.get({ key })).res); + + const merged = addAmounts(resolved, deltaAmounts); + if (Object.keys(merged).length === 0) return null; + return unflattenAmounts(merged); + } + + // -- Internals: flush --------------------------------------------- + + #scheduleFlush(): void { + if (this.#stopped) return; + const delay = + FLUSH_INTERVAL_MS + (Math.random() * 2 - 1) * FLUSH_JITTER_MS; + this.#flushTimer = setTimeout(() => { + this.flushCycle() + .catch((e) => { + console.error('[metering] flush cycle failed', e); + }) + .finally(() => this.#scheduleFlush()); + }, delay); + this.#flushTimer.unref?.(); + } + + /** + * Write one cycle's worth of buffered counters onward. Driven by the flush + * timer; returns how many counters it handled so a drain loop knows when + * there is nothing left. + */ + async flushCycle(): Promise { + const work: Array<() => Promise> = []; + let truncated = 0; + + const buckets = await Promise.all( + Array.from({ length: BUCKET_COUNT }, (_, bucket) => + this.#drainBucket(`m${bucket}`), + ), + ); + + for (const bucket of buckets) { + if (bucket.truncated) truncated++; + for (const key of bucket.keys) { + work.push(() => this.#flushOne(bucket.tag, key)); + } + for (const orphan of bucket.orphans) { + work.push(() => this.#settleOrphan(bucket.tag, orphan)); + } + } + + if (truncated > 0) { + console.warn( + `[metering] ${truncated} bucket(s) hit the per-cycle claim cap; the rest flush next cycle`, + ); + } + + for (let i = 0; i < work.length; i += SETTLE_CONCURRENCY) { + await Promise.all( + work.slice(i, i + SETTLE_CONCURRENCY).map((run) => + run().catch((e: Error) => { + // Leaving the claim in place is what makes this safe: + // the sweep re-drives it. + console.warn(`[metering] settle failed: ${e.message}`); + }), + ), + ); + } + + this.#reportCompression(work.length); + return work.length; + } + + /** + * How many increments each write onward stood in for. This is the number + * that says whether buffering is earning anything: a ratio near 1 means + * calls for the same counter arrive too far apart to batch. + */ + #reportCompression(flushed: number): void { + this.#flushedCount += flushed; + if (++this.#cyclesSinceReport < CYCLES_PER_COMPRESSION_REPORT) return; + + const absorbed = this.#absorbedCount; + const writes = this.#flushedCount; + this.#absorbedCount = 0; + this.#flushedCount = 0; + this.#cyclesSinceReport = 0; + + if (writes === 0) return; + console.log( + `[metering] buffer absorbed ${absorbed} increments into ${writes} writes (${(absorbed / writes).toFixed(2)}x)`, + ); + } + + async #drainBucket(tag: string): Promise<{ + tag: string; + keys: string[]; + orphans: Array<{ nonce: string; key: string }>; + truncated: boolean; + }> { + const redis = this.clients.redis; + const [popped, pending] = await Promise.all([ + redis.spop(dirtyKey(tag), CLAIMS_PER_BUCKET), + redis.hgetall(pendingIndexKey(tag)), + ]); + + // An empty pop can come back as nothing at all rather than an empty + // list, depending on the client. + const keys = popped ?? []; + + const cutoff = Date.now() - ORPHAN_AGE_MS; + const orphans: Array<{ nonce: string; key: string }> = []; + for (const [nonce, encoded] of Object.entries(pending ?? {})) { + const parsed = parsePendingEntry(encoded); + if (!parsed || parsed.claimedAt > cutoff) continue; + orphans.push({ nonce, key: parsed.key }); + } + + return { + tag, + keys, + orphans, + truncated: keys.length >= CLAIMS_PER_BUCKET, + }; + } + + async #flushOne(tag: string, key: string): Promise { + const nonce = randomUUID().replace(/-/g, ''); + const claimed = await this.#redis.meterClaim( + deltaKey(tag, key), + pendingKey(tag, nonce), + pendingIndexKey(tag), + nonce, + encodePendingEntry(Date.now(), key), + String(BUFFER_TTL_MS), + ); + // Nothing buffered for this counter — another flush already took it. + if (!claimed) return; + + await this.#settle(tag, key, nonce, pairsToAmounts(claimed)); + } + + async #settleOrphan( + tag: string, + orphan: { nonce: string; key: string }, + ): Promise { + const pairs = await this.clients.redis.hgetall( + pendingKey(tag, orphan.nonce), + ); + const amounts: FlatAmounts = {}; + for (const [path, amount] of Object.entries(pairs ?? {})) { + amounts[path] = Number(amount); + } + if (Object.keys(amounts).length === 0) { + // The claimed data is gone but its index entry outlived it. + await this.clients.redis.hdel(pendingIndexKey(tag), orphan.nonce); + return; + } + await this.#settle(tag, orphan.key, orphan.nonce, amounts); + } + + async #settle( + tag: string, + key: string, + nonce: string, + amounts: FlatAmounts, + ): Promise { + const { res } = await this.stores.kv.incr({ + key, + pathAndAmountMap: amounts, + }); + + const flat = flattenAmounts(res); + await this.#redis.meterSettle( + baseKey(tag, key), + pendingKey(tag, nonce), + pendingIndexKey(tag), + nonce, + String(BUFFER_TTL_MS), + flat['total'] === undefined ? '' : String(flat['total']), + ...toScriptArgs(flat), + ); + } +} + +// -- Pending index encoding ------------------------------------------- + +const encodePendingEntry = (claimedAt: number, key: string): string => + `${claimedAt}:${key}`; + +export const parsePendingEntry = ( + encoded: string, +): { claimedAt: number; key: string } | null => { + const separator = encoded.indexOf(':'); + if (separator < 0) return null; + + const claimedAt = Number(encoded.slice(0, separator)); + const key = encoded.slice(separator + 1); + if (!Number.isFinite(claimedAt) || !key) return null; + + return { claimedAt, key }; +}; diff --git a/src/backend/stores/systemKv/SystemKVStore.ts b/src/backend/stores/systemKv/SystemKVStore.ts index 27346dbc7..be65acd2f 100644 --- a/src/backend/stores/systemKv/SystemKVStore.ts +++ b/src/backend/stores/systemKv/SystemKVStore.ts @@ -267,7 +267,10 @@ export class SystemKVStore extends PuterStore { // -- Public API --------------------------------------------------- async get( - { key }: { key: string | string[] }, + { + key, + consistentRead, + }: { key: string | string[]; consistentRead?: boolean }, opts?: KVOpts, ): Promise> { const actor = ensureActor(opts); @@ -289,10 +292,11 @@ export class SystemKVStore extends PuterStore { kvEntries = entries; usage = u; } else { - const response = await this.clients.dynamo.get(this.tableName, { - namespace, - key, - }); + const response = await this.clients.dynamo.get( + this.tableName, + { namespace, key }, + consistentRead, + ); kvEntries = response.Item ? [response.Item as (typeof kvEntries)[number]] : [];