Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 39 additions & 3 deletions src/backend/services/metering/MeteringService.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import {
import type { Actor } from '../../core/actor.ts';
import { SYSTEM_ACTOR } from '../../core/actor.ts';
import { PuterServer } from '../../server.ts';
import { bucketTag } from '../../stores/metering/MeteringBufferStore.ts';
import { setupTestServer } from '../../testUtil.ts';
import {
DEFAULT_FREE_SUBSCRIPTION,
Expand Down Expand Up @@ -1245,6 +1246,43 @@ describe('MeteringService', () => {
);
});

it('stays set once the adjustment has been written onward', async () => {
await target.incrementUsage(actor, 'kv:read', 1, 10_000);
await server.stores.meteringBuffer.flushCycle();

await target.setActorCurrentMonthUsageTotal(actor, 0);
// The adjustment is buffered like any other amount, so the read
// that matters is the one after it has settled — a correction that
// only holds until then is a correction nobody keeps.
await server.stores.meteringBuffer.flushCycle();

const { usage } =
await target.getActorCurrentMonthUsageDetails(actor);
expect(usage.total).toBe(0);
expect(usage.allowanceUsed).toBe(0);
});

it('repairs a cached view that has drifted from the record', async () => {
await target.incrementUsage(actor, 'kv:read', 1, 10_000);
await server.stores.meteringBuffer.flushCycle();

// Whatever the drift came from, re-applying the total the record
// already holds is the support-facing repair for it, so it has to
// take even though there is nothing to write.
const key = `${METRICS_PREFIX}:actor:${actor.user.uuid}:${new Date().toISOString().slice(0, 7)}`;
await server.clients.redis.hset(
`meter:b:{${bucketTag(key)}}:${key}`,
'total',
'999999',
);

await target.setActorCurrentMonthUsageTotal(actor, 10_000);

const { usage } =
await target.getActorCurrentMonthUsageDetails(actor);
expect(usage.total).toBe(10_000);
});

it('rejects a negative total', async () => {
await expect(
target.setActorCurrentMonthUsageTotal(actor, -1),
Expand Down Expand Up @@ -1448,9 +1486,7 @@ describe('MeteringService', () => {
});

const allowed = await target.getAllowedUsage(actor);
expect(allowed.remaining).toBe(
sub.monthUsageAllowance - 1_000_000,
);
expect(allowed.remaining).toBe(sub.monthUsageAllowance - 1_000_000);
});

it('folds the legacy baseline in exactly once under concurrent increments', async () => {
Expand Down
8 changes: 8 additions & 0 deletions src/backend/services/metering/MeteringService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -944,7 +944,15 @@ export class MeteringService extends PuterService {
subscription.monthUsageAllowance,
);

// The record already reads as asked, so there is nothing to write — but
// an adjustment is also how a cached view that has drifted from the
// record gets repaired, and answering "already correct" from the record
// while readers keep being told something else is how that drift
// survives being corrected at all. Drop the view either way.
await this.stores.meteringBuffer.forgetBase(actorUsageKey);

if (delta === 0 && allowanceUsedDelta === 0) {
this.invalidateActorCredits(userId);
return (current as UsageByType) || ({ total: 0 } as UsageByType);
}

Expand Down
66 changes: 64 additions & 2 deletions src/backend/stores/metering/MeteringBufferStore.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -503,13 +503,14 @@ describe('MeteringBufferStore', () => {
).toBe('10');
});

it('does not let the base total move backwards', async () => {
it('adopts the view of a re-driven claim it wrote onward', 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.
// A claim another deployment left behind, re-driven here: its
// amounts reach the store, so the base takes what came back.
await server.clients.redis.hset(
`meter:p:{${tag}}:stalenonce`,
'total',
Expand All @@ -527,6 +528,67 @@ describe('MeteringBufferStore', () => {
);
});

it('does not let a settle replace a base laid down after it', async () => {
await target.incr({ key, pathAndAmountMap: { total: 100 } });
await target.flushCycle();

const tag = bucketTag(key);
const base = `meter:b:{${tag}}:${key}`;
// Stand in for a settle that finished later: this one's view of the
// store is the older of the two however large its total is.
await server.clients.redis.set(
`meter:bq:{${tag}}:${key}`,
'1000000000',
);

await target.incr({ key, pathAndAmountMap: { total: 5 } });
await target.flushCycle();

expect(await storedTotal(key)).toBe(105);
expect(Number(await server.clients.redis.hget(base, 'total'))).toBe(
100,
);
});

it('seeds a forgotten base from the store, keeping what is buffered', async () => {
await target.incr({ key, pathAndAmountMap: { total: 100 } });
await target.flushCycle();
await target.incr({ key, pathAndAmountMap: { total: 5 } });

const tag = bucketTag(key);
// Stand in for a cached view that no longer matches the record.
await server.clients.redis.hset(
`meter:b:{${tag}}:${key}`,
'total',
'999',
);
await target.forgetBase(key);

const { res } = await target.get({ key });
expect((res as { total: number }).total).toBe(105);
});

it('lets a correction take the base down', async () => {
// The counter is authoritative in both directions: an amount can be
// corrected downwards, and reads have to follow it down rather than
// answer with the number the correction replaced.
await target.incr({
key,
pathAndAmountMap: { total: 100, allowanceUsed: 100 },
});
await target.flushCycle();

await target.incr({
key,
pathAndAmountMap: { total: -100, allowanceUsed: -100 },
});
await target.flushCycle();

expect(await storedTotal(key)).toBe(0);
const { res } = await target.get({ key });
expect(res).toEqual({ total: 0, allowanceUsed: 0 });
});

it('flushes many counters in one cycle', async () => {
const keys = Array.from(
{ length: 25 },
Expand Down
84 changes: 74 additions & 10 deletions src/backend/stores/metering/MeteringBufferStore.ts
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,27 @@ export const bucketTag = (key: string): string =>
const deltaKey = (tag: string, key: string): string =>
`meter:d:{${tag}}:${key}`;
const baseKey = (tag: string, key: string): string => `meter:b:{${tag}}:${key}`;
/**
* Which settle a counter's base came from — see `seqKey`. Kept beside the base
* rather than in it, so the base holds amounts and nothing else: a bookkeeping
* field inside it would read back as a counter path of its own.
*/
const baseSeqKey = (tag: string, key: string): string =>
`meter:bq:{${tag}}:${key}`;
/**
* Hands out the ordering token a settle carries, one sequence per bucket.
*
* A settle replaces the base with what the store returned, so of two settles
* for the same counter the one that finished writing last is the one holding
* the newer view — and that is the only thing the base has to be ordered by.
* The token is taken the moment the write comes back, so the order the tokens
* are in is the order the writes completed in.
*
* Deliberately without a TTL: it orders every settle the bucket will ever do,
* and one that started over would hand out tokens the stamps already written
* are ahead of, leaving those bases in place until it caught up again.
*/
const seqKey = (tag: string): string => `meter:q:{${tag}}`;
const dirtyKey = (tag: string): string => `meter:dirty:{${tag}}`;
const pendingKey = (tag: string, nonce: string): string =>
`meter:p:{${tag}}:${nonce}`;
Expand Down Expand Up @@ -398,30 +419,37 @@ return { 1, redis.call('HGETALL', KEYS[2]) }
`;

/**
* KEYS: base, pending, pending index, delta, tracked set. ARGV: nonce, ttl, new
* total (or ''), member, then the authoritative path/value pairs.
* KEYS: base, pending, pending index, delta, tracked set, base sequence. ARGV:
* nonce, ttl, sequence, member, 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.
* picks up other deployments' contributions without a separate read. Two
* flushes settling out of order must not leave the older view in place, so a
* settle only replaces a base laid down by a settle that finished before it —
* ordered by the sequence each took when its write came back.
*
* Ordering by sequence rather than by which view holds the larger total is what
* lets a counter go down at all: an amount can be corrected downwards, and
* comparing totals reads that correction as the stale view it must refuse,
* pinning the base to the pre-correction number for as long as it lives.
*
* The counter leaves the tracked set here, but only if nothing has started a
* fresh delta for it in the meantime — checking and removing in the same step
* is what stops an increment that landed mid-settle from being forgotten.
*/
const SETTLE_SCRIPT = `
local current = redis.call('HGET', KEYS[1], 'total')
local current = redis.call('GET', KEYS[6])
local replace = true
if ARGV[3] ~= '' and current then
if tonumber(current) > tonumber(ARGV[3]) then replace = false end
if current and tonumber(current) > tonumber(ARGV[3]) then
replace = false
end
if replace then
redis.call('DEL', KEYS[1])
for i = 5, #ARGV, 2 do
redis.call('HSET', KEYS[1], ARGV[i], ARGV[i + 1])
end
redis.call('PEXPIRE', KEYS[1], ARGV[2])
redis.call('SET', KEYS[6], ARGV[3], 'PX', ARGV[2])
end
redis.call('DEL', KEYS[2])
redis.call('HDEL', KEYS[3], ARGV[1])
Expand Down Expand Up @@ -505,6 +533,7 @@ type ScriptRunner = {
meterRetire(...args: string[]): Promise<number>;
meterReconcile(...args: string[]): Promise<number>;
meterSeed(...args: string[]): Promise<string[]>;
incr(key: string): Promise<number>;
};

// -- MeteringBufferStore ----------------------------------------------
Expand Down Expand Up @@ -650,6 +679,35 @@ export class MeteringBufferStore extends PuterStore {
return this.stores.kv.get({ key, consistentRead: true });
}

/**
* Forget the cached view of what the store holds for a counter, so the next
* read seeds it from the store again.
*
* For a counter corrected outside this buffer's own accounting — an
* adjustment applied to the record itself rather than metered onto it —
* where the cached view would otherwise keep answering with what it had
* until the correction settles. Buffered increments are deliberately left
* alone: they are amounts the store hasn't seen yet, and the seeded view is
* what they are added to.
*
* Never throws. The correction is in the store either way; failing the call
* that made it over a cache that is about to be replaced anyway would be
* the worse outcome.
*/
async forgetBase(key: string): Promise<void> {
const tag = bucketTag(key);
try {
await this.clients.redis.del(
baseKey(tag, key),
baseSeqKey(tag, key),
);
} catch (e) {
console.warn(
`[metering] cached base not dropped for ${key}: ${(e as Error).message}`,
);
}
}

// -- Internals: client & scripts ----------------------------------

get #redis(): ScriptRunner {
Expand Down Expand Up @@ -682,7 +740,7 @@ export class MeteringBufferStore extends PuterStore {
lua: RECLAIM_SCRIPT,
});
client.defineCommand('meterSettle', {
numberOfKeys: 5,
numberOfKeys: 6,
lua: SETTLE_SCRIPT,
});
client.defineCommand('meterRetire', {
Expand Down Expand Up @@ -1113,16 +1171,22 @@ export class MeteringBufferStore extends PuterStore {
return;
}

// Taken here rather than before the writes: what the base has to be
// ordered by is which settle came away with the newer view of the
// store, and that is decided by the write that just returned.
const seq = await this.#redis.incr(seqKey(tag));

const flat = flattenAmounts(settled);
await this.#redis.meterSettle(
baseKey(tag, key),
pendingKey(tag, nonce),
pendingIndexKey(tag),
deltaKey(tag, key),
trackedKey(tag),
baseSeqKey(tag, key),
nonce,
String(BUFFER_TTL_MS),
flat['total'] === undefined ? '' : String(flat['total']),
String(seq),
key,
...toScriptArgs(flat),
);
Expand Down
Loading