diff --git a/keys.mjs b/keys.mjs index c597fe7..7d6d3f9 100644 --- a/keys.mjs +++ b/keys.mjs @@ -371,6 +371,36 @@ export function getCacheStats() { return { entries: total, totalHits, sizeBytes }; } +// ── Singleflight stampede protection ── + +// In-memory singleflight Map: hash → { promise, requesters } +// Deduplicates concurrent identical cache-miss flows so only one upstream call runs. +// Per ADR 0005 / spec D4: in-process scope only (single Node process per host). +const inflightMap = new Map(); + +export function singleflight(hash, fn) { + const existing = inflightMap.get(hash); + if (existing) { + existing.requesters++; + return existing.promise; + } + // Wrap fn() in Promise.resolve().then() so synchronous throws don't escape. + const promise = Promise.resolve().then(fn).finally(() => { + inflightMap.delete(hash); + }); + inflightMap.set(hash, { promise, requesters: 1 }); + return promise; +} + +export function getInflightStats() { + let totalRequesters = 0; + for (const entry of inflightMap.values()) totalRequesters += entry.requesters; + return { + inflight: inflightMap.size, + requesters: totalRequesters, + }; +} + // Find a key by id or name (returns { id, name } or null) export function findKey(idOrName) { const d = getDb(); diff --git a/server.mjs b/server.mjs index 72c936d..11b0aae 100644 --- a/server.mjs +++ b/server.mjs @@ -34,7 +34,7 @@ import { readFileSync, accessSync, existsSync, constants } from "node:fs"; import { fileURLToPath } from "node:url"; import { dirname, join } from "node:path"; import { homedir } from "node:os"; -import { validateKey, recordUsage, getUsageByKey, getUsageTimeline, getRecentUsage, createKey, listKeys, revokeKey, closeDb, checkQuota, updateKeyQuota, getKeyQuota, findKey, cacheHash, getCachedResponse, setCachedResponse, clearCache, getCacheStats, hasCacheControl } from "./keys.mjs"; +import { validateKey, recordUsage, getUsageByKey, getUsageTimeline, getRecentUsage, createKey, listKeys, revokeKey, closeDb, checkQuota, updateKeyQuota, getKeyQuota, findKey, cacheHash, getCachedResponse, setCachedResponse, clearCache, getCacheStats, hasCacheControl, singleflight, getInflightStats } from "./keys.mjs"; const __dirname = dirname(fileURLToPath(import.meta.url)); const _pkg = JSON.parse(readFileSync(join(__dirname, "package.json"), "utf8")); @@ -590,6 +590,7 @@ function startHeartbeat(res, intervalMs, sessionId) { // ── Call claude CLI (real streaming) ───────────────────────────────────── // Pipes stdout from the claude process directly to SSE chunks as they arrive. // Each data chunk becomes a proper SSE event with delta content in real time. +// TODO(cache-singleflight-stream): streaming-path singleflight is out of scope for v3.13.0; see spec D4 streaming caveat. function callClaudeStreaming(model, messages, conversationId, res, authInfo = {}) { const id = `chatcmpl-${randomUUID()}`; const created = Math.floor(Date.now() / 1000); @@ -1265,14 +1266,45 @@ async function handleChatCompletions(req, res) { const t0Usage = Date.now(); const promptChars = messages.reduce((a, m) => a + (typeof m.content === "string" ? m.content.length : JSON.stringify(m.content).length), 0); + + // Non-streaming path with stampede protection: wrap the upstream call in singleflight + // when cache is enabled and a hash is present. Concurrent identical requests share + // one upstream spawn; followers receive the same promise. Streaming-path dedup is + // explicitly out of scope (see TODO comment above callClaudeStreaming). + if (CACHE_TTL > 0 && req._cacheHash) { + try { + const content = await singleflight(req._cacheHash, async () => { + // Re-check cache inside the singleflight: a follower that enters before the + // leader finishes will wait on the shared promise (not reach here), but a + // request that races in just after the previous singleflight cleared the map + // will re-read the freshly-populated cache entry here rather than spawning. + const recheck = getCachedResponse(req._cacheHash, CACHE_TTL); + if (recheck) return recheck.response; + const c = await callClaude(model, messages, conversationId); + try { setCachedResponse(req._cacheHash, model, c); } catch (e) { logEvent("error", "cache_write_failed", { error: e.message }); } + return c; + }); + const id = `chatcmpl-${randomUUID()}`; + completionResponse(res, id, model, content); + try { recordUsage({ keyId: req._authKeyId, keyName: req._authKeyName, model, promptChars, responseChars: content.length, elapsedMs: Date.now() - t0Usage, success: true }); } catch (e) { logEvent("error", "usage_record_failed", { error: e.message }); } + return; + } catch (err) { + try { recordUsage({ keyId: req._authKeyId, keyName: req._authKeyName, model, promptChars, responseChars: 0, elapsedMs: Date.now() - t0Usage, success: false }); } catch (e) { logEvent("error", "usage_record_failed", { error: e.message }); } + console.error(`[proxy] error: ${err.message}`); + if (res.headersSent || res.writableEnded || res.destroyed) { + try { res.end(); } catch {} + return; + } + const safeMessage = (err.message || "Internal error").replace(/\/[\w/.\-]+/g, "[path]"); + return jsonResponse(res, 500, { error: { message: safeMessage, type: "proxy_error" } }); + } + } + + // Fallback: cache disabled (CACHE_TTL=0) or no _cacheHash — original path untouched. try { const content = await callClaude(model, messages, conversationId); const id = `chatcmpl-${randomUUID()}`; completionResponse(res, id, model, content); - // Write to cache - if (CACHE_TTL > 0 && req._cacheHash) { - try { setCachedResponse(req._cacheHash, model, content); } catch (e) { logEvent("error", "cache_write_failed", { error: e.message }); } - } try { recordUsage({ keyId: req._authKeyId, keyName: req._authKeyName, model, promptChars, responseChars: content.length, elapsedMs: Date.now() - t0Usage, success: true }); } catch (e) { logEvent("error", "usage_record_failed", { error: e.message }); } } catch (err) { try { recordUsage({ keyId: req._authKeyId, keyName: req._authKeyName, model, promptChars, responseChars: 0, elapsedMs: Date.now() - t0Usage, success: false }); } catch (e) { logEvent("error", "usage_record_failed", { error: e.message }); } @@ -1551,10 +1583,10 @@ const server = createServer(async (req, res) => { }); } - // GET /cache/stats — cache statistics + // GET /cache/stats — cache statistics (entries, hits, size, inflight singleflight count) if (pathname === "/cache/stats" && req.method === "GET") { if (!isAdmin) return jsonResponse(res, 403, { error: "Admin access required" }); - return jsonResponse(res, 200, getCacheStats()); + return jsonResponse(res, 200, { ...getCacheStats(), ...getInflightStats() }); } // DELETE /cache — clear cache diff --git a/test-features.mjs b/test-features.mjs index 5eccf93..63be8dc 100644 --- a/test-features.mjs +++ b/test-features.mjs @@ -3,7 +3,7 @@ * Integration test for Quota + Cache features. * Tests database layer functions directly — no server needed. */ -import { getDb, createKey, listKeys, validateKey, recordUsage, checkQuota, updateKeyQuota, getKeyQuota, findKey, cacheHash, getCachedResponse, setCachedResponse, clearCache, getCacheStats, closeDb, hasCacheControl } from "./keys.mjs"; +import { getDb, createKey, listKeys, validateKey, recordUsage, checkQuota, updateKeyQuota, getKeyQuota, findKey, cacheHash, getCachedResponse, setCachedResponse, clearCache, getCacheStats, closeDb, hasCacheControl, singleflight, getInflightStats } from "./keys.mjs"; import { createHash } from "node:crypto"; import { strict as assert } from "node:assert"; import { unlinkSync } from "node:fs"; @@ -354,6 +354,106 @@ test("D3: chunked replay uses Array.from — multibyte codepoints stay intact", } }); +// ── PR-B Singleflight tests (async) ── +async function asyncTest(name, fn) { + try { + await fn(); + passed++; + console.log(` ✓ ${name}`); + } catch (e) { + failed++; + console.log(` ✗ ${name}: ${e.message}`); + } +} + +async function runSingleflightTests() { + console.log("\nPR-B Singleflight:"); + + // 1. Basic dedup: 10 concurrent calls with same hash execute fn only once. + await asyncTest("basic dedup: 10 concurrent callers execute fn only once", async () => { + let callCount = 0; + const fn = () => new Promise(resolve => { + callCount++; + setTimeout(() => resolve("result-A"), 20); + }); + const results = await Promise.all(Array.from({ length: 10 }, () => singleflight("sf-dedup-1", fn))); + assert.equal(callCount, 1, `fn called ${callCount} times, expected 1`); + assert.ok(results.every(r => r === "result-A"), "all 10 callers should receive the same return value"); + }); + + // 2. Failure fan-out: all followers reject when leader rejects. + await asyncTest("failure fan-out: all followers reject with leader error", async () => { + let callCount = 0; + const fn = () => new Promise((_, reject) => { + callCount++; + setTimeout(() => reject(new Error("upstream-fail")), 20); + }); + const promises = Array.from({ length: 10 }, () => singleflight("sf-fail-1", fn)); + const results = await Promise.allSettled(promises); + assert.equal(callCount, 1, `fn called ${callCount} times, expected 1`); + assert.ok(results.every(r => r.status === "rejected"), "all 10 should be rejected"); + assert.ok(results.every(r => r.reason?.message === "upstream-fail"), "all should share the same error message"); + }); + + // 3a. Map cleanup after success: inflight count returns to 0 after promise resolves. + await asyncTest("map cleanup after success: inflight=0 after promise settles", async () => { + const fn = () => new Promise(resolve => setTimeout(() => resolve("done"), 10)); + await singleflight("sf-cleanup-ok", fn); + const stats = getInflightStats(); + assert.equal(stats.inflight, 0, `expected inflight=0 after settlement, got ${stats.inflight}`); + }); + + // 3b. Map cleanup after failure: inflight count returns to 0 after promise rejects. + await asyncTest("map cleanup after failure: inflight=0 after promise rejects", async () => { + const fn = () => new Promise((_, reject) => setTimeout(() => reject(new Error("fail")), 10)); + try { await singleflight("sf-cleanup-fail", fn); } catch {} + const stats = getInflightStats(); + assert.equal(stats.inflight, 0, `expected inflight=0 after rejection, got ${stats.inflight}`); + }); + + // 4. Different hashes don't share: two parallel calls with distinct hashes both execute. + await asyncTest("different hashes do not share a singleflight entry", async () => { + let countA = 0; + let countB = 0; + const fnA = () => new Promise(resolve => { countA++; setTimeout(() => resolve("A"), 20); }); + const fnB = () => new Promise(resolve => { countB++; setTimeout(() => resolve("B"), 20); }); + const [rA, rB] = await Promise.all([singleflight("sf-hash-A", fnA), singleflight("sf-hash-B", fnB)]); + assert.equal(countA, 1); + assert.equal(countB, 1); + assert.equal(rA, "A"); + assert.equal(rB, "B"); + }); + + // 5. getInflightStats shape: returns { inflight: number, requesters: number }. + await asyncTest("getInflightStats returns correct shape", async () => { + // Verify shape against a settled state (inflight=0 is still the right shape). + const stats = getInflightStats(); + assert.equal(typeof stats.inflight, "number", "inflight should be a number"); + assert.equal(typeof stats.requesters, "number", "requesters should be a number"); + // Also verify live counts: start a pending fn, check inflight>0, then resolve. + const { promise: blocker, resolve: resolveBlocker } = Promise.withResolvers(); + const fn = () => blocker; + const p = singleflight("sf-stats-shape", fn); + const liveStats = getInflightStats(); + assert.ok(liveStats.inflight >= 1, `expected inflight>=1, got ${liveStats.inflight}`); + resolveBlocker("ok"); + await p; + }); + + // 6. Sequential calls don't share: singleflight is for concurrent dedup only. + await asyncTest("sequential calls with same hash each execute fn independently", async () => { + let callCount = 0; + const fn = () => new Promise(resolve => { callCount++; setTimeout(() => resolve(callCount), 10); }); + const r1 = await singleflight("sf-sequential", fn); + const r2 = await singleflight("sf-sequential", fn); + assert.equal(callCount, 2, `fn should have been called twice, got ${callCount}`); + assert.equal(r1, 1); + assert.equal(r2, 2); + }); +} + +await runSingleflightTests(); + // ── Cleanup ── closeDb();