diff --git a/apps/api/src/common/utils/__tests__/pinned-request.spec.ts b/apps/api/src/common/utils/__tests__/pinned-request.spec.ts new file mode 100644 index 000000000..58f2e71b5 --- /dev/null +++ b/apps/api/src/common/utils/__tests__/pinned-request.spec.ts @@ -0,0 +1,58 @@ +import { createServer, type Server } from 'http'; +import type { AddressInfo, LookupFunction } from 'net'; +import { pinnedRequest } from '../pinned-request'; + +const toLoopback: LookupFunction = (_hostname, options, callback) => { + if (options.all) { + (callback as unknown as (error: null, addresses: { address: string; family: number }[]) => void)(null, [ + { address: '127.0.0.1', family: 4 }, + ]); + return; + } + callback(null, '127.0.0.1', 4); +}; + +describe('pinnedRequest', () => { + let server: Server; + let port: number; + + beforeAll(async () => { + server = createServer((request, response) => { + if (request.url === '/reset') { + response.writeHead(205); + response.end('ignored'); + return; + } + if (request.url === '/redirect') { + response.writeHead(302, { location: '/metrics' }); + response.end(); + return; + } + response.writeHead(200, { 'content-type': 'text/plain', 'x-auth': request.headers.authorization ?? '' }); + response.end('lmcache:a_total 1'); + }); + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)); + port = (server.address() as AddressInfo).port; + }); + + afterAll(() => new Promise((resolve) => server.close(() => resolve()))); + + const send = (path: string, headers: Record = {}) => + pinnedRequest(`http://pinned.example:${port}${path}`, { headers, signal: AbortSignal.timeout(5_000), lookup: toLoopback }); + + it('connects through the given lookup and does not follow redirects', async () => { + const ok = await send('/metrics', { Authorization: 'Bearer t' }); + expect(ok.status).toBe(200); + expect(ok.headers.get('x-auth')).toBe('Bearer t'); + expect(await ok.text()).toBe('lmcache:a_total 1'); + const redirect = await send('/redirect'); + expect(redirect.status).toBe(302); + await redirect.body?.cancel(); + }); + + it('returns a null body for 205', async () => { + const reset = await send('/reset'); + expect(reset.status).toBe(205); + expect(reset.body).toBeNull(); + }); +}); diff --git a/apps/web/src/pages/KvCache.tsx b/apps/web/src/pages/KvCache.tsx index f3ab72ca9..11c53f0cd 100644 --- a/apps/web/src/pages/KvCache.tsx +++ b/apps/web/src/pages/KvCache.tsx @@ -150,10 +150,17 @@ export function KvCache() { recentBuckets={recent.data?.buckets ?? []} /> {latest && } - {samplesSince !== null && ( -

- Showing engine data since {new Date(samplesSince).toLocaleString()} + {samples.data?.truncated ? ( +

+ This range holds more samples than can be shown, so the charts and table cover only its most recent + part. Pick a shorter range to see all of it.

+ ) : ( + samplesSince !== null && ( +

+ Showing engine data since {new Date(samplesSince).toLocaleString()} +

+ ) )} { }); it('computes the eviction delta per node and treats a lower counter as a restart', () => { - expect(evictedKeysDelta(null, { a: 5 })).toBeNull(); - expect(evictedKeysDelta({ a: 20 }, { a: 50 })).toBe(30); - expect(evictedKeysDelta({ a: 20 }, { a: 5 })).toBe(5); - expect(evictedKeysDelta({ a: 10, b: 1000 }, { a: 12 })).toBe(2); - expect(evictedKeysDelta({ a: 10 }, { a: 12, b: 1000 })).toBe(2); - expect(evictedKeysDelta({ a: 10 }, { b: 1000 })).toBeNull(); - }); - - it('passes the eviction delta through to the snapshot', () => { - const node: NodeObservation = { scannedKeys: 0, matchedKeys: [], scanComplete: true, dbSize: 0, samples: [] }; - expect(buildSnapshot({ ...base, nodes: [node] }).evictedKeysDelta).toBeNull(); - expect(buildSnapshot({ ...base, nodes: [node], evictedKeysDelta: 7 }).evictedKeysDelta).toBe(7); + expect(evictedKeysDelta(undefined, new Map([['a', 5]]))).toBeNull(); + expect(evictedKeysDelta(new Map([['a', 20]]), new Map([['a', 50]]))).toBe(30); + expect(evictedKeysDelta(new Map([['a', 20]]), new Map([['a', 5]]))).toBe(5); + expect(evictedKeysDelta(new Map([['a', 20], ['b', 1000]]), new Map([['a', 25]]))).toBe(5); + expect(evictedKeysDelta(new Map([['a', 20]]), new Map([['a', 20], ['b', 1000]]))).toBe(0); + expect(evictedKeysDelta(new Map([['a', 20]]), new Map([['b', 1000]]))).toBeNull(); }); it('detects through the GLIDE client even with no keys', () => { @@ -351,4 +345,16 @@ describe('buildSnapshot', () => { const s = buildSnapshot({ ...base, nodes: [n1, n2] }); expect(s).toMatchObject({ scannedKeys: 10, matchedKeys: 3, chunksEst: 3, bytesEst: 50 }); }); + + it('weights each node sample by the keys it stands for when splitting per model', () => { + const a = Array.from({ length: 100 }, (_, i) => key(i)); + const b = key(0, '', 'other', 'half'); + const n1: NodeObservation = { scannedKeys: 100, matchedKeys: parse(a), scanComplete: true, dbSize: 100, samples: [sample(a[0], { memoryBytes: 100 })] }; + const n2: NodeObservation = { scannedKeys: 1, matchedKeys: parse([b]), scanComplete: true, dbSize: 1, samples: [sample(b, { memoryBytes: 100 })] }; + const s = buildSnapshot({ ...base, nodes: [n1, n2] }); + expect(s.perModel).toEqual([ + { model: M, dtype: 'bfloat16', chunksEst: 100, bytesEst: 10_000 }, + { model: 'other', dtype: 'half', chunksEst: 1, bytesEst: 100 }, + ]); + }); }); diff --git a/proprietary/kv-cache/__tests__/footprint-scanner.spec.ts b/proprietary/kv-cache/__tests__/footprint-scanner.spec.ts index 641de1736..0e4bc57a8 100644 --- a/proprietary/kv-cache/__tests__/footprint-scanner.spec.ts +++ b/proprietary/kv-cache/__tests__/footprint-scanner.spec.ts @@ -59,19 +59,46 @@ describe('scanLmcacheKeys', () => { expect(client.call).toHaveBeenCalledTimes(2); expect(result.scanComplete).toBe(false); }); + + it('gives up after the probe pages when no LMCache key is found', async () => { + const pages: Array<[string, string[]]> = [['7', ['user:1']], ['8', ['a@b@c@d@e']], ['0', [k(0)]]]; + const probing = scanClient(pages); + await scanLmcacheKeys(probing, { maxScanned: 200_000, maxMatched: 2000 }, 10_000, true); + expect(probing.call).toHaveBeenCalledTimes(2); + + const full = scanClient(pages); + const result = await scanLmcacheKeys(full, { maxScanned: 200_000, maxMatched: 2000 }, 10_000); + expect(result.matchedKeys).toEqual(p(k(0))); + }); + + it('keeps scanning past the probe pages once an LMCache key is found', async () => { + const client = scanClient([['7', [k(1)]], ['8', []], ['0', [k(0)]]]); + const result = await scanLmcacheKeys(client, { maxScanned: 200_000, maxMatched: 2000 }, 3000, true); + expect(result.matchedKeys).toEqual(p(k(1), k(0))); + }); }); +function sampleClient(replies: Record) { + const call = jest.fn(async (command: string, args: string[]) => { + const reply = replies[`${command} ${args.join(' ')}`]; + if (reply instanceof Error) throw reply; + return reply; + }); + return { call, getClient: () => { throw new Error('Raw client access not available through agent connections'); } } as any; +} + describe('sampleLmcacheKeys', () => { - it('reads memory, ttl and sibling existence in one pipeline', async () => { - const pipeline = { call: jest.fn().mockReturnThis(), ttl: jest.fn().mockReturnThis(), exists: jest.fn().mockReturnThis(), exec: jest.fn() }; - pipeline.exec.mockResolvedValue([ - [null, 3670112], [null, -1], [null, 0], - [null, 120], [null, 50], - [null, null], [null, -2], - ]); - const client = { getClient: () => ({ pipeline: () => pipeline }) } as any; + it('reads memory, ttl and sibling existence through call only', async () => { + const client = sampleClient({ + [`MEMORY USAGE ${k(0, 'kv_bytes')}`]: 3670112, + [`TTL ${k(0, 'kv_bytes')}`]: -1, + [`EXISTS ${k(0, 'metadata')}`]: 0, + [`MEMORY USAGE ${k(1)}`]: 120, + [`TTL ${k(1)}`]: 50, + [`MEMORY USAGE ${k(2)}`]: null, + [`TTL ${k(2)}`]: -2, + }); const samples = await sampleLmcacheKeys(client, p(k(0, 'kv_bytes'), k(1), k(2))); - expect(pipeline.exists).toHaveBeenCalledWith(k(0, 'metadata')); expect(samples.map(({ memoryBytes, ttl, siblingExists }) => ({ memoryBytes, ttl, siblingExists }))).toEqual([ { memoryBytes: 3670112, ttl: -1, siblingExists: false }, { memoryBytes: 120, ttl: 50, siblingExists: null }, @@ -80,42 +107,78 @@ describe('sampleLmcacheKeys', () => { }); it('treats a cross-slot sibling probe rejected with MOVED as unknown, not orphaned', async () => { - const pipeline = { call: jest.fn().mockReturnThis(), ttl: jest.fn().mockReturnThis(), exists: jest.fn().mockReturnThis(), exec: jest.fn() }; - pipeline.exec.mockResolvedValue([ - [null, 3670112], [null, -1], [new Error('MOVED 1234 10.0.0.2:6379'), null], - [null, 3670112], [null, -1], [null, 1], - [null, 3670112], [null, -1], [null, 0], - ]); - const client = { getClient: () => ({ pipeline: () => pipeline }) } as any; - const samples = await sampleLmcacheKeys(client, p(k(0, 'kv_bytes'), k(1, 'kv_bytes'), k(2, 'kv_bytes'))); - expect(samples.map((s) => s.siblingExists)).toEqual([null, true, false]); + const client = sampleClient({ + [`MEMORY USAGE ${k(0, 'kv_bytes')}`]: 3670112, + [`TTL ${k(0, 'kv_bytes')}`]: -1, + [`EXISTS ${k(0, 'metadata')}`]: new Error('MOVED 1234 10.0.0.2:6379'), + [`MEMORY USAGE ${k(1, 'kv_bytes')}`]: 3670112, + [`TTL ${k(1, 'kv_bytes')}`]: -1, + [`EXISTS ${k(1, 'metadata')}`]: 1, + }); + const samples = await sampleLmcacheKeys(client, p(k(0, 'kv_bytes'), k(1, 'kv_bytes'))); + expect(samples.map((s) => s.siblingExists)).toEqual([null, true]); + }); + + it('reports a key as unsized when MEMORY is unavailable', async () => { + const client = sampleClient({ + [`MEMORY USAGE ${k(0)}`]: new Error("ERR unknown command 'MEMORY'"), + [`TTL ${k(0)}`]: -1, + }); + const [sample] = await sampleLmcacheKeys(client, p(k(0))); + expect(sample.memoryBytes).toBeNull(); + }); + + it('fails the read when every command for a key is rejected', async () => { + const closed = new Error('Connection is closed.'); + const client = sampleClient({ + [`MEMORY USAGE ${k(0)}`]: 120, + [`TTL ${k(0)}`]: 50, + [`MEMORY USAGE ${k(1, 'kv_bytes')}`]: closed, + [`TTL ${k(1, 'kv_bytes')}`]: closed, + [`EXISTS ${k(1, 'metadata')}`]: closed, + }); + await expect(sampleLmcacheKeys(client, p(k(0), k(1, 'kv_bytes')))).rejects.toThrow('Connection is closed.'); + }); + + it('keeps the ttl of a key whose MEMORY and sibling probes are rejected', async () => { + const client = sampleClient({ + [`MEMORY USAGE ${k(0, 'kv_bytes')}`]: new Error("ERR unknown command 'MEMORY'"), + [`TTL ${k(0, 'kv_bytes')}`]: 30, + [`EXISTS ${k(0, 'metadata')}`]: new Error('MOVED 1234 10.0.0.2:6379'), + }); + const [sample] = await sampleLmcacheKeys(client, p(k(0, 'kv_bytes'))); + expect(sample).toMatchObject({ memoryBytes: null, ttl: 30, siblingExists: null }); + }); + + it('issues no commands for no keys', async () => { + const client = sampleClient({}); + expect(await sampleLmcacheKeys(client, [])).toEqual([]); + expect(client.call).not.toHaveBeenCalled(); }); it('skips the sibling probe when the node does not own the sibling slot', async () => { - const pipeline = { call: jest.fn().mockReturnThis(), ttl: jest.fn().mockReturnThis(), exists: jest.fn().mockReturnThis(), exec: jest.fn() }; - pipeline.exec.mockResolvedValue([ - [null, 3670112], [null, -1], - [null, 3670112], [null, -1], [null, 1], - ]); - const client = { getClient: () => ({ pipeline: () => pipeline }) } as any; + const client = sampleClient({ + [`MEMORY USAGE ${k(0, 'kv_bytes')}`]: 3670112, + [`TTL ${k(0, 'kv_bytes')}`]: -1, + [`MEMORY USAGE ${k(1, 'kv_bytes')}`]: 3670112, + [`TTL ${k(1, 'kv_bytes')}`]: -1, + [`EXISTS ${k(1, 'metadata')}`]: 1, + }); const owned = keySlot(k(1, 'metadata')); const samples = await sampleLmcacheKeys(client, p(k(0, 'kv_bytes'), k(1, 'kv_bytes')), (slot) => slot === owned); - expect(pipeline.exists).toHaveBeenCalledTimes(1); - expect(pipeline.exists).toHaveBeenCalledWith(k(1, 'metadata')); + const exists = client.call.mock.calls.filter(([command]: [string]) => command === 'EXISTS'); + expect(exists).toEqual([['EXISTS', [k(1, 'metadata')]]]); expect(samples.map((s) => s.siblingExists)).toEqual([null, true]); }); it('reports an unknown TTL when the TTL call fails', async () => { - const pipeline = { call: jest.fn().mockReturnThis(), ttl: jest.fn().mockReturnThis(), exists: jest.fn().mockReturnThis(), exec: jest.fn() }; - pipeline.exec.mockResolvedValue([[null, 120], [new Error('NOPERM'), null]]); - const client = { getClient: () => ({ pipeline: () => pipeline }) } as any; + const client = sampleClient({ + [`MEMORY USAGE ${k(0)}`]: 120, + [`TTL ${k(0)}`]: new Error('NOPERM'), + }); const [sample] = await sampleLmcacheKeys(client, p(k(0))); expect(sample).toMatchObject({ memoryBytes: 120, ttl: null }); }); - - it('skips the pipeline for no keys', async () => { - expect(await sampleLmcacheKeys({ getClient: () => { throw new Error('unused'); } } as any, [])).toEqual([]); - }); }); describe('helpers', () => { @@ -144,6 +207,11 @@ describe('helpers', () => { expect(await hasLmcacheGlideClient(client, 3)).toBe(true); }); + it('ignores a client name that imitates the GLIDE library field', async () => { + const list = 'id=1 name=lib-name=GlidePySync(lmcache:0.5.5) lib-name=redis-py lib-ver=5.0\n'; + expect(await hasLmcacheGlideClient({ call: jest.fn().mockResolvedValue(list) } as any, 0)).toBe(false); + }); + it('reads memory stats', () => { const info = { memory: { used_memory: '100', maxmemory: '200', maxmemory_policy: 'volatile-lru' }, diff --git a/proprietary/kv-cache/__tests__/kv-cache-alerts.service.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-alerts.service.spec.ts index 4e1c61dbb..778e6fc41 100644 --- a/proprietary/kv-cache/__tests__/kv-cache-alerts.service.spec.ts +++ b/proprietary/kv-cache/__tests__/kv-cache-alerts.service.spec.ts @@ -148,6 +148,70 @@ describe('KvCacheAlertsService eviction risk', () => { expect(evictingStates()).toEqual([true, true, true, true, true, false, false]); }); + it('keeps the evicting state when a snapshot only read some nodes', async () => { + const { service, pro } = setup(); + const evictingStates = () => + pro.dispatchKvCacheEvictionRisk.mock.calls.filter(([e]) => e.reason === 'evicting').map(([e]) => e.active); + await service.onSnapshot(snapshot({ evictedKeysDelta: 50 }), { partial: true }); + await service.onSnapshot(snapshot({ evictedKeysDelta: 50 })); + for (let i = 0; i < 3; i += 1) await service.onSnapshot(snapshot({ evictedKeysDelta: 0 }), { partial: true }); + await service.onSnapshot(snapshot({ evictedKeysDelta: 0 })); + expect(evictingStates()).toEqual([false, true, true, true, true, true]); + }); + + it('keeps the evicting state when a partial snapshot does not detect LMCache', async () => { + const { service, pro } = setup(); + await service.onSnapshot(snapshot({ evictedKeysDelta: 50 })); + await service.onSnapshot(snapshot({ detected: false, evictedKeysDelta: 0 }), { partial: true }); + expect(pro.dispatchKvCacheEvictionRisk).toHaveBeenLastCalledWith(expect.objectContaining({ reason: 'evicting', active: true })); + }); + + it('keeps the unevictable state when a snapshot only read some nodes', async () => { + const { service, pro } = setup(); + const unevictableStates = () => + pro.dispatchKvCacheEvictionRisk.mock.calls.filter(([e]) => e.reason === 'unevictable').map(([e]) => e.active); + await service.onSnapshot(snapshot({ noTtlRatio: 0.95 }), { partial: true }); + await service.onSnapshot(snapshot({ noTtlRatio: 0.95 })); + await service.onSnapshot(snapshot({ noTtlRatio: 0.1 }), { partial: true }); + await service.onSnapshot(snapshot({ noTtlRatio: null }), { partial: true }); + await service.onSnapshot(snapshot({ detected: false }), { partial: true }); + await service.onSnapshot(snapshot({ noTtlRatio: 0.1 })); + await service.onSnapshot(snapshot({ noTtlRatio: 0.95 }), { partial: true }); + expect(unevictableStates()).toEqual([false, true, true, true, true, false, false]); + }); + + it('forgets the unevictable state when the connection is removed or alerts are disabled', async () => { + const { service, pro, remove, storage } = setup(); + const lastUnevictable = () => + pro.dispatchKvCacheEvictionRisk.mock.calls.filter(([e]) => e.reason === 'unevictable').at(-1)?.[0].active; + service.onModuleInit(); + await service.onSnapshot(snapshot()); + remove('c1'); + await service.onSnapshot(snapshot(), { partial: true }); + expect(lastUnevictable()).toBe(false); + await service.onSnapshot(snapshot()); + storage.getKvCacheSettings.mockResolvedValueOnce({ + connectionId: 'c1', + hitRateAlertEnabled: true, + hitRateThreshold: 0.2, + evictionAlertEnabled: false, + updatedAt: 1, + }); + await service.onSnapshot(snapshot()); + await service.onSnapshot(snapshot(), { partial: true }); + expect(lastUnevictable()).toBe(false); + service.onModuleDestroy(); + }); + + it('passes the partial flag from footprint snapshots through', async () => { + const { service, pro, footprint } = setup(); + service.onModuleInit(); + const listener = footprint.onSnapshot.mock.calls[0][0]; + await listener(snapshot({ evictedKeysDelta: 50 }), { partial: true }); + expect(pro.dispatchKvCacheEvictionRisk).toHaveBeenCalledWith(expect.objectContaining({ reason: 'evicting', active: false })); + service.onModuleDestroy(); + }); + it('forgets the evicting hold when the connection is removed', async () => { const { service, pro, remove } = setup(); service.onModuleInit(); diff --git a/proprietary/kv-cache/__tests__/kv-cache-engine-registry.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-engine-registry.spec.ts index 8a3aa7e65..0dcee5de5 100644 --- a/proprietary/kv-cache/__tests__/kv-cache-engine-registry.spec.ts +++ b/proprietary/kv-cache/__tests__/kv-cache-engine-registry.spec.ts @@ -13,10 +13,15 @@ function setup(initial = [engine('a', 'c1', 'lmc-a'), engine('b', 'c2')]) { deleteKvCacheEngine: jest.fn().mockResolvedValue(true), } as any; let removal: (id: string) => unknown = () => undefined; - const footprint = { onConnectionRemoval: jest.fn((l) => (removal = l)) } as any; - const samples = { forgetEngine: jest.fn(), idle: jest.fn().mockResolvedValue(undefined) } as any; + let restore: (id: string) => unknown = () => undefined; + const footprint = { + onConnectionRemoval: jest.fn((l) => (removal = l)), + onConnectionRestore: jest.fn((l) => (restore = l)), + hasConnection: jest.fn().mockReturnValue(false), + } as any; + const samples = { forgetEngine: jest.fn(), settled: jest.fn().mockResolvedValue(undefined) } as any; const registry = new KvCacheEngineRegistry(storage, footprint, samples); - return { registry, storage, samples, removal: (id: string) => removal(id) }; + return { registry, storage, samples, footprint, removal: (id: string) => removal(id), restore: (id: string) => restore(id) }; } describe('KvCacheEngineRegistry', () => { @@ -277,9 +282,83 @@ describe('KvCacheEngineRegistry', () => { it('drops only the removed connection engines on broadcast', async () => { const { registry, samples, removal } = setup(); await registry.onModuleInit(); - removal('c1'); + const removing = removal('c1'); expect(registry.list().map((e) => e.id)).toEqual(['b']); expect(samples.forgetEngine).toHaveBeenCalledTimes(1); expect(samples.forgetEngine).toHaveBeenCalledWith('a'); + await removing; + expect(samples.settled).toHaveBeenCalled(); + }); + + it('reloads the engines of a connection that was registered again before cleanup', async () => { + const { registry, storage, removal, restore } = setup(); + await registry.onModuleInit(); + await removal('c1'); + expect(registry.get('a')).toBeNull(); + await restore('c1'); + expect(registry.list('c1').map((e) => e.id)).toEqual(['a']); + registry.recordResult('a', { lastSeenAt: 9, lastError: null }); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.saveKvCacheEngine).toHaveBeenCalledWith(expect.objectContaining({ id: 'a', lastSeenAt: 9 })); + }); + + it('does not restore engines when the connection is removed again during the reload', async () => { + const { registry, storage, removal, restore } = setup(); + await registry.onModuleInit(); + await removal('c1'); + let finishLoad: (engines: unknown) => void = () => undefined; + storage.getKvCacheEngines.mockImplementationOnce(() => new Promise((resolve) => (finishLoad = resolve))); + const restoring = restore('c1'); + await removal('c1'); + finishLoad([engine('a', 'c1', 'lmc-a')]); + await restoring; + expect(registry.get('a')).toBeNull(); + }); + + it('keeps a link created during cleanup and deletes the stored engine it supersedes', async () => { + const { registry, storage, footprint, removal, restore } = setup([engine('a', 'c1', 'lmc-a'), engine('old', 'c1')]); + await registry.onModuleInit(); + await removal('c1'); + footprint.hasConnection.mockReturnValue(true); + await registry.save({ ...engine('new', 'c1', 'lmc-a'), name: 'fresh' }); + await registry.save({ ...engine('dup', 'c1'), name: 'old' }); + await restore('c1'); + expect(registry.list('c1').map((e) => e.id).sort()).toEqual(['dup', 'new']); + expect(storage.deleteKvCacheEngine).toHaveBeenCalledWith('a'); + expect(storage.deleteKvCacheEngine).toHaveBeenCalledWith('old'); + }); + + it('does not reload engines after a removal that lands while a superseded engine is deleted', async () => { + const { registry, storage, footprint, removal, restore } = setup([engine('a', 'c1', 'lmc-a'), engine('b', 'c1')]); + await registry.onModuleInit(); + await removal('c1'); + footprint.hasConnection.mockReturnValue(true); + await registry.save(engine('new', 'c1', 'lmc-a')); + let finishDelete: () => void = () => undefined; + storage.deleteKvCacheEngine.mockImplementationOnce(() => new Promise((resolve) => (finishDelete = () => resolve(true)))); + const restoring = restore('c1'); + await new Promise((resolve) => setImmediate(resolve)); + await removal('c1'); + finishDelete(); + await restoring; + expect(registry.list('c1')).toEqual([]); + }); + + it('skips deleting a superseded engine once the connection is removed again', async () => { + const { registry, storage, footprint, removal, restore } = setup([engine('a', 'c1', 'lmc-a')]); + await registry.onModuleInit(); + await removal('c1'); + footprint.hasConnection.mockReturnValue(true); + await registry.save(engine('new', 'c1', 'lmc-a')); + let finishPending: () => void = () => undefined; + storage.deleteKvCacheEngine.mockImplementationOnce(() => new Promise((resolve) => (finishPending = () => resolve(true)))); + const pending = registry.remove('a'); + const restoring = restore('c1'); + await new Promise((resolve) => setImmediate(resolve)); + await removal('c1'); + finishPending(); + await pending; + await restoring; + expect(storage.deleteKvCacheEngine).toHaveBeenCalledTimes(1); }); }); diff --git a/proprietary/kv-cache/__tests__/kv-cache-engines.service.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-engines.service.spec.ts index 80cb8b1be..90ea481c6 100644 --- a/proprietary/kv-cache/__tests__/kv-cache-engines.service.spec.ts +++ b/proprietary/kv-cache/__tests__/kv-cache-engines.service.spec.ts @@ -303,7 +303,7 @@ describe('KvCacheEnginesService', () => { saveKvCacheEngine: jest.fn(async (e: any) => e), deleteKvCacheEngine: jest.fn(() => new Promise((resolve) => (finishDelete = () => resolve(true)))), } as any; - const footprint = { onConnectionRemoval: jest.fn() } as any; + const footprint = { onConnectionRemoval: jest.fn(), onConnectionRestore: jest.fn() } as any; const samples = { forgetEngine: jest.fn(), idle: jest.fn(), deleteEngine: jest.fn().mockResolvedValue(undefined), resetBaselines: jest.fn() } as any; const registry = new KvCacheEngineRegistry(storage, footprint, samples); await registry.onModuleInit(); diff --git a/proprietary/kv-cache/__tests__/kv-cache-footprint.service.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-footprint.service.spec.ts index 7fb71d5d2..600d480e9 100644 --- a/proprietary/kv-cache/__tests__/kv-cache-footprint.service.spec.ts +++ b/proprietary/kv-cache/__tests__/kv-cache-footprint.service.spec.ts @@ -1,4 +1,4 @@ -import { BadRequestException, Logger, NotFoundException } from '@nestjs/common'; +import { BadRequestException, Logger, NotFoundException, ServiceUnavailableException } from '@nestjs/common'; import { Feature } from '@betterdb/shared'; import * as classifier from '../key-classifier'; import * as scanner from '../footprint-scanner'; @@ -12,13 +12,15 @@ import { const KEY = (i: number) => `m/x@1@0@${(i + 1).toString(16)}@bfloat16`; +const REPLIES: Record = { MEMORY: 100, TTL: -1, EXISTS: 0, CLIENT: 'id=1 lib-name=redis-py\n' }; + function fakeClient(keys: string[]) { - const pipeline = { call: jest.fn().mockReturnThis(), ttl: jest.fn().mockReturnThis(), exists: jest.fn().mockReturnThis(), exec: jest.fn() }; - pipeline.exec.mockImplementation(async () => keys.flatMap(() => [[null, 100], [null, -1]])); return { getDbSize: jest.fn().mockResolvedValue(keys.length), - call: jest.fn(async (cmd: string) => (cmd === 'SCAN' ? ['0', keys] : 'id=1 lib-name=redis-py\n')), - getClient: () => ({ pipeline: () => pipeline }), + call: jest.fn(async (cmd: string) => (cmd === 'SCAN' ? ['0', keys] : REPLIES[cmd])), + getClient: () => { + throw new Error('Raw client access not available through agent connections'); + }, getInfoParsed: jest.fn().mockResolvedValue({ memory: { used_memory: '1000', maxmemory: '0', maxmemory_policy: 'noeviction' }, stats: { evicted_keys: '3' }, @@ -59,10 +61,50 @@ describe('KvCacheFootprintService', () => { const snapshot = await service.triggerCollection('c1'); expect(snapshot).toMatchObject({ connectionId: 'c1', detected: true, matchedKeys: 6, chunksEst: 6, evictedKeysDelta: null }); expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledWith(snapshot); - expect(listener).toHaveBeenCalledWith(snapshot); + expect(listener).toHaveBeenCalledWith(snapshot, { partial: false }); expect(service.getSampleKey('c1')).toBe('m/x@1@0@1…@bfloat16'); }); + it('reports the eviction delta on the second run', async () => { + const { service } = setup(); + await service.triggerCollection('c1'); + expect((await service.triggerCollection('c1'))?.evictedKeysDelta).toBe(0); + }); + + it('returns the in-flight collection to a manual trigger', async () => { + const { service, client, storage } = setup(); + let finishRead: (size: number) => void = () => undefined; + client.getDbSize.mockImplementationOnce(() => new Promise((resolve) => (finishRead = resolve))); + const first = service.triggerCollection('c1'); + await new Promise((resolve) => setImmediate(resolve)); + const second = service.triggerCollection('c1'); + finishRead(6); + const [a, b] = await Promise.all([first, second]); + expect(b).not.toBeNull(); + expect(b).toBe(a); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(1); + }); + + it('returns the in-flight collection to a trigger that starts in the same tick', async () => { + const { service, storage } = setup(); + const first = service.triggerCollection('c1'); + const second = service.triggerCollection('c1'); + const [a, b] = await Promise.all([first, second]); + expect(a).not.toBeNull(); + expect(b).toBe(a); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(1); + }); + + it('returns the in-flight collection to a poll that overlaps a manual trigger', async () => { + const { service, client } = setup(); + const ctx = { connectionId: 'c1', connectionName: 'one', client, host: 'h', port: 6379 }; + const manual = service.triggerCollection('c1'); + const polled = service.collect(ctx); + const [a, b] = await Promise.all([manual, polled]); + expect(a).not.toBeNull(); + expect(b).toBe(a); + }); + it('classifies each scanned key once per collection', async () => { const { service } = setup(); const classify = jest.spyOn(classifier, 'classifyKey'); @@ -74,10 +116,39 @@ describe('KvCacheFootprintService', () => { } }); - it('reports the eviction delta on the second run', async () => { - const { service } = setup(); - await service.triggerCollection('c1'); - expect((await service.triggerCollection('c1'))?.evictedKeysDelta).toBe(0); + it('rejects a manual trigger for a cluster member, which its seed collects', async () => { + const { service, client, registry } = setup(); + registry.list.mockReturnValue([ + { id: 'seed', name: 'seed', isConnected: true, connectionType: 'direct' }, + { + id: 'c1', + name: 'one', + isConnected: true, + connectionType: 'direct', + membership: { seedId: 'seed', nodeId: 'n', origin: 'auto', source: 'cluster' }, + }, + ]); + await expect(service.triggerCollection('c1')).rejects.toThrow(/seed/); + await expect(service.triggerCollection('c1')).rejects.toBeInstanceOf(BadRequestException); + expect(client.getDbSize).not.toHaveBeenCalled(); + }); + + it('rejects a manual trigger for an external connection', async () => { + const { service, client, registry } = setup(); + registry.list.mockReturnValue([{ id: 'c1', name: 'one', isConnected: true, connectionType: 'external' }]); + registry.getConfig.mockReturnValue({ id: 'c1', name: 'one', host: 'h', port: 6379, connectionType: 'external' }); + registry.get.mockImplementation(() => { + throw new Error('external connections have no client'); + }); + await expect(service.triggerCollection('c1')).rejects.toBeInstanceOf(BadRequestException); + expect(client.getDbSize).not.toHaveBeenCalled(); + }); + + it('rejects a manual trigger for a sentinel connection', async () => { + const { service, client } = setup(); + client.getCapabilities = () => ({ isSentinel: true }); + await expect(service.triggerCollection('c1')).rejects.toBeInstanceOf(BadRequestException); + expect(client.getDbSize).not.toHaveBeenCalled(); }); it('keeps the eviction baseline when the connection is renamed', async () => { @@ -204,8 +275,8 @@ describe('KvCacheFootprintService', () => { expect((await service.triggerCollection('c1'))?.evictedKeysDelta).toBe(1); }); - it('keeps the readable primaries and marks the scan incomplete when one primary fails', async () => { - const { service, client, storage } = setup(); + it('keeps the readable nodes and marks the scan incomplete when one node fails', async () => { + const { service, client } = setup(); const broken = fakeClient([]); broken.getDbSize.mockRejectedValue(new Error('LOADING')); jest.spyOn(service as any, 'nodesFor').mockResolvedValue({ @@ -216,38 +287,27 @@ describe('KvCacheFootprintService', () => { skipped: 0, }); const snapshot = await service.triggerCollection('c1'); - expect(snapshot).toMatchObject({ matchedKeys: 6, detected: true, scanComplete: false }); - expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledWith(snapshot); + expect(snapshot).toMatchObject({ matchedKeys: 6, scanComplete: false }); }); - it('keeps the previous detection and sample key when a partial read finds too few keys', async () => { - const { service, storage } = setup(); - await service.triggerCollection('c1'); - storage.getKvCacheFootprintSnapshots.mockResolvedValue([{ connectionId: 'c1', detected: true }]); - const sparse = fakeClient([KEY(0)]); + it('tells snapshot listeners whether every node was read', async () => { + const { service, client } = setup(); + const listener = jest.fn(); + service.onSnapshot(listener); const broken = fakeClient([]); broken.getDbSize.mockRejectedValue(new Error('LOADING')); const nodes = jest.spyOn(service as any, 'nodesFor').mockResolvedValue({ nodes: [ - { name: 'a', client: sparse }, - { name: 'b', client: broken }, - ], - skipped: 0, - }); - const snapshot = await service.triggerCollection('c1'); - expect(snapshot).toMatchObject({ detected: true, scanComplete: false }); - expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenLastCalledWith(snapshot); - expect(storage.getKvCacheFootprintSnapshots).toHaveBeenCalledWith({ connectionId: 'c1', limit: 1 }); - expect(service.getSampleKey('c1')).toBe('m/x@1@0@1…@bfloat16'); - nodes.mockResolvedValue({ - nodes: [ - { name: 'a', client: fakeClient([]) }, + { name: 'a', client }, { name: 'b', client: broken }, ], skipped: 0, }); await service.triggerCollection('c1'); - expect(service.getSampleKey('c1')).toBe('m/x@1@0@1…@bfloat16'); + expect(listener).toHaveBeenLastCalledWith(expect.objectContaining({ connectionId: 'c1' }), { partial: true }); + nodes.mockResolvedValue({ nodes: [{ name: 'a', client }], skipped: 0 }); + await service.triggerCollection('c1'); + expect(listener).toHaveBeenLastCalledWith(expect.objectContaining({ connectionId: 'c1' }), { partial: false }); }); it('clears an earlier detection after consecutive partial runs that find no LMCache', async () => { @@ -293,6 +353,27 @@ describe('KvCacheFootprintService', () => { expect(await service.triggerCollection('c1')).toMatchObject({ detected: false, scanComplete: false }); }); + it('marks the snapshot partial when a node drops its connection while sampling', async () => { + const { service, client } = setup(); + const listener = jest.fn(); + service.onSnapshot(listener); + const dropped = fakeClient(Array.from({ length: 6 }, (_, i) => KEY(i + 10))); + dropped.call.mockImplementation(async (cmd: string) => { + if (cmd === 'SCAN') return ['0', Array.from({ length: 6 }, (_, i) => KEY(i + 10))]; + throw new Error('Connection is closed.'); + }); + jest.spyOn(service as any, 'nodesFor').mockResolvedValue({ + nodes: [ + { name: 'a', client }, + { name: 'b', client: dropped }, + ], + skipped: 0, + }); + const snapshot = await service.triggerCollection('c1'); + expect(snapshot).toMatchObject({ matchedKeys: 6, scanComplete: false }); + expect(listener).toHaveBeenLastCalledWith(expect.objectContaining({ connectionId: 'c1' }), { partial: true }); + }); + describe('cluster members that were never read', () => { const primary = (id: string, slots: number[][], flags: string[] = []) => ({ id, flags: [...flags, 'master'], slots }); const replica = (id: string) => ({ id, flags: ['slave'], slots: [] }); @@ -302,6 +383,8 @@ describe('KvCacheFootprintService', () => { view: unknown[] = [primary('n0', [[0, 8191]], ['myself']), member === 'replica' ? replica('n1') : primary('n1', [[8192, 16383]])], ) { const ctx = setup(); + const listener = jest.fn(); + ctx.service.onSnapshot(listener); const other = fakeClient([]); ctx.client.getClusterNodes = jest.fn().mockResolvedValue(view); if (member === 'role-failed') other.getRole.mockRejectedValue(new Error('timeout')); @@ -320,19 +403,21 @@ describe('KvCacheFootprintService', () => { membership: { seedId: 'c1', nodeId: 'n1', origin: 'auto', source: 'cluster' }, }, ]); - return ctx; + return { ...ctx, listener }; } - it.each(['deregistered', 'role-failed'] as const)('marks the snapshot incomplete when a member is %s', async (member) => { - const { service, storage } = clusterSetup(member); + it.each(['deregistered', 'role-failed'] as const)('marks the snapshot partial when a member is %s', async (member) => { + const { service, listener } = clusterSetup(member); const snapshot = await service.triggerCollection('c1'); - expect(snapshot).toMatchObject({ matchedKeys: 6, scanComplete: false }); - expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledWith(snapshot); + expect(snapshot).toMatchObject({ scanComplete: false }); + expect(listener).toHaveBeenLastCalledWith(expect.objectContaining({ connectionId: 'c1' }), { partial: true }); }); it('does not treat a replica as unread', async () => { - const { service } = clusterSetup('replica'); - expect(await service.triggerCollection('c1')).toMatchObject({ scanComplete: true }); + const { service, listener } = clusterSetup('replica'); + const snapshot = await service.triggerCollection('c1'); + expect(snapshot).toMatchObject({ scanComplete: true }); + expect(listener).toHaveBeenLastCalledWith(expect.objectContaining({ connectionId: 'c1' }), { partial: false }); }); it('keeps the snapshot complete when every slot owner was scanned', async () => { @@ -497,6 +582,8 @@ describe('KvCacheFootprintService', () => { describe('disconnected cluster members', () => { function offlineSetup(clusterNodes: () => Promise) { const ctx = setup(); + const listener = jest.fn(); + ctx.service.onSnapshot(listener); ctx.registry.list.mockReturnValue([ { id: 'c1', name: 'one', isConnected: true, connectionType: 'direct' }, { @@ -508,35 +595,129 @@ describe('KvCacheFootprintService', () => { }, ]); ctx.client.getClusterNodes = jest.fn(clusterNodes); - return ctx; + return { ...ctx, listener }; } - it('marks the snapshot incomplete when a disconnected member still owns slots as a primary', async () => { - const { service } = offlineSetup(async () => [{ id: 'n1', flags: ['master', 'fail?'], slots: [[0, 5460]] }]); - expect(await service.triggerCollection('c1')).toMatchObject({ scanComplete: false }); + it('marks the snapshot partial when a disconnected member still owns slots as a primary', async () => { + const { service, listener } = offlineSetup(async () => [{ id: 'n1', flags: ['master', 'fail?'], slots: [[0, 5460]] }]); + const snapshot = await service.triggerCollection('c1'); + expect(snapshot).toMatchObject({ scanComplete: false }); + expect(listener).toHaveBeenLastCalledWith(expect.objectContaining({ connectionId: 'c1' }), { partial: true }); }); - it('marks the snapshot incomplete when the cluster view cannot be read', async () => { - const { service } = offlineSetup(async () => { + it('marks the snapshot partial when the cluster view cannot be read', async () => { + const { service, listener } = offlineSetup(async () => { throw new Error('timeout'); }); - expect(await service.triggerCollection('c1')).toMatchObject({ scanComplete: false }); + await service.triggerCollection('c1'); + expect(listener).toHaveBeenLastCalledWith(expect.objectContaining({ connectionId: 'c1' }), { partial: true }); }); it.each([ ['a replica', { id: 'n1', flags: ['slave'], slots: [] }], ['a failed-over primary without slots', { id: 'n1', flags: ['master', 'fail'], slots: [] }], ])('does not treat %s as unread', async (_label, node) => { - const { service } = offlineSetup(async () => [node]); - expect(await service.triggerCollection('c1')).toMatchObject({ scanComplete: true }); + const { service, listener } = offlineSetup(async () => [node]); + const snapshot = await service.triggerCollection('c1'); + expect(snapshot).toMatchObject({ scanComplete: true }); + expect(listener).toHaveBeenLastCalledWith(expect.objectContaining({ connectionId: 'c1' }), { partial: false }); }); }); - it('fails the collection when every primary fails', async () => { - const { service, client, storage } = setup(); + describe('partial reads with too few keys', () => { + function partialSetup() { + const ctx = setup([KEY(0)]); + const broken = fakeClient([]); + broken.getDbSize.mockRejectedValue(new Error('LOADING')); + jest.spyOn(ctx.service as any, 'nodesFor').mockResolvedValue({ + nodes: [ + { name: 'a', client: ctx.client }, + { name: 'b', client: broken }, + ], + skipped: 0, + }); + return ctx; + } + + it('keeps the previous detection instead of flipping to not detected', async () => { + const { service, storage } = partialSetup(); + storage.getKvCacheFootprintSnapshots.mockResolvedValue([{ connectionId: 'c1', detected: true }]); + const snapshot = await service.triggerCollection('c1'); + expect(snapshot).toMatchObject({ detected: true, scanComplete: false }); + expect(storage.getKvCacheFootprintSnapshots).toHaveBeenCalledWith({ connectionId: 'c1', limit: 1 }); + }); + + it('keeps the sample key when the readable nodes hold no LMCache keys', async () => { + const { service, storage } = partialSetup(); + storage.getKvCacheFootprintSnapshots.mockResolvedValue([{ connectionId: 'c1', detected: true }]); + await service.triggerCollection('c1'); + expect(service.getSampleKey('c1')).toBe('m/x@1@0@1…@bfloat16'); + const broken = fakeClient([]); + broken.getDbSize.mockRejectedValue(new Error('LOADING')); + (service as any).nodesFor.mockResolvedValue({ + nodes: [ + { name: 'a', client: fakeClient([]) }, + { name: 'b', client: broken }, + ], + skipped: 0, + }); + await service.triggerCollection('c1'); + expect(service.getSampleKey('c1')).toBe('m/x@1@0@1…@bfloat16'); + }); + + it('neither backs off polling nor switches the next scan to a probe', async () => { + const { service, storage, client } = partialSetup(); + await service.triggerCollection('c1'); + expect(storage.saveKvCacheFootprintSnapshot.mock.calls[0][0].detected).toBe(false); + expect((service as any).undetected.has('c1')).toBe(false); + client.getDbSize.mockClear(); + await (service as any).tick(); + expect(client.getDbSize).toHaveBeenCalledTimes(1); + }); + }); + + it('fails the collection when every node fails', async () => { + const { service, client } = setup(); client.getDbSize.mockRejectedValue(new Error('LOADING')); await expect(service.triggerCollection('c1')).rejects.toThrow('LOADING'); - expect(storage.saveKvCacheFootprintSnapshot).not.toHaveBeenCalled(); + }); + + it('scans fully first and only probes a connection whose last snapshot was undetected', async () => { + const { service, client } = setup(); + client.call.mockImplementation(async (cmd: string) => (cmd === 'SCAN' ? ['7', []] : REPLIES[cmd])); + client.getDbSize.mockResolvedValue(1_000_000); + const scans = () => client.call.mock.calls.filter(([cmd]: [string]) => cmd === 'SCAN').length; + await service.triggerCollection('c1'); + expect(scans()).toBe(200); + + client.call.mockClear(); + await service.triggerCollection('c1'); + expect(scans()).toBe(2); + + (service as any).forget('c1'); + client.call.mockClear(); + await service.triggerCollection('c1'); + expect(scans()).toBe(200); + }); + + it('deletes the data of a retired member once it is deleted, even with no connected connections', async () => { + jest.useFakeTimers({ doNotFake: ['setImmediate', 'nextTick'] }); + const { service, storage, registry } = setup(); + const member = { id: 'm1', name: 'member', isConnected: false, connectionType: 'direct' }; + registry.list.mockImplementation((options?: { includeRetired?: boolean }) => (options?.includeRetired ? [member] : [])); + try { + await service.onModuleInit(); + await jest.advanceTimersByTimeAsync(300000); + expect(storage.deleteKvCacheConnectionData).not.toHaveBeenCalled(); + + registry.list.mockImplementation(() => []); + await jest.advanceTimersByTimeAsync(300000); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledWith('m1'); + } finally { + await service.onModuleDestroy(); + jest.useRealTimers(); + } }); it('lists other databases relative to the database the client actually selected', async () => { @@ -585,6 +766,7 @@ describe('KvCacheFootprintService', () => { jest.spyOn(service as any, 'start').mockImplementation(() => undefined); await service.onModuleInit(); await (service as any).tick(); + registry.list.mockReturnValue([]); registry.removed('c1'); await new Promise((resolve) => setImmediate(resolve)); expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledWith('c1'); @@ -595,8 +777,10 @@ describe('KvCacheFootprintService', () => { it('stops polling a retired member but keeps its data', async () => { const { service, storage, client, registry } = setup(); - const listener = jest.fn(); - service.onConnectionRemoval(listener); + const removed = jest.fn(); + const restored = jest.fn(); + service.onConnectionRemoval(removed); + service.onConnectionRestore(restored); jest.spyOn(service as any, 'start').mockImplementation(() => undefined); await service.onModuleInit(); await (service as any).tick(); @@ -608,7 +792,8 @@ describe('KvCacheFootprintService', () => { await new Promise((resolve) => setImmediate(resolve)); expect(client.getDbSize).toHaveBeenCalledTimes(1); expect(storage.deleteKvCacheConnectionData).not.toHaveBeenCalled(); - expect(listener).not.toHaveBeenCalled(); + expect(removed).not.toHaveBeenCalled(); + expect(restored).not.toHaveBeenCalled(); expect(service.getSampleKey('c1')).toBe('m/x@1@0@1…@bfloat16'); await service.onModuleDestroy(); }); @@ -643,6 +828,25 @@ describe('KvCacheFootprintService', () => { return ctx; } + it('does not let a discarded run touch the miss count of a re-registered connection', async () => { + const { service, storage, client, registry } = partialSetup(); + storage.getKvCacheFootprintSnapshots.mockResolvedValue([{ connectionId: 'c1', detected: true }]); + const listed = registry.list(); + let finishRead: (size: number) => void = () => undefined; + client.getDbSize.mockImplementationOnce(() => new Promise((resolve) => (finishRead = resolve))); + const stale = service.triggerCollection('c1'); + await new Promise((resolve) => setImmediate(resolve)); + registry.list.mockReturnValue([]); + (service as any).forget('c1'); + registry.list.mockReturnValue(listed); + await new Promise((resolve) => setImmediate(resolve)); + finishRead(0); + await expect(stale).resolves.toBeNull(); + for (let i = 0; i < PARTIAL_MISSES_BEFORE_CLEAR; i += 1) await service.triggerCollection('c1'); + const saved = storage.saveKvCacheFootprintSnapshot.mock.calls.map(([snapshot]: [any]) => snapshot.detected); + expect(saved).toEqual([true, true, false]); + }); + it('keeps polling when slot owners were left unscanned and nothing was detected', async () => { const { service, storage } = partialSetup(); await (service as any).tick(); @@ -702,8 +906,7 @@ describe('KvCacheFootprintService', () => { await (service as any).tick(); expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(1); client.getDbSize.mockResolvedValue(keys.length); - client.call.mockImplementation(async (cmd: string) => (cmd === 'SCAN' ? ['0', keys] : '')); - client.getClient().pipeline().exec.mockImplementation(async () => keys.flatMap(() => [[null, 100], [null, -1]])); + client.call.mockImplementation(async (cmd: string) => (cmd === 'SCAN' ? ['0', keys] : REPLIES[cmd])); expect((await service.triggerCollection('c1'))?.detected).toBe(true); await (service as any).tick(); expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(3); @@ -718,111 +921,281 @@ describe('KvCacheFootprintService', () => { it('deletes stored data when a connection is removed', async () => { const { service, storage, registry } = setup(); - jest.spyOn(service as any, 'start').mockImplementation(() => undefined); - await service.onModuleInit(); - registry.removed('c1'); + registry.list.mockReturnValue([]); + (service as any).forget('c1'); await new Promise((resolve) => setImmediate(resolve)); expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledWith('c1'); - await service.onModuleDestroy(); }); - it('notifies removal listeners with the connection id', async () => { - const { service, registry } = setup(); - jest.spyOn(service as any, 'start').mockImplementation(() => undefined); - await service.onModuleInit(); + it('deletes stored data only after async removal listeners settle', async () => { + const { service, storage, registry } = setup(); + registry.list.mockReturnValue([]); + let release: () => void = () => undefined; + service.onConnectionRemoval(() => new Promise((resolve) => (release = resolve))); + (service as any).forget('c1'); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.deleteKvCacheConnectionData).not.toHaveBeenCalled(); + release(); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledWith('c1'); + }); + + it('discards an in-flight collection and deletes data after it settles when the connection is removed', async () => { + const { service, storage, client, registry } = setup(); + let finishRead: (size: number) => void = () => undefined; + client.getDbSize.mockImplementationOnce(() => new Promise((resolve) => (finishRead = resolve))); const listener = jest.fn(); - const failing = jest.fn(() => { throw new Error('x'); }); - service.onConnectionRemoval(failing); - service.onConnectionRemoval(listener); - registry.removed('c1'); + service.onSnapshot(listener); + const collecting = service.triggerCollection('c1'); await new Promise((resolve) => setImmediate(resolve)); - expect(listener).toHaveBeenCalledWith('c1'); - await service.onModuleDestroy(); + registry.list.mockReturnValue([]); + (service as any).forget('c1'); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.deleteKvCacheConnectionData).not.toHaveBeenCalled(); + finishRead(6); + await expect(collecting).resolves.toBeNull(); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.saveKvCacheFootprintSnapshot).not.toHaveBeenCalled(); + expect(listener).not.toHaveBeenCalled(); + expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledWith('c1'); + }); + + it('keeps fresh data when the connection is re-registered before an in-flight run settles', async () => { + const { service, storage, client, registry } = setup(); + const listed = registry.list(); + let finishRead: (size: number) => void = () => undefined; + client.getDbSize.mockImplementationOnce(() => new Promise((resolve) => (finishRead = resolve))); + const stale = service.triggerCollection('c1'); + await new Promise((resolve) => setImmediate(resolve)); + registry.list.mockReturnValue([]); + (service as any).forget('c1'); + registry.list.mockReturnValue(listed); + const fresh = await service.triggerCollection('c1'); + expect(fresh).not.toBeNull(); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(1); + finishRead(6); + await expect(stale).resolves.toBeNull(); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(1); + expect(storage.deleteKvCacheConnectionData).not.toHaveBeenCalled(); + expect(await service.triggerCollection('c1')).not.toBeNull(); + }); + + it('announces a restore instead of deleting when the connection is registered again', async () => { + const { service, storage } = setup(); + const restored = jest.fn(); + service.onConnectionRestore(restored); + (service as any).forget('c1'); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.deleteKvCacheConnectionData).not.toHaveBeenCalled(); + expect(restored).toHaveBeenCalledWith('c1'); }); - it('finishes removal listeners before deleting stored data', async () => { + it('announces a restore instead of deleting when the connection was only retired', async () => { const { service, storage, registry } = setup(); - jest.spyOn(service as any, 'start').mockImplementation(() => undefined); - await service.onModuleInit(); - let finish: () => void = () => undefined; - const listener = jest.fn(() => new Promise((resolve) => (finish = resolve))); - service.onConnectionRemoval(listener); - registry.removed('c1'); + const restored = jest.fn(); + service.onConnectionRestore(restored); + registry.list.mockImplementation((options?: { includeRetired?: boolean }) => + options?.includeRetired ? [{ id: 'c1', name: 'one', isConnected: true, connectionType: 'direct' }] : [], + ); + (service as any).forget('c1'); await new Promise((resolve) => setImmediate(resolve)); - expect(listener).toHaveBeenCalledWith('c1'); expect(storage.deleteKvCacheConnectionData).not.toHaveBeenCalled(); - finish(); + expect(restored).toHaveBeenCalledWith('c1'); + }); + + it('does not announce a restore when the connection stays removed', async () => { + const { service, storage, registry } = setup(); + const restored = jest.fn(); + service.onConnectionRestore(restored); + registry.list.mockReturnValue([]); + (service as any).forget('c1'); await new Promise((resolve) => setImmediate(resolve)); expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledWith('c1'); - await service.onModuleDestroy(); + expect(restored).not.toHaveBeenCalled(); }); - it('deletes data for a connection removed from the registry before the poller saw it', async () => { + it('waits for an in-flight removal delete before saving for a re-registered connection', async () => { const { service, storage, registry } = setup(); - jest.spyOn(service as any, 'start').mockImplementation(() => undefined); - await service.onModuleInit(); - registry.removed('c2'); + const listed = registry.list(); + const order: string[] = []; + let finishDelete: () => void = () => undefined; + storage.deleteKvCacheConnectionData.mockImplementationOnce( + () => new Promise((resolve) => (finishDelete = () => { order.push('delete'); resolve(); })), + ); + storage.saveKvCacheFootprintSnapshot.mockImplementation(async () => { order.push('save'); }); + registry.list.mockReturnValue([]); + (service as any).forget('c1'); await new Promise((resolve) => setImmediate(resolve)); - expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledWith('c2'); - await service.onModuleDestroy(); + expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledWith('c1'); + registry.list.mockReturnValue(listed); + const collecting = service.triggerCollection('c1'); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.saveKvCacheFootprintSnapshot).not.toHaveBeenCalled(); + finishDelete(); + await expect(collecting).resolves.not.toBeNull(); + expect(order).toEqual(['delete', 'save']); }); - it('waits for an in-flight collection before deleting a removed connection', async () => { + it('holds a save that races an in-flight removal delete until the delete completes', async () => { const { service, storage, client, registry } = setup(); - jest.spyOn(service as any, 'start').mockImplementation(() => undefined); - await service.onModuleInit(); - let release!: () => void; - const info = client.getInfoParsed.getMockImplementation(); - client.getInfoParsed.mockImplementationOnce( - () => new Promise((resolve) => { release = () => resolve(info!()); }), + const listed = registry.list(); + const order: string[] = []; + let release: () => void = () => undefined; + let finishRead: (size: number) => void = () => undefined; + let finishDelete: () => void = () => undefined; + service.onConnectionRemoval(() => new Promise((resolve) => (release = resolve))); + client.getDbSize.mockImplementationOnce(() => new Promise((resolve) => (finishRead = resolve))); + storage.deleteKvCacheConnectionData.mockImplementationOnce( + () => new Promise((resolve) => (finishDelete = () => { order.push('delete'); resolve(); })), ); - const collecting = service.triggerCollection('c1'); + storage.saveKvCacheFootprintSnapshot.mockImplementation(async () => { order.push('save'); }); + registry.list.mockReturnValue([]); + (service as any).forget('c1'); + const collecting = service.collect({ connectionId: 'c1', connectionName: 'one', client, host: 'h', port: 6379 }); await new Promise((resolve) => setImmediate(resolve)); - registry.removed('c1'); + release(); await new Promise((resolve) => setImmediate(resolve)); - expect(storage.deleteKvCacheConnectionData).not.toHaveBeenCalled(); + expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledWith('c1'); + registry.list.mockReturnValue(listed); + finishRead(6); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.saveKvCacheFootprintSnapshot).not.toHaveBeenCalled(); + finishDelete(); + await expect(collecting).resolves.not.toBeNull(); + expect(order).toEqual(['delete', 'save']); + }); + + it('drops a poll that waited on the removal delete when the connection stays removed', async () => { + const { service, storage, client, registry } = setup(); + const listener = jest.fn(); + service.onSnapshot(listener); + let finishDelete: () => void = () => undefined; + storage.deleteKvCacheConnectionData.mockImplementationOnce( + () => new Promise((resolve) => (finishDelete = resolve)), + ); + registry.list.mockReturnValue([]); + registry.getConfig.mockReturnValue(null); + (service as any).forget('c1'); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledWith('c1'); + const collecting = service.collect({ connectionId: 'c1', connectionName: 'one', client, host: 'h', port: 6379 }); + await new Promise((resolve) => setImmediate(resolve)); + finishDelete(); + await expect(collecting).resolves.toBeNull(); + expect(storage.saveKvCacheFootprintSnapshot).not.toHaveBeenCalled(); + expect(listener).not.toHaveBeenCalled(); + }); + + it('does not save or announce when the connection disappears mid-collect', async () => { + const { service, storage, client, registry } = setup(); + const listener = jest.fn(); + service.onSnapshot(listener); + let finishRead: (size: number) => void = () => undefined; + client.getDbSize.mockImplementationOnce(() => new Promise((resolve) => (finishRead = resolve))); + const collecting = service.collect({ connectionId: 'c1', connectionName: 'one', client, host: 'h', port: 6379 }); + await new Promise((resolve) => setImmediate(resolve)); + registry.getConfig.mockReturnValue(null); + finishRead(6); + await expect(collecting).resolves.toBeNull(); + expect(storage.saveKvCacheFootprintSnapshot).not.toHaveBeenCalled(); + expect(listener).not.toHaveBeenCalled(); + }); + + it('deletes once for the latest removal when a connection is removed again before data is deleted', async () => { + const { service, storage, registry } = setup(); + registry.list.mockReturnValue([]); + const releases: Array<() => void> = []; + service.onConnectionRemoval(() => new Promise((resolve) => releases.push(resolve))); + (service as any).forget('c1'); + (service as any).forget('c1'); + releases.forEach((release) => release()); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledTimes(1); + }); + + it('forgets the removal generation once the data is deleted', async () => { + const { service, storage, registry } = setup(); + registry.list.mockReturnValue([]); + (service as any).forget('c1'); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledWith('c1'); + expect((service as any).generations.has('c1')).toBe(false); + }); + + it('notifies removal listeners with the connection id', () => { + const { service } = setup(); + const listener = jest.fn(); + const failing = jest.fn(() => { throw new Error('x'); }); + service.onConnectionRemoval(failing); + service.onConnectionRemoval(listener); + (service as any).forget('c1'); + expect(listener).toHaveBeenCalledWith('c1'); + }); + + it('runs removal listeners only after an in-flight snapshot announce finishes', async () => { + const { service, registry } = setup(); + const order: string[] = []; + let release: () => void = () => undefined; + service.onSnapshot(async () => { + order.push('announce:start'); + await new Promise((resolve) => (release = resolve)); + order.push('announce:end'); + }); + service.onConnectionRemoval(() => { + order.push('removed'); + }); + const collecting = service.triggerCollection('c1'); + while (!order.includes('announce:start')) await new Promise((resolve) => setImmediate(resolve)); + registry.list.mockReturnValue([]); + (service as any).forget('c1'); + await new Promise((resolve) => setImmediate(resolve)); + expect(order).toEqual(['announce:start']); release(); await collecting; await new Promise((resolve) => setImmediate(resolve)); - const saved = storage.saveKvCacheFootprintSnapshot.mock.invocationCallOrder[0]; - const deleted = storage.deleteKvCacheConnectionData.mock.invocationCallOrder[0]; - expect(deleted).toBeGreaterThan(saved); - expect(service.getSampleKey('c1')).toBeNull(); - await service.onModuleDestroy(); + expect(order).toEqual(['announce:start', 'announce:end', 'removed']); }); - it('returns the in-flight collection to a refresh', async () => { - const { service, client, storage } = setup(); - let release!: () => void; - const info = client.getInfoParsed.getMockImplementation(); - client.getInfoParsed.mockImplementationOnce( - () => new Promise((resolve) => { release = () => resolve(info!()); }), - ); - const first = service.triggerCollection('c1'); + it('keeps the snapshot of a re-registered connection when the older announce finishes after it', async () => { + const { service } = setup(); + const order: string[] = []; + let release: () => void = () => undefined; + let calls = 0; + service.onSnapshot(async () => { + calls += 1; + if (calls > 1) { + order.push('snapshot:new'); + return; + } + order.push('announce:start'); + await new Promise((resolve) => (release = resolve)); + order.push('announce:end'); + }); + service.onConnectionRemoval(() => { + order.push('removed'); + }); + const stale = service.triggerCollection('c1'); + while (!order.includes('announce:start')) await new Promise((resolve) => setImmediate(resolve)); + (service as any).forget('c1'); + const fresh = service.triggerCollection('c1'); await new Promise((resolve) => setImmediate(resolve)); - const second = service.triggerCollection('c1'); release(); - const [a, b] = await Promise.all([first, second]); - expect(b).not.toBeNull(); - expect(b).toBe(a); - expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(1); + await stale; + await expect(fresh).resolves.not.toBeNull(); + await new Promise((resolve) => setImmediate(resolve)); + expect(order).toEqual(['announce:start', 'announce:end', 'removed', 'snapshot:new']); }); - it('rejects a refresh for a cluster member, which its seed collects', async () => { - const { service, client, registry } = setup(); - registry.list.mockReturnValue([ - { id: 'seed', name: 'seed', isConnected: true, connectionType: 'direct' }, - { - id: 'c1', - name: 'one', - isConnected: true, - connectionType: 'direct', - membership: { seedId: 'seed', nodeId: 'n', origin: 'auto', source: 'cluster' }, - }, - ]); - await expect(service.triggerCollection('c1')).rejects.toThrow(/seed/); - await expect(service.triggerCollection('c1')).rejects.toBeInstanceOf(BadRequestException); - expect(client.getDbSize).not.toHaveBeenCalled(); + it('deletes data for a connection removed from the registry before the poller saw it', async () => { + const { service, storage, registry } = setup(); + jest.spyOn(service as any, 'start').mockImplementation(() => undefined); + await service.onModuleInit(); + registry.list.mockReturnValue([]); + registry.removed('c2'); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledWith('c2'); + await service.onModuleDestroy(); }); it('rejects a refresh for a retired member without scanning it', async () => { @@ -849,7 +1222,7 @@ describe('KvCacheFootprintService', () => { it('rejects a refresh for a disconnected connection without scanning it', async () => { const { service, client, registry } = setup(); registry.list.mockReturnValue([{ id: 'c1', name: 'one', isConnected: false, connectionType: 'direct' }]); - await expect(service.triggerCollection('c1')).rejects.toBeInstanceOf(BadRequestException); + await expect(service.triggerCollection('c1')).rejects.toBeInstanceOf(ServiceUnavailableException); expect(client.getDbSize).not.toHaveBeenCalled(); }); diff --git a/proprietary/kv-cache/__tests__/kv-cache-samples.service.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-samples.service.spec.ts index c33570675..27087a67c 100644 --- a/proprietary/kv-cache/__tests__/kv-cache-samples.service.spec.ts +++ b/proprietary/kv-cache/__tests__/kv-cache-samples.service.spec.ts @@ -166,7 +166,7 @@ describe('KvCacheSamplesService', () => { connectionId: 'c1', from: to - SAMPLES_MAX_RANGE_MS, to, - limit: SAMPLES_LIMIT, + limit: SAMPLES_LIMIT + 1, bucketMs: expect.any(Number), engineId: 'e1', modelName: 'm', @@ -207,9 +207,45 @@ describe('KvCacheSamplesService', () => { expect((await service.getSamples('c1', { from: 0, to: 1000 })).rangeHitRate).toBe(0.75); }); + it('flags a result only when rows beyond the limit exist', async () => { + const { service, storage } = setup(); + storage.getKvCacheEngineSamples.mockResolvedValue(Array.from({ length: SAMPLES_LIMIT }, (_, i) => sample({ timestamp: i }))); + expect((await service.getSamples('c1', { from: 0, to: 1000 })).truncated).toBe(false); + storage.getKvCacheEngineSamples.mockResolvedValue(Array.from({ length: SAMPLES_LIMIT + 1 }, (_, i) => sample({ timestamp: i }))); + const result = await service.getSamples('c1', { from: 0, to: 1000 }); + expect(result.truncated).toBe(true); + expect(result.buckets).toHaveLength(SAMPLES_LIMIT); + expect(result.buckets[0].timestamp).toBe(1); + }); + + it('drops the oldest bucket when the row limit splits it', async () => { + const { service, storage } = setup(); + const boundary = [sample({ timestamp: 0, requestedTokens: 1, hitTokens: 1 }), sample({ timestamp: 0, engineId: 'e2', requestedTokens: 1 })]; + const newer = Array.from({ length: SAMPLES_LIMIT - 1 }, (_, i) => sample({ timestamp: Math.floor(i / 2) + 1, engineId: i % 2 ? 'e2' : 'e1', requestedTokens: 2, hitTokens: 1 })); + storage.getKvCacheEngineSamples.mockResolvedValue([...boundary, ...newer]); + storage.getKvCacheSampleTotals.mockResolvedValue({ requestedTokens: 4, hitTokens: 3 }); + const result = await service.getSamples('c1', { from: 0, to: 1_000_000 }); + expect(result.truncated).toBe(true); + expect(result.buckets.some((bucket) => bucket.timestamp === 0)).toBe(false); + expect(result.buckets[0]).toMatchObject({ timestamp: 1, requestedTokens: 2 }); + expect(result.rangeHitRate).toBe(0.75); + }); + + it('keeps the newest rows when every fetched row shares the oldest timestamp', async () => { + const { service, storage } = setup(); + const rows = Array.from({ length: SAMPLES_LIMIT + 1 }, (_, i) => sample({ timestamp: 5, engineId: `e${i}`, requestedTokens: 2, hitTokens: 1 })); + storage.getKvCacheEngineSamples.mockResolvedValue(rows); + storage.getKvCacheSampleTotals.mockResolvedValue({ requestedTokens: 4, hitTokens: 3 }); + const result = await service.getSamples('c1', { from: 0, to: 1000 }); + expect(result.truncated).toBe(true); + expect(result.buckets).toHaveLength(SAMPLES_LIMIT); + expect(result.buckets.some((bucket) => bucket.engineId === 'e0')).toBe(false); + expect(result.rangeHitRate).toBe(0.75); + }); + it('returns a null range hit rate when nothing was requested', async () => { const { service } = setup(); - expect(await service.getSamples('c1', { from: 0, to: 1000 })).toEqual({ buckets: [], rangeHitRate: null, bucketMs: 60_000 }); + expect(await service.getSamples('c1', { from: 0, to: 1000 })).toEqual({ buckets: [], rangeHitRate: null, truncated: false, bucketMs: 60_000 }); }); it('waits for an in-flight save even when there is nothing new to drain', async () => { diff --git a/proprietary/kv-cache/__tests__/kv-cache-scrape.service.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-scrape.service.spec.ts index 5ac9a988d..f58e07367 100644 --- a/proprietary/kv-cache/__tests__/kv-cache-scrape.service.spec.ts +++ b/proprietary/kv-cache/__tests__/kv-cache-scrape.service.spec.ts @@ -20,7 +20,7 @@ const engine = (over: Record = {}) => ({ const body = 'lmcache:num_hit_tokens_total{model_name="m",worker_id="0",role="worker"} 42\n'; -function setup(engines: unknown[], licensed = true) { +function setup(engines: unknown[], licensed = true, connectionIds = ['c1']) { const registry = { list: jest.fn().mockReturnValue(engines), get: jest.fn((id: string) => (engines as { id: string }[]).find((e) => e.id === id) ?? null), @@ -29,8 +29,9 @@ function setup(engines: unknown[], licensed = true) { const samples = { observe: jest.fn() }; const engineService = { authHeaderFor: jest.fn().mockReturnValue(null) }; const license = { hasFeature: jest.fn().mockReturnValue(licensed) }; - const service = new KvCacheScrapeService(registry as any, samples as any, engineService as any, license as any); - return { registry, samples, engineService, license, service }; + const connections = { getConfig: jest.fn((id: string) => (connectionIds.includes(id) ? { id } : null)) }; + const service = new KvCacheScrapeService(registry as any, samples as any, engineService as any, license as any, connections as any); + return { registry, samples, engineService, license, connections, service }; } describe('KvCacheScrapeService', () => { @@ -48,6 +49,24 @@ describe('KvCacheScrapeService', () => { expect(fetchMock).not.toHaveBeenCalled(); }); + it('skips engines whose connection was removed', async () => { + const { service } = setup([engine()], true, []); + await service.tick(1000); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it('drops a result when the connection is removed mid-fetch', async () => { + fetchMock.mockResolvedValue(body); + const { service, samples, registry, connections } = setup([engine()]); + fetchMock.mockImplementationOnce(async () => { + connections.getConfig.mockReturnValue(null); + return body; + }); + await service.tick(1000); + expect(samples.observe).not.toHaveBeenCalled(); + expect(registry.recordResult).not.toHaveBeenCalled(); + }); + it('observes cumulative samples and records success', async () => { fetchMock.mockResolvedValue(body); const { service, samples, registry } = setup([engine()]); diff --git a/proprietary/kv-cache/__tests__/metrics-fetch.spec.ts b/proprietary/kv-cache/__tests__/metrics-fetch.spec.ts index 1b02ad8cb..ef18c4ff4 100644 --- a/proprietary/kv-cache/__tests__/metrics-fetch.spec.ts +++ b/proprietary/kv-cache/__tests__/metrics-fetch.spec.ts @@ -2,7 +2,7 @@ import { BadRequestException, Logger } from '@nestjs/common'; import { promises as dns } from 'dns'; import { createServer, type IncomingHttpHeaders, type Server } from 'http'; import { createServer as createTcpServer, type AddressInfo } from 'net'; -import { assertSafeOutboundUrl } from '@app/common/utils/outbound-url-guard'; +import { assertSafeOutboundUrl, BlockedAddressError } from '@app/common/utils/outbound-url-guard'; import { fetchMetricsText, scrapeAllowsPrivateNetworks, ScrapeError } from '../metrics-fetch'; jest.mock('@app/common/utils/outbound-url-guard', () => ({ @@ -47,6 +47,16 @@ describe('fetchMetricsText', () => { } }); + it('surfaces a connect-time blocked address', async () => { + const blocked = jest.fn().mockRejectedValue(new BlockedAddressError('Metrics URL resolves to a blocked address')); + await expect(fetchMetricsText('http://x', null, blocked)).rejects.toThrow(new ScrapeError('Metrics URL resolves to a blocked address')); + }); + + it('reports a DNS failure', async () => { + const missing = jest.fn().mockRejectedValue(Object.assign(new Error('getaddrinfo ENOTFOUND x'), { code: 'ENOTFOUND' })); + await expect(fetchMetricsText('http://x', null, missing)).rejects.toThrow(new ScrapeError('DNS lookup failed (ENOTFOUND)')); + }); + it('maps fetch failures to timeout and connection failed', async () => { const timeout = jest.fn().mockRejectedValue(Object.assign(new Error('t'), { name: 'TimeoutError' })); await expect(fetchMetricsText('http://x', null, timeout)).rejects.toThrow('timeout'); diff --git a/proprietary/kv-cache/footprint-estimate.ts b/proprietary/kv-cache/footprint-estimate.ts index 674cd8271..9d5236dd4 100644 --- a/proprietary/kv-cache/footprint-estimate.ts +++ b/proprietary/kv-cache/footprint-estimate.ts @@ -23,20 +23,6 @@ export interface MemoryStats { evictedKeys: number; } -export function evictedKeysDelta( - previous: Record | null, - current: Record, -): number | null { - if (previous === null) return null; - const shared = Object.keys(current).filter((node) => node in previous); - if (shared.length === 0) return null; - return shared.reduce((sum, node) => { - const before = previous[node]; - const now = current[node]; - return sum + (now >= before ? now - before : now); - }, 0); -} - export interface ScanCoverage { primariesScanned: number; primariesTotal: number | null; @@ -56,6 +42,17 @@ export interface SnapshotInput { export const DETECTION_MIN_MATCHES = 5; +export function evictedKeysDelta(previous: Map | undefined, current: Map): number | null { + if (!previous) return null; + let delta: number | null = null; + for (const [node, evicted] of current) { + const before = previous.get(node); + if (before === undefined) continue; + delta = (delta ?? 0) + (evicted >= before ? evicted - before : evicted); + } + return delta; +} + const ratio = (value: number) => Math.round(value * 1e4) / 1e4; export function layoutOf(keys: ParsedLmcacheKey[]): KvCacheLayout | null { diff --git a/proprietary/kv-cache/footprint-scanner.ts b/proprietary/kv-cache/footprint-scanner.ts index f37078d25..5ad214ce0 100644 --- a/proprietary/kv-cache/footprint-scanner.ts +++ b/proprietary/kv-cache/footprint-scanner.ts @@ -16,6 +16,8 @@ export interface ScanResult { } const SCAN_COUNT = 1000; +const PROBE_PAGES = 2; +const SAMPLE_BATCH = 100; const GLIDE_LMCACHE = /(?:^| )lib-name=GlidePySync\(lmcache(?::|\))/; const CLIENT_DB = /(?:^| )db=(\d+)(?: |$)/; @@ -25,7 +27,12 @@ function isLmcacheClientOn(line: string, db: number): boolean { return clientDb === undefined || Number(clientDb) === db; } -export async function scanLmcacheKeys(client: DatabasePort, budgets: ScanBudgets, dbSize: number): Promise { +export async function scanLmcacheKeys( + client: DatabasePort, + budgets: ScanBudgets, + dbSize: number, + probe = false, +): Promise { let cursor = '0'; let iterations = 0; let returned = 0; @@ -39,6 +46,7 @@ export async function scanLmcacheKeys(client: DatabasePort, budgets: ScanBudgets const parsed = classifyKey(key); if (parsed) matched.push(parsed); } + if (probe && iterations >= PROBE_PAGES && matched.length === 0) break; } while ( cursor !== '0' && returned < budgets.maxScanned && @@ -55,38 +63,43 @@ export async function scanLmcacheKeys(client: DatabasePort, budgets: ScanBudgets }; } +type Settled = { ok: true; value: unknown } | { ok: false; error: unknown }; + +const settle = (promise: Promise): Promise => + promise.then( + (value) => ({ ok: true as const, value }), + (error: unknown) => ({ ok: false as const, error }), + ); + +async function sampleKey(client: DatabasePort, parsed: ParsedLmcacheKey, ownsSlot?: SlotOwnership): Promise { + const { key } = parsed; + const candidate = siblingKey(parsed); + const sibling = candidate && (!ownsSlot || ownsSlot(keySlot(candidate))) ? candidate : null; + const [memory, ttl, exists] = await Promise.all([ + settle(client.call('MEMORY', ['USAGE', key])), + settle(client.call('TTL', [key])), + sibling ? settle(client.call('EXISTS', [sibling])) : Promise.resolve(null), + ]); + if (!memory.ok && !ttl.ok && exists?.ok !== true) throw ttl.error; + return { + parsed, + memoryBytes: memory.ok && memory.value !== null ? Number(memory.value) : null, + ttl: ttl.ok ? Number(ttl.value) : null, + siblingExists: exists?.ok ? Number(exists.value) === 1 : null, + }; +} + export async function sampleLmcacheKeys( client: DatabasePort, keys: ParsedLmcacheKey[], ownsSlot?: SlotOwnership, ): Promise { - if (keys.length === 0) return []; - const pipeline = client.getClient().pipeline(); - const plan = keys.map((parsed) => { - const candidate = siblingKey(parsed); - const sibling = candidate && (!ownsSlot || ownsSlot(keySlot(candidate))) ? candidate : null; - pipeline.call('MEMORY', 'USAGE', parsed.key); - pipeline.ttl(parsed.key); - if (sibling) pipeline.exists(sibling); - return { parsed, sibling }; - }); - const replies = ((await pipeline.exec()) ?? []) as Array<[Error | null, unknown]>; - let index = 0; - return plan.map(({ parsed, sibling }) => { - const [memoryErr, memory] = replies[index++] ?? [new Error('missing'), null]; - const [ttlErr, ttl] = replies[index++] ?? [new Error('missing'), null]; - let siblingExists: boolean | null = null; - if (sibling) { - const [existsErr, exists] = replies[index++] ?? [new Error('missing'), null]; - siblingExists = existsErr ? null : Number(exists) === 1; - } - return { - parsed, - memoryBytes: memoryErr || memory === null ? null : Number(memory), - ttl: ttlErr ? null : Number(ttl), - siblingExists, - }; - }); + const samples: KeySample[] = []; + for (let start = 0; start < keys.length; start += SAMPLE_BATCH) { + const batch = keys.slice(start, start + SAMPLE_BATCH); + samples.push(...(await Promise.all(batch.map((parsed) => sampleKey(client, parsed, ownsSlot))))); + } + return samples; } export async function hasLmcacheGlideClient(client: DatabasePort, db: number): Promise { diff --git a/proprietary/kv-cache/kv-cache-alerts.service.ts b/proprietary/kv-cache/kv-cache-alerts.service.ts index 1caa1639f..900c077ac 100644 --- a/proprietary/kv-cache/kv-cache-alerts.service.ts +++ b/proprietary/kv-cache/kv-cache-alerts.service.ts @@ -16,7 +16,7 @@ import { PrometheusService } from '@app/prometheus/prometheus.service'; import { LicenseService } from '@proprietary/licenses/license.service'; import { evictionConditions, shouldCheckHitRate, windowHitRates } from './kv-cache-alert-rules'; import { KvCacheEngineRegistry } from './kv-cache-engine-registry'; -import { KvCacheFootprintService } from './kv-cache-footprint.service'; +import { KvCacheFootprintService, type SnapshotContext } from './kv-cache-footprint.service'; const EVALUATE_INTERVAL_MS = 60_000; const EVICTION_REASONS: KvCacheEvictionReason[] = ['unevictable', 'evicting']; @@ -35,7 +35,7 @@ export class KvCacheAlertsService implements OnModuleInit, OnModuleDestroy { private readonly exportedConnections = new Set(); private readonly checkedWindows = new Map>(); private readonly evictingCleanStreaks = new Map(); - private readonly unevictableActive = new Map(); + private readonly unevictableActive = new Set(); private evaluating = false; constructor( @@ -51,7 +51,7 @@ export class KvCacheAlertsService implements OnModuleInit, OnModuleDestroy { ) {} onModuleInit(): void { - this.footprint.onSnapshot((snapshot) => this.onSnapshot(snapshot)); + this.footprint.onSnapshot((snapshot, context) => this.onSnapshot(snapshot, context)); this.footprint.onConnectionRemoval((connectionId) => { this.exportedConnections.delete(connectionId); this.checkedWindows.delete(connectionId); @@ -87,13 +87,13 @@ export class KvCacheAlertsService implements OnModuleInit, OnModuleDestroy { return this.storage.saveKvCacheSettings({ ...current, ...definedUpdate, connectionId, updatedAt: Date.now() }); } - async onSnapshot(snapshot: KvCacheFootprintSnapshot): Promise { + async onSnapshot(snapshot: KvCacheFootprintSnapshot, context: SnapshotContext = { partial: false }): Promise { try { if (!this.licensed()) return; if (snapshot.scanComplete) this.prometheus?.setKvCacheFootprint(snapshot.connectionId, snapshot.perModel); if (!this.webhookEventsPro) return; const settings = await this.getSettings(snapshot.connectionId); - const conditions = this.evictionState(snapshot, settings.evictionAlertEnabled); + const conditions = this.evictionState(snapshot, settings.evictionAlertEnabled, context.partial); const instance = this.instanceFor(snapshot.connectionId); for (const reason of EVICTION_REASONS) { void this.webhookEventsPro @@ -121,14 +121,24 @@ export class KvCacheAlertsService implements OnModuleInit, OnModuleDestroy { this.logger.warn(`Eviction risk alert failed for ${connectionId}: ${error instanceof Error ? error.message : error}`); } - private evictionState(snapshot: KvCacheFootprintSnapshot, enabled: boolean): Record { - if (!enabled || !snapshot.detected) { + private evictionState( + snapshot: KvCacheFootprintSnapshot, + enabled: boolean, + partial: boolean, + ): Record { + if (!enabled || (!partial && !snapshot.detected)) { this.evictingCleanStreaks.delete(snapshot.connectionId); this.unevictableActive.delete(snapshot.connectionId); return { unevictable: false, evicting: false }; } - const conditions = evictionConditions(snapshot, this.unevictableActive.get(snapshot.connectionId)); - if (conditions.unevictable) this.unevictableActive.set(snapshot.connectionId, true); + if (partial) { + return { + unevictable: this.unevictableActive.has(snapshot.connectionId), + evicting: this.evictingCleanStreaks.has(snapshot.connectionId), + }; + } + const conditions = evictionConditions(snapshot, this.unevictableActive.has(snapshot.connectionId)); + if (conditions.unevictable) this.unevictableActive.add(snapshot.connectionId); else this.unevictableActive.delete(snapshot.connectionId); conditions.evicting = this.holdEvicting(snapshot.connectionId, conditions.evicting); return conditions; diff --git a/proprietary/kv-cache/kv-cache-engine-registry.ts b/proprietary/kv-cache/kv-cache-engine-registry.ts index e7c37b988..8c96a8d90 100644 --- a/proprietary/kv-cache/kv-cache-engine-registry.ts +++ b/proprietary/kv-cache/kv-cache-engine-registry.ts @@ -13,6 +13,7 @@ export class KvCacheEngineRegistry implements OnModuleInit { private readonly writes = new Map>(); private readonly statusVersions = new Map(); private readonly removed = new Map(); + private readonly removals = new Map(); private readonly removedConnections = new Map(); private readonly persistedSeenAt = new Map(); @@ -28,6 +29,7 @@ export class KvCacheEngineRegistry implements OnModuleInit { this.persistedSeenAt.set(engine.id, engine.lastSeenAt); } this.footprint.onConnectionRemoval(async (connectionId) => { + this.removals.set(connectionId, (this.removals.get(connectionId) ?? 0) + 1); this.tombstone(this.removedConnections, connectionId); const ids = this.list(connectionId).map((engine) => engine.id); for (const id of ids) { @@ -37,8 +39,46 @@ export class KvCacheEngineRegistry implements OnModuleInit { this.persistedSeenAt.delete(id); this.samples.forgetEngine(id); } - await Promise.all([...ids.map((id) => this.writes.get(id)), this.samples.idle()]); + await Promise.all(ids.map((id) => this.writes.get(id))); + return this.samples.settled(); }); + this.footprint.onConnectionRestore((connectionId) => this.restore(connectionId)); + } + + private async restore(connectionId: string): Promise { + this.removedConnections.delete(connectionId); + const removals = this.removals.get(connectionId); + const stored = (await this.storage.getKvCacheEngines()).filter((engine) => engine.connectionId === connectionId); + if (this.removals.get(connectionId) !== removals) return; + const superseded: string[] = []; + for (const engine of stored) { + if (this.engines.has(engine.id)) continue; + if (this.clashes(engine)) { + superseded.push(engine.id); + continue; + } + this.removed.delete(engine.id); + this.engines.set(engine.id, engine); + this.persistedSeenAt.set(engine.id, engine.lastSeenAt); + } + await Promise.all( + superseded.map((id) => + this.serialize(id, async () => { + if (this.removals.get(connectionId) !== removals || this.engines.has(id)) return; + await this.storage.deleteKvCacheEngine(id); + }).catch((error) => + this.logger.warn(`Could not delete superseded KV cache engine ${id}: ${error instanceof Error ? error.message : String(error)}`), + ), + ), + ); + } + + private clashes(engine: StoredKvCacheEngine): boolean { + return [...this.engines.values()].some( + (other) => + (other.connectionId === engine.connectionId && other.name === engine.name) || + (engine.otlpEngineId !== null && other.otlpEngineId === engine.otlpEngineId), + ); } list(connectionId?: string): StoredKvCacheEngine[] { @@ -62,10 +102,11 @@ export class KvCacheEngineRegistry implements OnModuleInit { return this.serialize(engine.id, async () => { if (this.removed.has(engine.id)) throw new NotFoundException(`KV cache engine ${engine.id} not found`); this.assertConnectionLive(engine.connectionId); + const removals = this.removals.get(engine.connectionId); const saved = await this.storage.saveKvCacheEngine(engine); - if (this.removedConnections.has(saved.connectionId)) { + if (this.removals.get(saved.connectionId) !== removals) { await this.storage.deleteKvCacheEngine(saved.id); - this.assertConnectionLive(saved.connectionId); + throw new NotFoundException(`Connection '${saved.connectionId}' not found`); } if (this.removed.has(saved.id)) throw new NotFoundException(`KV cache engine ${saved.id} not found`); const latest = this.engines.get(saved.id); @@ -130,7 +171,9 @@ export class KvCacheEngineRegistry implements OnModuleInit { } private assertConnectionLive(connectionId: string): void { - if (this.removedConnections.has(connectionId)) throw new NotFoundException(`Connection '${connectionId}' not found`); + if (this.removedConnections.has(connectionId) && !this.footprint.hasConnection(connectionId)) { + throw new NotFoundException(`Connection '${connectionId}' not found`); + } } private serialize(id: string, task: () => Promise): Promise { diff --git a/proprietary/kv-cache/kv-cache-footprint.service.ts b/proprietary/kv-cache/kv-cache-footprint.service.ts index a394f9d9a..b34dff921 100644 --- a/proprietary/kv-cache/kv-cache-footprint.service.ts +++ b/proprietary/kv-cache/kv-cache-footprint.service.ts @@ -1,4 +1,12 @@ -import { BadRequestException, Inject, Injectable, Logger, NotFoundException, OnModuleInit } from '@nestjs/common'; +import { + BadRequestException, + Inject, + Injectable, + Logger, + NotFoundException, + OnModuleInit, + ServiceUnavailableException, +} from '@nestjs/common'; import { Feature, type ConnectionStatus, @@ -32,7 +40,11 @@ import { import { hasSingleKeyEvidence, hasTwoKeyEvidence, maskKey, resolveTagSuffixes } from './key-classifier'; import { slotOwner, type SlotOwnership } from './key-slot'; -export type SnapshotListener = (snapshot: KvCacheFootprintSnapshot) => void | Promise; +export interface SnapshotContext { + partial: boolean; +} + +export type SnapshotListener = (snapshot: KvCacheFootprintSnapshot, context: SnapshotContext) => void | Promise; export type RemovalListener = (connectionId: string) => void | Promise; interface NodeReading { @@ -105,7 +117,7 @@ export class KvCacheFootprintService extends MultiConnectionPoller implements On ); private readonly budgets = { maxScanned: this.setting('KV_CACHE_SCAN_MAX_KEYS', 200000, 1), - maxMatched: this.setting('KV_CACHE_MATCH_MAX_KEYS', 2000, 1), + maxMatched: this.setting('KV_CACHE_MATCH_MAX_KEYS', 2000, DETECTION_MIN_MATCHES), }; private readonly sampleSize = this.setting('KV_CACHE_SAMPLE_KEYS', 500, 1); private readonly undetectedRecheckMs = this.setting( @@ -113,14 +125,23 @@ export class KvCacheFootprintService extends MultiConnectionPoller implements On 1800000, MIN_FOOTPRINT_INTERVAL_MS, ); - private readonly inFlight = new Map>(); - private readonly lastEvicted = new Map>(); - private unsubscribeRemoval: (() => void) | null = null; + private readonly running = new Map>(); + private readonly generations = new Map(); + private readonly deleting = new Map>(); + private readonly removing = new Map>(); + private readonly announcing = new Map>(); + private readonly lastEvicted = new Map>(); private readonly sampleKeys = new Map(); private readonly undetectedUntil = new Map(); + private readonly undetected = new Set(); private readonly partialMisses = new Map(); + private retired = new Set(); + private sweepTimer: NodeJS.Timeout | null = null; + private unsubscribeRemoval: (() => void) | null = null; + private readonly memoryBlocked = new Set(); private readonly listeners: SnapshotListener[] = []; private readonly removalListeners: RemovalListener[] = []; + private readonly restoreListeners: RemovalListener[] = []; constructor( connectionRegistry: ConnectionRegistry, @@ -133,9 +154,14 @@ export class KvCacheFootprintService extends MultiConnectionPoller implements On async onModuleInit(): Promise { this.unsubscribeRemoval = this.connectionRegistry.addConnectionRemovalListener((id) => this.forget(id)); this.start(); + this.sweepDeletedRetirees(); + this.sweepTimer = setInterval(() => this.sweepDeletedRetirees(), this.intervalMs); + this.sweepTimer.unref?.(); } async onModuleDestroy(): Promise { + if (this.sweepTimer) clearInterval(this.sweepTimer); + this.sweepTimer = null; this.unsubscribeRemoval?.(); this.unsubscribeRemoval = null; await super.onModuleDestroy(); @@ -162,28 +188,101 @@ export class KvCacheFootprintService extends MultiConnectionPoller implements On await this.collect(ctx); } - private async forget(connectionId: string): Promise { - await this.inFlight.get(connectionId)?.catch(() => null); + private forget(connectionId: string): void { + const active = this.running.get(connectionId); + this.running.delete(connectionId); + const generation = this.generationOf(connectionId) + 1; + this.generations.set(connectionId, generation); this.lastEvicted.delete(connectionId); this.sampleKeys.delete(connectionId); this.undetectedUntil.delete(connectionId); + this.undetected.delete(connectionId); this.partialMisses.delete(connectionId); - for (const listener of this.removalListeners) { + const settled = active?.then( + () => undefined, + () => undefined, + ); + const announcing = this.announcing.get(connectionId); + const removed = announcing + ? announcing.then(() => this.announceRemoval(connectionId)) + : this.announceRemoval(connectionId); + const removal = removed.then( + () => undefined, + () => undefined, + ); + this.removing.set(connectionId, removal); + void removal.then(() => { + if (this.removing.get(connectionId) === removal) this.removing.delete(connectionId); + }); + void Promise.all([removed, settled]) + .then(() => this.deleteRemovedData(connectionId, generation)) + .catch((error) => this.logger.warn(`Could not delete KV cache data for ${connectionId}: ${errorMessage(error)}`)); + } + + private async announceRemoval(connectionId: string): Promise { + await Promise.all( + this.removalListeners.map(async (listener) => { + try { + await listener(connectionId); + } catch (error) { + this.logger.warn(`KV cache removal listener failed: ${errorMessage(error)}`); + } + }), + ); + } + + private sweepDeletedRetirees(): void { + const all = this.connectionRegistry.list({ includeRetired: true }).map((connection) => connection.id); + const live = new Set(this.connectionRegistry.list().map((connection) => connection.id)); + const known = new Set(all); + const deleted = [...this.retired].filter((id) => !known.has(id)); + this.retired = new Set(all.filter((id) => !live.has(id))); + for (const id of deleted) this.forget(id); + } + + private async deleteRemovedData(connectionId: string, generation: number): Promise { + if (this.generationOf(connectionId) !== generation) return; + if (this.isRegistered(connectionId)) return this.announceRestore(connectionId); + if (!this.running.has(connectionId)) this.generations.delete(connectionId); + const deletion = this.storage.deleteKvCacheConnectionData(connectionId); + const settled = deletion.then( + () => undefined, + () => undefined, + ); + this.deleting.set(connectionId, settled); + try { + await deletion; + } finally { + if (this.deleting.get(connectionId) === settled) this.deleting.delete(connectionId); + } + } + + private async announceRestore(connectionId: string): Promise { + for (const listener of this.restoreListeners) { try { await listener(connectionId); } catch (error) { - this.logger.warn(`KV cache removal listener failed: ${errorMessage(error)}`); + this.logger.warn(`KV cache restore listener failed: ${errorMessage(error)}`); } } - await this.storage - .deleteKvCacheConnectionData(connectionId) - .catch((error) => this.logger.warn(`Could not delete KV cache data for ${connectionId}: ${errorMessage(error)}`)); + } + + private isRegistered(connectionId: string): boolean { + return this.connectionRegistry.list({ includeRetired: true }).some((connection) => connection.id === connectionId); + } + + hasConnection(connectionId: string): boolean { + return Boolean(this.connectionRegistry.getConfig(connectionId)); } onConnectionRemoval(listener: RemovalListener): void { this.removalListeners.push(listener); } + onConnectionRestore(listener: RemovalListener): void { + this.restoreListeners.push(listener); + } + onSnapshot(listener: SnapshotListener): void { this.listeners.push(listener); } @@ -200,7 +299,10 @@ export class KvCacheFootprintService extends MultiConnectionPoller implements On const target = this.target(connectionId); if (!target) throw new NotFoundException(`Connection '${connectionId}' not found`); if (target.rejection) throw new BadRequestException(target.rejection); - const running = this.inFlight.get(connectionId); + if (target.connection && !target.connection.isConnected) { + throw new ServiceUnavailableException(`Connection '${connectionId}' is not connected`); + } + const running = this.running.get(connectionId); if (running) return running; const recent = await this.recentSnapshot(connectionId); if (recent) return recent; @@ -213,6 +315,7 @@ export class KvCacheFootprintService extends MultiConnectionPoller implements On } async collect(ctx: ConnectionContext): Promise { + await this.deleting.get(ctx.connectionId); const target = this.target(ctx.connectionId); if (!target || target.rejection) return null; return this.collectWith(ctx, target.plan); @@ -220,11 +323,16 @@ export class KvCacheFootprintService extends MultiConnectionPoller implements On private target( connectionId: string, - ): { config: DatabaseConnectionConfig; plan: CollectionPlan; rejection: string | null } | null { + ): { + config: DatabaseConnectionConfig; + connection: ConnectionStatus | undefined; + plan: CollectionPlan; + rejection: string | null; + } | null { const config = this.connectionRegistry.getConfig(connectionId); if (!config) return null; const { connection, plan } = resolvePlan(this.connectionRegistry, connectionId); - return { config, plan, rejection: this.rejection(connectionId, config, connection, plan) }; + return { config, connection, plan, rejection: this.rejection(connectionId, config, connection, plan) }; } private async recentSnapshot(connectionId: string): Promise { @@ -238,66 +346,94 @@ export class KvCacheFootprintService extends MultiConnectionPoller implements On connection: ConnectionStatus | undefined, plan: CollectionPlan, ): string | null { + if (this.isSentinelConnection(connectionId)) return 'Sentinel connections have no keyspace to scan'; if (config.host === 'agent') return 'KV cache footprint is not available for agent connections'; if (config.connectionType === 'external') return 'KV cache footprint is not available for external connections'; - if (connection?.isConnected === false) return 'The connection is not connected'; - if (this.isSentinelConnection(connectionId)) return 'Sentinel connections have no keyspace to scan'; if (connection?.membership?.retiredAt !== undefined) return 'This node has been retired from its topology'; if (plan.kind === 'skip') return 'KV cache footprint for this node is collected through its seed or primary connection'; return null; } private async collectWith(ctx: ConnectionContext, plan: CollectionPlan): Promise { - const running = this.inFlight.get(ctx.connectionId); + await this.deleting.get(ctx.connectionId); + if (!this.connectionRegistry.getConfig(ctx.connectionId)) return null; + const running = this.running.get(ctx.connectionId); if (running) return running; - const run = this.run(ctx, plan); - this.inFlight.set(ctx.connectionId, run); + const run = this.runCollection(ctx, plan); + this.running.set(ctx.connectionId, run); try { return await run; } finally { - this.inFlight.delete(ctx.connectionId); + if (this.running.get(ctx.connectionId) === run) this.running.delete(ctx.connectionId); } } - private async run(ctx: ConnectionContext, plan: CollectionPlan): Promise { + private generationOf(connectionId: string): number { + return this.generations.get(connectionId) ?? 0; + } + + private async runCollection(ctx: ConnectionContext, plan: CollectionPlan): Promise { + const generation = this.generationOf(ctx.connectionId); + const current = () => this.generationOf(ctx.connectionId) === generation; const targets = await this.nodesFor(ctx, plan); const { nodes, skipped } = targets; - const settled = await this.readAll(nodes, targets.slots?.rangesByNode.map(slotOwner)); + const probe = this.undetected.has(ctx.connectionId); + const settled = await this.readAll(nodes, probe, targets.slots?.rangesByNode.map(slotOwner)); const readings: NodeReading[] = []; const readIndexes: number[] = []; + const evicted = new Map(); settled.forEach((result, index) => { if (result.status === 'fulfilled') { readings.push(result.value); readIndexes.push(index); - } else this.logger.warn(`KV cache read failed on ${nodes[index].name}: ${errorMessage(result.reason)}`); + evicted.set(result.value.nodeKey, result.value.memory.evictedKeys); + } else { + this.logger.warn(`KV cache read failed on ${nodes[index].name}: ${errorMessage(result.reason)}`); + } }); const failed = settled.find((result): result is PromiseRejectedResult => result.status === 'rejected'); if (readings.length === 0 && failed) throw failed.reason; - const partial = failed !== undefined || skipped > 0; const memory = this.sumMemory(readings.map((reading) => reading.memory)); - const evicted = Object.fromEntries(readings.map((reading) => [reading.nodeKey, reading.memory.evictedKeys])); const built = buildSnapshot({ connectionId: ctx.connectionId, timestamp: Date.now(), nodes: readings.map((reading) => reading.observation), memory, - evictedKeysDelta: evictedKeysDelta(this.lastEvicted.get(ctx.connectionId) ?? null, evicted), + evictedKeysDelta: evictedKeysDelta(this.lastEvicted.get(ctx.connectionId), evicted), otherDbs: readings[0]?.otherDbs ?? [], clientDetected: readings.some((reading) => reading.glide), coverage: this.coverage(targets, readIndexes), }); - const detected = partial && !built.detected ? await this.carriedDetection(ctx.connectionId) : built.detected; - if (!partial || built.detected) this.partialMisses.delete(ctx.connectionId); + const partial = failed !== undefined || skipped > 0; + const missed = partial && !built.detected; + const misses = missed ? (this.partialMisses.get(ctx.connectionId) ?? 0) + 1 : 0; + const detected = missed ? await this.carriedDetection(ctx.connectionId, misses) : built.detected; const keepsEarlier = partial && detected; const snapshot = partial ? { ...built, scanComplete: false, detected } : built; + await this.deleting.get(ctx.connectionId); + if (current()) await this.removing.get(ctx.connectionId); + const live = () => current() && Boolean(this.connectionRegistry.getConfig(ctx.connectionId)); + if (!live()) return null; await this.storage.saveKvCacheFootprintSnapshot(snapshot); + if (!live()) return null; this.lastEvicted.set(ctx.connectionId, evicted); - if (built.detected) this.undetectedUntil.delete(ctx.connectionId); - else if (!snapshot.detected && (!partial || this.partialMissesExhausted(ctx.connectionId))) { + if (missed) this.partialMisses.set(ctx.connectionId, misses); + else this.partialMisses.delete(ctx.connectionId); + if (built.detected) { + this.undetectedUntil.delete(ctx.connectionId); + this.undetected.delete(ctx.connectionId); + } else if (!snapshot.detected && (!partial || misses >= PARTIAL_MISSES_BEFORE_CLEAR)) { this.undetectedUntil.set(ctx.connectionId, snapshot.timestamp + this.undetectedRecheckMs); + this.undetected.add(ctx.connectionId); } this.rememberSampleKey(ctx.connectionId, readings, keepsEarlier); - await this.announce(snapshot); + const announced = this.announce(snapshot, { partial }); + this.announcing.set(ctx.connectionId, announced); + try { + await announced; + } finally { + if (this.announcing.get(ctx.connectionId) === announced) this.announcing.delete(ctx.connectionId); + } return snapshot; } @@ -378,9 +514,10 @@ export class KvCacheFootprintService extends MultiConnectionPoller implements On private async readAll( nodes: ScanNode[], + probe: boolean, ownership?: SlotOwnership[], ): Promise[]> { - const scanned = await Promise.allSettled(nodes.map((node) => this.scanNode(node))); + const scanned = await Promise.allSettled(nodes.map((node) => this.scanNode(node, probe))); const runKeys = scanned.flatMap((result) => (result.status === 'fulfilled' ? result.value.scan.matchedKeys : [])); const evidence = { twoKey: hasTwoKeyEvidence(runKeys), singleKey: hasSingleKeyEvidence(runKeys) }; const checkClients = runKeys.length < DETECTION_MIN_MATCHES; @@ -393,8 +530,17 @@ export class KvCacheFootprintService extends MultiConnectionPoller implements On ); } - private async scanNode(node: ScanNode): Promise { - const [scan, info] = await Promise.all([this.scanKeys(node), node.client.getInfoParsed(['memory', 'stats', 'keyspace'])]); + private async carriedDetection(connectionId: string, misses: number): Promise { + if (misses >= PARTIAL_MISSES_BEFORE_CLEAR) return false; + const [previous] = await this.storage.getKvCacheFootprintSnapshots({ connectionId, limit: 1 }); + return previous?.detected === true; + } + + private async scanNode(node: ScanNode, probe: boolean): Promise { + const [scan, info] = await Promise.all([ + this.scanKeys(node, probe), + node.client.getInfoParsed(['memory', 'stats', 'keyspace']), + ]); const ownDb = node.client.getActiveDbIndex?.() ?? 0; return { node, @@ -406,9 +552,9 @@ export class KvCacheFootprintService extends MultiConnectionPoller implements On }; } - private async scanKeys(node: ScanNode): Promise { + private async scanKeys(node: ScanNode, probe: boolean): Promise { const dbSize = await node.client.getDbSize(); - return { ...(await scanLmcacheKeys(node.client, this.budgets, dbSize)), dbSize }; + return { ...(await scanLmcacheKeys(node.client, this.budgets, dbSize, probe)), dbSize }; } private async sampleNode( @@ -422,9 +568,20 @@ export class KvCacheFootprintService extends MultiConnectionPoller implements On sampleLmcacheKeys(node.client, downsampleSeries(matchedKeys, this.sampleSize), ownsSlot), checkClients ? hasLmcacheGlideClient(node.client, node.client.getActiveDbIndex?.() ?? 0) : false, ]); + this.noteMemoryAccess(node.name, samples.length > 0 && samples.every((sample) => sample.memoryBytes === null)); return { ...rest, glide, nodeKey: node.connectionId ?? node.name, observation: { ...scan, matchedKeys, samples } }; } + private noteMemoryAccess(nodeName: string, blocked: boolean): void { + if (!blocked) { + this.memoryBlocked.delete(nodeName); + return; + } + if (this.memoryBlocked.has(nodeName)) return; + this.memoryBlocked.add(nodeName); + this.logger.warn(`MEMORY USAGE returned nothing on ${nodeName}; KV cache byte estimates will read as unknown`); + } + private sumMemory(stats: MemoryStats[]): MemoryStats { const policies = new Set(stats.map((s) => s.maxmemoryPolicy)); const unlimited = stats.some((s) => s.maxmemory === 0); @@ -436,28 +593,16 @@ export class KvCacheFootprintService extends MultiConnectionPoller implements On }; } - private async carriedDetection(connectionId: string): Promise { - const misses = (this.partialMisses.get(connectionId) ?? 0) + 1; - this.partialMisses.set(connectionId, misses); - if (misses >= PARTIAL_MISSES_BEFORE_CLEAR) return false; - const [previous] = await this.storage.getKvCacheFootprintSnapshots({ connectionId, limit: 1 }); - return previous?.detected === true; - } - - private partialMissesExhausted(connectionId: string): boolean { - return (this.partialMisses.get(connectionId) ?? 0) >= PARTIAL_MISSES_BEFORE_CLEAR; - } - private rememberSampleKey(connectionId: string, readings: NodeReading[], keepsEarlier: boolean): void { const parsed = readings.flatMap((reading) => reading.observation.matchedKeys)[0]; if (parsed) this.sampleKeys.set(connectionId, maskKey(parsed)); else if (!keepsEarlier) this.sampleKeys.delete(connectionId); } - private async announce(snapshot: KvCacheFootprintSnapshot): Promise { + private async announce(snapshot: KvCacheFootprintSnapshot, context: SnapshotContext): Promise { for (const listener of this.listeners) { try { - await listener(snapshot); + await listener(snapshot, context); } catch (error) { this.logger.warn(`KV cache snapshot listener failed: ${errorMessage(error)}`); } diff --git a/proprietary/kv-cache/kv-cache-samples.service.ts b/proprietary/kv-cache/kv-cache-samples.service.ts index 5628f449b..2f08d3a9b 100644 --- a/proprietary/kv-cache/kv-cache-samples.service.ts +++ b/proprietary/kv-cache/kv-cache-samples.service.ts @@ -34,6 +34,12 @@ export function mergeBuckets(samples: KvCacheEngineSample[]): KvCacheSampleBucke .map((row) => ({ ...row, hitRate: hitRate(row.requestedTokens, row.hitTokens) })); } +function dropOldestTimestamp(samples: KvCacheEngineSample[]): KvCacheEngineSample[] { + const oldest = samples.reduce((min, row) => Math.min(min, row.timestamp), Infinity); + const newer = samples.filter((row) => row.timestamp > oldest); + return newer.length > 0 ? newer : samples.slice(-SAMPLES_LIMIT); +} + @Injectable() export class KvCacheSamplesService implements OnModuleInit, OnModuleDestroy { private readonly logger = new Logger(KvCacheSamplesService.name); @@ -66,7 +72,7 @@ export class KvCacheSamplesService implements OnModuleInit, OnModuleDestroy { this.tracker.forgetEngine(engineId); } - idle(): Promise { + settled(): Promise { return this.flushing; } @@ -105,10 +111,12 @@ export class KvCacheSamplesService implements OnModuleInit, OnModuleDestroy { modelName: query.model, }; const bucketMs = chartBucketMs(filter.from, filter.to); - const [rows, totals] = await Promise.all([ - this.storage.getKvCacheEngineSamples({ ...filter, limit: SAMPLES_LIMIT, bucketMs }), + const [fetched, totals] = await Promise.all([ + this.storage.getKvCacheEngineSamples({ ...filter, limit: SAMPLES_LIMIT + 1, bucketMs }), this.storage.getKvCacheSampleTotals(filter), ]); - return { buckets: mergeBuckets(rows), rangeHitRate: hitRate(totals.requestedTokens, totals.hitTokens), bucketMs }; + const truncated = fetched.length > SAMPLES_LIMIT; + const rows = truncated ? dropOldestTimestamp(fetched) : fetched; + return { buckets: mergeBuckets(rows), rangeHitRate: hitRate(totals.requestedTokens, totals.hitTokens), truncated, bucketMs }; } } diff --git a/proprietary/kv-cache/kv-cache-scrape.service.ts b/proprietary/kv-cache/kv-cache-scrape.service.ts index d8892fde6..cbecac9f5 100644 --- a/proprietary/kv-cache/kv-cache-scrape.service.ts +++ b/proprietary/kv-cache/kv-cache-scrape.service.ts @@ -1,6 +1,7 @@ import { Injectable, Logger, OnModuleDestroy, OnModuleInit } from '@nestjs/common'; import { Feature } from '@betterdb/shared'; import type { StoredKvCacheEngine } from '@app/common/interfaces/storage-port.interface'; +import { ConnectionRegistry } from '@app/connections/connection-registry.service'; import { LicenseService } from '@proprietary/licenses/license.service'; import { labelSeries } from './counter-deltas'; import { KvCacheEngineRegistry } from './kv-cache-engine-registry'; @@ -35,6 +36,7 @@ export class KvCacheScrapeService implements OnModuleInit, OnModuleDestroy { private readonly samples: KvCacheSamplesService, private readonly engines: KvCacheEnginesService, private readonly license: LicenseService, + private readonly connections: ConnectionRegistry, ) {} onModuleInit(): void { @@ -53,7 +55,9 @@ export class KvCacheScrapeService implements OnModuleInit, OnModuleDestroy { if (this.running || !this.license.hasFeature(Feature.KV_CACHE_MONITORING)) return; this.running = true; try { - const targets = this.registry.list().filter((e) => e.enabled && e.source === 'scrape' && e.scrapeUrl); + const targets = this.registry + .list() + .filter((e) => e.enabled && e.source === 'scrape' && e.scrapeUrl && this.connectionExists(e.connectionId)); await Promise.allSettled(targets.map((engine) => this.scrape(engine, now))); } finally { this.running = false; @@ -98,6 +102,10 @@ export class KvCacheScrapeService implements OnModuleInit, OnModuleDestroy { private stillCurrent(engine: StoredKvCacheEngine): boolean { const current = this.registry.get(engine.id); - return !!current?.enabled && current.scrapeUrl === engine.scrapeUrl; + return !!current?.enabled && current.scrapeUrl === engine.scrapeUrl && this.connectionExists(engine.connectionId); + } + + private connectionExists(connectionId: string): boolean { + return this.connections.getConfig(connectionId) !== null; } }