feat(quota): QuotaStore.poolConsumedTotal — pool aggregate per dimension

Add poolConsumedTotal(poolId, dim) to the QuotaStore interface and implement
it in both SqliteQuotaStore and RedisQuotaStore so the enforce path can read
the real pool-wide consumption (sum across all keys) rather than the per-key
saturation signal, which is 0 for countable units and therefore never blocks.

- db/quotaConsumption.ts: add sumPoolDimension() — single SQL COALESCE(SUM)
  for curr and prev buckets across all api_key_id rows for a dimensionKey.
- localDb.ts: re-export sumPoolDimension.
- quota/types.ts: add poolConsumedTotal to QuotaStore interface (with JSDoc).
- sqliteQuotaStore.ts: implement using sumPoolDimension + existing
  slidingWindowEffective helper — one consistent sliding-window read.
- redisQuotaStore.ts: implement by fetching pool allocations from SQLite (F2),
  then issuing a single MGET for all (key, curr-bucket) and (key, prev-bucket)
  Redis keys, summing raw values, and applying the sliding-window formula once.
- tests/unit/quota-store-pool-total.test.ts: 4 tests (sum, pool isolation,
  unit isolation, peek-no-regression) all passing.
This commit is contained in:
diegosouzapw
2026-05-31 17:57:02 -03:00
parent be77a03aa0
commit abb77879db
6 changed files with 246 additions and 0 deletions

View File

@@ -200,6 +200,58 @@ export function listConsumptionForPool(poolId: string, limit: number): Consumpti
});
}
// ---------------------------------------------------------------------------
// Pool-wide aggregate
// ---------------------------------------------------------------------------
interface BucketPairRow {
api_key_id: string;
curr: number;
prev: number;
}
/**
* Sum the consumed values for a given (dimensionKey, currentBucketIndex) and
* (dimensionKey, currentBucketIndex - 1) across ALL api_key_id values.
*
* Returns { currTotal, prevTotal } so the caller can apply the sliding-window
* formula once with the pool-wide totals instead of summing per-key.
*
* @param dimensionKey "<poolId>:<unit>:<window>" string — same format as consume/peek.
* @param currentBucket floor(nowMs / windowMs) — caller must pass the same value.
*/
export function sumPoolDimension(
dimensionKey: string,
currentBucket: number
): { currTotal: number; prevTotal: number } {
const prevBucket = currentBucket - 1;
interface SumRow {
total: number;
}
const currRow = getDb()
.prepare<SumRow>(
`SELECT COALESCE(SUM(consumed), 0) AS total
FROM quota_consumption
WHERE dimension_key = ? AND bucket_index = ?`
)
.get(dimensionKey, currentBucket);
const prevRow = getDb()
.prepare<SumRow>(
`SELECT COALESCE(SUM(consumed), 0) AS total
FROM quota_consumption
WHERE dimension_key = ? AND bucket_index = ?`
)
.get(dimensionKey, prevBucket);
return {
currTotal: currRow?.total ?? 0,
prevTotal: prevRow?.total ?? 0,
};
}
// ---------------------------------------------------------------------------
// GC
// ---------------------------------------------------------------------------

View File

