mirror of
https://github.com/HeyPuter/puter.git
synced 2026-08-24 06:58:21 +00:00
perf: improve metering perf + add extra compat for monthly charges (#3513)
closes PUT-1445 and PUT-1446
This commit is contained in:
@@ -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];
|
||||
|
||||
@@ -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<T extends keyof EventMap>(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.
|
||||
*
|
||||
|
||||
@@ -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 string> =
|
||||
K extends `${infer Head}.${infer Tail}`
|
||||
? | `${Head}.*`
|
||||
| (Tail extends `${string}.${string}`
|
||||
? `${Head}.${WildcardPrefixes<Tail>}`
|
||||
: never)
|
||||
?
|
||||
| `${Head}.*`
|
||||
| (Tail extends `${string}.${string}`
|
||||
? `${Head}.${WildcardPrefixes<Tail>}`
|
||||
: never)
|
||||
: never;
|
||||
|
||||
export type ListenKey = EventKey | WildcardPrefixes<EventKey>;
|
||||
|
||||
@@ -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<string, unknown>)
|
||||
.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<void>;
|
||||
|
||||
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<string> }
|
||||
).settledActors.clear();
|
||||
|
||||
const registered: ChargeListener[] = [];
|
||||
const listen = (fn: ChargeListener) => {
|
||||
server.clients.event.on(
|
||||
'metering.monthly.charges',
|
||||
fn as Parameters<typeof server.clients.event.on>[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<typeof server.clients.event.off>[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', () => {
|
||||
|
||||
@@ -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> | 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<typeof setInterval> | 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<string>();
|
||||
/**
|
||||
* 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<string>();
|
||||
|
||||
// -- 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<string, AppTotals> | 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<string, AppTotals> = {};
|
||||
@@ -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<UsageByType> {
|
||||
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<unknown>): 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<UsageByType | null> {
|
||||
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<UsageByType | null> {
|
||||
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<string, unknown>)?.[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,
|
||||
|
||||
@@ -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';
|
||||
|
||||
|
||||
@@ -32,9 +32,22 @@ export interface UsageRecord {
|
||||
units: number;
|
||||
}
|
||||
|
||||
export type UsageByType = { total: number } & Partial<
|
||||
Record<Exclude<string, 'total'>, 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<Record<Exclude<string, 'total'>, UsageRecord>>;
|
||||
|
||||
export interface AppTotals {
|
||||
total: number;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 <https://www.gnu.org/licenses/>.
|
||||
*/
|
||||
|
||||
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<number> => {
|
||||
const { res } = await kv.get({ key });
|
||||
return (res as { total?: number } | null)?.total ?? 0;
|
||||
};
|
||||
|
||||
const cacheKeys = async (): Promise<string[]> =>
|
||||
(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<never>;
|
||||
},
|
||||
'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<never>;
|
||||
},
|
||||
'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');
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -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 <https://www.gnu.org/licenses/>.
|
||||
*/
|
||||
|
||||
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<string, number>;
|
||||
}
|
||||
|
||||
type FlatAmounts = Record<string, number>;
|
||||
|
||||
// -- 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<number> => {
|
||||
const out: Record<string, unknown> = {};
|
||||
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<string, unknown>;
|
||||
}
|
||||
node[parts[parts.length - 1]!] = amount;
|
||||
}
|
||||
return out as RecursiveRecord<number>;
|
||||
};
|
||||
|
||||
/** 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, number>): 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<string[] | null>;
|
||||
meterSettle(...args: string[]): Promise<number>;
|
||||
meterSeed(...args: string[]): Promise<string[]>;
|
||||
};
|
||||
|
||||
// -- 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<typeof setTimeout> | null = null;
|
||||
#stopped = false;
|
||||
#definedScripts = false;
|
||||
#absorbedCount = 0;
|
||||
#flushedCount = 0;
|
||||
#cyclesSinceReport = 0;
|
||||
|
||||
// -- Lifecycle ----------------------------------------------------
|
||||
|
||||
override onServerStart(): void {
|
||||
this.#defineScripts();
|
||||
this.#scheduleFlush();
|
||||
}
|
||||
|
||||
override async onServerShutdown(): Promise<void> {
|
||||
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<number>; 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<void> {
|
||||
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<number> }> {
|
||||
return this.stores.kv.incr(input) as Promise<{
|
||||
res: RecursiveRecord<number>;
|
||||
}>;
|
||||
}
|
||||
|
||||
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<FlatAmounts> {
|
||||
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<unknown | null> {
|
||||
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<number> {
|
||||
const work: Array<() => Promise<void>> = [];
|
||||
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<void> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
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 };
|
||||
};
|
||||
@@ -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<KVResult<unknown | null | (unknown | null)[]>> {
|
||||
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]]
|
||||
: [];
|
||||
|
||||
Reference in New Issue
Block a user