diff --git a/.github/workflows/api-tests.yml b/.github/workflows/api-tests.yml index ef44ab4df..57edd0774 100644 --- a/.github/workflows/api-tests.yml +++ b/.github/workflows/api-tests.yml @@ -240,6 +240,8 @@ jobs: - name: Run API tests working-directory: apps/api + env: + KV_CACHE_IT_VALKEY_PORT: '6390' run: pnpm test - name: Run entitlement unit tests diff --git a/apps/api/src/__tests__/kv-cache-feature.spec.ts b/apps/api/src/__tests__/kv-cache-feature.spec.ts new file mode 100644 index 000000000..57f9c7e06 --- /dev/null +++ b/apps/api/src/__tests__/kv-cache-feature.spec.ts @@ -0,0 +1,13 @@ +import { Feature, TIER_FEATURES, Tier, DEFAULT_KV_CACHE_HIT_RATE_THRESHOLD } from '@betterdb/shared'; + +describe('kv cache feature', () => { + it('is a Pro feature', () => { + expect(Feature.KV_CACHE_MONITORING).toBe('kvCacheMonitoring'); + expect(TIER_FEATURES[Tier.pro]).toContain(Feature.KV_CACHE_MONITORING); + expect(TIER_FEATURES[Tier.community]).not.toContain(Feature.KV_CACHE_MONITORING); + }); + + it('defaults the hit rate threshold to 0.2', () => { + expect(DEFAULT_KV_CACHE_HIT_RATE_THRESHOLD).toBe(0.2); + }); +}); diff --git a/apps/api/src/app.module.ts b/apps/api/src/app.module.ts index 0583b7412..fcf6f728f 100644 --- a/apps/api/src/app.module.ts +++ b/apps/api/src/app.module.ts @@ -47,6 +47,7 @@ let LatencyRegressionModule: any = null; let WebhookProModule: any = null; let InferenceLatencyProModule: any = null; let ScalingReadinessProModule: any = null; +let KvCacheModule: any = null; let CacheProposalsModule: any = null; let MemoryProposalsModule: any = null; let AgentModule: any = null; @@ -126,6 +127,14 @@ try { // Proprietary module not available } +try { + const kvCacheModule = require('../../../proprietary/kv-cache/kv-cache.module'); + KvCacheModule = kvCacheModule.KvCacheModule; + console.log('[KvCache] Proprietary module loaded'); +} catch { + // Proprietary module not available +} + try { const cacheProposalsModule = require('../../../proprietary/cache-proposals/cache-proposals.module'); CacheProposalsModule = cacheProposalsModule.CacheProposalsModule; @@ -241,6 +250,7 @@ const proprietaryImports = [ WebhookProModule, InferenceLatencyProModule, ScalingReadinessProModule, + KvCacheModule, CacheProposalsModule, MemoryProposalsModule, AiModule, diff --git a/apps/api/src/common/interfaces/database-port.interface.ts b/apps/api/src/common/interfaces/database-port.interface.ts index a9a620c57..037bf7288 100644 --- a/apps/api/src/common/interfaces/database-port.interface.ts +++ b/apps/api/src/common/interfaces/database-port.interface.ts @@ -54,6 +54,7 @@ export interface DatabasePort { disconnect(): Promise; isConnected(): boolean; sampleVersion?(): number | null; + getActiveDbIndex?(): number; ping(): Promise; getInfo(sections?: string[]): Promise>; getCapabilities(): DatabaseCapabilities; diff --git a/apps/api/src/common/interfaces/storage-port.interface.ts b/apps/api/src/common/interfaces/storage-port.interface.ts index e6e92ff4d..6319f4c8c 100644 --- a/apps/api/src/common/interfaces/storage-port.interface.ts +++ b/apps/api/src/common/interfaces/storage-port.interface.ts @@ -39,6 +39,10 @@ import type { ReadinessDimension, ReadinessDimensionKey, ScalingReadinessSettings, + KvCacheEngineSample, + KvCacheEngineSource, + KvCacheFootprintSnapshot, + KvCacheSettings, } from '@betterdb/shared'; export type { MetricForecastSettings, MetricKind } from '@betterdb/shared'; export type { StoredCveDataset, CveScanResult } from '@betterdb/shared'; @@ -405,6 +409,37 @@ export interface ScalingReadinessScoreQuery { limit?: number; } +export interface StoredKvCacheEngine { + id: string; + connectionId: string; + name: string; + source: KvCacheEngineSource; + scrapeUrl: string | null; + scrapeAuthHeader: string | null; + scrapeAuthEncrypted: boolean; + otlpEngineId: string | null; + enabled: boolean; + createdAt: number; + lastSeenAt: number | null; + lastError: string | null; +} + +export interface KvCacheRangeQuery { + connectionId: string; + from?: number; + to?: number; + limit?: number; +} + +export interface KvCacheFootprintQuery extends KvCacheRangeQuery { + bucketMs?: number; +} + +export interface KvCacheSampleQuery extends KvCacheRangeQuery { + engineId?: string; + modelName?: string; +} + export interface MemorySnapshotQueryOptions { connectionId?: string; startTime?: number; @@ -669,6 +704,19 @@ export interface StoragePort { pruneOldScalingReadinessScores(cutoffTimestamp: number, connectionId?: string): Promise; getScalingReadinessSettings(connectionId: string): Promise; saveScalingReadinessSettings(settings: ScalingReadinessSettings): Promise; + saveKvCacheFootprintSnapshot(snapshot: KvCacheFootprintSnapshot): Promise; + getKvCacheFootprintSnapshots(query: KvCacheFootprintQuery): Promise; + pruneOldKvCacheFootprintSnapshots(cutoffTimestamp: number, connectionId?: string): Promise; + saveKvCacheEngineSamples(samples: KvCacheEngineSample[]): Promise; + getKvCacheEngineSamples(query: KvCacheSampleQuery): Promise; + pruneOldKvCacheEngineSamples(cutoffTimestamp: number, connectionId?: string): Promise; + saveKvCacheEngine(engine: StoredKvCacheEngine): Promise; + getKvCacheEngines(connectionId?: string): Promise; + getKvCacheEngine(id: string): Promise; + deleteKvCacheEngine(id: string): Promise; + getKvCacheSettings(connectionId: string): Promise; + saveKvCacheSettings(settings: KvCacheSettings): Promise; + deleteKvCacheConnectionData(connectionId: string): Promise; // Command Stats Sample Methods - connectionId required for writes saveCommandStatsSamples( diff --git a/apps/api/src/config/env.schema.kv-cache.spec.ts b/apps/api/src/config/env.schema.kv-cache.spec.ts new file mode 100644 index 000000000..dd255974a --- /dev/null +++ b/apps/api/src/config/env.schema.kv-cache.spec.ts @@ -0,0 +1,19 @@ +import { envSchema } from './env.schema'; + +const KV_CACHE_FOOTPRINT_VARS = [ + 'KV_CACHE_FOOTPRINT_INTERVAL_MS', + 'KV_CACHE_SCAN_MAX_KEYS', + 'KV_CACHE_MATCH_MAX_KEYS', + 'KV_CACHE_SAMPLE_KEYS', + 'KV_CACHE_UNDETECTED_RECHECK_MS', +] as const; + +describe('KV cache footprint env vars', () => { + it.each(KV_CACHE_FOOTPRINT_VARS)('registers %s', (name) => { + expect((envSchema.parse({ [name]: '600000' }) as Record)[name]).toBe('600000'); + }); + + it.each(KV_CACHE_FOOTPRINT_VARS)('leaves an invalid %s to the service fallback instead of failing startup', (name) => { + expect(envSchema.safeParse({ [name]: 'abc' }).success).toBe(true); + }); +}); diff --git a/apps/api/src/config/env.schema.ts b/apps/api/src/config/env.schema.ts index 00d66c2c9..195497476 100644 --- a/apps/api/src/config/env.schema.ts +++ b/apps/api/src/config/env.schema.ts @@ -169,6 +169,12 @@ export const envSchema = z ANOMALY_CACHE_TTL_MS: z.coerce.number().int().min(1000).default(3600000), ANOMALY_PROMETHEUS_INTERVAL_MS: z.coerce.number().int().min(1000).default(30000), + KV_CACHE_FOOTPRINT_INTERVAL_MS: z.string().optional(), + KV_CACHE_SCAN_MAX_KEYS: z.string().optional(), + KV_CACHE_MATCH_MAX_KEYS: z.string().optional(), + KV_CACHE_SAMPLE_KEYS: z.string().optional(), + KV_CACHE_UNDETECTED_RECHECK_MS: z.string().optional(), + // License configuration (optional) BETTERDB_LICENSE_KEY: z.string().optional(), // Offline (air-gapped) license: the signed token itself, or a path to it diff --git a/apps/api/src/connections/__tests__/connection-registry.membership.spec.ts b/apps/api/src/connections/__tests__/connection-registry.membership.spec.ts index 606e467ea..b03693c35 100644 --- a/apps/api/src/connections/__tests__/connection-registry.membership.spec.ts +++ b/apps/api/src/connections/__tests__/connection-registry.membership.spec.ts @@ -635,4 +635,35 @@ describe('ConnectionRegistry membership', () => { await registry.retireChild('adopted'); expect(registry.getDefaultId()).toBe('seed'); }); + + it('notifies removal listeners after a connection and its auto children are deleted', async () => { + const { registry, storage } = build(); + put(registry, seed); + put(registry, { id: 'child', membership: { seedId: 'seed', nodeId: 'c', origin: 'auto', source: 'cluster' } }); + const removed: string[] = []; + registry.addConnectionRemovalListener((id) => { + expect(storage.deleteConnection).toHaveBeenCalledWith(id); + removed.push(id); + }); + const unsubscribe = registry.addConnectionRemovalListener(() => { + throw new Error('boom'); + }); + await registry.removeConnection('seed'); + expect(removed).toEqual(['child', 'seed']); + unsubscribe(); + }); + + it('logs a removal listener that rejects asynchronously', async () => { + const { registry } = build(); + put(registry, seed); + const warn = jest + .spyOn((registry as unknown as { logger: { warn: (message: string) => void } }).logger, 'warn') + .mockImplementation(() => undefined); + registry.addConnectionRemovalListener(async () => { + throw new Error('async boom'); + }); + await registry.removeConnection('seed'); + await new Promise((resolve) => setImmediate(resolve)); + expect(warn).toHaveBeenCalledWith(expect.stringContaining('async boom')); + }); }); diff --git a/apps/api/src/connections/connection-registry.service.ts b/apps/api/src/connections/connection-registry.service.ts index afc071f04..f61441d27 100644 --- a/apps/api/src/connections/connection-registry.service.ts +++ b/apps/api/src/connections/connection-registry.service.ts @@ -20,6 +20,8 @@ import { ENV_DEFAULT_ID } from './connection.constants'; export const CHILD_CONNECT_TIMEOUT_MS = 10_000; +export type ConnectionRemovalListener = (id: string) => void | Promise; + function certificateHint(errorMsg: string, node: DesiredNode): string { if (node.hostname !== undefined || !/altnames|ERR_TLS_CERT_ALTNAME_INVALID/i.test(errorMsg)) { return ''; @@ -91,6 +93,7 @@ export class ConnectionRegistry implements OnModuleInit, OnModuleDestroy { private readonly encryption: EnvelopeEncryptionService | null; private startupConnectionErrors: Array<{ name: string; host: string; port: number; error: string }> = []; private readonly seedLocks = new Map>(); + private readonly removalListeners = new Set(); constructor( @Inject('STORAGE_CLIENT') private readonly storage: StoragePort, @@ -709,6 +712,27 @@ export class ConnectionRegistry implements OnModuleInit, OnModuleDestroy { this.discoveredInstances?.forgetPushed(discoveredHost, discoveredPort, Date.now()); } + addConnectionRemovalListener(listener: ConnectionRemovalListener): () => void { + this.removalListeners.add(listener); + return () => { + this.removalListeners.delete(listener); + }; + } + + private notifyRemoved(id: string): void { + const report = (error: unknown) => { + const reason = error instanceof Error ? error.message : String(error); + this.logger.warn(`Connection removal listener failed for ${id}: ${reason}`); + }; + for (const listener of this.removalListeners) { + try { + void Promise.resolve(listener(id)).catch(report); + } catch (error: unknown) { + report(error); + } + } + } + async removeConnection(id: string): Promise { if (id === ENV_DEFAULT_ID) { throw new Error('Cannot remove the default environment connection'); @@ -751,6 +775,7 @@ export class ConnectionRegistry implements OnModuleInit, OnModuleDestroy { this.runtimeCapabilityTracker.removeConnection(id); this.externalMetricsStore.clear(id); await this.storage.deleteConnection(id); + this.notifyRemoved(id); if (removedConfig) { this.usageTelemetry?.trackDbRemove({ diff --git a/apps/api/src/database/adapters/unified.adapter.ts b/apps/api/src/database/adapters/unified.adapter.ts index 3bb066731..015de0f8b 100644 --- a/apps/api/src/database/adapters/unified.adapter.ts +++ b/apps/api/src/database/adapters/unified.adapter.ts @@ -7,6 +7,7 @@ import { } from '../../common/interfaces/database-port.interface'; import { InfoParser } from '../parsers/info.parser'; import { MetricsParser } from '../parsers/metrics.parser'; +import { parseKeyspaceDatabases, type KeyspaceDatabase } from '../parsers/keyspace-databases'; import { CLUSTER_TOTAL_SLOTS } from '../../common/constants/cluster.constants'; import { InfoResponse, @@ -884,22 +885,13 @@ export class UnifiedDatabaseAdapter implements DatabasePort { * proxies that hide it), in which case the caller falls back to scanning * the connection's own database. */ - private async listKeyspaceDatabases(): Promise> { + private async listKeyspaceDatabases(): Promise { try { // Reuse the centralized INFO keyspace parser (MetricsParser via // getInfoParsed) rather than a second regex: it is the single definition // of a db entry and already drops malformed `db:` lines as raw strings. const { keyspace } = await this.getInfoParsed(['keyspace']); - const dbs: Array<{ db: number; keys: number }> = []; - for (const [key, info] of Object.entries(keyspace ?? {})) { - const match = /^db(\d+)$/.exec(key); - // A raw-string value is a malformed line the parser kept verbatim; only - // typed entries with a positive key count are non-empty databases. - if (match && typeof info === 'object' && info.keys > 0) { - dbs.push({ db: parseInt(match[1], 10), keys: info.keys }); - } - } - return dbs.sort((a, b) => a.db - b.db); + return parseKeyspaceDatabases(keyspace); } catch (error) { this.logger.debug( `Could not list keyspace databases: ${error instanceof Error ? error.message : error}`, diff --git a/apps/api/src/database/parsers/keyspace-databases.spec.ts b/apps/api/src/database/parsers/keyspace-databases.spec.ts new file mode 100644 index 000000000..8933ec7d5 --- /dev/null +++ b/apps/api/src/database/parsers/keyspace-databases.spec.ts @@ -0,0 +1,22 @@ +import { parseKeyspaceDatabases } from './keyspace-databases'; + +describe('parseKeyspaceDatabases', () => { + it('lists non-empty databases sorted by index', () => { + expect( + parseKeyspaceDatabases({ + db3: { keys: 2, expires: 0, avg_ttl: 0 }, + db0: { keys: 5, expires: 0, avg_ttl: 0 }, + db4: { keys: 0, expires: 0, avg_ttl: 0 }, + junk: 'x', + db7: 'malformed', + }), + ).toEqual([ + { db: 0, keys: 5 }, + { db: 3, keys: 2 }, + ]); + }); + + it('returns an empty list without a keyspace section', () => { + expect(parseKeyspaceDatabases(undefined)).toEqual([]); + }); +}); diff --git a/apps/api/src/database/parsers/keyspace-databases.ts b/apps/api/src/database/parsers/keyspace-databases.ts new file mode 100644 index 000000000..5d761feba --- /dev/null +++ b/apps/api/src/database/parsers/keyspace-databases.ts @@ -0,0 +1,17 @@ +import type { KeyspaceInfo } from '../../common/types/metrics.types'; + +export interface KeyspaceDatabase { + db: number; + keys: number; +} + +export function parseKeyspaceDatabases(keyspace: KeyspaceInfo | undefined): KeyspaceDatabase[] { + const dbs: KeyspaceDatabase[] = []; + for (const [key, info] of Object.entries(keyspace ?? {})) { + const match = /^db(\d+)$/.exec(key); + if (match && typeof info === 'object' && info.keys > 0) { + dbs.push({ db: parseInt(match[1], 10), keys: info.keys }); + } + } + return dbs.sort((a, b) => a.db - b.db); +} diff --git a/apps/api/src/retention/__tests__/local-retention.service.spec.ts b/apps/api/src/retention/__tests__/local-retention.service.spec.ts index 6e3b35d01..a1fcbcae8 100644 --- a/apps/api/src/retention/__tests__/local-retention.service.spec.ts +++ b/apps/api/src/retention/__tests__/local-retention.service.spec.ts @@ -29,6 +29,8 @@ describe('LocalRetentionService', () => { pruneOldLatencyHistograms: jest.fn().mockResolvedValue(0), pruneOldMemorySnapshots: jest.fn().mockResolvedValue(0), pruneOldScalingReadinessScores: jest.fn().mockResolvedValue(3), + pruneOldKvCacheFootprintSnapshots: jest.fn().mockResolvedValue(4), + pruneOldKvCacheEngineSamples: jest.fn().mockResolvedValue(5), pruneOldCaptureChunks: jest.fn().mockResolvedValue(0), pruneOldCaptureSessions: jest.fn().mockResolvedValue(0), pruneOldCaptureTriggers: jest.fn().mockResolvedValue(0), @@ -84,6 +86,15 @@ describe('LocalRetentionService', () => { expect(results.scaling_readiness_scores).toBe(3); }); + it('reports pruned kv cache rows under their own names', async () => { + const results = await runRetentionSweep(storage, 1_234, { error: jest.fn() } as any); + + expect(storage.pruneOldKvCacheFootprintSnapshots).toHaveBeenCalledWith(1_234); + expect(storage.pruneOldKvCacheEngineSamples).toHaveBeenCalledWith(1_234); + expect(results.kv_cache_footprint_snapshots).toBe(4); + expect(results.kv_cache_engine_samples).toBe(5); + }); + it('sweeps client analytics rollups so removed connections do not keep them forever', async () => { storage.pruneOldClientAnalyticsRollups.mockResolvedValue(4); const results = await runRetentionSweep(storage, 5_678, { error: jest.fn() } as any); diff --git a/apps/api/src/retention/retention-sweep.ts b/apps/api/src/retention/retention-sweep.ts index f58f90fba..798fbcc2d 100644 --- a/apps/api/src/retention/retention-sweep.ts +++ b/apps/api/src/retention/retention-sweep.ts @@ -8,9 +8,9 @@ import { StoragePort } from '../common/interfaces/storage-port.interface'; * self-hosted local sweep so the two never drift on which tables are covered. * * `sampleCutoff` applies to the high-volume sample stores (command/latency - * stats samples, vector index snapshots), which in cloud keep a tighter 7-day - * cap than the tier's analytics window. Self-hosted passes one cutoff for - * everything, so it defaults to `cutoff`. + * stats samples, vector index snapshots, KV cache engine samples), which in + * cloud keep a tighter 7-day cap than the tier's analytics window. + * Self-hosted passes one cutoff for everything, so it defaults to `cutoff`. * * A store that fails to prune is reported as -1 and does not stop the sweep. */ @@ -46,6 +46,8 @@ export async function runRetentionSweep( { name: 'latency_histograms', fn: () => storage.pruneOldLatencyHistograms(cutoff) }, { name: 'memory_snapshots', fn: () => storage.pruneOldMemorySnapshots(cutoff) }, { name: 'scaling_readiness_scores', fn: () => storage.pruneOldScalingReadinessScores(cutoff) }, + { name: 'kv_cache_footprint_snapshots', fn: () => storage.pruneOldKvCacheFootprintSnapshots(cutoff) }, + { name: 'kv_cache_engine_samples', fn: () => storage.pruneOldKvCacheEngineSamples(sampleCutoff) }, { name: 'capture_chunks', fn: () => storage.pruneOldCaptureChunks(cutoff) }, { name: 'capture_sessions', fn: () => storage.pruneOldCaptureSessions(cutoff) }, { name: 'capture_triggers', fn: () => storage.pruneOldCaptureTriggers(cutoff) }, diff --git a/apps/api/src/storage/adapters/__tests__/kv-cache-storage.spec.ts b/apps/api/src/storage/adapters/__tests__/kv-cache-storage.spec.ts new file mode 100644 index 000000000..62cfd20cb --- /dev/null +++ b/apps/api/src/storage/adapters/__tests__/kv-cache-storage.spec.ts @@ -0,0 +1,308 @@ +import { execFileSync } from 'child_process'; +import { randomUUID } from 'crypto'; +import * as fs from 'fs'; +import * as os from 'os'; +import * as path from 'path'; +import { Pool } from 'pg'; +import { SqliteAdapter } from '../sqlite.adapter'; +import * as sqliteChunkedDelete from '../sqlite-chunked-delete'; +import { MemoryAdapter } from '../memory.adapter'; +import { PostgresAdapter } from '../postgres.adapter'; +import type { KvCacheEngineSample, KvCacheFootprintSnapshot } from '@betterdb/shared'; +import type { StoragePort, StoredKvCacheEngine } from '../../../common/interfaces/storage-port.interface'; + +const snapshot = (o: Partial = {}): KvCacheFootprintSnapshot => ({ + connectionId: 'conn-a', + timestamp: 1_000, + detected: true, + layout: 'single_key', + scannedKeys: 100, + matchedKeys: 10, + sampledKeys: 10, + scanComplete: true, + primariesScanned: 1, + primariesTotal: 3, + chunksEst: 10, + bytesEst: 36_701_120, + usedMemory: 40_000_000, + maxmemory: 0, + maxmemoryPolicy: 'noeviction', + lmcacheMemoryShare: 0.92, + noTtlRatio: 1, + orphanRatio: null, + evictedKeysDelta: null, + evictedKeysIntervalMs: null, + otherDbs: [2], + perModel: [{ model: 'Qwen/Qwen2.5-0.5B-Instruct', dtype: 'bfloat16', chunksEst: 10, bytesEst: 36_701_120 }], + ...o, +}); + +const sample = (o: Partial = {}): KvCacheEngineSample => ({ + engineId: 'eng-1', + connectionId: 'conn-a', + modelName: 'm', + timestamp: 60_000, + requestedTokens: 1000, + hitTokens: 700, + lookupTokens: 1000, + lookupHits: 3, + remoteReadBytes: 10, + remoteWriteBytes: 20, + remoteReadRequests: 1, + remoteWriteRequests: 2, + remotePingErrors: 0, + ...o, +}); + +const engine = (o: Partial = {}): StoredKvCacheEngine => ({ + id: 'eng-1', + connectionId: 'conn-a', + name: 'vllm-0', + source: 'scrape', + scrapeUrl: 'http://vllm:8000/metrics', + scrapeAuthHeader: null, + scrapeAuthEncrypted: false, + otlpEngineId: null, + enabled: true, + createdAt: 1, + lastSeenAt: null, + lastError: null, + ...o, +}); + +const POSTGRES_DSN = + process.env.KV_CACHE_POSTGRES_TEST_DSN ?? 'postgres://betterdb:devpassword@localhost:5433/betterdb'; + +function isReachable(connectionString: string): boolean { + const { hostname, port } = new URL(connectionString); + const probe = ` + const socket = require('net').createConnection({ host: process.env.PROBE_HOST, port: Number(process.env.PROBE_PORT) }); + socket.setTimeout(1500); + socket.on('connect', () => { socket.destroy(); process.exit(0); }); + socket.on('timeout', () => { socket.destroy(); process.exit(1); }); + socket.on('error', () => { process.exit(1); }); + `; + try { + execFileSync(process.execPath, ['-e', probe], { + stdio: 'ignore', + timeout: 3000, + env: { ...process.env, PROBE_HOST: hostname, PROBE_PORT: port || '5432' }, + }); + return true; + } catch { + return false; + } +} + +type AdapterFactory = () => Promise<{ storage: StoragePort; cleanup: () => Promise }>; + +const postgresAdapters: Array<[string, AdapterFactory]> = isReachable(POSTGRES_DSN) + ? [ + [ + 'postgres', + async () => { + const schema = `kv_cache_${randomUUID().replace(/-/g, '')}`; + const storage = new PostgresAdapter({ connectionString: POSTGRES_DSN, schema }); + await storage.initialize(); + return { + storage: storage as StoragePort, + cleanup: async () => { + await storage.close(); + const pool = new Pool({ connectionString: POSTGRES_DSN }); + try { + await pool.query(`DROP SCHEMA IF EXISTS ${schema} CASCADE`); + } finally { + await pool.end(); + } + }, + }; + }, + ], + ] + : []; + +describe.each<[string, AdapterFactory]>([ + ['memory', async () => { + const storage = new MemoryAdapter(); + await storage.initialize(); + return { storage: storage as StoragePort, cleanup: async () => storage.close() }; + }], + ['sqlite', async () => { + const dir = fs.mkdtempSync(path.join(os.tmpdir(), 'kv-cache-')); + const storage = new SqliteAdapter({ filepath: path.join(dir, 'db.sqlite') }); + await storage.initialize(); + return { + storage: storage as StoragePort, + cleanup: async () => { + await storage.close(); + fs.rmSync(dir, { recursive: true, force: true }); + }, + }; + }], + ...postgresAdapters, +])('%s adapter kv cache storage', (_name, make) => { + let storage: StoragePort; + let cleanup: () => Promise; + + beforeEach(async () => { + ({ storage, cleanup } = await make()); + }); + + afterEach(async () => { + await cleanup(); + }); + + it('round-trips a footprint snapshot', async () => { + await storage.saveKvCacheFootprintSnapshot(snapshot()); + expect(await storage.getKvCacheFootprintSnapshots({ connectionId: 'conn-a' })).toEqual([snapshot()]); + }); + + it('round-trips an eviction delta with the interval it covers', async () => { + await storage.saveKvCacheFootprintSnapshot(snapshot({ evictedKeysDelta: 12, evictedKeysIntervalMs: 300_000 })); + expect((await storage.getKvCacheFootprintSnapshots({ connectionId: 'conn-a' }))[0]).toMatchObject({ + evictedKeysDelta: 12, + evictedKeysIntervalMs: 300_000, + }); + }); + + it('round-trips a snapshot without a TTL ratio', async () => { + await storage.saveKvCacheFootprintSnapshot(snapshot({ noTtlRatio: null })); + expect((await storage.getKvCacheFootprintSnapshots({ connectionId: 'conn-a' }))[0].noTtlRatio).toBeNull(); + }); + + it('round-trips a snapshot whose byte estimate is unknown', async () => { + const unknown = snapshot({ + bytesEst: null, + lmcacheMemoryShare: null, + primariesTotal: null, + perModel: [{ model: 'm', dtype: 'half', chunksEst: 10, bytesEst: null }], + }); + await storage.saveKvCacheFootprintSnapshot(unknown); + expect(await storage.getKvCacheFootprintSnapshots({ connectionId: 'conn-a' })).toEqual([unknown]); + }); + + it('returns the newest snapshots in ascending order when limited', async () => { + for (const timestamp of [1, 2, 3]) await storage.saveKvCacheFootprintSnapshot(snapshot({ timestamp })); + const rows = await storage.getKvCacheFootprintSnapshots({ connectionId: 'conn-a', limit: 2 }); + expect(rows.map((r) => r.timestamp)).toEqual([2, 3]); + }); + + it('keeps the latest snapshot per time bucket when bucketing', async () => { + for (const timestamp of [100, 105, 109, 110, 125, 139, 140, 200]) { + await storage.saveKvCacheFootprintSnapshot(snapshot({ timestamp })); + } + await storage.saveKvCacheFootprintSnapshot(snapshot({ timestamp: 150, connectionId: 'conn-b' })); + const rows = await storage.getKvCacheFootprintSnapshots({ connectionId: 'conn-a', from: 100, to: 199, bucketMs: 10 }); + expect(rows.map((r) => r.timestamp)).toEqual([109, 110, 125, 139, 140]); + expect(rows[0]).toEqual(snapshot({ timestamp: 109 })); + const limited = await storage.getKvCacheFootprintSnapshots({ connectionId: 'conn-a', from: 100, to: 199, bucketMs: 10, limit: 2 }); + expect(limited.map((r) => r.timestamp)).toEqual([139, 140]); + }); + + it('prunes snapshots by cutoff and optionally by connection', async () => { + await storage.saveKvCacheFootprintSnapshot(snapshot({ timestamp: 1 })); + await storage.saveKvCacheFootprintSnapshot(snapshot({ timestamp: 1, connectionId: 'conn-b' })); + await storage.saveKvCacheFootprintSnapshot(snapshot({ timestamp: 5 })); + expect(await storage.pruneOldKvCacheFootprintSnapshots(2, 'conn-a')).toBe(1); + expect(await storage.pruneOldKvCacheFootprintSnapshots(2)).toBe(1); + expect((await storage.getKvCacheFootprintSnapshots({ connectionId: 'conn-a' })).map((r) => r.timestamp)).toEqual([5]); + }); + + it('filters samples by engine, model and range', async () => { + await storage.saveKvCacheEngineSamples([ + sample(), + sample({ engineId: 'eng-2' }), + sample({ modelName: 'other' }), + sample({ timestamp: 999_999 }), + ]); + const rows = await storage.getKvCacheEngineSamples({ + connectionId: 'conn-a', + engineId: 'eng-1', + modelName: 'm', + from: 0, + to: 100_000, + }); + expect(rows).toEqual([sample()]); + }); + + it('prunes samples', async () => { + await storage.saveKvCacheEngineSamples([sample({ timestamp: 1 }), sample({ timestamp: 10 })]); + expect(await storage.pruneOldKvCacheEngineSamples(5)).toBe(1); + }); + + it('upserts, lists, reads and deletes engines', async () => { + await storage.saveKvCacheEngine(engine()); + await storage.saveKvCacheEngine(engine({ id: 'eng-2', source: 'otlp', scrapeUrl: null, otlpEngineId: 'lmc-1', createdAt: 2 })); + await storage.saveKvCacheEngine(engine({ name: 'renamed', lastSeenAt: 9, lastError: 'HTTP 401' })); + expect((await storage.getKvCacheEngines('conn-a')).map((e) => e.name)).toEqual(['renamed', 'vllm-0']); + expect(await storage.getKvCacheEngine('eng-2')).toEqual( + engine({ id: 'eng-2', source: 'otlp', scrapeUrl: null, otlpEngineId: 'lmc-1', createdAt: 2 }), + ); + expect(await storage.deleteKvCacheEngine('eng-2')).toBe(true); + expect(await storage.deleteKvCacheEngine('eng-2')).toBe(false); + expect(await storage.getKvCacheEngine('eng-2')).toBeNull(); + }); + + it('keeps the identity of an engine on re-save and returns the stored row', async () => { + await storage.saveKvCacheEngine(engine()); + const saved = await storage.saveKvCacheEngine( + engine({ connectionId: 'conn-b', source: 'otlp', createdAt: 99, name: 'renamed', lastSeenAt: 5 }), + ); + const expected = engine({ name: 'renamed', lastSeenAt: 5 }); + expect(saved).toEqual(expected); + expect(await storage.getKvCacheEngine('eng-1')).toEqual(expected); + expect(await storage.getKvCacheEngines('conn-b')).toEqual([]); + }); + + it('upserts settings', async () => { + expect(await storage.getKvCacheSettings('conn-a')).toBeNull(); + const settings = { connectionId: 'conn-a', hitRateAlertEnabled: false, hitRateThreshold: 0.35, evictionAlertEnabled: true, updatedAt: 7 }; + await storage.saveKvCacheSettings(settings); + await storage.saveKvCacheSettings({ ...settings, hitRateThreshold: 0.4 }); + expect(await storage.getKvCacheSettings('conn-a')).toEqual({ ...settings, hitRateThreshold: 0.4 }); + }); + + it('deletes everything for a connection', async () => { + await storage.saveKvCacheFootprintSnapshot(snapshot()); + await storage.saveKvCacheEngineSamples([sample()]); + await storage.saveKvCacheEngine(engine()); + await storage.saveKvCacheSettings({ connectionId: 'conn-a', hitRateAlertEnabled: true, hitRateThreshold: 0.2, evictionAlertEnabled: true, updatedAt: 1 }); + await storage.saveKvCacheFootprintSnapshot(snapshot({ connectionId: 'conn-b' })); + await storage.deleteKvCacheConnectionData('conn-a'); + expect(await storage.getKvCacheFootprintSnapshots({ connectionId: 'conn-a' })).toEqual([]); + expect(await storage.getKvCacheEngineSamples({ connectionId: 'conn-a' })).toEqual([]); + expect(await storage.getKvCacheEngines('conn-a')).toEqual([]); + expect(await storage.getKvCacheSettings('conn-a')).toBeNull(); + expect(await storage.getKvCacheFootprintSnapshots({ connectionId: 'conn-b' })).toHaveLength(1); + }); +}); + +describe('postgres adapter kv cache storage before initialize', () => { + it('rejects engine samples instead of dropping them', async () => { + const adapter = new PostgresAdapter({ connectionString: 'postgres://unused' }); + await expect(adapter.saveKvCacheEngineSamples([sample()])).rejects.toThrow('Database not initialized'); + }); +}); + +describe('sqlite adapter kv cache connection cleanup', () => { + it('deletes each table in bounded chunks', async () => { + const dir = fs.mkdtempSync(path.join(os.tmpdir(), 'kv-cache-')); + const storage = new SqliteAdapter({ filepath: path.join(dir, 'db.sqlite') }); + await storage.initialize(); + const chunked = jest.spyOn(sqliteChunkedDelete, 'chunkedSqliteDelete'); + try { + await storage.deleteKvCacheConnectionData('conn-a'); + expect(chunked.mock.calls.map(([, table]) => table)).toEqual([ + 'kv_cache_footprint_snapshots', + 'kv_cache_engine_samples', + 'kv_cache_engines', + 'kv_cache_settings', + ]); + expect(chunked.mock.calls.every(([, , where, params]) => where === 'connection_id = ?' && params[0] === 'conn-a')).toBe(true); + } finally { + chunked.mockRestore(); + await storage.close(); + fs.rmSync(dir, { recursive: true, force: true }); + } + }); +}); diff --git a/apps/api/src/storage/adapters/memory.adapter.ts b/apps/api/src/storage/adapters/memory.adapter.ts index 2e6fd5100..ad9a52529 100644 --- a/apps/api/src/storage/adapters/memory.adapter.ts +++ b/apps/api/src/storage/adapters/memory.adapter.ts @@ -102,12 +102,19 @@ import type { ActivityRepository } from '../../common/interfaces/activity-reposi import type { StoredScalingReadinessScore, ScalingReadinessScoreQuery, + StoredKvCacheEngine, + KvCacheRangeQuery, + KvCacheFootprintQuery, + KvCacheSampleQuery, } from '../../common/interfaces/storage-port.interface'; import type { ReadinessBand, ReadinessDimension, ReadinessDimensionKey, ScalingReadinessSettings, + KvCacheEngineSample, + KvCacheFootprintSnapshot, + KvCacheSettings, } from '@betterdb/shared'; const NULL_SUB_DISCRIMINATOR = '__betterdb_null__'; @@ -1228,6 +1235,113 @@ export class MemoryAdapter implements StoragePort, RawDatabaseHandleProvider { return settings; } + private kvCacheSnapshots: KvCacheFootprintSnapshot[] = []; + private kvCacheSamples: KvCacheEngineSample[] = []; + private kvCacheEngines = new Map(); + private kvCacheSettings = new Map(); + + private inKvCacheRange(row: { connectionId: string; timestamp: number }, query: KvCacheRangeQuery): boolean { + return ( + row.connectionId === query.connectionId && + (query.from === undefined || row.timestamp >= query.from) && + (query.to === undefined || row.timestamp <= query.to) + ); + } + + async saveKvCacheFootprintSnapshot(snapshot: KvCacheFootprintSnapshot): Promise { + this.kvCacheSnapshots.push(structuredClone(snapshot)); + } + + async getKvCacheFootprintSnapshots(query: KvCacheFootprintQuery): Promise { + const rows = this.kvCacheSnapshots + .filter((r) => this.inKvCacheRange(r, query)) + .sort((a, b) => a.timestamp - b.timestamp); + const bucketMs = query.bucketMs; + const bucketed = bucketMs + ? [ + ...new Map(rows.map((r) => [Math.floor((r.timestamp - (query.from ?? 0)) / bucketMs), r])).values(), + ] + : rows; + return bucketed.slice(-(query.limit ?? 2000)).map((r) => structuredClone(r)); + } + + async pruneOldKvCacheFootprintSnapshots(cutoffTimestamp: number, connectionId?: string): Promise { + const before = this.kvCacheSnapshots.length; + this.kvCacheSnapshots = this.kvCacheSnapshots.filter( + (r) => r.timestamp >= cutoffTimestamp || (connectionId !== undefined && r.connectionId !== connectionId), + ); + return before - this.kvCacheSnapshots.length; + } + + async saveKvCacheEngineSamples(samples: KvCacheEngineSample[]): Promise { + this.kvCacheSamples.push(...samples.map((s) => ({ ...s }))); + } + + async getKvCacheEngineSamples(query: KvCacheSampleQuery): Promise { + return this.kvCacheSamples + .filter( + (r) => + this.inKvCacheRange(r, query) && + (query.engineId === undefined || r.engineId === query.engineId) && + (query.modelName === undefined || r.modelName === query.modelName), + ) + .sort((a, b) => a.timestamp - b.timestamp) + .slice(-(query.limit ?? 50_000)) + .map((r) => ({ ...r })); + } + + async pruneOldKvCacheEngineSamples(cutoffTimestamp: number, connectionId?: string): Promise { + const before = this.kvCacheSamples.length; + this.kvCacheSamples = this.kvCacheSamples.filter( + (r) => r.timestamp >= cutoffTimestamp || (connectionId !== undefined && r.connectionId !== connectionId), + ); + return before - this.kvCacheSamples.length; + } + + async saveKvCacheEngine(engine: StoredKvCacheEngine): Promise { + const existing = this.kvCacheEngines.get(engine.id); + const stored = existing + ? { ...engine, connectionId: existing.connectionId, source: existing.source, createdAt: existing.createdAt } + : { ...engine }; + this.kvCacheEngines.set(engine.id, stored); + return { ...stored }; + } + + async getKvCacheEngines(connectionId?: string): Promise { + return [...this.kvCacheEngines.values()] + .filter((e) => connectionId === undefined || e.connectionId === connectionId) + .sort((a, b) => a.createdAt - b.createdAt) + .map((e) => ({ ...e })); + } + + async getKvCacheEngine(id: string): Promise { + const engine = this.kvCacheEngines.get(id); + return engine ? { ...engine } : null; + } + + async deleteKvCacheEngine(id: string): Promise { + return this.kvCacheEngines.delete(id); + } + + async getKvCacheSettings(connectionId: string): Promise { + const settings = this.kvCacheSettings.get(connectionId); + return settings ? { ...settings } : null; + } + + async saveKvCacheSettings(settings: KvCacheSettings): Promise { + this.kvCacheSettings.set(settings.connectionId, { ...settings }); + return settings; + } + + async deleteKvCacheConnectionData(connectionId: string): Promise { + this.kvCacheSnapshots = this.kvCacheSnapshots.filter((r) => r.connectionId !== connectionId); + this.kvCacheSamples = this.kvCacheSamples.filter((r) => r.connectionId !== connectionId); + for (const [id, engine] of this.kvCacheEngines) { + if (engine.connectionId === connectionId) this.kvCacheEngines.delete(id); + } + this.kvCacheSettings.delete(connectionId); + } + // Command Stats Sample Methods private commandStatsSamples: StoredCommandStatsSample[] = []; diff --git a/apps/api/src/storage/adapters/postgres.adapter.ts b/apps/api/src/storage/adapters/postgres.adapter.ts index 8a8c64d4a..3aa2fc008 100644 --- a/apps/api/src/storage/adapters/postgres.adapter.ts +++ b/apps/api/src/storage/adapters/postgres.adapter.ts @@ -11,10 +11,18 @@ import { parseSshTunnel, parseMembership } from '@betterdb/shared'; import type { StoredScalingReadinessScore, ScalingReadinessScoreQuery, + StoredKvCacheEngine, + KvCacheRangeQuery, + KvCacheFootprintQuery, + KvCacheSampleQuery, } from '../../common/interfaces/storage-port.interface'; import type { ReadinessDimension, ScalingReadinessSettings, + KvCacheEngineSample, + KvCacheFootprintSnapshot, + KvCacheModelFootprint, + KvCacheSettings, } from '@betterdb/shared'; import type { RawDatabaseHandle, RawDatabaseHandleProvider } from '../raw-database-handle'; import { @@ -1690,6 +1698,82 @@ export class PostgresAdapter implements StoragePort, RawDatabaseHandleProvider { updated_at BIGINT NOT NULL ); + CREATE TABLE IF NOT EXISTS kv_cache_engines ( + id TEXT PRIMARY KEY, + connection_id TEXT NOT NULL, + name TEXT NOT NULL, + source TEXT NOT NULL, + scrape_url TEXT, + scrape_auth_header TEXT, + scrape_auth_encrypted BOOLEAN NOT NULL DEFAULT FALSE, + otlp_engine_id TEXT, + enabled BOOLEAN NOT NULL DEFAULT TRUE, + created_at BIGINT NOT NULL, + last_seen_at BIGINT, + last_error TEXT + ); + CREATE INDEX IF NOT EXISTS idx_kv_cache_engines_conn ON kv_cache_engines(connection_id); + CREATE UNIQUE INDEX IF NOT EXISTS idx_kv_cache_engines_otlp ON kv_cache_engines(otlp_engine_id); + + CREATE TABLE IF NOT EXISTS kv_cache_engine_samples ( + id BIGSERIAL PRIMARY KEY, + engine_id TEXT NOT NULL, + connection_id TEXT NOT NULL, + model_name TEXT NOT NULL, + timestamp BIGINT NOT NULL, + requested_tokens DOUBLE PRECISION NOT NULL, + hit_tokens DOUBLE PRECISION NOT NULL, + lookup_tokens DOUBLE PRECISION NOT NULL, + lookup_hits DOUBLE PRECISION NOT NULL, + remote_read_bytes DOUBLE PRECISION NOT NULL, + remote_write_bytes DOUBLE PRECISION NOT NULL, + remote_read_requests DOUBLE PRECISION NOT NULL, + remote_write_requests DOUBLE PRECISION NOT NULL, + remote_ping_errors DOUBLE PRECISION NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_kv_cache_samples_conn_ts ON kv_cache_engine_samples(connection_id, timestamp); + CREATE INDEX IF NOT EXISTS idx_kv_cache_samples_engine_ts ON kv_cache_engine_samples(engine_id, timestamp); + CREATE INDEX IF NOT EXISTS idx_kv_cache_samples_ts ON kv_cache_engine_samples(timestamp); + + CREATE TABLE IF NOT EXISTS kv_cache_footprint_snapshots ( + id BIGSERIAL PRIMARY KEY, + connection_id TEXT NOT NULL, + timestamp BIGINT NOT NULL, + detected BOOLEAN NOT NULL, + layout TEXT, + scanned_keys BIGINT NOT NULL, + matched_keys BIGINT NOT NULL, + sampled_keys BIGINT NOT NULL, + scan_complete BOOLEAN NOT NULL, + primaries_scanned INTEGER NOT NULL DEFAULT 1, + primaries_total INTEGER, + chunks_est DOUBLE PRECISION NOT NULL, + bytes_est DOUBLE PRECISION, + used_memory DOUBLE PRECISION NOT NULL, + maxmemory DOUBLE PRECISION NOT NULL, + maxmemory_policy TEXT NOT NULL, + lmcache_memory_share DOUBLE PRECISION, + no_ttl_ratio DOUBLE PRECISION, + orphan_ratio DOUBLE PRECISION, + evicted_keys_delta DOUBLE PRECISION, + evicted_keys_interval_ms BIGINT, + other_dbs JSONB NOT NULL, + per_model JSONB NOT NULL + ); + ALTER TABLE kv_cache_footprint_snapshots ADD COLUMN IF NOT EXISTS primaries_scanned INTEGER NOT NULL DEFAULT 1; + ALTER TABLE kv_cache_footprint_snapshots ADD COLUMN IF NOT EXISTS primaries_total INTEGER; + ALTER TABLE kv_cache_footprint_snapshots ADD COLUMN IF NOT EXISTS evicted_keys_interval_ms BIGINT; + CREATE INDEX IF NOT EXISTS idx_kv_cache_footprint_conn_ts ON kv_cache_footprint_snapshots(connection_id, timestamp); + CREATE INDEX IF NOT EXISTS idx_kv_cache_footprint_ts ON kv_cache_footprint_snapshots(timestamp); + + CREATE TABLE IF NOT EXISTS kv_cache_settings ( + connection_id TEXT PRIMARY KEY, + hit_rate_alert_enabled BOOLEAN NOT NULL DEFAULT TRUE, + hit_rate_threshold DOUBLE PRECISION NOT NULL DEFAULT 0.2, + eviction_alert_enabled BOOLEAN NOT NULL DEFAULT TRUE, + updated_at BIGINT NOT NULL + ); + CREATE TABLE IF NOT EXISTS webhooks ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), name VARCHAR(255) NOT NULL, @@ -4101,6 +4185,265 @@ export class PostgresAdapter implements StoragePort, RawDatabaseHandleProvider { return settings; } + + private kvCacheSnapshotFromRow(r: Record): KvCacheFootprintSnapshot { + const json = (v: unknown): T => (typeof v === 'string' ? JSON.parse(v) : v) as T; + const nullableNumber = (v: unknown): number | null => (v === null || v === undefined ? null : Number(v)); + return { + connectionId: r.connection_id, + timestamp: Number(r.timestamp), + detected: r.detected, + layout: r.layout, + scannedKeys: Number(r.scanned_keys), + matchedKeys: Number(r.matched_keys), + sampledKeys: Number(r.sampled_keys), + scanComplete: r.scan_complete, + primariesScanned: Number(r.primaries_scanned), + primariesTotal: nullableNumber(r.primaries_total), + chunksEst: Number(r.chunks_est), + bytesEst: nullableNumber(r.bytes_est), + usedMemory: Number(r.used_memory), + maxmemory: Number(r.maxmemory), + maxmemoryPolicy: r.maxmemory_policy, + lmcacheMemoryShare: nullableNumber(r.lmcache_memory_share), + noTtlRatio: nullableNumber(r.no_ttl_ratio), + orphanRatio: nullableNumber(r.orphan_ratio), + evictedKeysDelta: nullableNumber(r.evicted_keys_delta), + evictedKeysIntervalMs: nullableNumber(r.evicted_keys_interval_ms), + otherDbs: json(r.other_dbs), + perModel: json(r.per_model), + }; + } + + private kvCacheSampleFromRow(r: Record): KvCacheEngineSample { + return { + engineId: r.engine_id, + connectionId: r.connection_id, + modelName: r.model_name, + timestamp: Number(r.timestamp), + requestedTokens: Number(r.requested_tokens), + hitTokens: Number(r.hit_tokens), + lookupTokens: Number(r.lookup_tokens), + lookupHits: Number(r.lookup_hits), + remoteReadBytes: Number(r.remote_read_bytes), + remoteWriteBytes: Number(r.remote_write_bytes), + remoteReadRequests: Number(r.remote_read_requests), + remoteWriteRequests: Number(r.remote_write_requests), + remotePingErrors: Number(r.remote_ping_errors), + }; + } + + private kvCacheEngineFromRow(r: Record): StoredKvCacheEngine { + return { + id: r.id, + connectionId: r.connection_id, + name: r.name, + source: r.source, + scrapeUrl: r.scrape_url, + scrapeAuthHeader: r.scrape_auth_header, + scrapeAuthEncrypted: r.scrape_auth_encrypted, + otlpEngineId: r.otlp_engine_id, + enabled: r.enabled, + createdAt: Number(r.created_at), + lastSeenAt: r.last_seen_at === null ? null : Number(r.last_seen_at), + lastError: r.last_error, + }; + } + + async saveKvCacheFootprintSnapshot(s: KvCacheFootprintSnapshot): Promise { + if (!this.pool) throw new Error('Database not initialized'); + await this.pool.query( + `INSERT INTO kv_cache_footprint_snapshots + (connection_id, timestamp, detected, layout, scanned_keys, matched_keys, sampled_keys, + scan_complete, primaries_scanned, primaries_total, chunks_est, bytes_est, used_memory, maxmemory, + maxmemory_policy, lmcache_memory_share, no_ttl_ratio, orphan_ratio, evicted_keys_delta, evicted_keys_interval_ms, + other_dbs, per_model) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22)`, + [ + s.connectionId, s.timestamp, s.detected, s.layout, s.scannedKeys, s.matchedKeys, + s.sampledKeys, s.scanComplete, s.primariesScanned, s.primariesTotal, s.chunksEst, s.bytesEst, s.usedMemory, s.maxmemory, + s.maxmemoryPolicy, s.lmcacheMemoryShare, s.noTtlRatio, s.orphanRatio, s.evictedKeysDelta, s.evictedKeysIntervalMs, + JSON.stringify(s.otherDbs), JSON.stringify(s.perModel), + ], + ); + } + + async getKvCacheFootprintSnapshots(query: KvCacheFootprintQuery): Promise { + if (!this.pool) throw new Error('Database not initialized'); + const params = [query.connectionId, query.from ?? 0, query.to ?? Number.MAX_SAFE_INTEGER, query.limit ?? 2000]; + const result = query.bucketMs + ? await this.pool.query( + `SELECT * FROM ( + SELECT * FROM ( + SELECT *, ROW_NUMBER() OVER ( + PARTITION BY (timestamp - $2::bigint) / $5::bigint ORDER BY timestamp DESC, id DESC + ) AS bucket_rank + FROM kv_cache_footprint_snapshots + WHERE connection_id = $1 AND timestamp >= $2 AND timestamp <= $3 + ) AS ranked WHERE bucket_rank = 1 + ORDER BY timestamp DESC, id DESC LIMIT $4 + ) AS recent ORDER BY timestamp ASC, id ASC`, + [...params, Math.ceil(query.bucketMs)], + ) + : await this.pool.query( + `SELECT * FROM ( + SELECT * FROM kv_cache_footprint_snapshots + WHERE connection_id = $1 AND timestamp >= $2 AND timestamp <= $3 + ORDER BY timestamp DESC, id DESC LIMIT $4 + ) AS recent ORDER BY timestamp ASC, id ASC`, + params, + ); + return result.rows.map((r) => this.kvCacheSnapshotFromRow(r)); + } + + async pruneOldKvCacheFootprintSnapshots(cutoffTimestamp: number, connectionId?: string): Promise { + if (!this.pool) throw new Error('Database not initialized'); + if (connectionId) { + return chunkedPostgresDelete(this.pool, 'kv_cache_footprint_snapshots', 'timestamp < $1 AND connection_id = $2', [cutoffTimestamp, connectionId]); + } + return chunkedPostgresDelete(this.pool, 'kv_cache_footprint_snapshots', 'timestamp < $1', [cutoffTimestamp]); + } + + async saveKvCacheEngineSamples(samples: KvCacheEngineSample[]): Promise { + if (!this.pool) throw new Error('Database not initialized'); + if (samples.length === 0) return; + const columns = 13; + for (let offset = 0; offset < samples.length; offset += 500) { + const chunk = samples.slice(offset, offset + 500); + const values: unknown[] = []; + const placeholders = chunk.map((s, i) => { + values.push( + s.engineId, s.connectionId, s.modelName, s.timestamp, s.requestedTokens, s.hitTokens, + s.lookupTokens, s.lookupHits, s.remoteReadBytes, s.remoteWriteBytes, s.remoteReadRequests, + s.remoteWriteRequests, s.remotePingErrors, + ); + const base = i * columns; + return `(${Array.from({ length: columns }, (_, c) => `$${base + c + 1}`).join(', ')})`; + }); + await this.pool.query( + `INSERT INTO kv_cache_engine_samples + (engine_id, connection_id, model_name, timestamp, requested_tokens, hit_tokens, lookup_tokens, + lookup_hits, remote_read_bytes, remote_write_bytes, remote_read_requests, remote_write_requests, + remote_ping_errors) + VALUES ${placeholders.join(', ')}`, + values, + ); + } + } + + async getKvCacheEngineSamples(query: KvCacheSampleQuery): Promise { + if (!this.pool) throw new Error('Database not initialized'); + const params: unknown[] = [query.connectionId, query.from ?? 0, query.to ?? Number.MAX_SAFE_INTEGER]; + const where = ['connection_id = $1', 'timestamp >= $2', 'timestamp <= $3']; + if (query.engineId !== undefined) { + params.push(query.engineId); + where.push(`engine_id = $${params.length}`); + } + if (query.modelName !== undefined) { + params.push(query.modelName); + where.push(`model_name = $${params.length}`); + } + params.push(query.limit ?? 50_000); + const result = await this.pool.query( + `SELECT * FROM ( + SELECT * FROM kv_cache_engine_samples WHERE ${where.join(' AND ')} + ORDER BY timestamp DESC, id DESC LIMIT $${params.length} + ) AS recent ORDER BY timestamp ASC, id ASC`, + params, + ); + return result.rows.map((r) => this.kvCacheSampleFromRow(r)); + } + + async pruneOldKvCacheEngineSamples(cutoffTimestamp: number, connectionId?: string): Promise { + if (!this.pool) throw new Error('Database not initialized'); + if (connectionId) { + return chunkedPostgresDelete(this.pool, 'kv_cache_engine_samples', 'timestamp < $1 AND connection_id = $2', [cutoffTimestamp, connectionId]); + } + return chunkedPostgresDelete(this.pool, 'kv_cache_engine_samples', 'timestamp < $1', [cutoffTimestamp]); + } + + async saveKvCacheEngine(e: StoredKvCacheEngine): Promise { + if (!this.pool) throw new Error('Database not initialized'); + const result = await this.pool.query( + `INSERT INTO kv_cache_engines + (id, connection_id, name, source, scrape_url, scrape_auth_header, scrape_auth_encrypted, + otlp_engine_id, enabled, created_at, last_seen_at, last_error) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) + ON CONFLICT (id) DO UPDATE SET + name = EXCLUDED.name, + scrape_url = EXCLUDED.scrape_url, + scrape_auth_header = EXCLUDED.scrape_auth_header, + scrape_auth_encrypted = EXCLUDED.scrape_auth_encrypted, + otlp_engine_id = EXCLUDED.otlp_engine_id, + enabled = EXCLUDED.enabled, + last_seen_at = EXCLUDED.last_seen_at, + last_error = EXCLUDED.last_error + RETURNING *`, + [ + e.id, e.connectionId, e.name, e.source, e.scrapeUrl, e.scrapeAuthHeader, + e.scrapeAuthEncrypted, e.otlpEngineId, e.enabled, e.createdAt, e.lastSeenAt, e.lastError, + ], + ); + return this.kvCacheEngineFromRow(result.rows[0]); + } + + async getKvCacheEngines(connectionId?: string): Promise { + if (!this.pool) throw new Error('Database not initialized'); + const result = + connectionId === undefined + ? await this.pool.query('SELECT * FROM kv_cache_engines ORDER BY created_at ASC') + : await this.pool.query('SELECT * FROM kv_cache_engines WHERE connection_id = $1 ORDER BY created_at ASC', [connectionId]); + return result.rows.map((r) => this.kvCacheEngineFromRow(r)); + } + + async getKvCacheEngine(id: string): Promise { + if (!this.pool) throw new Error('Database not initialized'); + const result = await this.pool.query('SELECT * FROM kv_cache_engines WHERE id = $1', [id]); + return result.rows.length === 0 ? null : this.kvCacheEngineFromRow(result.rows[0]); + } + + async deleteKvCacheEngine(id: string): Promise { + if (!this.pool) throw new Error('Database not initialized'); + const result = await this.pool.query('DELETE FROM kv_cache_engines WHERE id = $1', [id]); + return (result.rowCount ?? 0) > 0; + } + + async getKvCacheSettings(connectionId: string): Promise { + if (!this.pool) throw new Error('Database not initialized'); + const result = await this.pool.query('SELECT * FROM kv_cache_settings WHERE connection_id = $1', [connectionId]); + if (result.rows.length === 0) return null; + const row = result.rows[0]; + return { + connectionId: row.connection_id, + hitRateAlertEnabled: row.hit_rate_alert_enabled, + hitRateThreshold: Number(row.hit_rate_threshold), + evictionAlertEnabled: row.eviction_alert_enabled, + updatedAt: Number(row.updated_at), + }; + } + + async saveKvCacheSettings(settings: KvCacheSettings): Promise { + if (!this.pool) throw new Error('Database not initialized'); + await this.pool.query( + `INSERT INTO kv_cache_settings (connection_id, hit_rate_alert_enabled, hit_rate_threshold, eviction_alert_enabled, updated_at) + VALUES ($1, $2, $3, $4, $5) + ON CONFLICT (connection_id) DO UPDATE SET + hit_rate_alert_enabled = EXCLUDED.hit_rate_alert_enabled, + hit_rate_threshold = EXCLUDED.hit_rate_threshold, + eviction_alert_enabled = EXCLUDED.eviction_alert_enabled, + updated_at = EXCLUDED.updated_at`, + [settings.connectionId, settings.hitRateAlertEnabled, settings.hitRateThreshold, settings.evictionAlertEnabled, settings.updatedAt], + ); + return settings; + } + + async deleteKvCacheConnectionData(connectionId: string): Promise { + if (!this.pool) throw new Error('Database not initialized'); + for (const table of ['kv_cache_footprint_snapshots', 'kv_cache_engine_samples', 'kv_cache_engines', 'kv_cache_settings']) { + await chunkedPostgresDelete(this.pool, table, 'connection_id = $1', [connectionId]); + } + } + // Command Stats Sample Methods async saveCommandStatsSamples( samples: Omit[], diff --git a/apps/api/src/storage/adapters/sqlite.adapter.ts b/apps/api/src/storage/adapters/sqlite.adapter.ts index 8d07cdf83..0cd6aa0cc 100644 --- a/apps/api/src/storage/adapters/sqlite.adapter.ts +++ b/apps/api/src/storage/adapters/sqlite.adapter.ts @@ -12,12 +12,20 @@ import { parseSshTunnel, parseMembership } from '@betterdb/shared'; import type { StoredScalingReadinessScore, ScalingReadinessScoreQuery, + StoredKvCacheEngine, + KvCacheRangeQuery, + KvCacheFootprintQuery, + KvCacheSampleQuery, } from '../../common/interfaces/storage-port.interface'; import type { ReadinessBand, ReadinessDimension, ReadinessDimensionKey, ScalingReadinessSettings, + KvCacheEngineSample, + KvCacheFootprintSnapshot, + KvCacheModelFootprint, + KvCacheSettings, } from '@betterdb/shared'; import type { RawDatabaseHandle, RawDatabaseHandleProvider } from '../raw-database-handle'; import { @@ -272,6 +280,20 @@ function addMemoryProposalIntegrityColumns(db: Database.Database): void { * Idempotent migrations for capture_sessions columns landed across PR 14a + 14b. * Extracted so the call-site in createSchema stays readable. */ +function addKvCacheCoverageColumns(db: Database.Database): void { + const cols = db.prepare(`PRAGMA table_info(kv_cache_footprint_snapshots)`).all() as { name: string }[]; + if (cols.length === 0) return; + if (!cols.some((c) => c.name === 'primaries_scanned')) { + db.prepare(`ALTER TABLE kv_cache_footprint_snapshots ADD COLUMN primaries_scanned INTEGER NOT NULL DEFAULT 1`).run(); + } + if (!cols.some((c) => c.name === 'primaries_total')) { + db.prepare(`ALTER TABLE kv_cache_footprint_snapshots ADD COLUMN primaries_total INTEGER`).run(); + } + if (!cols.some((c) => c.name === 'evicted_keys_interval_ms')) { + db.prepare(`ALTER TABLE kv_cache_footprint_snapshots ADD COLUMN evicted_keys_interval_ms INTEGER`).run(); + } +} + function addCaptureSessionsTargetNodeColumn(db: Database.Database): void { try { const cols = db.prepare(`PRAGMA table_info(capture_sessions)`).all() as { name: string }[]; @@ -1892,6 +1914,79 @@ export class SqliteAdapter implements StoragePort, RawDatabaseHandleProvider { updated_at INTEGER NOT NULL ); + CREATE TABLE IF NOT EXISTS kv_cache_engines ( + id TEXT PRIMARY KEY, + connection_id TEXT NOT NULL, + name TEXT NOT NULL, + source TEXT NOT NULL, + scrape_url TEXT, + scrape_auth_header TEXT, + scrape_auth_encrypted INTEGER NOT NULL DEFAULT 0, + otlp_engine_id TEXT, + enabled INTEGER NOT NULL DEFAULT 1, + created_at INTEGER NOT NULL, + last_seen_at INTEGER, + last_error TEXT + ); + CREATE INDEX IF NOT EXISTS idx_kv_cache_engines_conn ON kv_cache_engines(connection_id); + CREATE UNIQUE INDEX IF NOT EXISTS idx_kv_cache_engines_otlp ON kv_cache_engines(otlp_engine_id); + + CREATE TABLE IF NOT EXISTS kv_cache_engine_samples ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + engine_id TEXT NOT NULL, + connection_id TEXT NOT NULL, + model_name TEXT NOT NULL, + timestamp INTEGER NOT NULL, + requested_tokens REAL NOT NULL, + hit_tokens REAL NOT NULL, + lookup_tokens REAL NOT NULL, + lookup_hits REAL NOT NULL, + remote_read_bytes REAL NOT NULL, + remote_write_bytes REAL NOT NULL, + remote_read_requests REAL NOT NULL, + remote_write_requests REAL NOT NULL, + remote_ping_errors REAL NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_kv_cache_samples_conn_ts ON kv_cache_engine_samples(connection_id, timestamp); + CREATE INDEX IF NOT EXISTS idx_kv_cache_samples_engine_ts ON kv_cache_engine_samples(engine_id, timestamp); + CREATE INDEX IF NOT EXISTS idx_kv_cache_samples_ts ON kv_cache_engine_samples(timestamp); + + CREATE TABLE IF NOT EXISTS kv_cache_footprint_snapshots ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + connection_id TEXT NOT NULL, + timestamp INTEGER NOT NULL, + detected INTEGER NOT NULL, + layout TEXT, + scanned_keys INTEGER NOT NULL, + matched_keys INTEGER NOT NULL, + sampled_keys INTEGER NOT NULL, + scan_complete INTEGER NOT NULL, + primaries_scanned INTEGER NOT NULL DEFAULT 1, + primaries_total INTEGER, + chunks_est REAL NOT NULL, + bytes_est REAL, + used_memory REAL NOT NULL, + maxmemory REAL NOT NULL, + maxmemory_policy TEXT NOT NULL, + lmcache_memory_share REAL, + no_ttl_ratio REAL, + orphan_ratio REAL, + evicted_keys_delta REAL, + evicted_keys_interval_ms INTEGER, + other_dbs TEXT NOT NULL, + per_model TEXT NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_kv_cache_footprint_conn_ts ON kv_cache_footprint_snapshots(connection_id, timestamp); + CREATE INDEX IF NOT EXISTS idx_kv_cache_footprint_ts ON kv_cache_footprint_snapshots(timestamp); + + CREATE TABLE IF NOT EXISTS kv_cache_settings ( + connection_id TEXT PRIMARY KEY, + hit_rate_alert_enabled INTEGER NOT NULL DEFAULT 1, + hit_rate_threshold REAL NOT NULL DEFAULT 0.2, + eviction_alert_enabled INTEGER NOT NULL DEFAULT 1, + updated_at INTEGER NOT NULL + ); + CREATE TABLE IF NOT EXISTS webhooks ( id TEXT PRIMARY KEY, name TEXT NOT NULL, @@ -2363,6 +2458,7 @@ export class SqliteAdapter implements StoragePort, RawDatabaseHandleProvider { ); addCaptureSessionsTargetNodeColumn(this.db!); addMemoryProposalIntegrityColumns(this.db!); + addKvCacheCoverageColumns(this.db!); // Agent Tokens Table (cloud-only, but created in all environments for interface compliance) this.db.exec(` @@ -4021,6 +4117,271 @@ export class SqliteAdapter implements StoragePort, RawDatabaseHandleProvider { return settings; } + + private kvCacheSnapshotFromRow(r: Record): KvCacheFootprintSnapshot { + return { + connectionId: r.connection_id, + timestamp: r.timestamp, + detected: r.detected === 1, + layout: r.layout, + scannedKeys: r.scanned_keys, + matchedKeys: r.matched_keys, + sampledKeys: r.sampled_keys, + scanComplete: r.scan_complete === 1, + primariesScanned: r.primaries_scanned, + primariesTotal: r.primaries_total, + chunksEst: r.chunks_est, + bytesEst: r.bytes_est, + usedMemory: r.used_memory, + maxmemory: r.maxmemory, + maxmemoryPolicy: r.maxmemory_policy, + lmcacheMemoryShare: r.lmcache_memory_share, + noTtlRatio: r.no_ttl_ratio, + orphanRatio: r.orphan_ratio, + evictedKeysDelta: r.evicted_keys_delta, + evictedKeysIntervalMs: r.evicted_keys_interval_ms ?? null, + otherDbs: JSON.parse(r.other_dbs) as number[], + perModel: JSON.parse(r.per_model) as KvCacheModelFootprint[], + }; + } + + private kvCacheSampleFromRow(r: Record): KvCacheEngineSample { + return { + engineId: r.engine_id, + connectionId: r.connection_id, + modelName: r.model_name, + timestamp: r.timestamp, + requestedTokens: r.requested_tokens, + hitTokens: r.hit_tokens, + lookupTokens: r.lookup_tokens, + lookupHits: r.lookup_hits, + remoteReadBytes: r.remote_read_bytes, + remoteWriteBytes: r.remote_write_bytes, + remoteReadRequests: r.remote_read_requests, + remoteWriteRequests: r.remote_write_requests, + remotePingErrors: r.remote_ping_errors, + }; + } + + private kvCacheEngineFromRow(r: Record): StoredKvCacheEngine { + return { + id: r.id, + connectionId: r.connection_id, + name: r.name, + source: r.source, + scrapeUrl: r.scrape_url, + scrapeAuthHeader: r.scrape_auth_header, + scrapeAuthEncrypted: r.scrape_auth_encrypted === 1, + otlpEngineId: r.otlp_engine_id, + enabled: r.enabled === 1, + createdAt: r.created_at, + lastSeenAt: r.last_seen_at, + lastError: r.last_error, + }; + } + + async saveKvCacheFootprintSnapshot(s: KvCacheFootprintSnapshot): Promise { + if (!this.db) throw new Error('Database not initialized'); + this.db + .prepare( + `INSERT INTO kv_cache_footprint_snapshots + (connection_id, timestamp, detected, layout, scanned_keys, matched_keys, sampled_keys, + scan_complete, primaries_scanned, primaries_total, chunks_est, bytes_est, used_memory, maxmemory, + maxmemory_policy, lmcache_memory_share, no_ttl_ratio, orphan_ratio, evicted_keys_delta, evicted_keys_interval_ms, + other_dbs, per_model) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + ) + .run( + s.connectionId, s.timestamp, s.detected ? 1 : 0, s.layout, s.scannedKeys, s.matchedKeys, + s.sampledKeys, s.scanComplete ? 1 : 0, s.primariesScanned, s.primariesTotal, s.chunksEst, s.bytesEst, s.usedMemory, s.maxmemory, + s.maxmemoryPolicy, s.lmcacheMemoryShare, s.noTtlRatio, s.orphanRatio, s.evictedKeysDelta, s.evictedKeysIntervalMs, + JSON.stringify(s.otherDbs), JSON.stringify(s.perModel), + ); + } + + async getKvCacheFootprintSnapshots(query: KvCacheFootprintQuery): Promise { + if (!this.db) throw new Error('Database not initialized'); + const from = query.from ?? 0; + const range = [query.connectionId, from, query.to ?? Number.MAX_SAFE_INTEGER]; + const limit = query.limit ?? 2000; + const rows = ( + query.bucketMs + ? this.db + .prepare( + `SELECT * FROM ( + SELECT * FROM ( + SELECT *, ROW_NUMBER() OVER ( + PARTITION BY CAST((timestamp - ?) / ? AS INTEGER) ORDER BY timestamp DESC, id DESC + ) AS bucket_rank + FROM kv_cache_footprint_snapshots + WHERE connection_id = ? AND timestamp >= ? AND timestamp <= ? + ) WHERE bucket_rank = 1 + ORDER BY timestamp DESC, id DESC LIMIT ? + ) ORDER BY timestamp ASC, id ASC`, + ) + .all(from, Math.ceil(query.bucketMs), ...range, limit) + : this.db + .prepare( + `SELECT * FROM ( + SELECT * FROM kv_cache_footprint_snapshots + WHERE connection_id = ? AND timestamp >= ? AND timestamp <= ? + ORDER BY timestamp DESC, id DESC LIMIT ? + ) ORDER BY timestamp ASC, id ASC`, + ) + .all(...range, limit) + ) as Record[]; + return rows.map((r) => this.kvCacheSnapshotFromRow(r)); + } + + async pruneOldKvCacheFootprintSnapshots(cutoffTimestamp: number, connectionId?: string): Promise { + if (!this.db) throw new Error('Database not initialized'); + if (connectionId) { + return chunkedSqliteDelete(this.db, 'kv_cache_footprint_snapshots', 'timestamp < ? AND connection_id = ?', [cutoffTimestamp, connectionId]); + } + return chunkedSqliteDelete(this.db, 'kv_cache_footprint_snapshots', 'timestamp < ?', [cutoffTimestamp]); + } + + async saveKvCacheEngineSamples(samples: KvCacheEngineSample[]): Promise { + if (!this.db) throw new Error('Database not initialized'); + const insert = this.db.prepare( + `INSERT INTO kv_cache_engine_samples + (engine_id, connection_id, model_name, timestamp, requested_tokens, hit_tokens, lookup_tokens, + lookup_hits, remote_read_bytes, remote_write_bytes, remote_read_requests, remote_write_requests, + remote_ping_errors) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + ); + const insertAll = this.db.transaction((rows: KvCacheEngineSample[]) => { + for (const s of rows) { + insert.run( + s.engineId, s.connectionId, s.modelName, s.timestamp, s.requestedTokens, s.hitTokens, + s.lookupTokens, s.lookupHits, s.remoteReadBytes, s.remoteWriteBytes, s.remoteReadRequests, + s.remoteWriteRequests, s.remotePingErrors, + ); + } + }); + insertAll(samples); + } + + async getKvCacheEngineSamples(query: KvCacheSampleQuery): Promise { + if (!this.db) throw new Error('Database not initialized'); + const where = ['connection_id = ?', 'timestamp >= ?', 'timestamp <= ?']; + const params: unknown[] = [query.connectionId, query.from ?? 0, query.to ?? Number.MAX_SAFE_INTEGER]; + if (query.engineId !== undefined) { + where.push('engine_id = ?'); + params.push(query.engineId); + } + if (query.modelName !== undefined) { + where.push('model_name = ?'); + params.push(query.modelName); + } + params.push(query.limit ?? 50_000); + const rows = this.db + .prepare( + `SELECT * FROM ( + SELECT * FROM kv_cache_engine_samples WHERE ${where.join(' AND ')} + ORDER BY timestamp DESC, id DESC LIMIT ? + ) ORDER BY timestamp ASC, id ASC`, + ) + .all(...params) as Record[]; + return rows.map((r) => this.kvCacheSampleFromRow(r)); + } + + async pruneOldKvCacheEngineSamples(cutoffTimestamp: number, connectionId?: string): Promise { + if (!this.db) throw new Error('Database not initialized'); + if (connectionId) { + return chunkedSqliteDelete(this.db, 'kv_cache_engine_samples', 'timestamp < ? AND connection_id = ?', [cutoffTimestamp, connectionId]); + } + return chunkedSqliteDelete(this.db, 'kv_cache_engine_samples', 'timestamp < ?', [cutoffTimestamp]); + } + + async saveKvCacheEngine(e: StoredKvCacheEngine): Promise { + if (!this.db) throw new Error('Database not initialized'); + const row = this.db + .prepare( + `INSERT INTO kv_cache_engines + (id, connection_id, name, source, scrape_url, scrape_auth_header, scrape_auth_encrypted, + otlp_engine_id, enabled, created_at, last_seen_at, last_error) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET + name = excluded.name, + scrape_url = excluded.scrape_url, + scrape_auth_header = excluded.scrape_auth_header, + scrape_auth_encrypted = excluded.scrape_auth_encrypted, + otlp_engine_id = excluded.otlp_engine_id, + enabled = excluded.enabled, + last_seen_at = excluded.last_seen_at, + last_error = excluded.last_error + RETURNING *`, + ) + .get( + e.id, e.connectionId, e.name, e.source, e.scrapeUrl, e.scrapeAuthHeader, + e.scrapeAuthEncrypted ? 1 : 0, e.otlpEngineId, e.enabled ? 1 : 0, e.createdAt, + e.lastSeenAt, e.lastError, + ) as Record; + return this.kvCacheEngineFromRow(row); + } + + async getKvCacheEngines(connectionId?: string): Promise { + if (!this.db) throw new Error('Database not initialized'); + const rows = (connectionId === undefined + ? this.db.prepare('SELECT * FROM kv_cache_engines ORDER BY created_at ASC').all() + : this.db.prepare('SELECT * FROM kv_cache_engines WHERE connection_id = ? ORDER BY created_at ASC').all(connectionId)) as Record[]; + return rows.map((r) => this.kvCacheEngineFromRow(r)); + } + + async getKvCacheEngine(id: string): Promise { + if (!this.db) throw new Error('Database not initialized'); + const row = this.db.prepare('SELECT * FROM kv_cache_engines WHERE id = ?').get(id) as Record | undefined; + return row ? this.kvCacheEngineFromRow(row) : null; + } + + async deleteKvCacheEngine(id: string): Promise { + if (!this.db) throw new Error('Database not initialized'); + return this.db.prepare('DELETE FROM kv_cache_engines WHERE id = ?').run(id).changes > 0; + } + + async getKvCacheSettings(connectionId: string): Promise { + if (!this.db) throw new Error('Database not initialized'); + const row = this.db.prepare('SELECT * FROM kv_cache_settings WHERE connection_id = ?').get(connectionId) as Record | undefined; + if (!row) return null; + return { + connectionId: row.connection_id, + hitRateAlertEnabled: row.hit_rate_alert_enabled === 1, + hitRateThreshold: row.hit_rate_threshold, + evictionAlertEnabled: row.eviction_alert_enabled === 1, + updatedAt: row.updated_at, + }; + } + + async saveKvCacheSettings(settings: KvCacheSettings): Promise { + if (!this.db) throw new Error('Database not initialized'); + this.db + .prepare( + `INSERT INTO kv_cache_settings (connection_id, hit_rate_alert_enabled, hit_rate_threshold, eviction_alert_enabled, updated_at) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT(connection_id) DO UPDATE SET + hit_rate_alert_enabled = excluded.hit_rate_alert_enabled, + hit_rate_threshold = excluded.hit_rate_threshold, + eviction_alert_enabled = excluded.eviction_alert_enabled, + updated_at = excluded.updated_at`, + ) + .run( + settings.connectionId, + settings.hitRateAlertEnabled ? 1 : 0, + settings.hitRateThreshold, + settings.evictionAlertEnabled ? 1 : 0, + settings.updatedAt, + ); + return settings; + } + + async deleteKvCacheConnectionData(connectionId: string): Promise { + if (!this.db) throw new Error('Database not initialized'); + for (const table of ['kv_cache_footprint_snapshots', 'kv_cache_engine_samples', 'kv_cache_engines', 'kv_cache_settings']) { + await chunkedSqliteDelete(this.db, table, 'connection_id = ?', [connectionId]); + } + } + // Command Stats Sample Methods async saveCommandStatsSamples( samples: Omit[], diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index e3d75da7d..d8a2c31f8 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -20,6 +20,7 @@ export * from './types/vector-index-snapshots'; export * from './types/migration'; export * from './types/metric-forecasting.types'; export * from './types/scaling-readiness.types'; +export * from './types/kv-cache.types'; export * from './types/telemetry'; export * from './types/cli.types'; export * from './types/command-safety'; diff --git a/packages/shared/src/license/types.ts b/packages/shared/src/license/types.ts index d17cf7036..9771f133f 100644 --- a/packages/shared/src/license/types.ts +++ b/packages/shared/src/license/types.ts @@ -42,6 +42,7 @@ export enum Feature { MONITOR_SCHEDULED_CAPTURES = 'monitorScheduledCaptures', MONITOR_CAPTURE_DIFF = 'monitorCaptureDiff', SCALING_READINESS_HISTORY = 'scalingReadinessHistory', + KV_CACHE_MONITORING = 'kvCacheMonitoring', // Enterprise-only features SSO_SAML = 'ssoSaml', COMPLIANCE_EXPORT = 'complianceExport', @@ -72,6 +73,7 @@ export const TIER_FEATURES: Record = { Feature.MONITOR_SCHEDULED_CAPTURES, Feature.MONITOR_CAPTURE_DIFF, Feature.SCALING_READINESS_HISTORY, + Feature.KV_CACHE_MONITORING, ], [Tier.enterprise]: Object.values(Feature), }; diff --git a/packages/shared/src/types/kv-cache.types.ts b/packages/shared/src/types/kv-cache.types.ts new file mode 100644 index 000000000..02af0e2ff --- /dev/null +++ b/packages/shared/src/types/kv-cache.types.ts @@ -0,0 +1,146 @@ +export type KvCacheLayout = 'two_key' | 'single_key' | 'mixed'; +export type KvCacheEngineSource = 'scrape' | 'otlp'; +export type KvCacheEvictionReason = 'unevictable' | 'evicting'; + +export interface KvCacheModelFootprint { + model: string; + dtype: string; + chunksEst: number; + bytesEst: number | null; +} + +export interface KvCacheFootprintSnapshot { + connectionId: string; + timestamp: number; + detected: boolean; + layout: KvCacheLayout | null; + scannedKeys: number; + matchedKeys: number; + sampledKeys: number; + scanComplete: boolean; + primariesScanned: number; + primariesTotal: number | null; + chunksEst: number; + bytesEst: number | null; + usedMemory: number; + maxmemory: number; + maxmemoryPolicy: string; + lmcacheMemoryShare: number | null; + noTtlRatio: number | null; + orphanRatio: number | null; + evictedKeysDelta: number | null; + evictedKeysIntervalMs: number | null; + otherDbs: number[]; + perModel: KvCacheModelFootprint[]; +} + +export interface KvCacheEngine { + id: string; + connectionId: string; + name: string; + source: KvCacheEngineSource; + scrapeUrl: string | null; + hasScrapeAuth: boolean; + otlpEngineId: string | null; + enabled: boolean; + createdAt: number; + lastSeenAt: number | null; + lastError: string | null; +} + +export interface KvCacheEngineCreate { + name: string; + source: KvCacheEngineSource; + scrapeUrl?: string; + scrapeAuthHeader?: string; + otlpEngineId?: string; + enabled?: boolean; +} + +export interface KvCacheEngineUpdate { + name?: string; + scrapeUrl?: string; + scrapeAuthHeader?: string | null; + enabled?: boolean; +} + +export interface KvCacheSampleCounters { + requestedTokens: number; + hitTokens: number; + lookupTokens: number; + lookupHits: number; + remoteReadBytes: number; + remoteWriteBytes: number; + remoteReadRequests: number; + remoteWriteRequests: number; + remotePingErrors: number; +} + +export interface KvCacheEngineSample extends KvCacheSampleCounters { + engineId: string; + connectionId: string; + modelName: string; + timestamp: number; +} + +export interface KvCacheSampleBucket extends KvCacheEngineSample { + hitRate: number | null; +} + +export interface KvCacheSamplesResponse { + buckets: KvCacheSampleBucket[]; + rangeHitRate: number | null; +} + +export interface KvCacheStatus { + hasLmcache: boolean; + latest: KvCacheFootprintSnapshot | null; + sampleKey: string | null; + engines: KvCacheEngine[]; +} + +export interface KvCacheSettings { + connectionId: string; + hitRateAlertEnabled: boolean; + hitRateThreshold: number; + evictionAlertEnabled: boolean; + updatedAt: number; +} + +export interface KvCacheSettingsUpdate { + hitRateAlertEnabled?: boolean; + hitRateThreshold?: number; + evictionAlertEnabled?: boolean; +} + +export interface KvCacheHitRateLowData { + connectionId: string; + engineId: string; + engineName: string; + model: string; + hitRate: number; + threshold: number; + requestedTokens: number; + windowMs: number; + timestamp: number; + instance?: { host: string; port: number }; +} + +export interface KvCacheEvictionRiskData { + connectionId: string; + reason: KvCacheEvictionReason; + active: boolean; + policy: string; + usedMemory: number; + maxmemory: number; + lmcacheMemoryShare: number | null; + noTtlRatio: number | null; + evictedKeysDelta: number | null; + timestamp: number; + instance?: { host: string; port: number }; +} + +export const DEFAULT_KV_CACHE_HIT_RATE_THRESHOLD = 0.2; +export const MAX_KV_CACHE_HIT_RATE_THRESHOLD = 0.9; +export const KV_CACHE_HIT_RATE_WINDOW_MS = 15 * 60_000; +export const KV_CACHE_HIT_RATE_MIN_TOKENS = 10_000; diff --git a/proprietary/data-retention/__tests__/data-retention.service.spec.ts b/proprietary/data-retention/__tests__/data-retention.service.spec.ts index 412e1cdb2..7ac7b67cf 100644 --- a/proprietary/data-retention/__tests__/data-retention.service.spec.ts +++ b/proprietary/data-retention/__tests__/data-retention.service.spec.ts @@ -42,6 +42,8 @@ describe('DataRetentionService', () => { pruneOldCommandStatsSamples: jest.fn().mockResolvedValue(18), pruneOldLatencyStatsSamples: jest.fn().mockResolvedValue(19), pruneOldVectorIndexSnapshots: jest.fn().mockResolvedValue(20), + pruneOldKvCacheFootprintSnapshots: jest.fn().mockResolvedValue(22), + pruneOldKvCacheEngineSamples: jest.fn().mockResolvedValue(23), } as any; licenseService = { @@ -97,6 +99,7 @@ describe('DataRetentionService', () => { 'pruneOldScheduledCaptures', 'pruneOldAiCacheSamples', 'pruneOldOtelSpans', + 'pruneOldKvCacheFootprintSnapshots', ] as const; // Rollups compare bucket_start, so they get the cutoff aligned down to its @@ -115,6 +118,7 @@ describe('DataRetentionService', () => { 'pruneOldCommandStatsSamples', 'pruneOldLatencyStatsSamples', 'pruneOldVectorIndexSnapshots', + 'pruneOldKvCacheEngineSamples', ] as const; it('community tier uses 7-day cutoff and calls all prune methods', async () => { diff --git a/proprietary/key-analytics/__tests__/key-analytics-topology.spec.ts b/proprietary/key-analytics/__tests__/key-analytics-topology.spec.ts index ece1df517..754f5b0af 100644 --- a/proprietary/key-analytics/__tests__/key-analytics-topology.spec.ts +++ b/proprietary/key-analytics/__tests__/key-analytics-topology.spec.ts @@ -11,7 +11,7 @@ import type { ConnectionRegistry } from '@app/connections/connection-registry.se import type { StoragePort } from '@app/common/interfaces/storage-port.interface'; import type { LicenseService } from '@proprietary/licenses'; import { KeyAnalyticsService } from '../key-analytics.service'; -import { collectionPlan } from '../key-analytics-plan'; +import { collectionPlan, planForConnection, resolvePlan } from '../key-analytics-plan'; import { buildPatternSnapshots, mergePatternSnapshots } from '../key-pattern-snapshots'; function connection(id: string, over: Partial = {}): ConnectionStatus { @@ -105,6 +105,30 @@ describe('collectionPlan', () => { }); }); +describe('planForConnection', () => { + const seed = connection('seed'); + const node = connection('node', { membership: member() }); + const registry = { list: () => [seed, node] } as unknown as ConnectionRegistry; + + it('plans a listed connection against the registry', () => { + expect(planForConnection(registry, 'seed')).toEqual({ kind: 'cluster', members: [node] }); + expect(planForConnection(registry, 'node')).toEqual({ kind: 'skip' }); + }); + + it('scans an unlisted connection on its own', () => { + expect(planForConnection(registry, 'missing')).toEqual({ kind: 'single' }); + }); + + it('skips a retired member that only the retired listing includes', () => { + const retired = connection('retired', { membership: member({ retiredAt: 1 }) }); + const withRetired = { + list: (options?: { includeRetired?: boolean }) => (options?.includeRetired ? [seed, node, retired] : [seed, node]), + } as unknown as ConnectionRegistry; + expect(planForConnection(withRetired, 'retired')).toEqual({ kind: 'skip' }); + expect(resolvePlan(withRetired, 'retired')).toEqual({ connection: retired, plan: { kind: 'skip' } }); + }); +}); + describe('mergePatternSnapshots', () => { it('returns a single node untouched', () => { const only = [snapshot({ keyCount: 5 })]; diff --git a/proprietary/key-analytics/__tests__/scan-nodes.spec.ts b/proprietary/key-analytics/__tests__/scan-nodes.spec.ts new file mode 100644 index 000000000..cf1266511 --- /dev/null +++ b/proprietary/key-analytics/__tests__/scan-nodes.spec.ts @@ -0,0 +1,41 @@ +import { Logger } from '@nestjs/common'; +import { clusterScanNodes, primaryScanNodes } from '../scan-nodes'; + +const logger = { warn: jest.fn() } as unknown as Logger; +const client = (role: string) => ({ getRole: jest.fn().mockResolvedValue({ role }) }) as any; + +describe('scan nodes', () => { + it('adds registered members after the seed and skips unregistered ones', () => { + const registry = { get: jest.fn((id: string) => { if (id === 'gone') throw new Error('x'); return client('master'); }) } as any; + const nodes = clusterScanNodes(registry, { name: 'seed', client: client('master') }, [{ id: 'a', name: 'A', membership: { nodeId: 'na' } }, { id: 'gone', name: 'G' }] as any, logger, 'KV cache'); + expect(nodes.map((n) => [n.name, n.connectionId, n.nodeId])).toEqual([['seed', undefined, undefined], ['A', 'a', 'na']]); + }); + + it('keeps primaries only and drops nodes whose role fails', async () => { + const failing = { getRole: jest.fn().mockRejectedValue(new Error('down')) } as any; + const nodes = await primaryScanNodes( + [{ name: 'p', client: client('master') }, { name: 'r', client: client('slave') }, { name: 'f', client: failing }], + logger, + 'KV cache', + ); + expect(nodes.map((n) => n.name)).toEqual(['p']); + }); + + it('drops a node whose cluster id an earlier node already reported', async () => { + const id = 'a'.repeat(40); + const other = 'b'.repeat(40); + const withId = (myId: unknown) => ({ getRole: jest.fn().mockResolvedValue({ role: 'master' }), call: jest.fn().mockResolvedValue(myId) }) as any; + const nodes = await primaryScanNodes( + [ + { name: 'dns-seed', client: withId(id) }, + { name: 'ip-member', client: withId(id), connectionId: 'm1', nodeId: 'stale' }, + { name: 'other', client: withId(other), connectionId: 'm2' }, + { name: 'no-id', client: withId('junk') }, + { name: 'no-call', client: client('master') }, + ], + logger, + 'KV cache', + ); + expect(nodes.map((n) => [n.name, n.nodeId])).toEqual([['dns-seed', id], ['other', other], ['no-id', undefined], ['no-call', undefined]]); + }); +}); diff --git a/proprietary/key-analytics/key-analytics-plan.ts b/proprietary/key-analytics/key-analytics-plan.ts index 8696fb855..a2dea0112 100644 --- a/proprietary/key-analytics/key-analytics-plan.ts +++ b/proprietary/key-analytics/key-analytics-plan.ts @@ -1,4 +1,5 @@ import { hasOwnKeyAnalytics, type ConnectionStatus } from '@betterdb/shared'; +import type { ConnectionRegistry } from '@app/connections/connection-registry.service'; export type CollectionPlan = | { kind: 'skip' } @@ -22,3 +23,18 @@ export function collectionPlan(connection: ConnectionStatus, all: ConnectionStat }); return nodes.length > 0 ? { kind: 'cluster', members: nodes.filter(isScannable) } : { kind: 'single' }; } + +export interface ResolvedPlan { + connection: ConnectionStatus | undefined; + plan: CollectionPlan; +} + +export function resolvePlan(registry: Pick, connectionId: string): ResolvedPlan { + const connections = registry.list({ includeRetired: true }); + const connection = connections.find((candidate) => candidate.id === connectionId); + return { connection, plan: connection ? collectionPlan(connection, connections) : { kind: 'single' } }; +} + +export function planForConnection(registry: Pick, connectionId: string): CollectionPlan { + return resolvePlan(registry, connectionId).plan; +} diff --git a/proprietary/key-analytics/key-analytics.service.ts b/proprietary/key-analytics/key-analytics.service.ts index 3cc479b0e..c4606c1c9 100644 --- a/proprietary/key-analytics/key-analytics.service.ts +++ b/proprietary/key-analytics/key-analytics.service.ts @@ -12,7 +12,6 @@ import { import { ConnectionRegistry } from '@app/connections/connection-registry.service'; import { LicenseService } from '@proprietary/licenses/license.service'; import { - ConnectionStatus, KeyAnalyticsResult, KeyDetail, KeySizeDistribution, @@ -27,8 +26,9 @@ import { } from '@betterdb/shared'; import type { DatabasePort } from '@app/common/interfaces/database-port.interface'; import { rankCompositeKeys } from './composite-key-ranker'; -import { CollectionPlan, collectionPlan, isScannable } from './key-analytics-plan'; +import { CollectionPlan, collectionPlan, isScannable, planForConnection } from './key-analytics-plan'; import { buildPatternSnapshots, mergePatternSnapshots } from './key-pattern-snapshots'; +import { clusterScanNodes, primaryScanNodes } from './scan-nodes'; import { randomUUID } from 'crypto'; // The collectors prune keyDetails mid-scan to the top KEY_DETAILS_TOP_N keys per @@ -49,11 +49,6 @@ const COMPOSITE_TOP_N = KEY_DETAILS_TOP_N; const NO_KEY_SIZES: KeySizeDistribution = { databases: {}, available: false }; -interface ScanNode { - name: string; - client: DatabasePort; -} - function dedupeByKeyMaxMemory(entries: HotKeyEntry[]): HotKeyEntry[] { const byKey = new Map(); for (const entry of entries) { @@ -134,12 +129,6 @@ export class KeyAnalyticsService extends MultiConnectionPoller implements OnModu return this.collect(ctx, false); } - private planFor(connectionId: string): CollectionPlan { - const connections = this.connectionRegistry.list(); - const connection = connections.find((candidate) => candidate.id === connectionId); - return connection ? collectionPlan(connection, connections) : { kind: 'single' }; - } - private hidesAnalytics(connectionId?: string): boolean { if (!connectionId) { return false; @@ -148,36 +137,12 @@ export class KeyAnalyticsService extends MultiConnectionPoller implements OnModu return connection !== undefined && !hasOwnKeyAnalytics(connection.membership); } - private clusterNodes(seed: ScanNode, members: ConnectionStatus[]): ScanNode[] { - const nodes = [seed]; - for (const member of members) { - try { - nodes.push({ name: member.name, client: this.connectionRegistry.get(member.id) }); - } catch { - this.logger.warn(`Key analytics skipped ${member.name}: the connection is no longer registered`); - } - } - return nodes; - } - - private async primaries(nodes: ScanNode[]): Promise { - const roles = await Promise.allSettled(nodes.map((node) => node.client.getRole())); - return nodes.filter((node, index) => { - const role = roles[index]; - if (role.status === 'rejected') { - this.logger.warn(`Key analytics could not read the role of ${node.name}: ${role.reason instanceof Error ? role.reason.message : role.reason}`); - return false; - } - return role.value.role === 'master'; - }); - } - private async scan(ctx: ConnectionContext, plan: CollectionPlan, fullScan: boolean): Promise { const options = { sampleSize: this.sampleSize, scanBatchSize: this.scanBatchSize, fullScan }; if (plan.kind !== 'cluster') { return await ctx.client.collectKeyAnalytics(options); } - const nodes = await this.primaries(this.clusterNodes({ name: ctx.connectionName, client: ctx.client }, plan.members)); + const nodes = await primaryScanNodes(clusterScanNodes(this.connectionRegistry, { name: ctx.connectionName, client: ctx.client }, plan.members, this.logger, 'Key analytics'), this.logger, 'Key analytics'); if (nodes.length === 0) { throw new Error(`No reachable primary found for ${ctx.connectionName}`); } @@ -201,7 +166,7 @@ export class KeyAnalyticsService extends MultiConnectionPoller implements OnModu } private async collect(ctx: ConnectionContext, fullScan: boolean): Promise { - const plan = this.planFor(ctx.connectionId); + const plan = planForConnection(this.connectionRegistry, ctx.connectionId); if (plan.kind === 'skip') { return; } @@ -678,7 +643,7 @@ export class KeyAnalyticsService extends MultiConnectionPoller implements OnModu return this.keySizesOf(client); } - const nodes = await this.primaries(this.clusterNodes({ name: target?.name ?? targetId, client }, plan.members)); + const nodes = await primaryScanNodes(clusterScanNodes(this.connectionRegistry, { name: target?.name ?? targetId, client }, plan.members, this.logger, 'Key analytics'), this.logger, 'Key analytics'); return mergeKeySizeDistributions(await Promise.all(nodes.map((node) => this.keySizesOf(node.client)))); } diff --git a/proprietary/key-analytics/scan-nodes.ts b/proprietary/key-analytics/scan-nodes.ts new file mode 100644 index 000000000..ae851e02f --- /dev/null +++ b/proprietary/key-analytics/scan-nodes.ts @@ -0,0 +1,75 @@ +import type { Logger } from '@nestjs/common'; +import type { DatabasePort } from '@app/common/interfaces/database-port.interface'; +import type { ConnectionRegistry } from '@app/connections/connection-registry.service'; +import type { ConnectionStatus } from '@betterdb/shared'; + +export interface ScanNode { + name: string; + client: DatabasePort; + connectionId?: string; + nodeId?: string; +} + +export function clusterScanNodes( + registry: ConnectionRegistry, + seed: ScanNode, + members: ConnectionStatus[], + logger: Logger, + label: string, +): ScanNode[] { + const nodes = [seed]; + for (const member of members) { + try { + nodes.push({ + name: member.name, + client: registry.get(member.id), + connectionId: member.id, + nodeId: member.membership?.nodeId, + }); + } catch { + logger.warn(`${label} skipped ${member.name}: the connection is no longer registered`); + } + } + return nodes; +} + +const CLUSTER_NODE_ID = /^[0-9a-f]{40}$/; + +async function clusterNodeId(node: ScanNode): Promise { + try { + const reply = await node.client.call('CLUSTER', ['MYID']); + const id = Buffer.isBuffer(reply) ? reply.toString() : reply; + return typeof id === 'string' && CLUSTER_NODE_ID.test(id) ? id : undefined; + } catch { + return undefined; + } +} + +export async function primaryScanNodes(nodes: ScanNode[], logger: Logger, label: string): Promise { + const roles = await Promise.allSettled(nodes.map((node) => node.client.getRole())); + const masters = nodes.filter((node, index) => { + const role = roles[index]; + if (role.status === 'rejected') { + logger.warn(`${label} could not read the role of ${node.name}: ${role.reason instanceof Error ? role.reason.message : role.reason}`); + return false; + } + return role.value.role === 'master'; + }); + const ids = await Promise.all(masters.map((node) => clusterNodeId(node))); + const seen = new Set(); + const primaries: ScanNode[] = []; + masters.forEach((node, index) => { + const id = ids[index]; + if (id === undefined) { + primaries.push(node); + return; + } + if (seen.has(id)) { + logger.warn(`${label} skipped ${node.name}: it is the same cluster node as an earlier connection`); + return; + } + seen.add(id); + primaries.push({ ...node, nodeId: id }); + }); + return primaries; +} diff --git a/proprietary/kv-cache/__tests__/footprint-estimate.spec.ts b/proprietary/kv-cache/__tests__/footprint-estimate.spec.ts new file mode 100644 index 000000000..657deb9c1 --- /dev/null +++ b/proprietary/kv-cache/__tests__/footprint-estimate.spec.ts @@ -0,0 +1,360 @@ +import { buildSnapshot, evictedKeysDelta, layoutOf, type KeySample, type NodeObservation } from '../footprint-estimate'; +import { classifyKey } from '../key-classifier'; + +const M = 'Qwen/Qwen2.5-0.5B-Instruct'; +const key = (i: number, suffix = '', model = M, dtype = 'bfloat16') => `${model}@1@0@${(i + 1).toString(16)}@${dtype}${suffix}`; +const sample = (k: string, o: Partial = {}): KeySample => ({ + parsed: classifyKey(k)!, + memoryBytes: 3_670_112, + ttl: -1, + siblingExists: null, + ...o, +}); +const parse = (keys: string[]) => keys.map((k) => classifyKey(k)!); +const memory = { usedMemory: 100_000_000, maxmemory: 0, maxmemoryPolicy: 'noeviction', evictedKeys: 0 }; +const base = { connectionId: 'c', timestamp: 1, memory, evictedKeysDelta: null, evictedKeysIntervalMs: null, otherDbs: [], clientDetected: false }; + +describe('layoutOf', () => { + it('detects each layout', () => { + expect(layoutOf([])).toBeNull(); + expect(layoutOf([classifyKey(key(0))!])).toBe('single_key'); + expect(layoutOf([classifyKey(key(0, 'kv_bytes'))!, classifyKey(key(0, 'metadata'))!])).toBe('two_key'); + expect(layoutOf([classifyKey(key(0))!, classifyKey(key(1, 'kv_bytes'))!])).toBe('mixed'); + }); +}); + +describe('buildSnapshot', () => { + it('counts exactly when the scan completed', () => { + const keys = Array.from({ length: 10 }, (_, i) => key(i)); + const node: NodeObservation = { scannedKeys: 12, matchedKeys: parse(keys), scanComplete: true, dbSize: 12, samples: keys.map((k) => sample(k)) }; + const s = buildSnapshot({ ...base, nodes: [node] }); + expect(s).toMatchObject({ + detected: true, + layout: 'single_key', + scannedKeys: 12, + matchedKeys: 10, + sampledKeys: 10, + scanComplete: true, + chunksEst: 10, + bytesEst: 36_701_120, + lmcacheMemoryShare: 0.367, + noTtlRatio: 1, + orphanRatio: null, + }); + expect(s.perModel).toEqual([{ model: M, dtype: 'bfloat16', chunksEst: 10, bytesEst: 36_701_120 }]); + }); + + it('reports every node as scanned when no coverage is given', () => { + const node: NodeObservation = { scannedKeys: 1, matchedKeys: [], scanComplete: true, dbSize: 1, samples: [] }; + expect(buildSnapshot({ ...base, nodes: [node, node] })).toMatchObject({ primariesScanned: 2, primariesTotal: 2 }); + }); + + it('scales the footprint by the share of slots scanned but keeps the memory share of the scanned nodes', () => { + const keys = Array.from({ length: 10 }, (_, i) => key(i)); + const node: NodeObservation = { scannedKeys: 10, matchedKeys: parse(keys), scanComplete: true, dbSize: 10, samples: keys.map((k) => sample(k)) }; + const s = buildSnapshot({ ...base, nodes: [node], coverage: { primariesScanned: 1, primariesTotal: 4, slotShare: 0.25 } }); + expect(s).toMatchObject({ + primariesScanned: 1, + primariesTotal: 4, + chunksEst: 40, + bytesEst: 146_804_480, + lmcacheMemoryShare: 0.367, + }); + expect(s.perModel).toEqual([{ model: M, dtype: 'bfloat16', chunksEst: 40, bytesEst: 146_804_480 }]); + }); + + it('does not scale when the scanned slot share is unknown', () => { + const keys = Array.from({ length: 10 }, (_, i) => key(i)); + const node: NodeObservation = { scannedKeys: 10, matchedKeys: parse(keys), scanComplete: true, dbSize: 10, samples: [] }; + const s = buildSnapshot({ ...base, nodes: [node], coverage: { primariesScanned: 1, primariesTotal: null, slotShare: null } }); + expect(s).toMatchObject({ chunksEst: 10, primariesScanned: 1, primariesTotal: null }); + }); + + it('extrapolates from DBSIZE when the scan stopped early', () => { + const keys = Array.from({ length: 4 }, (_, i) => key(i)); + const node: NodeObservation = { scannedKeys: 100, matchedKeys: parse(keys), scanComplete: false, dbSize: 1000, samples: [sample(keys[0], { memoryBytes: 1000 })] }; + const s = buildSnapshot({ ...base, nodes: [node] }); + expect(s.chunksEst).toBe(40); + expect(s.bytesEst).toBe(40_000); + expect(s.scanComplete).toBe(false); + }); + + it('counts a two-key chunk once and measures orphans', () => { + const keys = [key(0, 'kv_bytes'), key(0, 'metadata'), key(1, 'kv_bytes')]; + const node: NodeObservation = { + scannedKeys: 3, + matchedKeys: parse(keys), + scanComplete: true, + dbSize: 3, + samples: [ + sample(keys[0], { siblingExists: true }), + sample(keys[1], { memoryBytes: 128, siblingExists: true }), + sample(keys[2], { siblingExists: false }), + ], + }; + const s = buildSnapshot({ ...base, nodes: [node] }); + expect(s.layout).toBe('two_key'); + expect(s.chunksEst).toBe(2); + expect(s.orphanRatio).toBe(0.3333); + expect(s.detected).toBe(false); + }); + + it('counts metadata whose kv_bytes was evicted as an orphan', () => { + const keys = [key(0, 'metadata'), key(1, 'metadata'), key(2, 'kv_bytes')]; + const node: NodeObservation = { + scannedKeys: 3, + matchedKeys: parse(keys), + scanComplete: true, + dbSize: 3, + samples: [ + sample(keys[0], { memoryBytes: 128, siblingExists: false }), + sample(keys[1], { memoryBytes: null, siblingExists: false }), + sample(keys[2], { siblingExists: true }), + ], + }; + expect(buildSnapshot({ ...base, nodes: [node] }).orphanRatio).toBe(0.6667); + }); + + it('measures orphans only over suffixed keys whose sibling probe answered', () => { + const keys = [key(0, 'kv_bytes'), key(1, 'kv_bytes'), key(2, 'kv_bytes'), key(3, 'kv_bytes')]; + const node: NodeObservation = { + scannedKeys: 4, + matchedKeys: parse(keys), + scanComplete: true, + dbSize: 4, + samples: [ + sample(keys[0], { siblingExists: true }), + sample(keys[1], { siblingExists: null }), + sample(keys[2], { siblingExists: null }), + sample(keys[3], { siblingExists: null }), + ], + }; + expect(buildSnapshot({ ...base, nodes: [node] }).orphanRatio).toBe(0); + }); + + it('reports no orphan ratio when no sibling probe answered', () => { + const keys = [key(0, 'kv_bytes'), key(1, 'kv_bytes')]; + const node: NodeObservation = { + scannedKeys: 2, + matchedKeys: parse(keys), + scanComplete: true, + dbSize: 2, + samples: keys.map((k) => sample(k, { siblingExists: null })), + }; + expect(buildSnapshot({ ...base, nodes: [node] }).orphanRatio).toBeNull(); + }); + + it('ignores vanished keys and measures TTL coverage', () => { + const keys = [key(0), key(1), key(2)]; + const node: NodeObservation = { + scannedKeys: 3, + matchedKeys: parse(keys), + scanComplete: true, + dbSize: 3, + samples: [sample(keys[0]), sample(keys[1], { ttl: 3600 }), sample(keys[2], { ttl: -2, memoryBytes: null })], + }; + const s = buildSnapshot({ ...base, nodes: [node] }); + expect(s.sampledKeys).toBe(2); + expect(s.noTtlRatio).toBe(0.5); + }); + + it('keeps samples whose MEMORY USAGE failed for TTL coverage', () => { + const keys = [key(0), key(1), key(2)]; + const node: NodeObservation = { + scannedKeys: 3, + matchedKeys: parse(keys), + scanComplete: true, + dbSize: 3, + samples: [ + sample(keys[0], { memoryBytes: null }), + sample(keys[1], { memoryBytes: null, ttl: 3600 }), + sample(keys[2], { memoryBytes: null, ttl: -2 }), + ], + }; + const s = buildSnapshot({ ...base, nodes: [node] }); + expect(s.noTtlRatio).toBe(0.5); + expect(s.bytesEst).toBeNull(); + expect(s.lmcacheMemoryShare).toBeNull(); + expect(s.chunksEst).toBe(3); + }); + + it('sizes a node without readable samples from the other nodes', () => { + const n1: NodeObservation = { scannedKeys: 2, matchedKeys: parse([key(0), key(1)]), scanComplete: true, dbSize: 2, samples: [sample(key(0), { memoryBytes: 10 }), sample(key(1), { memoryBytes: 30 })] }; + const n2: NodeObservation = { scannedKeys: 2, matchedKeys: parse([key(2), key(3)]), scanComplete: true, dbSize: 2, samples: [sample(key(2), { memoryBytes: null }), sample(key(3), { memoryBytes: null })] }; + expect(buildSnapshot({ ...base, nodes: [n1, n2] }).bytesEst).toBe(80); + }); + + it('keeps the bytes of samples whose TTL is unknown but leaves them out of TTL coverage', () => { + const keys = [key(0), key(1)]; + const node: NodeObservation = { + scannedKeys: 2, + matchedKeys: parse(keys), + scanComplete: true, + dbSize: 2, + samples: [sample(keys[0], { memoryBytes: 100, ttl: null }), sample(keys[1], { memoryBytes: 300, ttl: 3600 })], + }; + const s = buildSnapshot({ ...base, nodes: [node] }); + expect(s.sampledKeys).toBe(2); + expect(s.bytesEst).toBe(400); + expect(s.noTtlRatio).toBe(0); + }); + + it('reports no TTL ratio when no sampled key answered TTL', () => { + const keys = [key(0)]; + const empty: NodeObservation = { scannedKeys: 0, matchedKeys: [], scanComplete: true, dbSize: 0, samples: [] }; + const vanished: NodeObservation = { + scannedKeys: 1, + matchedKeys: parse(keys), + scanComplete: true, + dbSize: 1, + samples: [sample(keys[0], { ttl: -2, memoryBytes: null })], + }; + expect(buildSnapshot({ ...base, nodes: [empty] }).noTtlRatio).toBeNull(); + expect(buildSnapshot({ ...base, nodes: [vanished] }).noTtlRatio).toBeNull(); + }); + + it('rounds the no-TTL ratio to four decimals', () => { + const keys = Array.from({ length: 12 }, (_, i) => key(i)); + const node: NodeObservation = { + scannedKeys: 12, + matchedKeys: parse(keys), + scanComplete: true, + dbSize: 12, + samples: keys.map((k, i) => sample(k, { ttl: i === 0 ? 3600 : -1 })), + }; + expect(buildSnapshot({ ...base, nodes: [node] }).noTtlRatio).toBe(0.9167); + }); + + it('splits per model by sample share', () => { + const a = [key(0), key(1), key(2)]; + const b = [key(3, '', 'other', 'half')]; + const node: NodeObservation = { scannedKeys: 4, matchedKeys: parse([...a, ...b]), scanComplete: true, dbSize: 4, samples: [...a, ...b].map((k) => sample(k, { memoryBytes: 100 })) }; + const s = buildSnapshot({ ...base, nodes: [node] }); + expect(s.perModel).toEqual([ + { model: M, dtype: 'bfloat16', chunksEst: 3, bytesEst: 300 }, + { model: 'other', dtype: 'half', chunksEst: 1, bytesEst: 100 }, + ]); + }); + + it('scales each model by the sampling rate of the node it lives on', () => { + const full: NodeObservation = { + scannedKeys: 2, + matchedKeys: parse([key(0), key(1)]), + scanComplete: true, + dbSize: 2, + samples: [key(0), key(1)].map((k) => sample(k, { memoryBytes: 10 })), + }; + const other = [key(2, '', 'other'), key(3, '', 'other')]; + const sampled: NodeObservation = { + scannedKeys: 2, + matchedKeys: parse(other), + scanComplete: false, + dbSize: 1000, + samples: other.map((k) => sample(k, { memoryBytes: 10 })), + }; + const s = buildSnapshot({ ...base, nodes: [full, sampled] }); + expect(s.chunksEst).toBe(1002); + expect(s.perModel).toEqual([ + { model: 'other', dtype: 'bfloat16', chunksEst: 1000, bytesEst: 10_000 }, + { model: M, dtype: 'bfloat16', chunksEst: 2, bytesEst: 20 }, + ]); + }); + + it('splits chunks per model from the matched keys when MEMORY USAGE is denied', () => { + const a = [key(0), key(1), key(2)]; + const b = [key(3, '', 'other', 'half')]; + const node: NodeObservation = { + scannedKeys: 4, + matchedKeys: parse([...a, ...b]), + scanComplete: true, + dbSize: 4, + samples: [...a, ...b].map((k) => sample(k, { memoryBytes: null })), + }; + expect(buildSnapshot({ ...base, nodes: [node] }).perModel).toEqual([ + { model: M, dtype: 'bfloat16', chunksEst: 3, bytesEst: null }, + { model: 'other', dtype: 'half', chunksEst: 1, bytesEst: null }, + ]); + }); + + it('counts a layerwise chunk once rather than once per layer', () => { + const layers = 32; + const layered = Array.from({ length: 4 }, (_, chunk) => + Array.from({ length: layers }, (_, layer) => `${key(chunk)}@${layer}`), + ).flat(); + const node: NodeObservation = { + scannedKeys: layered.length, + matchedKeys: parse(layered), + scanComplete: true, + dbSize: layered.length, + samples: layered.map((k) => sample(k, { memoryBytes: 10 })), + }; + const s = buildSnapshot({ ...base, nodes: [node] }); + expect(s.chunksEst).toBe(4); + expect(s.bytesEst).toBe(layered.length * 10); + expect(s.perModel).toEqual([{ model: M, dtype: 'bfloat16', chunksEst: 4, bytesEst: layered.length * 10 }]); + }); + + it('weighs layered keys per model when models mix layerwise and whole-chunk storage', () => { + const layered = [0, 1].flatMap((chunk) => [0, 1, 2, 3].map((layer) => `${key(chunk)}@${layer}`)); + const whole = [key(10, '', 'other'), key(11, '', 'other')]; + const all = [...layered, ...whole]; + const node: NodeObservation = { + scannedKeys: all.length, + matchedKeys: parse(all), + scanComplete: true, + dbSize: all.length, + samples: all.map((k) => sample(k, { memoryBytes: 10 })), + }; + const s = buildSnapshot({ ...base, nodes: [node] }); + expect(s.chunksEst).toBe(4); + expect(s.perModel.map(({ model, chunksEst }) => ({ model, chunksEst }))).toEqual([ + { model: M, chunksEst: 2 }, + { model: 'other', chunksEst: 2 }, + ]); + }); + + it('reports multiprocess adapter keys per model with an unknown dtype', () => { + const hash = (i: number) => (i + 1).toString(16).padStart(64, '0'); + const keys = Array.from({ length: 6 }, (_, i) => `${M}@0000000${i % 2}@0@${hash(i)}`); + const node: NodeObservation = { scannedKeys: 6, matchedKeys: parse(keys), scanComplete: true, dbSize: 6, samples: keys.map((k) => sample(k)) }; + const s = buildSnapshot({ ...base, nodes: [node] }); + expect(s).toMatchObject({ detected: true, layout: 'single_key', chunksEst: 6, orphanRatio: null }); + expect(s.perModel).toEqual([{ model: M, dtype: 'unknown', chunksEst: 6, bytesEst: 6 * 3_670_112 }]); + }); + + 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); + }); + + it('passes the eviction interval through only alongside a delta', () => { + const node: NodeObservation = { scannedKeys: 0, matchedKeys: [], scanComplete: true, dbSize: 0, samples: [] }; + expect(buildSnapshot({ ...base, nodes: [node], evictedKeysDelta: 7, evictedKeysIntervalMs: 60_000 }).evictedKeysIntervalMs).toBe(60_000); + expect(buildSnapshot({ ...base, nodes: [node], evictedKeysIntervalMs: 60_000 }).evictedKeysIntervalMs).toBeNull(); + }); + + it('detects through the GLIDE client even with no keys', () => { + const node: NodeObservation = { scannedKeys: 0, matchedKeys: [], scanComplete: true, dbSize: 0, samples: [] }; + const s = buildSnapshot({ ...base, nodes: [node], clientDetected: true }); + expect(s.detected).toBe(true); + expect(s.layout).toBeNull(); + expect(s.bytesEst).toBe(0); + }); + + it('sums cluster nodes', () => { + const n1: NodeObservation = { scannedKeys: 5, matchedKeys: parse([key(0), key(1)]), scanComplete: true, dbSize: 5, samples: [sample(key(0), { memoryBytes: 10 })] }; + const n2: NodeObservation = { scannedKeys: 5, matchedKeys: parse([key(2)]), scanComplete: true, dbSize: 5, samples: [sample(key(2), { memoryBytes: 30 })] }; + const s = buildSnapshot({ ...base, nodes: [n1, n2] }); + expect(s).toMatchObject({ scannedKeys: 10, matchedKeys: 3, chunksEst: 3, bytesEst: 50 }); + }); +}); diff --git a/proprietary/kv-cache/__tests__/footprint-scanner.spec.ts b/proprietary/kv-cache/__tests__/footprint-scanner.spec.ts new file mode 100644 index 000000000..4852de0db --- /dev/null +++ b/proprietary/kv-cache/__tests__/footprint-scanner.spec.ts @@ -0,0 +1,160 @@ +import { classifyKey } from '../key-classifier'; +import { hasLmcacheGlideClient, readMemoryStats, sampleLmcacheKeys, scanLmcacheKeys } from '../footprint-scanner'; +import { keySlot } from '../key-slot'; + +const M = 'm/x'; +const k = (i: number, s = '') => `${M}@1@0@${(i + 1).toString(16)}@bfloat16${s}`; +const p = (...keys: string[]) => keys.map((key) => classifyKey(key)!); + +function scanClient(pages: Array<[string, string[]]>) { + const call = jest.fn(); + for (const page of pages) call.mockResolvedValueOnce(page); + return { call } as any; +} + +describe('scanLmcacheKeys', () => { + it('scans to cursor 0 and keeps only LMCache keys', async () => { + const client = scanClient([['7', [k(0), 'a@b@c@d@e', 'user:1']], ['0', [k(1)]]]); + const result = await scanLmcacheKeys(client, { maxScanned: 200_000, maxMatched: 2000 }, 1500); + expect(client.call).toHaveBeenCalledWith('SCAN', ['0', 'COUNT', '1000']); + expect(result).toEqual({ scannedKeys: 1500, matchedKeys: p(k(0), k(1)), scanComplete: true }); + }); + + it('counts a key SCAN returns more than once only once', async () => { + const client = scanClient([['7', [k(0), k(1)]], ['0', [k(1), k(0), k(2)]]]); + const result = await scanLmcacheKeys(client, { maxScanned: 200_000, maxMatched: 2000 }, 3); + expect(result.matchedKeys.map((key) => key.key)).toEqual([k(0), k(1), k(2)]); + }); + + it('leaves tag/suffix ties to the caller, which sees every node of the run', async () => { + const ambiguous = `${M}@1@0@1@bfloat16@kind%metadata`; + const client = scanClient([['0', [ambiguous, `${M}@1@0@2@bfloat16@kind%decode`]]]); + const result = await scanLmcacheKeys(client, { maxScanned: 200_000, maxMatched: 2000 }, 2); + expect(result.matchedKeys.map((key) => key.suffix)).toEqual(['metadata', null]); + }); + + it('stops at the matched budget', async () => { + const client = scanClient([['7', [k(0), k(1), k(2)]]]); + const result = await scanLmcacheKeys(client, { maxScanned: 200_000, maxMatched: 2 }, 10_000); + expect(result).toEqual({ scannedKeys: 2, matchedKeys: p(k(0), k(1)), scanComplete: false }); + }); + + it('shrinks the scanned count in proportion when a completed scan is truncated', async () => { + const client = scanClient([['0', [k(0), k(1), k(2), k(3)]]]); + const result = await scanLmcacheKeys(client, { maxScanned: 200_000, maxMatched: 2 }, 1500); + expect(result).toEqual({ scannedKeys: 750, matchedKeys: p(k(0), k(1)), scanComplete: false }); + }); + + it('measures coverage from the keys SCAN returned when it stops early', async () => { + const filler = (n: number) => Array.from({ length: n }, (_, i) => `user:${i}`); + const client = scanClient([['7', [k(0), ...filler(1499)]], ['8', [k(1), ...filler(1499)]], ['9', [k(2)]]]); + const result = await scanLmcacheKeys(client, { maxScanned: 2500, maxMatched: 2000 }, 100_000); + expect(client.call).toHaveBeenCalledTimes(2); + expect(result).toEqual({ scannedKeys: 3000, matchedKeys: p(k(0), k(1)), scanComplete: false }); + }); + + it('does not inflate coverage when SCAN returns fewer keys than COUNT', async () => { + const client = scanClient([['7', [k(0)]], ['8', [k(1), 'user:1']]]); + const result = await scanLmcacheKeys(client, { maxScanned: 2000, maxMatched: 2000 }, 10_000); + expect(result).toEqual({ scannedKeys: 3, matchedKeys: p(k(0), k(1)), scanComplete: false }); + }); + + it('stops at the scanned budget', async () => { + const client = scanClient([['7', []], ['8', []], ['9', []]]); + const result = await scanLmcacheKeys(client, { maxScanned: 2000, maxMatched: 2000 }, 10_000); + expect(client.call).toHaveBeenCalledTimes(2); + expect(result.scanComplete).toBe(false); + }); +}); + +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; + 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 }, + { memoryBytes: null, ttl: -2, siblingExists: null }, + ]); + }); + + 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]); + }); + + 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 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')); + 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 [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', () => { + it('detects a GLIDE LMCache client and tolerates CLIENT LIST failures', async () => { + expect(await hasLmcacheGlideClient({ call: jest.fn().mockResolvedValue('id=1 db=0 lib-name=GlidePySync(lmcache:0.5.5) lib-ver=2.5.3\n') } as any, 0)).toBe(true); + expect(await hasLmcacheGlideClient({ call: jest.fn().mockResolvedValue('id=1 db=0 lib-name=redis-py\n') } as any, 0)).toBe(false); + expect(await hasLmcacheGlideClient({ call: jest.fn().mockRejectedValue(new Error('NOPERM')) } as any, 0)).toBe(false); + }); + + it('detects the bare lmcache tag LMCache falls back to', async () => { + const list = (libName: string) => ({ call: jest.fn().mockResolvedValue(`id=1 db=0 lib-name=${libName} lib-ver=2.5.3\n`) }) as any; + expect(await hasLmcacheGlideClient(list('GlidePySync(lmcache)'), 0)).toBe(true); + expect(await hasLmcacheGlideClient(list('GlidePySync(lmcache:0.5.5)'), 0)).toBe(true); + expect(await hasLmcacheGlideClient(list('GlidePySync(lmcachex)'), 0)).toBe(false); + expect(await hasLmcacheGlideClient(list('GlidePySync(other)'), 0)).toBe(false); + }); + + it('only counts LMCache clients on the scanned database', async () => { + const reply = [ + 'id=1 addr=1.1.1.1:1 db=3 lib-name=GlidePySync(lmcache:0.5.5) lib-ver=2.5.3', + 'id=2 addr=1.1.1.1:2 db=0 lib-name=redis-py lib-ver=5', + '', + ].join('\n'); + const client = { call: jest.fn().mockResolvedValue(reply) } as any; + expect(await hasLmcacheGlideClient(client, 0)).toBe(false); + expect(await hasLmcacheGlideClient(client, 3)).toBe(true); + }); + + it('reads memory stats', () => { + const info = { + memory: { used_memory: '100', maxmemory: '200', maxmemory_policy: 'volatile-lru' }, + stats: { evicted_keys: '7' }, + } as any; + expect(readMemoryStats(info)).toEqual({ usedMemory: 100, maxmemory: 200, maxmemoryPolicy: 'volatile-lru', evictedKeys: 7 }); + }); +}); diff --git a/proprietary/kv-cache/__tests__/key-classifier.spec.ts b/proprietary/kv-cache/__tests__/key-classifier.spec.ts new file mode 100644 index 000000000..8b21f044c --- /dev/null +++ b/proprietary/kv-cache/__tests__/key-classifier.spec.ts @@ -0,0 +1,232 @@ +import { classifyKey, isChunkKey, maskKey, resolveTagSuffixes, siblingKey } from '../key-classifier'; + +const MODEL = 'Qwen/Qwen2.5-0.5B-Instruct'; +const HASH = '9e3779b97f4a7c15'; + +describe('classifyKey', () => { + it('parses a single-key layout key from the spike', () => { + expect(classifyKey(`${MODEL}@1@0@${HASH}@bfloat16`)).toEqual({ + key: `${MODEL}@1@0@${HASH}@bfloat16`, + model: MODEL, + worldSize: 1, + workerId: 0, + chunkHash: HASH, + dtype: 'bfloat16', + layer: null, + tagged: false, + suffix: null, + format: 'classic', + }); + }); + + it('parses redis:// two-key suffixes without a separator', () => { + expect(classifyKey(`${MODEL}@1@0@${HASH}@bfloat16kv_bytes`)?.suffix).toBe('kv_bytes'); + expect(classifyKey(`${MODEL}@1@0@${HASH}@bfloat16metadata`)?.suffix).toBe('metadata'); + }); + + it('parses layerwise and tagged keys', () => { + const parsed = classifyKey(`${MODEL}@2@1@${HASH}@fp8_e4m3fn@12@k%v`); + expect(parsed).toMatchObject({ worldSize: 2, workerId: 1, dtype: 'fp8_e4m3fn', layer: 12, tagged: true }); + }); + + it.each([ + ['half'], + ['float'], + ['double'], + ['bfloat16'], + ['int8'], + ['uint8'], + ['int16'], + ['int32'], + ['int64'], + ['bool'], + ['fp8_e4m3fn'], + ['fp8_e4m3fnuz'], + ['fp8_e5m2'], + ['fp8_e5m2fnuz'], + ])('accepts the LMCache dtype string %s', (dtype) => { + expect(classifyKey(`${MODEL}@1@0@${HASH}@${dtype}`)?.dtype).toBe(dtype); + }); + + it.each([['float16'], ['float32'], ['float64'], ['float8_e4m3fn'], ['float8_e5m2']])( + 'rejects the torch attribute name %s', + (dtype) => { + expect(classifyKey(`${MODEL}@1@0@${HASH}@${dtype}`)).toBeNull(); + }, + ); + + it.each([ + ['@user%alice', null], + ['@user%alice@tenant%acme', null], + ['@user%alice@tenant%acme@session%a%b', null], + ['@user%', null], + ['@user%alice', 'kv_bytes'], + ['@user%alice', 'metadata'], + ['@user%alice@tenant%acme', 'kv_bytes'], + ['@user%alice@tenant%acme', 'metadata'], + ])('parses arbitrary tags %s with suffix %s', (tags, suffix) => { + const key = `${MODEL}@1@0@${HASH}@half${tags}${suffix ?? ''}`; + expect(classifyKey(key)).toMatchObject({ dtype: 'half', layer: null, tagged: true, suffix }); + }); + + it.each([ + ['@user%alice', null], + ['@user%alice@tenant%acme', null], + ['@user%alice', 'kv_bytes'], + ['@user%alice@tenant%acme', 'metadata'], + ['', 'kv_bytes'], + ])('parses layered keys with tags %s and suffix %s', (tags, suffix) => { + const key = `${MODEL}@2@1@${HASH}@bfloat16@7${tags}${suffix ?? ''}`; + expect(classifyKey(key)).toMatchObject({ layer: 7, tagged: tags !== '', suffix }); + }); + + it('pairs tagged two-key siblings', () => { + const key = `${MODEL}@1@0@${HASH}@half@user%alice@tenant%acmekv_bytes`; + expect(siblingKey(classifyKey(key)!)).toBe(`${MODEL}@1@0@${HASH}@half@user%alice@tenant%acmemetadata`); + }); + + it('accepts model names containing @', () => { + expect(classifyKey(`org/model@v2@1@0@${HASH}@half@user%alice`)?.model).toBe('org/model@v2'); + }); + + it('accepts a negative hash', () => { + expect(classifyKey(`${MODEL}@1@0@-1a2b@int8`)?.chunkHash).toBe('-1a2b'); + }); + + it.each([ + [`${MODEL}@1@0@XYZ@bfloat16`], + [`${MODEL}@1@0@${HASH}@float64`], + [`${MODEL}@one@0@${HASH}@bfloat16`], + [`user:1@a@b@c@d`], + [`@1@0@${HASH}@bfloat16`], + [`${MODEL}@1@0@${HASH}@bfloat16extra`], + [`${MODEL}@1@0@${HASH}@half@user`], + [`${MODEL}@1@0@${HASH}@half@%alice`], + [`${MODEL}@1@0@${HASH}@half@user%alice@`], + ])('rejects %s', (key) => { + expect(classifyKey(key)).toBeNull(); + }); +}); + +describe('classifyKey with the multiprocess Valkey L2 adapter format', () => { + const MP_HASH = 'a3f1c2d4e5b60718293a4b5c6d7e8f90a1b2c3d4e5f60718293a4b5c6d7e8f90'; + const mp = `${MODEL}@00000003@0@${MP_HASH}`; + + it('parses an unsalted chunk object key', () => { + expect(classifyKey(mp)).toEqual({ + key: mp, + model: MODEL, + worldSize: null, + workerId: 3, + chunkHash: MP_HASH, + dtype: 'unknown', + layer: null, + tagged: false, + suffix: null, + format: 'mp', + }); + }); + + it('parses a salted key and a key_prefix', () => { + expect(classifyKey(`${mp}@tenant-a`)).toMatchObject({ model: MODEL, chunkHash: MP_HASH, format: 'mp' }); + expect(classifyKey(`lmc@${MODEL}@0000000a@1f@${MP_HASH}`)).toMatchObject({ + model: MODEL, + workerId: 10, + chunkHash: MP_HASH, + }); + expect(classifyKey(`lmc@${MODEL}@00000000@0@${MP_HASH}@salt`)).toMatchObject({ model: MODEL, format: 'mp' }); + }); + + it('counts each object key as a single-key chunk with no sibling', () => { + const parsed = classifyKey(mp)!; + expect(isChunkKey(parsed)).toBe(true); + expect(siblingKey(parsed)).toBeNull(); + }); + + it('masks the hash', () => { + expect(maskKey(classifyKey(`lmc@${mp}@salt`)!)).toBe(`${MODEL}@00000003@0@a3f1c2d4…`); + }); + + it.each([ + [`${MODEL}@3@0@${MP_HASH}`], + [`${MODEL}@0000000g@0@${MP_HASH}`], + [`${MODEL}@00000003@0@abc`], + [`${MODEL}@00000003@0@${MP_HASH}0`], + [`${MODEL}@00000003@@${MP_HASH}`], + [`${MODEL}@00000003@0@${MP_HASH}@`], + [`${MODEL}@00000003@0@${MP_HASH}@a/b`], + [`@00000003@0@${MP_HASH}`], + [`user@00000001@1@ab`], + ])('rejects %s', (key) => { + expect(classifyKey(key)).toBeNull(); + }); + + it('keeps classifying classic keys as classic', () => { + expect(classifyKey(`${MODEL}@1@0@${HASH}@bfloat16`)?.format).toBe('classic'); + expect(classifyKey(`${MODEL}@1@0@${MP_HASH}@half`)?.format).toBe('classic'); + expect(classifyKey(`${MODEL}@1@0@${HASH}@half@user%alice`)?.format).toBe('classic'); + }); +}); + +describe('helpers', () => { + it('finds the sibling of a two-key chunk', () => { + const key = `${MODEL}@1@0@${HASH}@bfloat16kv_bytes`; + expect(siblingKey(classifyKey(key)!)).toBe(`${MODEL}@1@0@${HASH}@bfloat16metadata`); + const single = `${MODEL}@1@0@${HASH}@bfloat16`; + expect(siblingKey(classifyKey(single)!)).toBeNull(); + }); + + it('counts kv_bytes and single keys as chunks but not metadata', () => { + expect(isChunkKey(classifyKey(`${MODEL}@1@0@${HASH}@bfloat16`)!)).toBe(true); + expect(isChunkKey(classifyKey(`${MODEL}@1@0@${HASH}@bfloat16kv_bytes`)!)).toBe(true); + expect(isChunkKey(classifyKey(`${MODEL}@1@0@${HASH}@bfloat16metadata`)!)).toBe(false); + }); + + it('masks the hash', () => { + expect(maskKey(classifyKey(`${MODEL}@1@0@${HASH}@bfloat16`)!)).toBe(`${MODEL}@1@0@9e3779b9…@bfloat16`); + }); +}); + +describe('resolveTagSuffixes', () => { + const tagged = (i: number, tail: string) => classifyKey(`${MODEL}@1@0@${i}@half@kind%${tail}`)!; + + it('reads a tag value ending in a suffix name as the tag when the node stores single keys', () => { + const [ambiguous, plain] = resolveTagSuffixes([tagged(1, 'metadata'), tagged(2, 'prefill')]); + expect(ambiguous).toMatchObject({ suffix: null, tagged: true, key: `${MODEL}@1@0@1@half@kind%metadata` }); + expect(plain.suffix).toBeNull(); + }); + + it('keeps the suffix when the node stores two keys per chunk', () => { + const untagged = classifyKey(`${MODEL}@1@0@3@halfkv_bytes`)!; + const keys = [tagged(1, 'metadata'), tagged(2, 'prefill'), untagged]; + expect(resolveTagSuffixes(keys)[0].suffix).toBe('metadata'); + }); + + it('keeps tagged two-key pairs on a mixed node that also stores single keys', () => { + const keys = [tagged(3, 'prefillkv_bytes'), tagged(3, 'prefillmetadata'), tagged(4, 'decode')]; + expect(resolveTagSuffixes(keys).map((k) => k.suffix)).toEqual(['kv_bytes', 'metadata', null]); + }); + + it('keeps an unpaired tagged suffix when a tagged pair shows the node stores two keys', () => { + const keys = [tagged(1, 'metadata'), tagged(3, 'prefillkv_bytes'), tagged(3, 'prefillmetadata'), tagged(4, 'decode')]; + expect(resolveTagSuffixes(keys)[0].suffix).toBe('metadata'); + }); + + it('reads tag values as tags when single-key evidence comes from another node', () => { + expect(resolveTagSuffixes([tagged(1, 'metadata'), tagged(2, 'kv_bytes')], false, true).map((k) => k.suffix)).toEqual([ + null, + null, + ]); + }); + + it('keeps the suffix when the run has no single-key evidence even if the node does', () => { + expect(resolveTagSuffixes([tagged(1, 'metadata'), tagged(2, 'prefill')], false, false)[0].suffix).toBe('metadata'); + }); + + it('keeps the suffix when nothing on the node breaks the tie', () => { + expect(resolveTagSuffixes([tagged(1, 'metadata'), tagged(2, 'kv_bytes')]).map((k) => k.suffix)).toEqual([ + 'metadata', + 'kv_bytes', + ]); + }); +}); diff --git a/proprietary/kv-cache/__tests__/key-slot.spec.ts b/proprietary/kv-cache/__tests__/key-slot.spec.ts new file mode 100644 index 000000000..f5ddc49bf --- /dev/null +++ b/proprietary/kv-cache/__tests__/key-slot.spec.ts @@ -0,0 +1,26 @@ +import { keySlot, slotOwner } from '../key-slot'; + +describe('keySlot', () => { + it.each([ + ['foo', 12182], + ['bar', 5061], + ['123456789', 12739], + ['', 0], + ])('hashes %p to slot %p', (key, slot) => { + expect(keySlot(key)).toBe(slot); + }); + + it('hashes only the hash tag when one is present', () => { + expect(keySlot('{user1000}.following')).toBe(keySlot('user1000')); + expect(keySlot('foo{}{bar}')).not.toBe(keySlot('bar')); + expect(keySlot('foo{{bar}}zap')).toBe(keySlot('{bar')); + }); +}); + +describe('slotOwner', () => { + it('owns the slots inside its ranges, including single-slot ranges', () => { + const owns = slotOwner([[0, 99], [500]]); + expect([0, 99, 500].map(owns)).toEqual([true, true, true]); + expect([100, 499, 501].map(owns)).toEqual([false, false, false]); + }); +}); diff --git a/proprietary/kv-cache/__tests__/kv-cache-footprint.integration.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-footprint.integration.spec.ts new file mode 100644 index 000000000..57e96c1f4 --- /dev/null +++ b/proprietary/kv-cache/__tests__/kv-cache-footprint.integration.spec.ts @@ -0,0 +1,92 @@ +import { Feature } from '@betterdb/shared'; +import { UnifiedDatabaseAdapter } from '@app/database/adapters/unified.adapter'; +import { KvCacheFootprintService } from '../kv-cache-footprint.service'; + +const PORT_SETTING = process.env.KV_CACHE_IT_VALKEY_PORT?.trim(); +const HOST = process.env.KV_CACHE_IT_VALKEY_HOST || 'localhost'; +const PORT = Number(PORT_SETTING); +const PASSWORD = process.env.KV_CACHE_IT_VALKEY_PASSWORD ?? 'devpassword'; +const TEST_DB = 9; +const CONNECTION_ID = 'kv-cache-integration'; +const configured = Boolean(PORT_SETTING); +const suite = configured ? describe : describe.skip; +const title = configured + ? `KvCacheFootprintService against a real Valkey on ${HOST}:${PORT_SETTING}` + : 'KvCacheFootprintService against a real Valkey (skipped: set KV_CACHE_IT_VALKEY_PORT to a disposable Valkey)'; + +const chunk = (model: string, i: number, suffix: string) => `${model}@1@0@${(i + 1).toString(16)}@${suffix}`; + +suite(title, () => { + let selected = false; + let adapter: UnifiedDatabaseAdapter | undefined; + let service: KvCacheFootprintService; + let saved: unknown[] = []; + + beforeAll(async () => { + expect(Number.isInteger(PORT) && PORT > 0).toBe(true); + const target = new UnifiedDatabaseAdapter({ + host: HOST, + port: PORT, + username: 'default', + password: PASSWORD, + connectionId: CONNECTION_ID, + dbIndex: TEST_DB, + }); + adapter = target; + await target.connect(); + expect(target.getActiveDbIndex()).toBe(TEST_DB); + selected = true; + const client = target.getClient(); + await client.flushdb(); + + const kv = Buffer.alloc(4096, 1); + const meta = Buffer.alloc(128, 2); + const pipeline = client.pipeline(); + for (let i = 0; i < 6; i++) pipeline.set(chunk('m/x', i, 'bfloat16'), kv); + for (let i = 0; i < 3; i++) { + pipeline.set(chunk('m/y', i, 'halfkv_bytes'), kv); + if (i < 2) pipeline.set(chunk('m/y', i, 'halfmetadata'), meta); + } + pipeline.set(chunk('m/x', 100, 'bfloat16'), kv, 'EX', 3600); + for (const key of ['user:1', 'a@b', 'x@y@z@w@v']) pipeline.set(key, 'v'); + for (let i = 0; i < 17; i++) pipeline.set(`user:other:${i}`, 'v'); + await pipeline.exec(); + + const registry = { + list: () => [{ id: CONNECTION_ID, name: 'integration', isConnected: true, connectionType: 'direct' }], + get: () => target, + getConfig: () => ({ id: CONNECTION_ID, name: 'integration', host: HOST, port: PORT, dbIndex: TEST_DB }), + } as any; + const storage = { + saveKvCacheFootprintSnapshot: async (snapshot: unknown) => { + saved.push(snapshot); + }, + getKvCacheFootprintSnapshots: async () => [], + deleteKvCacheConnectionData: async () => undefined, + } as any; + const license = { hasFeature: (f: string) => f === Feature.KV_CACHE_MONITORING } as any; + service = new KvCacheFootprintService(registry, storage, license); + }); + + afterAll(async () => { + if (selected) await adapter?.getClient().flushdb(); + await adapter?.disconnect(); + }); + + it('measures the seeded LMCache keys', async () => { + const snapshot = await service.triggerCollection(CONNECTION_ID); + expect(snapshot).toMatchObject({ + connectionId: CONNECTION_ID, + detected: true, + layout: 'mixed', + matchedKeys: 12, + scanComplete: true, + chunksEst: 10, + orphanRatio: 0.2, + noTtlRatio: 0.9167, + }); + expect(snapshot?.perModel.map((m) => m.model)).toEqual(expect.arrayContaining(['m/x', 'm/y'])); + expect(snapshot?.otherDbs).not.toContain(TEST_DB); + expect(saved).toHaveLength(1); + }); +}); diff --git a/proprietary/kv-cache/__tests__/kv-cache-footprint.service.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-footprint.service.spec.ts new file mode 100644 index 000000000..01a064625 --- /dev/null +++ b/proprietary/kv-cache/__tests__/kv-cache-footprint.service.spec.ts @@ -0,0 +1,1111 @@ +import { BadRequestException, Logger, NotFoundException } from '@nestjs/common'; +import { Feature } from '@betterdb/shared'; +import * as classifier from '../key-classifier'; +import * as scanner from '../footprint-scanner'; +import { + KvCacheFootprintService, + MAX_FOOTPRINT_INTERVAL_MS, + PARTIAL_MISSES_BEFORE_CLEAR, + REFRESH_REUSE_MS, + resolvePositiveInt, +} from '../kv-cache-footprint.service'; + +const KEY = (i: number) => `m/x@1@0@${(i + 1).toString(16)}@bfloat16`; + +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 }), + getInfoParsed: jest.fn().mockResolvedValue({ + memory: { used_memory: '1000', maxmemory: '0', maxmemory_policy: 'noeviction' }, + stats: { evicted_keys: '3' }, + keyspace: { db0: { keys: keys.length, expires: 0, avg_ttl: 0 } }, + }), + getRole: jest.fn().mockResolvedValue({ role: 'master' }), + } as any; +} + +function setup(keys = Array.from({ length: 6 }, (_, i) => KEY(i)), licensed = true) { + const client = fakeClient(keys); + const removalListeners: Array<(id: string) => void> = []; + const registry = { + list: jest.fn().mockReturnValue([{ id: 'c1', name: 'one', isConnected: true, connectionType: 'direct' }]), + get: jest.fn().mockReturnValue(client), + getConfig: jest.fn().mockReturnValue({ id: 'c1', name: 'one', host: 'h', port: 6379, dbIndex: 0 }), + addConnectionRemovalListener: jest.fn((listener: (id: string) => void) => { + removalListeners.push(listener); + return () => undefined; + }), + removed: (id: string) => removalListeners.forEach((listener) => listener(id)), + } as any; + const storage = { + saveKvCacheFootprintSnapshot: jest.fn(), + getKvCacheFootprintSnapshots: jest.fn().mockResolvedValue([]), + deleteKvCacheConnectionData: jest.fn().mockResolvedValue(undefined), + } as any; + const license = { hasFeature: jest.fn((f: string) => licensed && f === Feature.KV_CACHE_MONITORING) } as any; + const service = new KvCacheFootprintService(registry, storage, license); + return { service, storage, client, registry, license }; +} + +describe('KvCacheFootprintService', () => { + it('collects, stores and announces a snapshot', async () => { + const { service, storage } = setup(); + const listener = jest.fn(); + service.onSnapshot(listener); + const snapshot = await service.triggerCollection('c1'); + expect(snapshot).toMatchObject({ connectionId: 'c1', detected: true, matchedKeys: 6, chunksEst: 6, evictedKeysDelta: null, evictedKeysIntervalMs: null }); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledWith(snapshot); + expect(listener).toHaveBeenCalledWith(snapshot); + expect(service.getSampleKey('c1')).toBe('m/x@1@0@1…@bfloat16'); + }); + + it('classifies each scanned key once per collection', async () => { + const { service } = setup(); + const classify = jest.spyOn(classifier, 'classifyKey'); + try { + await service.triggerCollection('c1'); + expect(classify).toHaveBeenCalledTimes(6); + } finally { + classify.mockRestore(); + } + }); + + 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('reports the time since the run the eviction delta is measured from', async () => { + const { service } = setup(); + const now = jest.spyOn(Date, 'now').mockReturnValue(1_000_000); + try { + await service.triggerCollection('c1'); + now.mockReturnValue(1_000_000 + 4_000); + expect(await service.triggerCollection('c1')).toMatchObject({ evictedKeysDelta: 0, evictedKeysIntervalMs: 4_000 }); + } finally { + now.mockRestore(); + } + }); + + it('keeps the eviction baseline when the connection is renamed', async () => { + const { service, registry } = setup(); + await service.triggerCollection('c1'); + registry.getConfig.mockReturnValue({ id: 'c1', name: 'renamed', host: 'h', port: 6379, dbIndex: 0 }); + expect((await service.triggerCollection('c1'))?.evictedKeysDelta).toBe(0); + }); + + it('keeps the eviction baseline of a cluster member when it is renamed', async () => { + const { service, client, registry } = setup(); + const member = fakeClient([]); + client.getClusterNodes = jest.fn().mockResolvedValue([ + { id: 'n0', flags: ['myself', 'master'], slots: [[0, 8191]] }, + { id: 'n1', flags: ['master'], slots: [[8192, 16383]] }, + ]); + registry.get.mockImplementation((id: string) => (id === 'm1' ? member : client)); + const listing = (name: string) => [ + { id: 'c1', name: 'one', isConnected: true, connectionType: 'direct' }, + { id: 'm1', name, isConnected: true, connectionType: 'direct', membership: { seedId: 'c1', nodeId: 'n1', origin: 'auto', source: 'cluster' } }, + ]; + const evicted = (count: number) => + member.getInfoParsed.mockResolvedValue({ + memory: { used_memory: '1000', maxmemory: '0', maxmemory_policy: 'noeviction' }, + stats: { evicted_keys: String(count) }, + keyspace: {}, + }); + registry.list.mockReturnValue(listing('m1')); + evicted(10); + await service.triggerCollection('c1'); + registry.list.mockReturnValue(listing('renamed')); + evicted(15); + expect((await service.triggerCollection('c1'))?.evictedKeysDelta).toBe(5); + }); + + describe('memory limits across primaries', () => { + const withMemory = (target: any, maxmemory: string, policy: string) => + target.getInfoParsed.mockResolvedValue({ + memory: { used_memory: '1000', maxmemory, maxmemory_policy: policy }, + stats: { evicted_keys: '0' }, + keyspace: {}, + }); + + function twoNodes(a: [string, string], b: [string, string]) { + const ctx = setup(); + const other = fakeClient([]); + withMemory(ctx.client, ...a); + withMemory(other, ...b); + jest.spyOn(ctx.service as any, 'nodesFor').mockResolvedValue({ + nodes: [{ name: 'a', client: ctx.client }, { name: 'b', client: other }], + skipped: 0, + }); + return ctx.service; + } + + it('sums the limits when every primary has one', async () => { + const service = twoNodes(['1000', 'allkeys-lru'], ['2000', 'allkeys-lru']); + expect(await service.triggerCollection('c1')).toMatchObject({ maxmemory: 3000, maxmemoryPolicy: 'allkeys-lru' }); + }); + + it('reports no limit when any primary is unlimited', async () => { + const service = twoNodes(['1000', 'allkeys-lru'], ['0', 'allkeys-lru']); + expect(await service.triggerCollection('c1')).toMatchObject({ maxmemory: 0 }); + }); + + it('reports a mixed policy when primaries disagree', async () => { + const service = twoNodes(['1000', 'allkeys-lru'], ['1000', 'volatile-lru']); + expect(await service.triggerCollection('c1')).toMatchObject({ maxmemoryPolicy: 'mixed' }); + }); + }); + + it('counts a tagged pair split across primaries as two-key evidence', async () => { + const pair = 'm/x@1@0@3@bfloat16@kind%prefill'; + const single = 'm/x@1@0@4@bfloat16@kind%decode'; + const { service } = setup(); + jest.spyOn(service as any, 'nodesFor').mockResolvedValue({ + nodes: [ + { name: 'a', client: fakeClient([`${pair}metadata`, single]) }, + { name: 'b', client: fakeClient([`${pair}kv_bytes`]) }, + ], + skipped: 0, + }); + expect(await service.triggerCollection('c1')).toMatchObject({ matchedKeys: 3, chunksEst: 2 }); + }); + + it('counts single keys on another primary as single-key evidence', async () => { + const { service } = setup(); + jest.spyOn(service as any, 'nodesFor').mockResolvedValue({ + nodes: [ + { name: 'a', client: fakeClient(['m/x@1@0@1@bfloat16', 'm/x@1@0@2@bfloat16']) }, + { name: 'b', client: fakeClient(['m/x@1@0@3@bfloat16@kind%metadata', 'm/x@1@0@4@bfloat16@kind%metadata']) }, + ], + skipped: 0, + }); + expect(await service.triggerCollection('c1')).toMatchObject({ matchedKeys: 4, chunksEst: 4 }); + }); + + it('reads a tag ending in a suffix name as a tag when no node shows two-key evidence', async () => { + const { service } = setup(['m/x@1@0@1@bfloat16@kind%metadata', 'm/x@1@0@2@bfloat16@kind%decode']); + expect(await service.triggerCollection('c1')).toMatchObject({ matchedKeys: 2, chunksEst: 2 }); + }); + + it('drops the eviction baseline of a node that misses a run', async () => { + const { service, client } = setup(); + const other = fakeClient([]); + const evicted = (target: any, count: number) => + target.getInfoParsed.mockResolvedValue({ + memory: { used_memory: '1000', maxmemory: '0', maxmemory_policy: 'noeviction' }, + stats: { evicted_keys: String(count) }, + keyspace: {}, + }); + const nodes = jest.spyOn(service as any, 'nodesFor'); + const both = { nodes: [{ name: 'a', client }, { name: 'b', client: other }], skipped: 0 }; + nodes.mockResolvedValue(both); + evicted(client, 10); + evicted(other, 100); + await service.triggerCollection('c1'); + nodes.mockResolvedValue({ nodes: [{ name: 'a', client }], skipped: 1 }); + evicted(client, 12); + expect((await service.triggerCollection('c1'))?.evictedKeysDelta).toBe(2); + nodes.mockResolvedValue(both); + evicted(client, 13); + evicted(other, 5000); + 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(); + const broken = fakeClient([]); + broken.getDbSize.mockRejectedValue(new Error('LOADING')); + jest.spyOn(service as any, 'nodesFor').mockResolvedValue({ + nodes: [ + { name: 'a', client }, + { name: 'b', client: broken }, + ], + skipped: 0, + }); + const snapshot = await service.triggerCollection('c1'); + expect(snapshot).toMatchObject({ matchedKeys: 6, detected: true, scanComplete: false }); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledWith(snapshot); + }); + + 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)]); + 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: 'b', client: broken }, + ], + skipped: 0, + }); + await service.triggerCollection('c1'); + expect(service.getSampleKey('c1')).toBe('m/x@1@0@1…@bfloat16'); + }); + + it('clears an earlier detection after consecutive partial runs that find no LMCache', async () => { + const { service, storage } = setup(); + await service.triggerCollection('c1'); + storage.getKvCacheFootprintSnapshots.mockImplementation(async () => [ + { ...storage.saveKvCacheFootprintSnapshot.mock.calls.at(-1)[0], timestamp: 0 }, + ]); + jest.spyOn(service as any, 'nodesFor').mockResolvedValue({ nodes: [{ name: 'a', client: fakeClient([]) }], skipped: 1 }); + const runs = []; + for (let i = 0; i < PARTIAL_MISSES_BEFORE_CLEAR; i += 1) runs.push(await service.triggerCollection('c1')); + expect(runs.slice(0, -1).every((run) => run?.detected === true)).toBe(true); + expect(runs.at(-1)).toMatchObject({ detected: false, scanComplete: false }); + expect(service.getSampleKey('c1')).toBeNull(); + expect(await service.triggerCollection('c1')).toMatchObject({ detected: false }); + }); + + it('restarts the partial miss count when a run finds LMCache again', async () => { + const { service, storage } = setup(); + await service.triggerCollection('c1'); + storage.getKvCacheFootprintSnapshots.mockResolvedValue([{ connectionId: 'c1', detected: true }]); + const nodes = jest.spyOn(service as any, 'nodesFor'); + const empty = { nodes: [{ name: 'a', client: fakeClient([]) }], skipped: 1 }; + nodes.mockResolvedValue(empty); + for (let i = 0; i < PARTIAL_MISSES_BEFORE_CLEAR - 1; i += 1) await service.triggerCollection('c1'); + nodes.mockRestore(); + await service.triggerCollection('c1'); + jest.spyOn(service as any, 'nodesFor').mockResolvedValue(empty); + expect(await service.triggerCollection('c1')).toMatchObject({ detected: true }); + }); + + it('stores not detected for a partial read when LMCache was not detected before', async () => { + const { service } = setup(); + const broken = fakeClient([]); + broken.getDbSize.mockRejectedValue(new Error('LOADING')); + jest.spyOn(service as any, 'nodesFor').mockResolvedValue({ + nodes: [ + { name: 'a', client: fakeClient([KEY(0)]) }, + { name: 'b', client: broken }, + ], + skipped: 0, + }); + expect(await service.triggerCollection('c1')).toMatchObject({ detected: false, scanComplete: false }); + }); + + 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: [] }); + + function clusterSetup( + member: 'deregistered' | 'role-failed' | 'replica' | 'primary', + view: unknown[] = [primary('n0', [[0, 8191]], ['myself']), member === 'replica' ? replica('n1') : primary('n1', [[8192, 16383]])], + ) { + const ctx = setup(); + const other = fakeClient([]); + ctx.client.getClusterNodes = jest.fn().mockResolvedValue(view); + if (member === 'role-failed') other.getRole.mockRejectedValue(new Error('timeout')); + if (member === 'replica') other.getRole.mockResolvedValue({ role: 'slave' }); + ctx.registry.get.mockImplementation((id: string) => { + if (id === 'm1' && member === 'deregistered') throw new Error('not registered'); + return id === 'm1' ? other : ctx.client; + }); + ctx.registry.list.mockReturnValue([ + { id: 'c1', name: 'one', isConnected: true, connectionType: 'direct' }, + { + id: 'm1', + name: 'm1', + isConnected: true, + connectionType: 'direct', + membership: { seedId: 'c1', nodeId: 'n1', origin: 'auto', source: 'cluster' }, + }, + ]); + return ctx; + } + + it.each(['deregistered', 'role-failed'] as const)('marks the snapshot incomplete when a member is %s', async (member) => { + const { service, storage } = clusterSetup(member); + const snapshot = await service.triggerCollection('c1'); + expect(snapshot).toMatchObject({ matchedKeys: 6, scanComplete: false }); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledWith(snapshot); + }); + + it('does not treat a replica as unread', async () => { + const { service } = clusterSetup('replica'); + expect(await service.triggerCollection('c1')).toMatchObject({ scanComplete: true }); + }); + + it('keeps the snapshot complete when every slot owner was scanned', async () => { + const { service } = clusterSetup('primary'); + expect(await service.triggerCollection('c1')).toMatchObject({ scanComplete: true }); + }); + + it('marks the snapshot incomplete when a slot owner has no member connection', async () => { + const { service } = clusterSetup('primary', [ + primary('n0', [[0, 5460]], ['myself']), + primary('n1', [[5461, 10922]]), + primary('n2', [[10923, 16383]]), + ]); + expect(await service.triggerCollection('c1')).toMatchObject({ scanComplete: false, primariesScanned: 2, primariesTotal: 3 }); + }); + + it('limits sibling probes on each primary to the slots it owns', async () => { + const { service, registry, client } = clusterSetup('primary'); + const sample = jest.spyOn(scanner, 'sampleLmcacheKeys'); + try { + await service.triggerCollection('c1'); + const ownership = new Map( + sample.mock.calls.map(([target, , owns]) => [target === client ? 'seed' : 'member', owns]), + ); + expect([0, 8191, 8192].map((slot) => ownership.get('seed')?.(slot))).toEqual([true, true, false]); + expect([8191, 8192, 16383].map((slot) => ownership.get('member')?.(slot))).toEqual([false, true, true]); + expect(registry.get).toHaveBeenCalledWith('m1'); + } finally { + sample.mockRestore(); + } + }); + + it('probes every sibling when the cluster view is unknown', async () => { + const { service, client } = clusterSetup('primary'); + client.getClusterNodes.mockRejectedValue(new Error('timeout')); + const sample = jest.spyOn(scanner, 'sampleLmcacheKeys'); + try { + await service.triggerCollection('c1'); + expect(sample.mock.calls.every(([, , owns]) => owns === undefined)).toBe(true); + } finally { + sample.mockRestore(); + } + }); + + it('scales by the slots of the primaries that were read when one read fails', async () => { + const { service, registry, client } = clusterSetup('primary'); + const broken = fakeClient([]); + broken.getDbSize.mockRejectedValue(new Error('LOADING')); + registry.get.mockImplementation((id: string) => (id === 'm1' ? broken : client)); + expect(await service.triggerCollection('c1')).toMatchObject({ + scanComplete: false, + matchedKeys: 6, + chunksEst: 12, + primariesScanned: 1, + primariesTotal: 2, + }); + }); + + it('keeps the snapshot complete when only a replica fails its role check', async () => { + const { service } = clusterSetup('role-failed', [primary('n0', [[0, 16383]], ['myself']), replica('n1')]); + expect(await service.triggerCollection('c1')).toMatchObject({ scanComplete: true }); + }); + + it.each([ + ['its IP', '10.0.0.2'], + ['its announced hostname', 'node-2.example'], + ])('finds a member whose stored node id is stale by %s', async (_label, host) => { + const { service, registry } = clusterSetup('primary', [ + primary('n0', [[0, 5460]], ['myself']), + { ...primary('n1-replaced', [[5461, 16383]]), address: '10.0.0.2:6379@16379', hostname: 'node-2.example' }, + ]); + registry.getConfig.mockImplementation((id: string) => + id === 'm1' + ? { id: 'm1', name: 'm1', host, port: 6379, dbIndex: 0 } + : { id: 'c1', name: 'one', host: 'h', port: 6379, dbIndex: 0 }, + ); + expect(await service.triggerCollection('c1')).toMatchObject({ + scanComplete: true, + chunksEst: 6, + primariesScanned: 2, + primariesTotal: 2, + }); + }); + + it('does not scale the estimate when a scanned member matches no slot owner', async () => { + const { service, registry, client } = clusterSetup('primary', [ + primary('n0', [[0, 5460]], ['myself']), + primary('n1-replaced', [[5461, 10922]]), + primary('n2', [[10923, 16383]]), + ]); + const unidentified = fakeClient(Array.from({ length: 6 }, (_, i) => KEY(i + 6))); + registry.get.mockImplementation((id: string) => (id === 'm1' ? unidentified : client)); + expect(await service.triggerCollection('c1')).toMatchObject({ + matchedKeys: 12, + chunksEst: 12, + primariesScanned: 1, + primariesTotal: 3, + scanComplete: false, + }); + }); + + it('scans a member once when it is the same node as the seed', async () => { + const seedId = 'a'.repeat(40); + const { service, registry, client } = clusterSetup('primary', [ + primary(seedId, [[0, 8191]], ['myself']), + primary('b'.repeat(40), [[8192, 16383]]), + ]); + const sameNode = fakeClient(Array.from({ length: 6 }, (_, i) => KEY(i))); + const myId = (target: any) => { + const call = target.call.getMockImplementation(); + target.call.mockImplementation(async (cmd: string, args: string[]) => (cmd === 'CLUSTER' && args[0] === 'MYID' ? seedId : call(cmd, args))); + }; + myId(client); + myId(sameNode); + registry.get.mockImplementation((id: string) => (id === 'm1' ? sameNode : client)); + const snapshot = await service.triggerCollection('c1'); + expect(snapshot).toMatchObject({ matchedKeys: 6, chunksEst: 12, primariesScanned: 1, primariesTotal: 2, scanComplete: false }); + expect(sameNode.getDbSize).not.toHaveBeenCalled(); + }); + + it('marks the snapshot incomplete when the cluster view cannot be read', async () => { + const { service, client } = clusterSetup('primary'); + client.getClusterNodes.mockRejectedValue(new Error('timeout')); + expect(await service.triggerCollection('c1')).toMatchObject({ scanComplete: false }); + }); + }); + + describe('cluster seed without registered members', () => { + function seedSetup(clusterNodes: () => Promise) { + const ctx = setup(); + ctx.client.getCapabilities = () => ({ clusterEnabled: true }); + ctx.client.getClusterNodes = jest.fn(clusterNodes); + return ctx; + } + + it('marks the snapshot incomplete when other primaries own slots', async () => { + const { service } = seedSetup(async () => [ + { id: 'a', flags: ['myself', 'master'], slots: [[0, 8191]] }, + { id: 'b', flags: ['master'], slots: [[8192, 16383]] }, + { id: 'c', flags: ['slave'], slots: [] }, + ]); + expect(await service.triggerCollection('c1')).toMatchObject({ matchedKeys: 6, scanComplete: false }); + }); + + it('reports the shards scanned and scales the footprint by the slots they own', async () => { + const { service } = seedSetup(async () => [ + { id: 'a', flags: ['myself', 'master'], slots: [[0, 4095]] }, + { id: 'b', flags: ['master'], slots: [[4096, 8191]] }, + { id: 'c', flags: ['master'], slots: [[8192, 16383]] }, + { id: 'd', flags: ['slave'], slots: [] }, + ]); + expect(await service.triggerCollection('c1')).toMatchObject({ + matchedKeys: 6, + chunksEst: 24, + primariesScanned: 1, + primariesTotal: 3, + }); + }); + + it('counts a replica seed as no primary scanned', async () => { + const { service } = seedSetup(async () => [ + { id: 'r', flags: ['myself', 'slave'], master: 'a', slots: [] }, + { id: 'a', flags: ['master'], slots: [[0, 16383]] }, + ]); + expect(await service.triggerCollection('c1')).toMatchObject({ + scanComplete: false, + primariesScanned: 0, + primariesTotal: 1, + }); + }); + + it('does not sample keys on a replica seed, which would answer MOVED', async () => { + const { service, client } = seedSetup(async () => [ + { id: 'r', flags: ['myself', 'slave'], master: 'a', slots: [] }, + { id: 'a', flags: ['master'], slots: [[0, 16383]] }, + ]); + const sample = jest.spyOn(scanner, 'sampleLmcacheKeys'); + try { + expect(await service.triggerCollection('c1')).toMatchObject({ matchedKeys: 6, primariesScanned: 0 }); + expect(sample.mock.calls.filter(([target]) => target === client)).toHaveLength(0); + } finally { + sample.mockRestore(); + } + }); + + it('samples keys on a primary seed', async () => { + const { service, client } = seedSetup(async () => [{ id: 'a', flags: ['myself', 'master'], slots: [[0, 16383]] }]); + const sample = jest.spyOn(scanner, 'sampleLmcacheKeys'); + try { + await service.triggerCollection('c1'); + expect(sample.mock.calls.filter(([target]) => target === client)).toHaveLength(1); + } finally { + sample.mockRestore(); + } + }); + + it('keeps the snapshot complete when the seed is the only primary', async () => { + const { service } = seedSetup(async () => [{ id: 'a', flags: ['myself', 'master'], slots: [[0, 16383]] }]); + expect(await service.triggerCollection('c1')).toMatchObject({ scanComplete: true, primariesScanned: 1, primariesTotal: 1 }); + }); + + it('marks the snapshot incomplete when the cluster view cannot be read', async () => { + const { service } = seedSetup(async () => { + throw new Error('timeout'); + }); + expect(await service.triggerCollection('c1')).toMatchObject({ + scanComplete: false, + chunksEst: 6, + primariesScanned: 1, + primariesTotal: null, + }); + }); + + it('does not read the cluster view of a standalone server', async () => { + const { service, client } = setup(); + client.getCapabilities = () => ({ clusterEnabled: false }); + client.getClusterNodes = jest.fn(); + expect(await service.triggerCollection('c1')).toMatchObject({ scanComplete: true, primariesScanned: 1, primariesTotal: 1 }); + expect(client.getClusterNodes).not.toHaveBeenCalled(); + }); + }); + + describe('disconnected cluster members', () => { + function offlineSetup(clusterNodes: () => Promise) { + const ctx = setup(); + ctx.registry.list.mockReturnValue([ + { id: 'c1', name: 'one', isConnected: true, connectionType: 'direct' }, + { + id: 'm1', + name: 'm1', + isConnected: false, + connectionType: 'direct', + membership: { seedId: 'c1', nodeId: 'n1', origin: 'auto', source: 'cluster' }, + }, + ]); + ctx.client.getClusterNodes = jest.fn(clusterNodes); + return ctx; + } + + 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 incomplete when the cluster view cannot be read', async () => { + const { service } = offlineSetup(async () => { + throw new Error('timeout'); + }); + expect(await service.triggerCollection('c1')).toMatchObject({ scanComplete: false }); + }); + + 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 }); + }); + }); + + it('fails the collection when every primary fails', async () => { + const { service, client, storage } = setup(); + client.getDbSize.mockRejectedValue(new Error('LOADING')); + await expect(service.triggerCollection('c1')).rejects.toThrow('LOADING'); + expect(storage.saveKvCacheFootprintSnapshot).not.toHaveBeenCalled(); + }); + + it('lists other databases relative to the database the client actually selected', async () => { + const { service, client, registry } = setup(); + registry.getConfig.mockReturnValue({ id: 'c1', name: 'one', host: 'h', port: 6379, dbIndex: 2 }); + client.getInfoParsed.mockResolvedValue({ + memory: { used_memory: '1000', maxmemory: '0', maxmemory_policy: 'noeviction' }, + stats: { evicted_keys: '0' }, + keyspace: { db0: { keys: 6, expires: 0, avg_ttl: 0 }, db2: { keys: 9, expires: 0, avg_ttl: 0 } }, + }); + client.getActiveDbIndex = jest.fn().mockReturnValue(0); + expect((await service.triggerCollection('c1'))?.otherDbs).toEqual([2]); + client.getActiveDbIndex.mockReturnValue(2); + expect((await service.triggerCollection('c1'))?.otherDbs).toEqual([0]); + }); + + it('keeps going when a listener throws', async () => { + const { service, storage } = setup(); + service.onSnapshot(() => { throw new Error('boom'); }); + await expect(service.triggerCollection('c1')).resolves.not.toBeNull(); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalled(); + }); + + it('starts the poller even when the licence has not resolved at init', async () => { + const { service } = setup(undefined, false); + const start = jest.spyOn(service as any, 'start').mockImplementation(() => undefined); + await service.onModuleInit(); + expect(start).toHaveBeenCalled(); + }); + + it('skips collection while unlicensed and collects once the licence resolves', async () => { + const { service, storage, license } = setup(undefined, false); + jest.spyOn(service as any, 'start').mockImplementation(() => undefined); + await service.onModuleInit(); + await (service as any).tick(); + expect(storage.saveKvCacheFootprintSnapshot).not.toHaveBeenCalled(); + license.hasFeature.mockImplementation((f: string) => f === Feature.KV_CACHE_MONITORING); + await (service as any).tick(); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(1); + }); + + it('runs connection removal cleanup while unlicensed', async () => { + const { service, storage, registry } = setup(undefined, false); + jest.spyOn(service as any, 'start').mockImplementation(() => undefined); + await service.onModuleInit(); + await (service as any).tick(); + registry.removed('c1'); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledWith('c1'); + expect(storage.saveKvCacheFootprintSnapshot).not.toHaveBeenCalled(); + await service.onModuleDestroy(); + }); + + it('stops polling a retired member but keeps its data', async () => { + const { service, storage, client, registry } = setup(); + jest.spyOn(service as any, 'start').mockImplementation(() => undefined); + await service.onModuleInit(); + await (service as any).tick(); + expect(client.getDbSize).toHaveBeenCalledTimes(1); + registry.list.mockImplementation((options?: { includeRetired?: boolean }) => + options?.includeRetired ? [{ id: 'c1', name: 'one', isConnected: false, connectionType: 'direct' }] : [], + ); + await (service as any).tick(); + await new Promise((resolve) => setImmediate(resolve)); + expect(client.getDbSize).toHaveBeenCalledTimes(1); + expect(storage.deleteKvCacheConnectionData).not.toHaveBeenCalled(); + expect(service.getSampleKey('c1')).toBe('m/x@1@0@1…@bfloat16'); + await service.onModuleDestroy(); + }); + + it('backs off full collection after a not-detected result until the recheck window passes', async () => { + const { service, storage, client } = setup([]); + const now = jest.spyOn(Date, 'now').mockReturnValue(1_000_000); + try { + await (service as any).tick(); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(1); + expect(storage.saveKvCacheFootprintSnapshot.mock.calls[0][0].detected).toBe(false); + now.mockReturnValue(1_000_000 + 1_800_000 - 1); + await (service as any).tick(); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(1); + expect(client.getDbSize).toHaveBeenCalledTimes(1); + now.mockReturnValue(1_000_000 + 1_800_000); + await (service as any).tick(); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(2); + } finally { + now.mockRestore(); + } + }); + + describe('back-off on partial runs', () => { + function partialSetup() { + const ctx = setup([]); + ctx.client.getCapabilities = () => ({ clusterEnabled: true }); + ctx.client.getClusterNodes = jest.fn().mockResolvedValue([ + { id: 'a', flags: ['myself', 'master'], slots: [[0, 8191]] }, + { id: 'b', flags: ['master'], slots: [[8192, 16383]] }, + ]); + return ctx; + } + + it('keeps polling when slot owners were left unscanned and nothing was detected', async () => { + const { service, storage } = partialSetup(); + await (service as any).tick(); + await (service as any).tick(); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(2); + expect(storage.saveKvCacheFootprintSnapshot.mock.calls[0][0]).toMatchObject({ detected: false, scanComplete: false }); + }); + + it('backs off once consecutive partial misses clear an earlier detection', async () => { + const { service, storage } = partialSetup(); + storage.getKvCacheFootprintSnapshots.mockResolvedValue([{ connectionId: 'c1', detected: true }]); + for (let i = 0; i < PARTIAL_MISSES_BEFORE_CLEAR + 1; i += 1) await (service as any).tick(); + const saved = storage.saveKvCacheFootprintSnapshot.mock.calls.map(([snapshot]: [any]) => snapshot.detected); + expect(saved).toEqual([true, true, false]); + }); + + it('backs off a cluster that is partial on every run after consecutive misses', async () => { + const { service, storage } = partialSetup(); + const now = jest.spyOn(Date, 'now').mockReturnValue(1_000_000); + try { + for (let i = 0; i < PARTIAL_MISSES_BEFORE_CLEAR + 2; i += 1) await (service as any).tick(); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(PARTIAL_MISSES_BEFORE_CLEAR); + now.mockReturnValue(1_000_000 + 1_800_000); + await (service as any).tick(); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(PARTIAL_MISSES_BEFORE_CLEAR + 1); + } finally { + now.mockRestore(); + } + }); + + it('keeps polling when a scanned primary could not be read', async () => { + const { service, storage } = partialSetup(); + const broken = fakeClient([]); + broken.getDbSize.mockRejectedValue(new Error('LOADING')); + jest.spyOn(service as any, 'nodesFor').mockResolvedValue({ + nodes: [{ name: 'a', client: fakeClient([]) }, { name: 'b', client: broken }], + skipped: 0, + }); + await (service as any).tick(); + await (service as any).tick(); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(2); + }); + + it('keeps polling when a partial run keeps an earlier detection', async () => { + const { service, storage } = partialSetup(); + storage.getKvCacheFootprintSnapshots.mockResolvedValue([{ connectionId: 'c1', detected: true }]); + await (service as any).tick(); + await (service as any).tick(); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(2); + }); + }); + + it('collects on a manual trigger during the back-off and resumes polling once detected', async () => { + const keys = Array.from({ length: 6 }, (_, i) => KEY(i)); + const { service, storage, client } = setup([]); + await (service as any).tick(); + 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]])); + expect((await service.triggerCollection('c1'))?.detected).toBe(true); + await (service as any).tick(); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(3); + }); + + it('rejects a manual trigger for an unknown connection with NotFoundException', async () => { + const { service, registry, storage } = setup(); + registry.getConfig.mockReturnValue(null); + await expect(service.triggerCollection('missing')).rejects.toBeInstanceOf(NotFoundException); + expect(storage.saveKvCacheFootprintSnapshot).not.toHaveBeenCalled(); + }); + + 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'); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledWith('c1'); + await service.onModuleDestroy(); + }); + + 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.removed('c2'); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.deleteKvCacheConnectionData).toHaveBeenCalledWith('c2'); + await service.onModuleDestroy(); + }); + + it('waits for an in-flight collection before deleting a removed connection', 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 collecting = service.triggerCollection('c1'); + await new Promise((resolve) => setImmediate(resolve)); + registry.removed('c1'); + await new Promise((resolve) => setImmediate(resolve)); + expect(storage.deleteKvCacheConnectionData).not.toHaveBeenCalled(); + 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(); + }); + + 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'); + 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); + }); + + 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('rejects a refresh for a retired member without scanning it', async () => { + const { service, client, registry } = setup(); + const retired = { + id: 'c1', + name: 'one', + isConnected: true, + connectionType: 'direct', + membership: { seedId: 'seed', nodeId: 'n', origin: 'auto', source: 'cluster', retiredAt: 1 }, + }; + registry.list.mockImplementation((options?: { includeRetired?: boolean }) => (options?.includeRetired ? [retired] : [])); + await expect(service.triggerCollection('c1')).rejects.toThrow(/retired/); + await expect(service.triggerCollection('c1')).rejects.toBeInstanceOf(BadRequestException); + expect(client.getDbSize).not.toHaveBeenCalled(); + }); + + it('reads the registry listing once per refresh', async () => { + const { service, registry } = setup(); + await service.triggerCollection('c1'); + expect(registry.list).toHaveBeenCalledTimes(1); + }); + + 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); + expect(client.getDbSize).not.toHaveBeenCalled(); + }); + + it('returns the latest snapshot to a refresh when it is only seconds old', async () => { + const { service, client, storage } = setup(); + const recent = { connectionId: 'c1', timestamp: Date.now() - 1000, detected: true }; + storage.getKvCacheFootprintSnapshots.mockResolvedValue([recent]); + expect(await service.triggerCollection('c1')).toBe(recent); + expect(client.getDbSize).not.toHaveBeenCalled(); + expect(storage.saveKvCacheFootprintSnapshot).not.toHaveBeenCalled(); + }); + + it('collects on a refresh when the latest snapshot is older than the reuse window', async () => { + const { service, storage } = setup(); + storage.getKvCacheFootprintSnapshots.mockResolvedValue([ + { connectionId: 'c1', timestamp: Date.now() - REFRESH_REUSE_MS - 1, detected: true }, + ]); + await service.triggerCollection('c1'); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(1); + }); + + it('rejects a refresh 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 refresh 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('reads INFO while the scan runs and CLIENT LIST only after it', async () => { + const { service, client } = setup([KEY(0)]); + let release!: () => void; + const scanned = new Promise((resolve) => { release = resolve; }); + const call = client.call.getMockImplementation(); + client.call.mockImplementation(async (cmd: string, args: string[]) => { + if (cmd === 'SCAN') await scanned; + return call!(cmd, args); + }); + const collecting = service.triggerCollection('c1'); + await new Promise((resolve) => setImmediate(resolve)); + expect(client.getInfoParsed).toHaveBeenCalled(); + expect(client.call).not.toHaveBeenCalledWith('CLIENT', ['LIST']); + release(); + await collecting; + expect(client.call).toHaveBeenCalledWith('CLIENT', ['LIST']); + }); + + it('skips CLIENT LIST when the scan finds enough LMCache keys', async () => { + const { service, client } = setup(); + const clientLists = () => client.call.mock.calls.filter(([cmd]: [string]) => cmd === 'CLIENT').length; + await service.triggerCollection('c1'); + await service.triggerCollection('c1'); + expect(clientLists()).toBe(0); + }); + + it('skips CLIENT LIST when the run has enough keys though each primary has few', async () => { + const { service } = setup(); + const a = fakeClient([KEY(0), KEY(1), KEY(2)]); + const b = fakeClient([KEY(3), KEY(4), KEY(5)]); + jest.spyOn(service as any, 'nodesFor').mockResolvedValue({ nodes: [{ name: 'a', client: a }, { name: 'b', client: b }], skipped: 0 }); + expect(await service.triggerCollection('c1')).toMatchObject({ detected: true, matchedKeys: 6 }); + expect(a.call).not.toHaveBeenCalledWith('CLIENT', ['LIST']); + expect(b.call).not.toHaveBeenCalledWith('CLIENT', ['LIST']); + }); + + it('checks CLIENT LIST again when a previously detected cache drops below the key threshold', async () => { + const { service, client, storage } = setup(); + await service.triggerCollection('c1'); + client.getDbSize.mockResolvedValue(1); + client.call.mockImplementation(async (cmd: string) => + cmd === 'SCAN' ? ['0', [KEY(0)]] : 'id=1 lib-name=GlidePySync(lmcache:0.3.0) lib-ver=1\n', + ); + const snapshot = await service.triggerCollection('c1'); + expect(client.call).toHaveBeenCalledWith('CLIENT', ['LIST']); + expect(snapshot).toMatchObject({ detected: true, matchedKeys: 1 }); + await (service as any).tick(); + expect(storage.saveKvCacheFootprintSnapshot).toHaveBeenCalledTimes(3); + }); + + it('ignores an LMCache client connected to another database', async () => { + const { service, client } = setup([KEY(0)]); + client.getActiveDbIndex = jest.fn().mockReturnValue(2); + client.call.mockImplementation(async (cmd: string) => + cmd === 'SCAN' ? ['0', [KEY(0)]] : 'id=1 db=0 lib-name=GlidePySync(lmcache:0.3.0) lib-ver=1\n', + ); + expect(await service.triggerCollection('c1')).toMatchObject({ detected: false, matchedKeys: 1 }); + client.getActiveDbIndex.mockReturnValue(0); + expect(await service.triggerCollection('c1')).toMatchObject({ detected: true, matchedKeys: 1 }); + }); + + it('skips agent connections, which cannot run SCAN or MEMORY USAGE', async () => { + const { service, storage, client, registry } = setup(); + registry.list.mockReturnValue([{ id: 'c1', name: 'one', isConnected: true, connectionType: 'agent' }]); + registry.getConfig.mockReturnValue({ id: 'c1', name: 'one', host: 'agent', port: 0, dbIndex: 0 }); + await (service as any).tick(); + await expect(service.triggerCollection('c1')).rejects.toBeInstanceOf(BadRequestException); + expect(client.getDbSize).not.toHaveBeenCalled(); + expect(storage.saveKvCacheFootprintSnapshot).not.toHaveBeenCalled(); + }); +}); + +describe('resolvePositiveInt', () => { + it('uses the default when unset', () => { + expect(resolvePositiveInt(undefined, 300000, 60000)).toBe(300000); + }); + + it('uses the default when the value is not a number', () => { + expect(resolvePositiveInt('abc', 300000, 60000)).toBe(300000); + expect(resolvePositiveInt('', 500, 1)).toBe(500); + }); + + it('uses the default when the value is below the floor', () => { + expect(resolvePositiveInt('5m', 300000, 60000)).toBe(300000); + expect(resolvePositiveInt('59999', 300000, 60000)).toBe(300000); + expect(resolvePositiveInt('0', 2000, 1)).toBe(2000); + expect(resolvePositiveInt('-3', 500, 1)).toBe(500); + }); + + it('rejects numeric prefixes instead of truncating them', () => { + const onInvalid = jest.fn(); + expect(resolvePositiveInt('500k', 500, 1, onInvalid)).toBe(500); + expect(resolvePositiveInt('1.5', 500, 1, onInvalid)).toBe(500); + expect(onInvalid.mock.calls).toEqual([['500k'], ['1.5']]); + }); + + it('accepts exponent notation as the number it spells', () => { + expect(resolvePositiveInt('1e6', 200000, 1)).toBe(1_000_000); + }); + + it('reports a value below the floor but not an unset one', () => { + const onInvalid = jest.fn(); + expect(resolvePositiveInt('0', 2000, 1, onInvalid)).toBe(2000); + expect(resolvePositiveInt(undefined, 2000, 1, onInvalid)).toBe(2000); + expect(onInvalid.mock.calls).toEqual([['0']]); + }); + + it('uses the default when the value is above the ceiling', () => { + const onInvalid = jest.fn(); + expect(resolvePositiveInt('86400001', 300000, 60000, onInvalid, 86400000)).toBe(300000); + expect(resolvePositiveInt('86400000', 300000, 60000, onInvalid, 86400000)).toBe(86400000); + expect(onInvalid.mock.calls).toEqual([['86400001']]); + }); + + it('accepts a value at or above the floor', () => { + expect(resolvePositiveInt('60000', 300000, 60000)).toBe(60000); + expect(resolvePositiveInt('1', 2000, 1)).toBe(1); + expect(resolvePositiveInt('750', 500, 1)).toBe(750); + }); +}); + +describe('KvCacheFootprintService settings', () => { + const keys = ['KV_CACHE_FOOTPRINT_INTERVAL_MS', 'KV_CACHE_SCAN_MAX_KEYS', 'KV_CACHE_MATCH_MAX_KEYS', 'KV_CACHE_SAMPLE_KEYS', 'KV_CACHE_UNDETECTED_RECHECK_MS']; + const saved = Object.fromEntries(keys.map((key) => [key, process.env[key]])); + + afterEach(() => { + for (const key of keys) { + if (saved[key] === undefined) delete process.env[key]; + else process.env[key] = saved[key]; + } + }); + + it('falls back to the defaults for invalid settings', () => { + process.env.KV_CACHE_FOOTPRINT_INTERVAL_MS = 'abc'; + process.env.KV_CACHE_SCAN_MAX_KEYS = '0'; + process.env.KV_CACHE_MATCH_MAX_KEYS = 'NaN'; + process.env.KV_CACHE_SAMPLE_KEYS = '-1'; + process.env.KV_CACHE_UNDETECTED_RECHECK_MS = '1000'; + const { service } = setup(); + expect((service as any).getIntervalMs()).toBe(300000); + expect((service as any).budgets).toEqual({ maxScanned: 200000, maxMatched: 2000 }); + expect((service as any).sampleSize).toBe(500); + expect((service as any).undetectedRecheckMs).toBe(1800000); + }); + + it('caps the footprint interval at a day so setTimeout cannot overflow', () => { + process.env.KV_CACHE_FOOTPRINT_INTERVAL_MS = String(2 ** 31); + const warn = jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined); + try { + const { service } = setup(); + expect((service as any).getIntervalMs()).toBe(300000); + expect(warn).toHaveBeenCalledWith(expect.stringContaining('KV_CACHE_FOOTPRINT_INTERVAL_MS')); + process.env.KV_CACHE_FOOTPRINT_INTERVAL_MS = String(MAX_FOOTPRINT_INTERVAL_MS); + expect((setup().service as any).getIntervalMs()).toBe(86_400_000); + } finally { + warn.mockRestore(); + } + }); + + it('logs the setting it ignores', () => { + process.env.KV_CACHE_SAMPLE_KEYS = '500k'; + const warn = jest.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined); + try { + const { service } = setup(); + expect((service as any).sampleSize).toBe(500); + expect(warn).toHaveBeenCalledWith(expect.stringContaining('KV_CACHE_SAMPLE_KEYS')); + } finally { + warn.mockRestore(); + } + }); + + it('keeps valid settings', () => { + process.env.KV_CACHE_FOOTPRINT_INTERVAL_MS = '120000'; + process.env.KV_CACHE_SCAN_MAX_KEYS = '1000'; + process.env.KV_CACHE_MATCH_MAX_KEYS = '10'; + process.env.KV_CACHE_SAMPLE_KEYS = '5'; + process.env.KV_CACHE_UNDETECTED_RECHECK_MS = '600000'; + const { service } = setup(); + expect((service as any).getIntervalMs()).toBe(120000); + expect((service as any).budgets).toEqual({ maxScanned: 1000, maxMatched: 10 }); + expect((service as any).sampleSize).toBe(5); + expect((service as any).undetectedRecheckMs).toBe(600000); + }); +}); diff --git a/proprietary/kv-cache/__tests__/kv-cache-status.service.spec.ts b/proprietary/kv-cache/__tests__/kv-cache-status.service.spec.ts new file mode 100644 index 000000000..9201a6a49 --- /dev/null +++ b/proprietary/kv-cache/__tests__/kv-cache-status.service.spec.ts @@ -0,0 +1,63 @@ +import { KvCacheStatusService, toPublicEngine } from '../kv-cache-status.service'; + +const engine = { + id: 'e1', connectionId: 'c1', name: 'vllm', source: 'scrape', scrapeUrl: 'http://x/metrics', + scrapeAuthHeader: 'enc', scrapeAuthEncrypted: true, otlpEngineId: null, enabled: true, + createdAt: 1, lastSeenAt: 2, lastError: null, +} as const; + +describe('KvCacheStatusService', () => { + const storage = { + getKvCacheFootprintSnapshots: jest.fn(), + getKvCacheEngines: jest.fn().mockResolvedValue([engine]), + } as any; + const footprint = { getSampleKey: jest.fn().mockReturnValue('m@1@0@abc…@bfloat16') } as any; + const service = new KvCacheStatusService(storage, footprint); + + it('never exposes the auth header', () => { + const pub = toPublicEngine(engine as any); + expect(pub).not.toHaveProperty('scrapeAuthHeader'); + expect(pub).not.toHaveProperty('scrapeAuthEncrypted'); + expect(pub.hasScrapeAuth).toBe(true); + }); + + it('builds status from the latest snapshot', async () => { + storage.getKvCacheFootprintSnapshots.mockResolvedValue([{ detected: true }]); + const status = await service.getStatus('c1'); + expect(storage.getKvCacheFootprintSnapshots).toHaveBeenCalledWith({ connectionId: 'c1', limit: 1 }); + expect(status).toEqual({ hasLmcache: true, latest: { detected: true }, sampleKey: 'm@1@0@abc…@bfloat16', engines: [toPublicEngine(engine as any)] }); + }); + + it('reads the snapshot and the engines in parallel', async () => { + storage.getKvCacheFootprintSnapshots.mockReturnValue(new Promise(() => undefined)); + storage.getKvCacheEngines.mockClear(); + void service.getStatus('c1'); + await Promise.resolve(); + expect(storage.getKvCacheEngines).toHaveBeenCalledWith('c1'); + }); + + it('reports not detected without a snapshot', async () => { + storage.getKvCacheFootprintSnapshots.mockResolvedValue([]); + expect((await service.getStatus('c1')).hasLmcache).toBe(false); + }); + + it('caps the history range and buckets it in storage', async () => { + storage.getKvCacheFootprintSnapshots.mockResolvedValue([]); + const to = 100 * 24 * 60 * 60_000; + const range = 31 * 24 * 60 * 60_000; + await service.getFootprintHistory('c1', 0, to); + expect(storage.getKvCacheFootprintSnapshots).toHaveBeenLastCalledWith({ + connectionId: 'c1', + from: to - range, + to, + bucketMs: Math.ceil(range / 1000), + limit: 1001, + }); + }); + + it('returns the bucketed rows as stored', async () => { + const snapshots = [{ timestamp: 1 }, { timestamp: 2 }]; + storage.getKvCacheFootprintSnapshots.mockResolvedValue(snapshots); + expect(await service.getFootprintHistory('c1', 0, 60_000)).toEqual(snapshots); + }); +}); diff --git a/proprietary/kv-cache/__tests__/kv-cache.controller.spec.ts b/proprietary/kv-cache/__tests__/kv-cache.controller.spec.ts new file mode 100644 index 000000000..d1be571cb --- /dev/null +++ b/proprietary/kv-cache/__tests__/kv-cache.controller.spec.ts @@ -0,0 +1,66 @@ +import { BadRequestException, NotFoundException } from '@nestjs/common'; +import { Reflector } from '@nestjs/core'; +import { Feature } from '@betterdb/shared'; +import { KvCacheController } from '../kv-cache.controller'; + +describe('KvCacheController', () => { + const status = { getStatus: jest.fn().mockResolvedValue({}), getFootprintHistory: jest.fn().mockResolvedValue([]) }; + const footprint = { triggerCollection: jest.fn().mockResolvedValue(null) }; + const registry = { + getDefaultId: jest.fn().mockReturnValue('ui-added'), + get: jest.fn((id: string) => { + if (id === 'nope') throw new NotFoundException(`Connection '${id}' not found`); + return {}; + }), + }; + const controller = new KvCacheController(status as any, footprint as any, registry as any); + + it('defaults to the registry default connection on every route', async () => { + await controller.getStatus(undefined); + expect(status.getStatus).toHaveBeenLastCalledWith('ui-added'); + await controller.refreshFootprint(undefined); + expect(footprint.triggerCollection).toHaveBeenLastCalledWith('ui-added'); + await controller.getFootprintHistory('100', '200', undefined); + expect(status.getFootprintHistory).toHaveBeenLastCalledWith('ui-added', 100, 200); + }); + + it('returns a 404 when no default connection exists', async () => { + registry.getDefaultId.mockReturnValueOnce(null); + await expect(controller.getStatus(undefined)).rejects.toBeInstanceOf(NotFoundException); + }); + + it('rejects an unknown explicit connection id', async () => { + await expect(controller.getStatus('nope')).rejects.toBeInstanceOf(NotFoundException); + }); + + it('refreshes the footprint for the connection', async () => { + await controller.refreshFootprint('c1'); + expect(footprint.triggerCollection).toHaveBeenCalledWith('c1'); + }); + + it('propagates a not-found connection as a 404', async () => { + footprint.triggerCollection.mockRejectedValueOnce(new NotFoundException("Connection 'nope' not found")); + await expect(controller.refreshFootprint('nope')).rejects.toBeInstanceOf(NotFoundException); + }); + + it('parses and validates the history range', async () => { + await controller.getFootprintHistory('100', '200', 'c1'); + expect(status.getFootprintHistory).toHaveBeenCalledWith('c1', 100, 200); + await expect(controller.getFootprintHistory('200', '100', 'c1')).rejects.toBeInstanceOf(BadRequestException); + await expect(controller.getFootprintHistory(undefined, '100', 'c1')).rejects.toBeInstanceOf(BadRequestException); + }); + + it.each([ + ['a fractional timestamp', '100', '1700000000000.5'], + ['a timestamp beyond the safe integer range', '100', '1e19'], + ])('rejects %s with a 400', async (_label, from, to) => { + await expect(controller.getFootprintHistory(from, to, 'c1')).rejects.toBeInstanceOf(BadRequestException); + }); + + it('requires the KV cache feature on every route', () => { + const reflector = new Reflector(); + for (const handler of ['getStatus', 'refreshFootprint', 'getFootprintHistory'] as const) { + expect(reflector.get('requiredFeature', KvCacheController.prototype[handler])).toEqual(Feature.KV_CACHE_MONITORING); + } + }); +}); diff --git a/proprietary/kv-cache/footprint-estimate.ts b/proprietary/kv-cache/footprint-estimate.ts new file mode 100644 index 000000000..00b9e4dfe --- /dev/null +++ b/proprietary/kv-cache/footprint-estimate.ts @@ -0,0 +1,219 @@ +import type { KvCacheFootprintSnapshot, KvCacheLayout, KvCacheModelFootprint } from '@betterdb/shared'; +import { isChunkKey, type ParsedLmcacheKey } from './key-classifier'; + +export interface KeySample { + parsed: ParsedLmcacheKey; + memoryBytes: number | null; + ttl: number | null; + siblingExists: boolean | null; +} + +export interface NodeObservation { + scannedKeys: number; + matchedKeys: ParsedLmcacheKey[]; + scanComplete: boolean; + dbSize: number; + samples: KeySample[]; +} + +export interface MemoryStats { + usedMemory: number; + maxmemory: number; + maxmemoryPolicy: string; + 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; + slotShare: number | null; +} + +export interface SnapshotInput { + connectionId: string; + timestamp: number; + nodes: NodeObservation[]; + memory: MemoryStats; + evictedKeysDelta: number | null; + evictedKeysIntervalMs: number | null; + otherDbs: number[]; + clientDetected: boolean; + coverage?: ScanCoverage; +} + +export const DETECTION_MIN_MATCHES = 5; + +const ratio = (value: number) => Math.round(value * 1e4) / 1e4; + +export function layoutOf(keys: ParsedLmcacheKey[]): KvCacheLayout | null { + if (keys.length === 0) return null; + const suffixed = keys.some((k) => k.suffix !== null); + const plain = keys.some((k) => k.suffix === null); + if (suffixed && plain) return 'mixed'; + return suffixed ? 'two_key' : 'single_key'; +} + +function isPresent(sample: KeySample): boolean { + return sample.ttl !== -2; +} + +function isLive(sample: KeySample): boolean { + return sample.memoryBytes !== null && isPresent(sample); +} + +type LayerCounts = Map; + +function layerCounts(keys: ParsedLmcacheKey[]): LayerCounts { + const layers = new Map>(); + for (const k of keys) { + if (k.layer === null) continue; + const seen = layers.get(k.model) ?? new Set(); + seen.add(k.layer); + layers.set(k.model, seen); + } + return new Map([...layers].map(([model, seen]) => [model, seen.size])); +} + +function chunkWeight(parsed: ParsedLmcacheKey, layers: LayerCounts): number { + if (!isChunkKey(parsed)) return 0; + return parsed.layer === null ? 1 : 1 / (layers.get(parsed.model) ?? 1); +} + +interface ModelTally { + model: string; + dtype: string; + keys: number; + chunks: number; + sampledBytes: number; + sampled: number; +} + +const meanBytes = (samples: KeySample[]): number | null => + samples.length === 0 ? null : samples.reduce((sum, s) => sum + (s.memoryBytes ?? 0), 0) / samples.length; + +const addBytes = (a: number | null, b: number | null) => (a === null || b === null ? null : a + b); + +function tallyFor(tallies: Map, parsed: ParsedLmcacheKey): ModelTally { + const id = `${parsed.model}|${parsed.dtype}`; + const tally = tallies.get(id) ?? { + model: parsed.model, + dtype: parsed.dtype, + keys: 0, + chunks: 0, + sampledBytes: 0, + sampled: 0, + }; + tallies.set(id, tally); + return tally; +} + +function nodeModels(node: NodeObservation, layers: LayerCounts, overallMean: number | null): KvCacheModelFootprint[] { + const scale = node.scanComplete || node.scannedKeys === 0 ? 1 : node.dbSize / node.scannedKeys; + const live = node.samples.filter(isLive); + const nodeMean = meanBytes(live) ?? overallMean; + const tallies = new Map(); + for (const parsed of node.matchedKeys) { + const tally = tallyFor(tallies, parsed); + tally.keys += 1; + tally.chunks += chunkWeight(parsed, layers); + } + for (const sample of live) { + const tally = tallyFor(tallies, sample.parsed); + tally.sampledBytes += sample.memoryBytes ?? 0; + tally.sampled += 1; + } + return [...tallies.values()].map((t) => { + const mean = t.sampled === 0 ? nodeMean : t.sampledBytes / t.sampled; + return { + model: t.model, + dtype: t.dtype, + chunksEst: t.chunks * scale, + bytesEst: mean === null ? null : t.keys * scale * mean, + }; + }); +} + +function scaleModel(m: KvCacheModelFootprint, factor: number): KvCacheModelFootprint { + return { ...m, chunksEst: m.chunksEst * factor, bytesEst: m.bytesEst === null ? null : m.bytesEst * factor }; +} + +function mergeModels(models: KvCacheModelFootprint[]): KvCacheModelFootprint[] { + const merged = new Map(); + for (const m of models) { + const id = `${m.model}|${m.dtype}`; + const total = merged.get(id) ?? { model: m.model, dtype: m.dtype, chunksEst: 0, bytesEst: 0 }; + total.chunksEst += m.chunksEst; + total.bytesEst = addBytes(total.bytesEst, m.bytesEst); + merged.set(id, total); + } + return [...merged.values()] + .map((m) => ({ ...m, chunksEst: Math.round(m.chunksEst), bytesEst: m.bytesEst === null ? null : Math.round(m.bytesEst) })) + .sort((a, b) => (b.bytesEst ?? 0) - (a.bytesEst ?? 0) || b.chunksEst - a.chunksEst); +} + +export function buildSnapshot(input: SnapshotInput): KvCacheFootprintSnapshot { + const allParsed = input.nodes.flatMap((node) => node.matchedKeys); + const allSamples = input.nodes.flatMap((node) => node.samples); + const live = allSamples.filter(isLive); + const present = allSamples.filter(isPresent); + const ttlKnown = present.filter((s) => s.ttl !== null); + const scanned = input.nodes.reduce((sum, node) => sum + node.scannedKeys, 0); + const matched = allParsed.length; + const layers = layerCounts(allParsed); + const scannedTallies = input.nodes.flatMap((node) => nodeModels(node, layers, meanBytes(live))); + const scannedBytes = scannedTallies.reduce((sum, m) => addBytes(sum, m.bytesEst), 0); + const coverage = input.coverage ?? { + primariesScanned: input.nodes.length, + primariesTotal: input.nodes.length, + slotShare: 1, + }; + const share = coverage.slotShare; + const factor = share !== null && share > 0 && share < 1 ? 1 / share : 1; + const nodeTallies = scannedTallies.map((m) => scaleModel(m, factor)); + const chunks = nodeTallies.reduce((sum, m) => sum + m.chunksEst, 0); + const bytes = nodeTallies.reduce((sum, m) => addBytes(sum, m.bytesEst), 0); + const probed = present.filter((s) => s.parsed.suffix !== null && s.siblingExists !== null); + const { usedMemory, maxmemory, maxmemoryPolicy } = input.memory; + return { + connectionId: input.connectionId, + timestamp: input.timestamp, + detected: matched >= DETECTION_MIN_MATCHES || input.clientDetected, + layout: layoutOf(allParsed), + scannedKeys: scanned, + matchedKeys: matched, + sampledKeys: live.length, + scanComplete: input.nodes.every((n) => n.scanComplete), + primariesScanned: coverage.primariesScanned, + primariesTotal: coverage.primariesTotal, + chunksEst: Math.round(chunks), + bytesEst: bytes === null ? null : Math.round(bytes), + usedMemory, + maxmemory, + maxmemoryPolicy, + lmcacheMemoryShare: + scannedBytes === null ? null : usedMemory > 0 ? ratio(Math.min(1, scannedBytes / usedMemory)) : 0, + noTtlRatio: ttlKnown.length === 0 ? null : ratio(ttlKnown.filter((s) => s.ttl === -1).length / ttlKnown.length), + orphanRatio: + probed.length === 0 + ? null + : ratio(probed.filter((s) => s.siblingExists === false).length / probed.length), + evictedKeysDelta: input.evictedKeysDelta, + evictedKeysIntervalMs: input.evictedKeysDelta === null ? null : input.evictedKeysIntervalMs, + otherDbs: input.otherDbs, + perModel: mergeModels(nodeTallies), + }; +} diff --git a/proprietary/kv-cache/footprint-scanner.ts b/proprietary/kv-cache/footprint-scanner.ts new file mode 100644 index 000000000..d0b2f2ffc --- /dev/null +++ b/proprietary/kv-cache/footprint-scanner.ts @@ -0,0 +1,117 @@ +import type { DatabasePort } from '@app/common/interfaces/database-port.interface'; +import type { InfoResponse } from '@app/common/types/metrics.types'; +import { classifyKey, siblingKey, type ParsedLmcacheKey } from './key-classifier'; +import type { KeySample, MemoryStats } from './footprint-estimate'; +import { keySlot, type SlotOwnership } from './key-slot'; + +export interface ScanBudgets { + maxScanned: number; + maxMatched: number; +} + +export interface ScanResult { + scannedKeys: number; + matchedKeys: ParsedLmcacheKey[]; + scanComplete: boolean; +} + +const SCAN_COUNT = 1000; +const GLIDE_LMCACHE = /(?:^| )lib-name=GlidePySync\(lmcache(?::|\))/; +const CLIENT_DB = /(?:^| )db=(\d+)(?: |$)/; + +function isLmcacheClientOn(line: string, db: number): boolean { + if (!GLIDE_LMCACHE.test(line)) return false; + const clientDb = CLIENT_DB.exec(line)?.[1]; + return clientDb === undefined || Number(clientDb) === db; +} + +export async function scanLmcacheKeys(client: DatabasePort, budgets: ScanBudgets, dbSize: number): Promise { + let cursor = '0'; + let iterations = 0; + let returned = 0; + const matched: ParsedLmcacheKey[] = []; + const seen = new Set(); + do { + const [next, keys] = (await client.call('SCAN', [cursor, 'COUNT', String(SCAN_COUNT)])) as [string, string[]]; + cursor = String(next); + iterations += 1; + returned += keys.length; + for (const key of keys) { + if (seen.has(key)) continue; + const parsed = classifyKey(key); + if (!parsed) continue; + seen.add(key); + matched.push(parsed); + } + } while ( + cursor !== '0' && + returned < budgets.maxScanned && + iterations * SCAN_COUNT < budgets.maxScanned && + matched.length < budgets.maxMatched + ); + const walked = cursor === '0' ? dbSize : Math.min(returned, dbSize); + const kept = matched.slice(0, budgets.maxMatched); + const scannedKeys = kept.length < matched.length ? Math.max(1, Math.floor((walked * kept.length) / matched.length)) : walked; + return { + scannedKeys, + matchedKeys: kept, + scanComplete: cursor === '0' && matched.length <= budgets.maxMatched, + }; +} + +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, + }; + }); +} + +export async function hasLmcacheGlideClient(client: DatabasePort, db: number): Promise { + try { + const lines = String(await client.call('CLIENT', ['LIST'])).split('\n'); + return lines.some((line) => isLmcacheClientOn(line.trim(), db)); + } catch { + return false; + } +} + +const num = (value: unknown) => { + const n = Number(value ?? 0); + return Number.isFinite(n) ? n : 0; +}; + +export function readMemoryStats(info: InfoResponse): MemoryStats { + return { + usedMemory: num(info.memory?.used_memory), + maxmemory: num(info.memory?.maxmemory), + maxmemoryPolicy: info.memory?.maxmemory_policy ?? 'unknown', + evictedKeys: num(info.stats?.evicted_keys), + }; +} diff --git a/proprietary/kv-cache/key-classifier.ts b/proprietary/kv-cache/key-classifier.ts new file mode 100644 index 000000000..3b9f555c7 --- /dev/null +++ b/proprietary/kv-cache/key-classifier.ts @@ -0,0 +1,117 @@ +export type LmcacheKeySuffix = 'kv_bytes' | 'metadata'; + +export type LmcacheKeyFormat = 'classic' | 'mp'; + +export interface ParsedLmcacheKey { + key: string; + model: string; + worldSize: number | null; + workerId: number; + chunkHash: string; + dtype: string; + layer: number | null; + tagged: boolean; + suffix: LmcacheKeySuffix | null; + format: LmcacheKeyFormat; +} + +export const MP_UNKNOWN_DTYPE = 'unknown'; + +export const LMCACHE_TORCH_DTYPE_STRINGS: readonly string[] = [ + 'half', + 'bfloat16', + 'float', + 'double', + 'int8', + 'uint8', + 'int16', + 'int32', + 'int64', + 'bool', + 'fp8_e4m3fn', + 'fp8_e4m3fnuz', + 'fp8_e5m2', + 'fp8_e5m2fnuz', +]; + +const KEY_PATTERN = new RegExp( + `^(?.+)@(?\\d+)@(?\\d+)@(?-?[0-9a-f]+)@(?${LMCACHE_TORCH_DTYPE_STRINGS.join('|')})` + + '(?:@(?\\d+))?(?(?:@[^@%]+%[^@]*?)*)(?kv_bytes|metadata)?$', +); + +const MP_KEY_PATTERN = + /^(?[^@].*)@(?[0-9a-f]{8})@(?[0-9a-f]+)@(?(?:[0-9a-f]{2}){8,})(?:@(?[^@/\\\0]{1,128}))?$/; + +function classifyMpKey(key: string): ParsedLmcacheKey | null { + const groups = MP_KEY_PATTERN.exec(key)?.groups; + if (!groups) return null; + const model = groups.lead.slice(groups.lead.lastIndexOf('@') + 1); + if (model === '') return null; + return { + key, + model, + worldSize: null, + workerId: parseInt(groups.rank, 16), + chunkHash: groups.hash, + dtype: MP_UNKNOWN_DTYPE, + layer: null, + tagged: false, + suffix: null, + format: 'mp', + }; +} + +export function classifyKey(key: string): ParsedLmcacheKey | null { + const groups = KEY_PATTERN.exec(key)?.groups; + if (!groups) return classifyMpKey(key); + return { + key, + model: groups.model, + worldSize: Number(groups.ws), + workerId: Number(groups.worker), + chunkHash: groups.hash, + dtype: groups.dtype, + layer: groups.layer === undefined ? null : Number(groups.layer), + tagged: groups.tags !== '', + suffix: (groups.suffix as LmcacheKeySuffix | undefined) ?? null, + format: 'classic', + }; +} + +const isTagOrSuffix = (parsed: ParsedLmcacheKey) => parsed.tagged && parsed.suffix !== null; + +export function siblingKey(parsed: ParsedLmcacheKey): string | null { + if (parsed.suffix === null) return null; + const other: LmcacheKeySuffix = parsed.suffix === 'kv_bytes' ? 'metadata' : 'kv_bytes'; + return parsed.key.slice(0, parsed.key.length - parsed.suffix.length) + other; +} + +export function hasTwoKeyEvidence(keys: ParsedLmcacheKey[]): boolean { + const present = new Set(keys.map((parsed) => parsed.key)); + return keys.some((parsed) => parsed.suffix !== null && (!parsed.tagged || present.has(siblingKey(parsed) as string))); +} + +export function hasSingleKeyEvidence(keys: ParsedLmcacheKey[]): boolean { + return keys.some((parsed) => parsed.suffix === null); +} + +export function resolveTagSuffixes( + keys: ParsedLmcacheKey[], + twoKey = hasTwoKeyEvidence(keys), + singleKey = hasSingleKeyEvidence(keys), +): ParsedLmcacheKey[] { + if (!singleKey || twoKey) return keys; + return keys.map((parsed) => (isTagOrSuffix(parsed) ? { ...parsed, suffix: null } : parsed)); +} + +export function isChunkKey(parsed: ParsedLmcacheKey): boolean { + return parsed.suffix !== 'metadata'; +} + +export function maskKey(parsed: ParsedLmcacheKey): string { + if (parsed.format === 'mp') { + const groups = MP_KEY_PATTERN.exec(parsed.key)?.groups; + return `${parsed.model}@${groups?.rank}@${groups?.group}@${parsed.chunkHash.slice(0, 8)}…`; + } + return `${parsed.model}@${parsed.worldSize}@${parsed.workerId}@${parsed.chunkHash.slice(0, 8)}…@${parsed.dtype}`; +} diff --git a/proprietary/kv-cache/key-slot.ts b/proprietary/kv-cache/key-slot.ts new file mode 100644 index 000000000..e6984770c --- /dev/null +++ b/proprietary/kv-cache/key-slot.ts @@ -0,0 +1,29 @@ +const SLOT_COUNT = 16384; + +function crc16(bytes: Uint8Array): number { + let crc = 0; + for (const byte of bytes) { + crc ^= byte << 8; + for (let bit = 0; bit < 8; bit += 1) { + crc = crc & 0x8000 ? ((crc << 1) ^ 0x1021) & 0xffff : (crc << 1) & 0xffff; + } + } + return crc; +} + +function hashTag(key: string): string { + const open = key.indexOf('{'); + if (open === -1) return key; + const close = key.indexOf('}', open + 1); + return close > open + 1 ? key.slice(open + 1, close) : key; +} + +export function keySlot(key: string): number { + return crc16(Buffer.from(hashTag(key), 'utf8')) % SLOT_COUNT; +} + +export type SlotOwnership = (slot: number) => boolean; + +export function slotOwner(ranges: number[][]): SlotOwnership { + return (slot) => ranges.some((range) => slot >= range[0] && slot <= range[range.length - 1]); +} diff --git a/proprietary/kv-cache/kv-cache-footprint.service.ts b/proprietary/kv-cache/kv-cache-footprint.service.ts new file mode 100644 index 000000000..3c6995b82 --- /dev/null +++ b/proprietary/kv-cache/kv-cache-footprint.service.ts @@ -0,0 +1,456 @@ +import { BadRequestException, Inject, Injectable, Logger, NotFoundException, OnModuleInit } from '@nestjs/common'; +import { + Feature, + type ConnectionStatus, + type DatabaseConnectionConfig, + type KvCacheFootprintSnapshot, +} from '@betterdb/shared'; +import { MultiConnectionPoller, type ConnectionContext } from '@app/common/services/multi-connection-poller'; +import type { StoragePort } from '@app/common/interfaces/storage-port.interface'; +import { downsampleSeries } from '@app/common/utils/downsample'; +import { ConnectionRegistry } from '@app/connections/connection-registry.service'; +import { parseKeyspaceDatabases } from '@app/database/parsers/keyspace-databases'; +import { LicenseService } from '@proprietary/licenses/license.service'; +import { resolvePlan, type CollectionPlan } from '../key-analytics/key-analytics-plan'; +import type { ClusterNode } from '@app/common/types/metrics.types'; +import { clusterScanNodes, primaryScanNodes, type ScanNode } from '../key-analytics/scan-nodes'; +import { + buildSnapshot, + DETECTION_MIN_MATCHES, + evictedKeysDelta, + type MemoryStats, + type NodeObservation, + type ScanCoverage, +} from './footprint-estimate'; +import { + hasLmcacheGlideClient, + readMemoryStats, + sampleLmcacheKeys, + scanLmcacheKeys, + type ScanResult, +} from './footprint-scanner'; +import { hasSingleKeyEvidence, hasTwoKeyEvidence, maskKey, resolveTagSuffixes } from './key-classifier'; +import { slotOwner, type SlotOwnership } from './key-slot'; + +export type SnapshotListener = (snapshot: KvCacheFootprintSnapshot) => void | Promise; + +interface NodeReading { + nodeKey: string; + observation: NodeObservation; + memory: MemoryStats; + otherDbs: number[]; + glide: boolean; +} + +interface NodeScan { + node: ScanNode; + scan: ScanResult & { dbSize: number }; + memory: MemoryStats; + otherDbs: number[]; +} + +interface SlotView { + owners: number; + ownedSlots: number; + slotsByNode: number[]; + rangesByNode: number[][][]; +} + +interface ScanTargets { + nodes: ScanNode[]; + skipped: number; + slots?: SlotView; + replicaSeed?: boolean; +} + +const slotCount = (peer: ClusterNode) => + peer.slots.reduce((sum, range) => sum + (range[range.length - 1] - range[0] + 1), 0); + +function peerAddresses(peer: ClusterNode): string[] { + const hostPort = (peer.address ?? '').split('@')[0]; + const port = hostPort.slice(hostPort.lastIndexOf(':') + 1); + return peer.hostname ? [hostPort, `${peer.hostname}:${port}`] : [hostPort]; +} + +const errorMessage = (error: unknown) => (error instanceof Error ? error.message : String(error)); + +export const MIN_FOOTPRINT_INTERVAL_MS = 60000; +export const MAX_FOOTPRINT_INTERVAL_MS = 86_400_000; +export const PARTIAL_MISSES_BEFORE_CLEAR = 3; +export const REFRESH_REUSE_MS = 10_000; + +export function resolvePositiveInt( + raw: string | undefined, + fallback: number, + floor: number, + onInvalid?: (raw: string) => void, + ceiling = Number.MAX_SAFE_INTEGER, +): number { + if (raw === undefined || raw.trim() === '') return fallback; + const parsed = Number(raw); + if (Number.isInteger(parsed) && parsed >= floor && parsed <= ceiling) return parsed; + onInvalid?.(raw); + return fallback; +} + +@Injectable() +export class KvCacheFootprintService extends MultiConnectionPoller implements OnModuleInit { + protected readonly logger = new Logger(KvCacheFootprintService.name); + private readonly intervalMs = this.setting( + 'KV_CACHE_FOOTPRINT_INTERVAL_MS', + 300000, + MIN_FOOTPRINT_INTERVAL_MS, + MAX_FOOTPRINT_INTERVAL_MS, + ); + private readonly budgets = { + maxScanned: this.setting('KV_CACHE_SCAN_MAX_KEYS', 200000, 1), + maxMatched: this.setting('KV_CACHE_MATCH_MAX_KEYS', 2000, 1), + }; + private readonly sampleSize = this.setting('KV_CACHE_SAMPLE_KEYS', 500, 1); + private readonly undetectedRecheckMs = this.setting( + 'KV_CACHE_UNDETECTED_RECHECK_MS', + 1800000, + MIN_FOOTPRINT_INTERVAL_MS, + ); + private readonly inFlight = new Map>(); + private readonly lastEvicted = new Map }>(); + private unsubscribeRemoval: (() => void) | null = null; + private readonly sampleKeys = new Map(); + private readonly undetectedUntil = new Map(); + private readonly partialMisses = new Map(); + private readonly listeners: SnapshotListener[] = []; + + constructor( + connectionRegistry: ConnectionRegistry, + @Inject('STORAGE_CLIENT') private readonly storage: StoragePort, + private readonly license: LicenseService, + ) { + super(connectionRegistry); + } + + async onModuleInit(): Promise { + this.unsubscribeRemoval = this.connectionRegistry.addConnectionRemovalListener((id) => this.forget(id)); + this.start(); + } + + async onModuleDestroy(): Promise { + this.unsubscribeRemoval?.(); + this.unsubscribeRemoval = null; + await super.onModuleDestroy(); + } + + private setting(name: string, fallback: number, floor: number, ceiling?: number): number { + const range = ceiling === undefined ? `>= ${floor}` : `between ${floor} and ${ceiling}`; + return resolvePositiveInt( + process.env[name], + fallback, + floor, + (raw) => this.logger.warn(`Ignoring ${name}='${raw}': expected an integer ${range}, using ${fallback}`), + ceiling, + ); + } + + protected getIntervalMs(): number { + return this.intervalMs; + } + + protected async pollConnection(ctx: ConnectionContext): Promise { + if (!this.license.hasFeature(Feature.KV_CACHE_MONITORING)) return; + if (Date.now() < (this.undetectedUntil.get(ctx.connectionId) ?? 0)) return; + await this.collect(ctx); + } + + private async forget(connectionId: string): Promise { + await this.inFlight.get(connectionId)?.catch(() => null); + this.lastEvicted.delete(connectionId); + this.sampleKeys.delete(connectionId); + this.undetectedUntil.delete(connectionId); + this.partialMisses.delete(connectionId); + await this.storage + .deleteKvCacheConnectionData(connectionId) + .catch((error) => this.logger.warn(`Could not delete KV cache data for ${connectionId}: ${errorMessage(error)}`)); + } + + onSnapshot(listener: SnapshotListener): void { + this.listeners.push(listener); + } + + getSampleKey(connectionId: string): string | null { + return this.sampleKeys.get(connectionId) ?? null; + } + + async triggerCollection(connectionId: string): Promise { + 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 (running) return running; + const recent = await this.recentSnapshot(connectionId); + if (recent) return recent; + const { config, plan } = target; + const client = this.connectionRegistry.get(connectionId); + return this.collectWith( + { connectionId, connectionName: config.name, client, host: config.host, port: config.port }, + plan, + ); + } + + async collect(ctx: ConnectionContext): Promise { + const target = this.target(ctx.connectionId); + if (!target || target.rejection) return null; + return this.collectWith(ctx, target.plan); + } + + private target( + connectionId: string, + ): { config: DatabaseConnectionConfig; 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) }; + } + + private async recentSnapshot(connectionId: string): Promise { + const [latest] = await this.storage.getKvCacheFootprintSnapshots({ connectionId, limit: 1 }); + return latest && Date.now() - latest.timestamp < REFRESH_REUSE_MS ? latest : null; + } + + private rejection( + connectionId: string, + config: DatabaseConnectionConfig, + connection: ConnectionStatus | undefined, + plan: CollectionPlan, + ): string | null { + 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); + if (running) return running; + const run = this.run(ctx, plan); + this.inFlight.set(ctx.connectionId, run); + try { + return await run; + } finally { + this.inFlight.delete(ctx.connectionId); + } + } + + private async run(ctx: ConnectionContext, plan: CollectionPlan): Promise { + const targets = await this.nodesFor(ctx, plan); + const { nodes, skipped } = targets; + const settled = await this.readAll(nodes, targets.slots?.rangesByNode.map(slotOwner), targets.replicaSeed === true); + const readings: NodeReading[] = []; + const readIndexes: number[] = []; + 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)}`); + }); + 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 timestamp = Date.now(); + const previousEvicted = this.lastEvicted.get(ctx.connectionId); + const built = buildSnapshot({ + connectionId: ctx.connectionId, + timestamp, + nodes: readings.map((reading) => reading.observation), + memory, + evictedKeysDelta: evictedKeysDelta(previousEvicted?.counts ?? null, evicted), + evictedKeysIntervalMs: previousEvicted ? timestamp - previousEvicted.timestamp : null, + 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 keepsEarlier = partial && detected; + const snapshot = partial ? { ...built, scanComplete: false, detected } : built; + await this.storage.saveKvCacheFootprintSnapshot(snapshot); + this.lastEvicted.set(ctx.connectionId, { timestamp, counts: evicted }); + if (built.detected) this.undetectedUntil.delete(ctx.connectionId); + else if (!snapshot.detected && (!partial || this.partialMissesExhausted(ctx.connectionId))) { + this.undetectedUntil.set(ctx.connectionId, snapshot.timestamp + this.undetectedRecheckMs); + } + this.rememberSampleKey(ctx.connectionId, readings, keepsEarlier); + await this.announce(snapshot); + return snapshot; + } + + private async nodesFor(ctx: ConnectionContext, plan: CollectionPlan): Promise { + const seed = { name: ctx.connectionName, client: ctx.client, connectionId: ctx.connectionId }; + if (plan.kind !== 'cluster') { + return this.isClusterMode(seed) ? this.withClusterView([seed]) : { nodes: [seed], skipped: 0 }; + } + const registered = clusterScanNodes(this.connectionRegistry, seed, plan.members, this.logger, 'KV cache'); + const nodes = await primaryScanNodes(registered, this.logger, 'KV cache'); + if (nodes.length === 0) throw new Error(`No reachable primary found for ${ctx.connectionName}`); + return this.withClusterView(nodes); + } + + private async withClusterView(nodes: ScanNode[]): Promise { + const [viewer] = nodes; + try { + const view = await viewer.client.getClusterNodes(); + const self = view.find((peer) => peer.flags.includes('myself')); + const myself = self?.id; + const owners = view.filter((peer) => this.ownsSlots(peer)); + const ids = nodes.map((node, index) => this.ownerId(node, owners, index === 0 ? myself : undefined)); + const scanned = new Set(ids); + const slotsOf = new Map(owners.map((peer) => [peer.id, slotCount(peer)])); + const rangesOf = new Map(owners.map((peer) => [peer.id, peer.slots])); + return { + nodes, + replicaSeed: self?.flags.includes('slave') === true, + skipped: owners.filter((peer) => !scanned.has(peer.id)).length, + slots: { + owners: owners.length, + ownedSlots: owners.reduce((sum, peer) => sum + slotCount(peer), 0), + slotsByNode: ids.map((id) => (id === undefined ? 0 : (slotsOf.get(id) ?? 0))), + rangesByNode: ids.map((id) => (id === undefined ? [] : (rangesOf.get(id) ?? []))), + }, + }; + } catch (error) { + this.logger.warn(`KV cache could not read the cluster view from ${viewer.name}: ${errorMessage(error)}`); + return { nodes, skipped: 1 }; + } + } + + private coverage({ nodes, skipped, slots }: ScanTargets, readIndexes: number[]): ScanCoverage { + const primariesScanned = readIndexes.length; + if (!slots) { + return { + primariesScanned, + primariesTotal: skipped === 0 ? nodes.length : null, + slotShare: skipped === 0 && primariesScanned === nodes.length ? 1 : null, + }; + } + const scannedSlots = readIndexes.reduce((sum, index) => sum + slots.slotsByNode[index], 0); + return { + primariesScanned: readIndexes.filter((index) => slots.slotsByNode[index] > 0).length, + primariesTotal: slots.owners, + slotShare: slots.ownedSlots > 0 && readIndexes.every((index) => slots.slotsByNode[index] > 0) ? scannedSlots / slots.ownedSlots : null, + }; + } + + private ownerId(node: ScanNode, owners: ClusterNode[], myself: string | undefined): string | undefined { + const stored = node.nodeId ?? myself; + const owned = [node.nodeId, myself].find((id) => id !== undefined && owners.some((peer) => peer.id === id)); + if (owned !== undefined) return owned; + const config = node.connectionId ? this.connectionRegistry.getConfig(node.connectionId) : undefined; + const byAddress = config && owners.find((peer) => peerAddresses(peer).includes(`${config.host}:${config.port}`)); + return byAddress?.id ?? stored; + } + + private ownsSlots(peer: ClusterNode): boolean { + return peer.flags.includes('master') && peer.slots.length > 0; + } + + private isClusterMode(node: ScanNode): boolean { + try { + return node.client.getCapabilities().clusterEnabled === true; + } catch { + return false; + } + } + + private async readAll( + nodes: ScanNode[], + ownership?: SlotOwnership[], + replicaSeed = false, + ): Promise[]> { + const scanned = await Promise.allSettled(nodes.map((node) => this.scanNode(node))); + 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; + return Promise.allSettled( + scanned.map((result, index) => + result.status === 'fulfilled' + ? this.sampleNode(result.value, evidence, checkClients, ownership?.[index], replicaSeed && index === 0) + : Promise.reject(result.reason), + ), + ); + } + + private async scanNode(node: ScanNode): Promise { + const [scan, info] = await Promise.all([this.scanKeys(node), node.client.getInfoParsed(['memory', 'stats', 'keyspace'])]); + const ownDb = node.client.getActiveDbIndex?.() ?? 0; + return { + node, + scan, + memory: readMemoryStats(info), + otherDbs: parseKeyspaceDatabases(info.keyspace) + .map(({ db }) => db) + .filter((db) => db !== ownDb), + }; + } + + private async scanKeys(node: ScanNode): Promise { + const dbSize = await node.client.getDbSize(); + return { ...(await scanLmcacheKeys(node.client, this.budgets, dbSize)), dbSize }; + } + + private async sampleNode( + { node, scan, ...rest }: NodeScan, + { twoKey, singleKey }: { twoKey: boolean; singleKey: boolean }, + checkClients: boolean, + ownsSlot?: SlotOwnership, + replica = false, + ): Promise { + const matchedKeys = resolveTagSuffixes(scan.matchedKeys, twoKey, singleKey); + const [samples, glide] = await Promise.all([ + replica ? [] : sampleLmcacheKeys(node.client, downsampleSeries(matchedKeys, this.sampleSize), ownsSlot), + checkClients ? hasLmcacheGlideClient(node.client, node.client.getActiveDbIndex?.() ?? 0) : false, + ]); + return { ...rest, glide, nodeKey: node.connectionId ?? node.name, observation: { ...scan, matchedKeys, samples } }; + } + + private sumMemory(stats: MemoryStats[]): MemoryStats { + const policies = new Set(stats.map((s) => s.maxmemoryPolicy)); + const unlimited = stats.some((s) => s.maxmemory === 0); + return { + usedMemory: stats.reduce((sum, s) => sum + s.usedMemory, 0), + maxmemory: unlimited ? 0 : stats.reduce((sum, s) => sum + s.maxmemory, 0), + maxmemoryPolicy: policies.size > 1 ? 'mixed' : (stats[0]?.maxmemoryPolicy ?? 'unknown'), + evictedKeys: stats.reduce((sum, s) => sum + s.evictedKeys, 0), + }; + } + + 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 { + for (const listener of this.listeners) { + try { + await listener(snapshot); + } catch (error) { + this.logger.warn(`KV cache snapshot listener failed: ${errorMessage(error)}`); + } + } + } +} diff --git a/proprietary/kv-cache/kv-cache-status.service.ts b/proprietary/kv-cache/kv-cache-status.service.ts new file mode 100644 index 000000000..d6739f607 --- /dev/null +++ b/proprietary/kv-cache/kv-cache-status.service.ts @@ -0,0 +1,44 @@ +import { Inject, Injectable } from '@nestjs/common'; +import type { KvCacheEngine, KvCacheFootprintSnapshot, KvCacheStatus } from '@betterdb/shared'; +import type { StoragePort, StoredKvCacheEngine } from '@app/common/interfaces/storage-port.interface'; +import { KvCacheFootprintService } from './kv-cache-footprint.service'; + +export const FOOTPRINT_HISTORY_MAX_RANGE_MS = 31 * 24 * 60 * 60_000; +export const FOOTPRINT_HISTORY_POINTS = 1000; + +export function toPublicEngine(engine: StoredKvCacheEngine): KvCacheEngine { + const { scrapeAuthHeader, scrapeAuthEncrypted: _encrypted, ...rest } = engine; + return { ...rest, hasScrapeAuth: scrapeAuthHeader !== null }; +} + +@Injectable() +export class KvCacheStatusService { + constructor( + @Inject('STORAGE_CLIENT') private readonly storage: StoragePort, + private readonly footprint: KvCacheFootprintService, + ) {} + + async getStatus(connectionId: string): Promise { + const [[latest], engines] = await Promise.all([ + this.storage.getKvCacheFootprintSnapshots({ connectionId, limit: 1 }), + this.storage.getKvCacheEngines(connectionId), + ]); + return { + hasLmcache: latest?.detected === true, + latest: latest ?? null, + sampleKey: this.footprint.getSampleKey(connectionId), + engines: engines.map(toPublicEngine), + }; + } + + async getFootprintHistory(connectionId: string, from: number, to: number): Promise { + const start = Math.max(from, to - FOOTPRINT_HISTORY_MAX_RANGE_MS); + return this.storage.getKvCacheFootprintSnapshots({ + connectionId, + from: start, + to, + bucketMs: Math.max(1, Math.ceil((to - start) / FOOTPRINT_HISTORY_POINTS)), + limit: FOOTPRINT_HISTORY_POINTS + 1, + }); + } +} diff --git a/proprietary/kv-cache/kv-cache.controller.ts b/proprietary/kv-cache/kv-cache.controller.ts new file mode 100644 index 000000000..67c30e90d --- /dev/null +++ b/proprietary/kv-cache/kv-cache.controller.ts @@ -0,0 +1,67 @@ +import { BadRequestException, Controller, Get, Post, Query, UseGuards } from '@nestjs/common'; +import { ApiHeader, ApiOperation, ApiQuery, ApiTags } from '@nestjs/swagger'; +import { Feature, type KvCacheFootprintSnapshot, type KvCacheStatus } from '@betterdb/shared'; +import { LicenseGuard } from '@proprietary/licenses'; +import { RequiresFeature } from '@proprietary/licenses/requires-feature.decorator'; +import { ConnectionId } from '@app/common/decorators'; +import { ConnectionRegistry } from '@app/connections/connection-registry.service'; +import { requireConnectionId } from '@app/connections/require-connection-id'; +import { KvCacheFootprintService } from './kv-cache-footprint.service'; +import { KvCacheStatusService } from './kv-cache-status.service'; + +export function parseRange(from: string | undefined, to: string | undefined): { from: number; to: number } { + const fromMs = Number(from); + const toMs = Number(to); + if (!from || !to || !Number.isSafeInteger(fromMs) || !Number.isSafeInteger(toMs) || fromMs > toMs) { + throw new BadRequestException('from and to must be millisecond timestamps with from <= to'); + } + return { from: fromMs, to: toMs }; +} + +@ApiTags('kv-cache') +@Controller('kv-cache') +export class KvCacheController { + constructor( + private readonly status: KvCacheStatusService, + private readonly footprint: KvCacheFootprintService, + private readonly connectionRegistry: ConnectionRegistry, + ) {} + + @Get('status') + @UseGuards(LicenseGuard) + @RequiresFeature(Feature.KV_CACHE_MONITORING) + @ApiOperation({ summary: 'LMCache detection, latest footprint and linked engines (Pro)' }) + @ApiHeader({ name: 'x-connection-id', required: false }) + async getStatus(@ConnectionId() connectionId?: string): Promise { + return this.status.getStatus(requireConnectionId(this.connectionRegistry, connectionId)); + } + + @Post('footprint/refresh') + @UseGuards(LicenseGuard) + @RequiresFeature(Feature.KV_CACHE_MONITORING) + @ApiOperation({ summary: 'Collect the LMCache footprint now (Pro)' }) + @ApiHeader({ name: 'x-connection-id', required: false }) + async refreshFootprint(@ConnectionId() connectionId?: string): Promise { + return this.footprint.triggerCollection(requireConnectionId(this.connectionRegistry, connectionId)); + } + + @Get('footprint/history') + @UseGuards(LicenseGuard) + @RequiresFeature(Feature.KV_CACHE_MONITORING) + @ApiOperation({ summary: 'Stored LMCache footprint snapshots (Pro)' }) + @ApiQuery({ name: 'from', required: true, type: Number }) + @ApiQuery({ name: 'to', required: true, type: Number }) + @ApiHeader({ name: 'x-connection-id', required: false }) + async getFootprintHistory( + @Query('from') from: string | undefined, + @Query('to') to: string | undefined, + @ConnectionId() connectionId?: string, + ): Promise { + const range = parseRange(from, to); + return this.status.getFootprintHistory( + requireConnectionId(this.connectionRegistry, connectionId), + range.from, + range.to, + ); + } +} diff --git a/proprietary/kv-cache/kv-cache.module.ts b/proprietary/kv-cache/kv-cache.module.ts new file mode 100644 index 000000000..5855da77a --- /dev/null +++ b/proprietary/kv-cache/kv-cache.module.ts @@ -0,0 +1,14 @@ +import { Module } from '@nestjs/common'; +import { ConnectionsModule } from '@app/connections/connections.module'; +import { StorageModule } from '@app/storage/storage.module'; +import { LicenseModule } from '@proprietary/licenses'; +import { KvCacheController } from './kv-cache.controller'; +import { KvCacheFootprintService } from './kv-cache-footprint.service'; +import { KvCacheStatusService } from './kv-cache-status.service'; + +@Module({ + imports: [StorageModule, ConnectionsModule, LicenseModule], + controllers: [KvCacheController], + providers: [KvCacheFootprintService, KvCacheStatusService], +}) +export class KvCacheModule {}