@@ -549,6 +549,7 @@ export {
getBucket,
incrementBucket,
getPair,
sumPoolDimension,
gcOlderThan as gcQuotaConsumption,
} from "./db/quotaConsumption";
export {

View File

@@ -165,6 +165,51 @@ export class RedisQuotaStore implements QuotaStore {
return slidingWindowEffective(curr, prev, nowMs, windowMs);
}
/**
* Return the real pool-wide consumption for a dimension in the current
* sliding window, summed across ALL apiKeyIds in the pool's allocations.
*
* Strategy: pool/allocation metadata lives in SQLite (F2), so we fetch the
* allocation list for the pool, then issue a single MGET for all
* (apiKeyId, curr-bucket) and (apiKeyId, prev-bucket) keys. The Redis keys
* are per-key counters — there is no pool-level aggregate key — so we read
* each key's two buckets and sum the raw values before applying the
* sliding-window formula once on the totals.
*
* If the pool does not exist or has no allocations, returns 0.
* Keys absent in Redis are treated as 0 (MGET returns null).
*/
async poolConsumedTotal(poolId: string, dim: DimensionKey): Promise<number> {
const pool = getPool(poolId);
if (!pool || pool.allocations.length === 0) return 0;
const nowMs = Date.now();
const dimKey = dimensionKeyToString(dim);
const windowMs = WINDOW_MS[dim.window];
const currentBucket = Math.floor(nowMs / windowMs);
const prevBucket = currentBucket - 1;
const client = await this.client();
// Build [currKey0, prevKey0, currKey1, prevKey1, ...] for all allocated keys
const redisKeys: string[] = [];
for (const alloc of pool.allocations) {
redisKeys.push(bucketKey(alloc.apiKeyId, dimKey, currentBucket));
redisKeys.push(bucketKey(alloc.apiKeyId, dimKey, prevBucket));
}
const values = await client.mget(...redisKeys);
let currTotal = 0;
let prevTotal = 0;
for (let i = 0; i < values.length; i += 2) {
currTotal += parseFloat(values[i] ?? "0") || 0;
prevTotal += parseFloat(values[i + 1] ?? "0") || 0;
}
return slidingWindowEffective(currTotal, prevTotal, nowMs, windowMs);
}
/**
* Aggregate pool usage. Pool and allocation metadata come from SQLite (F2);
* rolling counters come from Redis.

View File

@@ -21,6 +21,7 @@ import {
getBucket,
incrementBucket,
getPair,
sumPoolDimension,
} from "@/lib/localDb";
import { WINDOW_MS, dimensionKeyToString } from "./dimensions";
import type { DimensionKey } from "./dimensions";
@@ -112,6 +113,24 @@ export class SqliteQuotaStore implements QuotaStore {
return slidingWindowEffective(curr, prev, nowMs, windowMs);
}
/**
* Return the real pool-wide consumption for a dimension in the current
* sliding window, summed across ALL apiKeyIds that share the same
* dimensionKey (i.e. same poolId + unit + window).
*
* Uses the same 2-bucket sliding-window formula as peek(), applied once
* to the pool totals so the result is consistent with per-key semantics.
*/
async poolConsumedTotal(poolId: string, dim: DimensionKey): Promise<number> {
const nowMs = Date.now();
const dimKey = dimensionKeyToString(dim);
const windowMs = WINDOW_MS[dim.window];
const currentBucket = Math.floor(nowMs / windowMs);
const { currTotal, prevTotal } = sumPoolDimension(dimKey, currentBucket);
return slidingWindowEffective(currTotal, prevTotal, nowMs, windowMs);
}
/**
* Return a PoolUsageSnapshot for the given pool, aggregating per-key
* consumption across all dimensions and computing fairShare / deficit /

View File

@@ -34,6 +34,16 @@ export interface ConsumeResult {
export interface QuotaStore {
consume(apiKeyId: string, dim: DimensionKey, cost: number): Promise<number>;
peek(apiKeyId: string, dim: DimensionKey): Promise<number>;
/**
* Return the real pool-wide consumption for a dimension in the current
* sliding window — i.e. the sum of each key's effective consumption across
* ALL apiKeyIds that have contributed to (poolId, unit, window).
*
* Unlike the per-key saturation signal (which can be 0 for countable units
* whose hard-cap has never been set), this reflects actual spent units so
* the enforce path can block when the pool total hits the plan limit.
*/
poolConsumedTotal(poolId: string, dim: DimensionKey): Promise<number>;
poolUsage(poolId: string): Promise<PoolUsageSnapshot>;
clear(apiKeyId: string, dim: DimensionKey): Promise<void>;
}

View File

