diff --git a/src/__tests__/usage-flush.test.ts b/src/__tests__/usage-flush.test.ts new file mode 100644 index 0000000..886760c --- /dev/null +++ b/src/__tests__/usage-flush.test.ts @@ -0,0 +1,57 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; + +// Unmock usage middleware — we want to test the real implementation +vi.unmock("../middleware/usage.js"); + +import { connectWithRetry } from "../services/db.js"; + +describe("flushDirtyEntries – independent key flushing", () => { + let usageMod: typeof import("../middleware/usage.js"); + + beforeEach(async () => { + vi.clearAllMocks(); + vi.resetModules(); + // Re-import to get fresh state + usageMod = await import("../middleware/usage.js"); + }); + + it("should flush remaining keys even if one key fails", async () => { + // Set up two keys in the in-memory cache via the middleware + const next = vi.fn(); + const res = { status: vi.fn(() => ({ json: vi.fn() })) }; + + usageMod.usageMiddleware({ apiKeyInfo: { key: "key-good" } }, res, next); + usageMod.usageMiddleware({ apiKeyInfo: { key: "key-bad" } }, res, next); + + // Track which keys were successfully upserted + const flushedKeys: string[] = []; + let callCount = 0; + + const mockQuery = vi.fn().mockImplementation((sql: string, params?: any[]) => { + if (sql.includes("INSERT INTO usage")) { + callCount++; + if (params && params[0] === "key-bad") { + throw new Error("simulated constraint violation"); + } + flushedKeys.push(params![0]); + } + return { rows: [], rowCount: 0 }; + }); + + const mockRelease = vi.fn(); + + // Each call to connectWithRetry returns a fresh client + vi.mocked(connectWithRetry).mockImplementation(async () => ({ + query: mockQuery, + release: mockRelease, + }) as any); + + // Access the flush function (exported for testing) + await usageMod.flushDirtyEntries(); + + // The good key should have been flushed despite the bad key failing + expect(flushedKeys).toContain("key-good"); + // Release should be called for each key (independent clients) + expect(mockRelease.mock.calls.length).toBeGreaterThanOrEqual(2); + }); +}); diff --git a/src/middleware/usage.ts b/src/middleware/usage.ts index df11d79..3f9a428 100644 --- a/src/middleware/usage.ts +++ b/src/middleware/usage.ts @@ -36,45 +36,37 @@ export async function loadUsageData(): Promise { } // Batch flush dirty entries to DB (Audit #10 + #12) -async function flushDirtyEntries(): Promise { +export async function flushDirtyEntries(): Promise { if (dirtyKeys.size === 0) return; const keysToFlush = [...dirtyKeys]; - - const client = await connectWithRetry(); - try { - await client.query("BEGIN"); - for (const key of keysToFlush) { - const record = usage.get(key); - if (!record) continue; - try { - await client.query( - `INSERT INTO usage (key, count, month_key) VALUES ($1, $2, $3) - ON CONFLICT (key) DO UPDATE SET count = $2, month_key = $3`, - [key, record.count, record.monthKey] - ); + + for (const key of keysToFlush) { + const record = usage.get(key); + if (!record) continue; + const client = await connectWithRetry(); + try { + await client.query( + `INSERT INTO usage (key, count, month_key) VALUES ($1, $2, $3) + ON CONFLICT (key) DO UPDATE SET count = $2, month_key = $3`, + [key, record.count, record.monthKey] + ); + dirtyKeys.delete(key); + retryCount.delete(key); + } catch (error) { + // Audit #12: retry logic for failed writes + const retries = (retryCount.get(key) || 0) + 1; + if (retries >= MAX_RETRIES) { + logger.error({ key: key.slice(0, 8) + "...", retries }, "CRITICAL: Usage write failed after max retries, data may diverge"); dirtyKeys.delete(key); retryCount.delete(key); - } catch (error) { - // Audit #12: retry logic for failed writes - const retries = (retryCount.get(key) || 0) + 1; - if (retries >= MAX_RETRIES) { - logger.error({ key: key.slice(0, 8) + "...", retries }, "CRITICAL: Usage write failed after max retries, data may diverge"); - dirtyKeys.delete(key); - retryCount.delete(key); - } else { - retryCount.set(key, retries); - logger.warn({ key: key.slice(0, 8) + "...", retries }, "Usage write failed, will retry"); - } + } else { + retryCount.set(key, retries); + logger.warn({ key: key.slice(0, 8) + "...", retries }, "Usage write failed, will retry"); } + } finally { + client.release(); } - await client.query("COMMIT"); - } catch (error) { - await client.query("ROLLBACK").catch(() => {}); - logger.error({ err: error }, "Failed to flush usage batch"); - // Keep all keys dirty for retry - } finally { - client.release(); } }