mirror of
https://github.com/HeyPuter/puter.git
synced 2026-11-05 05:24:13 +00:00
fix(metering): count usage mid-flush in balance reads
A flush takes a counter's buffered amounts aside before writing them onward, and reads only looked at the live delta and the cached base. From the claim until the settle replaced the base, those amounts were missing from every read: one store write normally, up to the orphan sweep if the write failed. Balance checks, credit reservations among them, could see an account as having more than it does. Reads now add the counter's outstanding claim. A batch written onward leaves the claim and joins the cached base in one step, so it stays counted until the settle replaces the base with the store's view. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
1 parent
8416d9738d
commit
a64dcfe322
2 files changed
+140
-9
No files matched your search
@@ -420,6 +420,75 @@ describe('MeteringBufferStore', () => {
|
||||
expect(await storedTotal(key)).toBe(0);
|
||||
});
|
||||
|
||||
/** Holds `method` on `target` until the returned release is called. */
|
||||
const holdCalls = <T extends object>(on: T, method: keyof T) => {
|
||||
let release: () => void = () => {};
|
||||
const gate = new Promise<void>((r) => (release = r));
|
||||
const real = (on[method] as (...a: unknown[]) => unknown).bind(on);
|
||||
const spy = vi
|
||||
.spyOn(on, method as never)
|
||||
.mockImplementation((async (...args: unknown[]) => {
|
||||
await gate;
|
||||
return real(...args);
|
||||
}) as never);
|
||||
return { spy, release };
|
||||
};
|
||||
const readTotal = async () =>
|
||||
((await target.get({ key })).res as { total?: number } | null)
|
||||
?.total ?? 0;
|
||||
|
||||
it('counts usage a flush has claimed but not yet written onward', async () => {
|
||||
await target.incr({ key, pathAndAmountMap: { total: 10 } });
|
||||
const write = holdCalls(kv, 'incr');
|
||||
try {
|
||||
const flushing = target.flushCycle();
|
||||
await vi.waitFor(() => expect(write.spy).toHaveBeenCalled());
|
||||
expect(await readTotal()).toBe(10);
|
||||
write.release();
|
||||
await flushing;
|
||||
} finally {
|
||||
write.spy.mockRestore();
|
||||
}
|
||||
expect(await readTotal()).toBe(10);
|
||||
expect(await storedTotal(key)).toBe(10);
|
||||
});
|
||||
|
||||
it('counts usage written onward until its settle lands', async () => {
|
||||
await target.incr({ key, pathAndAmountMap: { total: 10 } });
|
||||
await target.flushCycle();
|
||||
await target.incr({ key, pathAndAmountMap: { total: 5 } });
|
||||
const redis = server.clients.redis as unknown as {
|
||||
meterSettle: (...args: string[]) => Promise<number>;
|
||||
};
|
||||
const settle = holdCalls(redis, 'meterSettle');
|
||||
try {
|
||||
const flushing = target.flushCycle();
|
||||
await vi.waitFor(() => expect(settle.spy).toHaveBeenCalled());
|
||||
expect(await storedTotal(key)).toBe(15);
|
||||
expect(await readTotal()).toBe(15);
|
||||
settle.release();
|
||||
await flushing;
|
||||
} finally {
|
||||
settle.spy.mockRestore();
|
||||
}
|
||||
expect(await readTotal()).toBe(15);
|
||||
});
|
||||
|
||||
it('counts a claim in array reads too', async () => {
|
||||
await target.incr({ key, pathAndAmountMap: { total: 4 } });
|
||||
const write = holdCalls(kv, 'incr');
|
||||
try {
|
||||
const flushing = target.flushCycle();
|
||||
await vi.waitFor(() => expect(write.spy).toHaveBeenCalled());
|
||||
const { res } = await target.get({ key: [key] });
|
||||
expect((res as [{ total: number }])[0].total).toBe(4);
|
||||
write.release();
|
||||
await flushing;
|
||||
} finally {
|
||||
write.spy.mockRestore();
|
||||
}
|
||||
});
|
||||
|
||||
it('reads an unknown counter as absent', async () => {
|
||||
const { res } = await target.get({ key });
|
||||
expect(res).toBeNull();
|
||||
@@ -906,12 +975,16 @@ describe('MeteringBufferStore', () => {
|
||||
'1000000000',
|
||||
);
|
||||
|
||||
// Another deployment's write, which only a replacing settle would
|
||||
// bring into the base.
|
||||
await kv.incr({ key, pathAndAmountMap: { total: 1000 } });
|
||||
await target.incr({ key, pathAndAmountMap: { total: 5 } });
|
||||
await target.flushCycle();
|
||||
|
||||
expect(await storedTotal(key)).toBe(105);
|
||||
expect(await storedTotal(key)).toBe(1105);
|
||||
// The newer base keeps its view, plus what this flush wrote.
|
||||
expect(Number(await server.clients.redis.hget(base, 'total'))).toBe(
|
||||
100,
|
||||
105,
|
||||
);
|
||||
});
|
||||
|
||||
@@ -1742,6 +1815,17 @@ describe('MeteringBufferStore', () => {
|
||||
order.push('dropFromClaim');
|
||||
return realHdel(...args);
|
||||
});
|
||||
// A written batch leaves the claim through this script.
|
||||
const redis = server.clients.redis as unknown as {
|
||||
meterApply: (...args: string[]) => Promise<number>;
|
||||
};
|
||||
const realApply = redis.meterApply.bind(redis);
|
||||
const applySpy = vi
|
||||
.spyOn(redis, 'meterApply')
|
||||
.mockImplementation((...args: string[]) => {
|
||||
order.push('dropFromClaim');
|
||||
return realApply(...args);
|
||||
});
|
||||
const alarmSpy = vi
|
||||
.spyOn(server.clients.alarm, 'create')
|
||||
.mockImplementation(() => {
|
||||
@@ -1751,6 +1835,7 @@ describe('MeteringBufferStore', () => {
|
||||
await target.flushCycle();
|
||||
incrSpy.mockRestore();
|
||||
hdelSpy.mockRestore();
|
||||
applySpy.mockRestore();
|
||||
alarmSpy.mockRestore();
|
||||
|
||||
expect(order).toEqual(['dropFromClaim', 'alarm']);
|
||||
|
||||
@@ -403,9 +403,39 @@ redis.call('PEXPIRE', KEYS[4], ARGV[1])
|
||||
return { redis.call('HGETALL', KEYS[1]), redis.call('HGETALL', KEYS[2]) }
|
||||
`;
|
||||
|
||||
/** KEYS: delta, base. Reads both without marking the counter for flushing. */
|
||||
/**
|
||||
* KEYS: delta, base, claimed. ARGV: counter, pending key prefix.
|
||||
*
|
||||
* Reads both without marking the counter for flushing, plus the counter's
|
||||
* outstanding claim: amounts a flush has taken but not yet written onward are
|
||||
* in neither the delta nor the base. The claim's key is built here, which is
|
||||
* safe for the reason given at `RECONCILE_SCRIPT`.
|
||||
*/
|
||||
const READ_SCRIPT = `
|
||||
return { redis.call('HGETALL', KEYS[1]), redis.call('HGETALL', KEYS[2]) }
|
||||
local claim = {}
|
||||
local nonce = redis.call('HGET', KEYS[3], ARGV[1])
|
||||
if nonce then claim = redis.call('HGETALL', ARGV[2] .. nonce) end
|
||||
return { redis.call('HGETALL', KEYS[1]), redis.call('HGETALL', KEYS[2]), claim }
|
||||
`;
|
||||
|
||||
/**
|
||||
* KEYS: pending, base. ARGV: path/amount pairs just written onward.
|
||||
*
|
||||
* Takes written paths off a claim and, when a base is cached, adds them to it
|
||||
* in the same step, so a read counts them all the way to the settle that
|
||||
* replaces the base with the store's own view. Without a base, reads go to the
|
||||
* store, which already has them.
|
||||
*/
|
||||
const APPLY_SCRIPT = `
|
||||
redis.call('HDEL', KEYS[1], '${INTENT_FIELD}')
|
||||
local based = redis.call('EXISTS', KEYS[2]) == 1
|
||||
for i = 1, #ARGV, 2 do
|
||||
redis.call('HDEL', KEYS[1], ARGV[i])
|
||||
if based then
|
||||
redis.call('HINCRBYFLOAT', KEYS[2], ARGV[i], ARGV[i + 1])
|
||||
end
|
||||
end
|
||||
return 1
|
||||
`;
|
||||
|
||||
/**
|
||||
@@ -653,7 +683,8 @@ return redis.call('HGETALL', KEYS[1])
|
||||
|
||||
type ScriptRunner = {
|
||||
meterIncr(...args: string[]): Promise<[string[], string[]]>;
|
||||
meterRead(...args: string[]): Promise<[string[], string[]]>;
|
||||
meterRead(...args: string[]): Promise<[string[], string[], string[]]>;
|
||||
meterApply(...args: string[]): Promise<number>;
|
||||
meterDrain(...args: string[]): Promise<[string[], number]>;
|
||||
meterClaim(...args: string[]): Promise<string[] | null>;
|
||||
/** `[1, amounts]` when re-claimed, `[0, hadIndexEntry]` when not. */
|
||||
@@ -1119,9 +1150,13 @@ export class MeteringBufferStore extends PuterStore {
|
||||
lua: INCR_SCRIPT,
|
||||
});
|
||||
client.defineCommand('meterRead', {
|
||||
numberOfKeys: 2,
|
||||
numberOfKeys: 3,
|
||||
lua: READ_SCRIPT,
|
||||
});
|
||||
client.defineCommand('meterApply', {
|
||||
numberOfKeys: 2,
|
||||
lua: APPLY_SCRIPT,
|
||||
});
|
||||
client.defineCommand('meterDrain', {
|
||||
numberOfKeys: 2,
|
||||
lua: DRAIN_SCRIPT,
|
||||
@@ -1216,11 +1251,18 @@ export class MeteringBufferStore extends PuterStore {
|
||||
key: string,
|
||||
): Promise<{ delta: FlatAmounts; base: FlatAmounts }> {
|
||||
const tag = bucketTag(key);
|
||||
const [delta, base] = await this.#redis.meterRead(
|
||||
const [delta, base, claim] = await this.#redis.meterRead(
|
||||
deltaKey(tag, key),
|
||||
baseKey(tag, key),
|
||||
claimedKey(tag),
|
||||
key,
|
||||
pendingKey(tag, ''),
|
||||
);
|
||||
return { delta: pairsToAmounts(delta), base: pairsToAmounts(base) };
|
||||
// Claimed amounts count like buffered ones until their settle.
|
||||
return {
|
||||
delta: addAmounts(pairsToAmounts(delta), parseClaim(claim).amounts),
|
||||
base: pairsToAmounts(base),
|
||||
};
|
||||
}
|
||||
|
||||
async #readThrough(key: string): Promise<unknown | null> {
|
||||
@@ -1618,7 +1660,11 @@ export class MeteringBufferStore extends PuterStore {
|
||||
// before anything else can fail: a re-drive then picks up only what
|
||||
// is still outstanding. The alarm below is the "anything else" — it
|
||||
// comes after.
|
||||
await this.#dropFromClaim(tag, nonce, paths);
|
||||
await this.#redis.meterApply(
|
||||
pendingKey(tag, nonce),
|
||||
baseKey(tag, key),
|
||||
...toScriptArgs(batch),
|
||||
);
|
||||
|
||||
if (usage.write >= LARGE_SETTLE_WRITE_UNITS) {
|
||||
this.#recordLargeSettle(key, usage.write);
|
||||
|
||||
Reference in new issue
Block a user