@@ -0,0 +1,119 @@
/**
* tests/unit/quota-store-pool-total.test.ts
*
* Coverage for QuotaStore.poolConsumedTotal():
* - Two keys in same pool → sum equals individual contributions.
* - Different pool or different dim → 0 (isolation).
* - peek() for a single key still returns only that key's own value (no regression).
*/
import test from "node:test";
import assert from "node:assert/strict";
import fs from "node:fs";
import os from "node:os";
import path from "node:path";
const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-pool-total-"));
process.env.DATA_DIR = TEST_DATA_DIR;
const core = await import("../../src/lib/db/core.ts");
async function resetStorage() {
core.resetDbInstance();
for (let attempt = 0; attempt < 10; attempt++) {
try {
if (fs.existsSync(TEST_DATA_DIR)) {
fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true });
}
break;
} catch (err: unknown) {
const e = err as { code?: string };
if ((e?.code === "EBUSY" || e?.code === "EPERM") && attempt < 9) {
await new Promise((resolve) => setTimeout(resolve, 50 * (attempt + 1)));
} else {
throw err;
}
}
}
fs.mkdirSync(TEST_DATA_DIR, { recursive: true });
}
test.beforeEach(async () => {
await resetStorage();
});
test.after(async () => {
core.resetDbInstance();
if (fs.existsSync(TEST_DATA_DIR)) {
fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true });
}
});
// ─── Main scenario: two keys in same pool ────────────────────────────────────
test("poolConsumedTotal: keyA(3) + keyB(2) in same pool → total === 5", async () => {
const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts");
const store = new SqliteQuotaStore();
const dim = { poolId: "p1", unit: "requests" as const, window: "5h" as const };
await store.consume("keyA", dim, 3);
await store.consume("keyB", dim, 2);
const total = await store.poolConsumedTotal("p1", dim);
// Both consumes happen in the same bucket (milliseconds apart) so
// prev=0, elapsed≈0 → total ≈ 5. Allow small delta for timing.
assert.ok(total > 4.9, `Expected >4.9, got ${total}`);
assert.ok(total <= 5, `Expected <=5, got ${total}`);
});
// ─── Isolation: different pool ───────────────────────────────────────────────
test("poolConsumedTotal: different poolId → returns 0 (isolation)", async () => {
const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts");
const store = new SqliteQuotaStore();
const dimP1 = { poolId: "p1", unit: "requests" as const, window: "5h" as const };
const dimP2 = { poolId: "p2", unit: "requests" as const, window: "5h" as const };
await store.consume("keyA", dimP1, 10);
// Pool p2 has no consumption → should be 0
const total = await store.poolConsumedTotal("p2", dimP2);
assert.equal(total, 0, `Expected 0 for different pool, got ${total}`);
});
// ─── Isolation: different dimension (unit) ───────────────────────────────────
test("poolConsumedTotal: different unit dim → returns 0 (isolation)", async () => {
const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts");
const store = new SqliteQuotaStore();
const dimRequests = { poolId: "p1", unit: "requests" as const, window: "5h" as const };
const dimTokens = { poolId: "p1", unit: "tokens" as const, window: "5h" as const };
await store.consume("keyA", dimRequests, 7);
// tokens dim has no consumption → should be 0
const total = await store.poolConsumedTotal("p1", dimTokens);
assert.equal(total, 0, `Expected 0 for different unit, got ${total}`);
});
// ─── No regression: peek still returns per-key value ─────────────────────────
test("poolConsumedTotal: peek(keyA) still returns only keyA's own consumption", async () => {
const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts");
const store = new SqliteQuotaStore();
const dim = { poolId: "p1", unit: "requests" as const, window: "5h" as const };
await store.consume("keyA", dim, 3);
await store.consume("keyB", dim, 2);
const peekA = await store.peek("keyA", dim);
// keyA's own peek should be ≈3, not 5
assert.ok(peekA > 2.9, `Expected >2.9, got ${peekA}`);
assert.ok(peekA <= 3, `Expected <=3, got ${peekA}`);
});