fix: metering buffer improvements for manual editing

Author: Salazareo

src/backend/services/metering/MeteringService.test.ts

@@ -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,
@@ -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),
@@ -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 () => {

src/backend/services/metering/MeteringService.ts

@@ -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);
         }
 

src/backend/stores/metering/MeteringBufferStore.test.ts

@@ -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',
@@ -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 },

src/backend/stores/metering/MeteringBufferStore.ts

@@ -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}`;
@@ -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])
@@ -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 ----------------------------------------------
@@ -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 {
@@ -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', {
@@ -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),
         );