diff --git a/apps/api/package.json b/apps/api/package.json index 37b33d3b5..f8176cf45 100644 --- a/apps/api/package.json +++ b/apps/api/package.json @@ -40,7 +40,9 @@ "docs:index:valkey": "SOURCE=valkey ts-node -r tsconfig-paths/register scripts/index-docs.ts", "docs:index:redis": "SOURCE=redis ts-node -r tsconfig-paths/register scripts/index-docs.ts", "docs:clean": "rm -rf data/docs data/chunks data/vector-store", - "seed:demo": "ts-node -r tsconfig-paths/register scripts/seed-demo-data.ts" + "seed:demo": "ts-node -r tsconfig-paths/register scripts/seed-demo-data.ts", + "cache-sim:seed": "ts-node -r tsconfig-paths/register scripts/cache-simulator/seed.ts", + "cache-sim:validate": "ts-node -r tsconfig-paths/register scripts/cache-simulator/validate.ts" }, "dependencies": { "@betterdb/agent-memory": "workspace:*", diff --git a/apps/api/scripts/cache-simulator/scenarios.ts b/apps/api/scripts/cache-simulator/scenarios.ts new file mode 100644 index 000000000..831505e76 --- /dev/null +++ b/apps/api/scripts/cache-simulator/scenarios.ts @@ -0,0 +1,108 @@ +export interface SeedCommand { + type: 'string' | 'hash' | 'list' | 'set' | 'zset'; + key: string; + value: string | string[]; + ttlSeconds?: number; +} + +export interface Scenario { + name: string; + tolerance: number | null; + runs: number; + generate(seed: number): Iterable; +} + +export function mulberry32(seed: number): () => number { + let state = seed >>> 0; + return () => { + state = (state + 0x6d2b79f5) >>> 0; + let t = state; + t = Math.imul(t ^ (t >>> 15), t | 1); + t ^= t + Math.imul(t ^ (t >>> 7), t | 61); + return ((t ^ (t >>> 14)) >>> 0) / 4294967296; + }; +} + +const between = (random: () => number, min: number, max: number): number => min + Math.floor(random() * (max - min + 1)); +const text = (length: number): string => 'x'.repeat(length); +const items = (count: number, length: number, prefix: string): string[] => Array.from({ length: count }, (_, index) => `${prefix}${index}:${text(length)}`); + +function* mixed(random: () => number, prefix: string, counts: { string: number; hash: number; list: number; set: number; zset: number }): Iterable { + for (let index = 0; index < counts.string; index++) yield { type: 'string', key: `${prefix}:s:${index}`, value: text(between(random, 32, 512)) }; + for (let index = 0; index < counts.hash; index++) yield { type: 'hash', key: `${prefix}:h:${index}`, value: items(between(random, 10, 50), 32, 'f') }; + for (let index = 0; index < counts.list; index++) yield { type: 'list', key: `${prefix}:l:${index}`, value: items(between(random, 20, 200), 64, 'i') }; + for (let index = 0; index < counts.set; index++) yield { type: 'set', key: `${prefix}:t:${index}`, value: items(between(random, 10, 100), 32, 'm') }; + for (let index = 0; index < counts.zset; index++) yield { type: 'zset', key: `${prefix}:z:${index}`, value: items(between(random, 10, 100), 32, 'm') }; +} + +export const SCENARIOS: Scenario[] = [ + { + name: 'homogeneous-small', + tolerance: 0.05, + runs: 10, + *generate(seed) { + for (let index = 0; index < 100_000; index++) yield { type: 'string', key: `hs:${seed}:${index}`, value: text(64) }; + }, + }, + { + name: 'homogeneous-large', + tolerance: 0.05, + runs: 10, + *generate(seed) { + for (let index = 0; index < 20_000; index++) yield { type: 'string', key: `hl:${seed}:${index}`, value: text(16 * 1024) }; + }, + }, + { + name: 'mixed', + tolerance: 0.05, + runs: 10, + generate: (seed) => mixed(mulberry32(seed), `mx:${seed}`, { string: 50_000, hash: 10_000, list: 5_000, set: 5_000, zset: 5_000 }), + }, + { + name: 'pareto', + tolerance: 0.1, + runs: 10, + *generate(seed) { + const random = mulberry32(seed); + for (let index = 0; index < 50_000; index++) { + const u = 1 - random(); + yield { type: 'string', key: `pa:${seed}:${index}`, value: text(Math.min(1024 * 1024, Math.floor(100 / u ** (1 / 1.5)))) }; + } + }, + }, + { + name: 'ttl-mix', + tolerance: 0.05, + runs: 10, + *generate(seed) { + const random = mulberry32(seed); + for (let index = 0; index < 50_000; index++) { + yield { type: 'string', key: `tt:${seed}:${index}`, value: text(256), ...(random() < 0.5 ? { ttlSeconds: 3600 } : {}) }; + } + }, + }, + { + name: 'rare-giant', + tolerance: null, + runs: 10, + *generate(seed) { + const giant = Math.floor(mulberry32(seed)() * 50_000); + for (let index = 0; index < 50_000; index++) { + yield { type: 'string', key: `rg:${seed}:${index}`, value: text(index === giant ? 64 * 1024 * 1024 : 100) }; + } + }, + }, + { + name: 'scale-500k', + tolerance: 0.05, + runs: 3, + *generate(seed) { + const random = mulberry32(seed); + for (let index = 0; index < 500_000; index++) { + yield random() < 0.8 + ? { type: 'string', key: `sc:${seed}:${index}`, value: text(between(random, 100, 300)) } + : { type: 'hash', key: `sc:${seed}:${index}`, value: items(between(random, 5, 20), 32, 'f') }; + } + }, + }, +]; diff --git a/apps/api/scripts/cache-simulator/seed.ts b/apps/api/scripts/cache-simulator/seed.ts new file mode 100644 index 000000000..326195fea --- /dev/null +++ b/apps/api/scripts/cache-simulator/seed.ts @@ -0,0 +1,76 @@ +import Valkey from 'iovalkey'; +import { SCENARIOS, type SeedCommand } from './scenarios'; + +export interface Target { + url: string; + db: number; +} + +export function parseTarget(argv: string[]): Target & { scenario?: string; seed: number } { + const value = (flag: string): string | undefined => { + const index = argv.indexOf(flag); + return index >= 0 ? argv[index + 1] : undefined; + }; + const url = value('--url'); + const db = Number(value('--db')); + if (!url || !Number.isInteger(db) || db <= 0) { + throw new Error('Pass --url and --db ; the target database must be empty'); + } + return { url, db, scenario: value('--scenario'), seed: Number(value('--seed') ?? 1) }; +} + +export async function connect(target: Target): Promise { + const client = new Valkey(target.url, { db: target.db }); + if ((await client.dbsize()) !== 0) { + client.disconnect(); + throw new Error(`Database ${target.db} is not empty; refusing to seed`); + } + return client; +} + +async function flush(pipeline: ReturnType): Promise { + for (const [error] of (await pipeline.exec()) ?? []) if (error) throw error; +} + +export async function seed(client: Valkey, commands: Iterable): Promise { + let pipeline = client.pipeline(); + let queued = 0; + let total = 0; + for (const command of commands) { + if (command.type === 'string') pipeline.set(command.key, command.value as string); + if (command.type === 'hash') pipeline.hset(command.key, Object.fromEntries((command.value as string[]).map((field) => [field, field]))); + if (command.type === 'list') pipeline.rpush(command.key, ...(command.value as string[])); + if (command.type === 'set') pipeline.sadd(command.key, ...(command.value as string[])); + if (command.type === 'zset') pipeline.zadd(command.key, ...(command.value as string[]).flatMap((member, index) => [String(index), member])); + if (command.ttlSeconds !== undefined) pipeline.expire(command.key, command.ttlSeconds); + queued += 1; + total += 1; + if (queued === 1000) { + await flush(pipeline); + pipeline = client.pipeline(); + queued = 0; + } + } + if (queued > 0) await flush(pipeline); + return total; +} + +async function main(): Promise { + const target = parseTarget(process.argv.slice(2)); + const scenario = SCENARIOS.find((candidate) => candidate.name === target.scenario); + if (!scenario) throw new Error(`Pass --scenario with one of: ${SCENARIOS.map((candidate) => candidate.name).join(', ')}`); + const client = await connect(target); + try { + const count = await seed(client, scenario.generate(target.seed)); + console.log(`Seeded ${count} keys for ${scenario.name} into db ${target.db}`); + } finally { + client.disconnect(); + } +} + +if (require.main === module) { + main().catch((error) => { + console.error(error instanceof Error ? error.message : error); + process.exit(1); + }); +} diff --git a/apps/api/scripts/cache-simulator/validate.ts b/apps/api/scripts/cache-simulator/validate.ts new file mode 100644 index 000000000..3bae9bb61 --- /dev/null +++ b/apps/api/scripts/cache-simulator/validate.ts @@ -0,0 +1,69 @@ +import type Valkey from 'iovalkey'; +import { estimate, keyMemoryCeiling } from '../../../../proprietary/cache-simulator/size-estimator'; +import { sampleNodes, SAMPLE_BUDGET, SAMPLE_DEADLINE_MS, SAMPLE_HARD_CAP } from '../../../../proprietary/cache-simulator/stratified-sampler'; +import type { SamplerClient } from '../../../../proprietary/cache-simulator/sample-types'; +import { SCENARIOS } from './scenarios'; +import { connect, parseTarget, seed } from './seed'; + +async function groundTruth(client: Valkey): Promise { + let cursor = '0'; + let total = 0; + do { + const [next, keys] = await client.scan(cursor, 'COUNT', 1000); + cursor = next; + const pipeline = client.pipeline(); + keys.forEach((key) => pipeline.call('MEMORY', 'USAGE', key, 'SAMPLES', '0')); + for (const [error, size] of (await pipeline.exec()) ?? []) if (!error && typeof size === 'number') total += size; + } while (cursor !== '0'); + return total; +} + +async function memoryCeiling(client: Valkey): Promise { + const match = /used_memory_dataset:(\d+)/.exec(await client.info('memory')); + return keyMemoryCeiling(match ? Number(match[1]) : 0, await client.dbsize()); +} + +function adapter(client: Valkey): SamplerClient { + return { + call: (command: string, args: string[]) => client.call(command, ...args), + getClient: () => client, + getDbSize: () => client.dbsize(), + } as unknown as SamplerClient; +} + +async function main(): Promise { + const target = parseTarget(process.argv.slice(2)); + const client = await connect(target); + let failed = false; + try { + for (const scenario of SCENARIOS.filter((candidate) => !target.scenario || candidate.name === target.scenario)) { + let covered = 0; + let worst = 0; + let widest = 0; + for (let run = 1; run <= scenario.runs; run++) { + await seed(client, scenario.generate(run)); + const truth = await groundTruth(client); + const ceiling = await memoryCeiling(client); + const result = estimate(await sampleNodes([{ name: 'validate', client: adapter(client) }], { budget: SAMPLE_BUDGET, deadlineMs: SAMPLE_DEADLINE_MS, hardCap: SAMPLE_HARD_CAP }), Date.now(), ceiling); + const error = Math.abs(result.totalBytes - truth) / truth; + worst = Math.max(worst, error); + widest = Math.max(widest, result.ci95Bytes / truth); + if (Math.abs(result.totalBytes - truth) <= result.ci95Bytes + 0.001 * truth) covered += 1; + await client.flushdb(); + } + const required = scenario.runs - Math.floor(scenario.runs / 10); + const pass = (scenario.tolerance === null || worst <= scenario.tolerance) && covered >= required; + failed ||= !pass; + const limit = scenario.tolerance === null ? 'coverage only' : `limit ${scenario.tolerance * 100}%`; + console.log(`${pass ? 'PASS' : 'FAIL'} ${scenario.name}: worst error ${(worst * 100).toFixed(2)}% (${limit}), CI coverage ${covered}/${scenario.runs} (need ${required}), widest CI ±${(widest * 100).toFixed(2)}%`); + } + } finally { + client.disconnect(); + } + if (failed) process.exit(1); +} + +main().catch((error) => { + console.error(error instanceof Error ? error.message : error); + process.exit(1); +}); diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index dfba8f87b..1f2dc633c 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -32,3 +32,5 @@ export * from './types/otel-traces'; export * from './types/cve'; export * from './utils/cve-modules'; export * from './types/workspace'; +export * from './types/cache-simulator.types'; +export * from './utils/cache-simulator'; diff --git a/packages/shared/src/types/cache-simulator.types.ts b/packages/shared/src/types/cache-simulator.types.ts new file mode 100644 index 000000000..58606cd52 --- /dev/null +++ b/packages/shared/src/types/cache-simulator.types.ts @@ -0,0 +1,42 @@ +export type CacheSimulatorStatus = 'sampling' | 'ready' | 'failed'; + +export interface CacheSimulatorInstant { + usedMemory: number; + maxmemory: number; + maxmemoryPolicy: string; + totalKeys: number; +} + +export interface CacheSimulatorTypeEstimate { + type: string; + keys: number; + avgBytes: number; + stdBytes: number; + sampled: number; +} + +export interface CacheSimulatorEstimate { + totalBytes: number; + ci95Bytes: number; + avgKeyBytes: number; + stdKeyBytes: number; + sampleCount: number; + byType: CacheSimulatorTypeEstimate[]; + sampledAt: number; + partial: boolean; +} + +export interface CacheSimulatorGrowth { + keysPerDay: number; + bytesPerDay: number; + days: number; +} + +export interface CacheSimulatorSnapshot { + status: CacheSimulatorStatus; + progress: { sampled: number; target: number } | null; + instant: CacheSimulatorInstant; + estimate: CacheSimulatorEstimate | null; + growth: CacheSimulatorGrowth | null; + error?: string; +} diff --git a/packages/shared/src/utils/cache-simulator/eviction-copy.ts b/packages/shared/src/utils/cache-simulator/eviction-copy.ts new file mode 100644 index 000000000..d26049421 --- /dev/null +++ b/packages/shared/src/utils/cache-simulator/eviction-copy.ts @@ -0,0 +1,21 @@ +const COPY: Record = { + noeviction: 'Writes start failing with OOM errors once memory is full; reads keep working.', + 'allkeys-lru': 'The least recently used keys are evicted to make room, whether or not they have a TTL.', + 'allkeys-lfu': 'The least frequently used keys are evicted to make room, whether or not they have a TTL.', + 'allkeys-random': 'Random keys are evicted to make room, whether or not they have a TTL.', + 'allkeys-lrm': 'The least recently modified keys are evicted to make room, whether or not they have a TTL.', + 'volatile-lru': 'The least recently used keys that have a TTL are evicted; once none are left, writes fail with OOM errors.', + 'volatile-lfu': 'The least frequently used keys that have a TTL are evicted; once none are left, writes fail with OOM errors.', + 'volatile-random': 'Random keys that have a TTL are evicted; once none are left, writes fail with OOM errors.', + 'volatile-ttl': 'Keys closest to expiring are evicted first; once no keys with a TTL are left, writes fail with OOM errors.', + 'volatile-lrm': 'The least recently modified keys that have a TTL are evicted; once none are left, writes fail with OOM errors.', +}; + +export const EVICTION_POLICIES = Object.keys(COPY); + +export const UNKNOWN_EVICTION_COPY = 'This eviction policy is not recognised, so its behaviour at maxmemory cannot be described.'; + +export function evictionConsequence(policy: string): string { + const key = policy.toLowerCase(); + return Object.prototype.hasOwnProperty.call(COPY, key) ? COPY[key] : UNKNOWN_EVICTION_COPY; +} diff --git a/packages/shared/src/utils/cache-simulator/index.ts b/packages/shared/src/utils/cache-simulator/index.ts new file mode 100644 index 000000000..e5406c04d --- /dev/null +++ b/packages/shared/src/utils/cache-simulator/index.ts @@ -0,0 +1,3 @@ +export * from './projection'; +export * from './node-catalog'; +export * from './eviction-copy'; diff --git a/packages/shared/src/utils/cache-simulator/node-catalog.ts b/packages/shared/src/utils/cache-simulator/node-catalog.ts new file mode 100644 index 000000000..bb8eab6e7 --- /dev/null +++ b/packages/shared/src/utils/cache-simulator/node-catalog.ts @@ -0,0 +1,142 @@ +import { HEADROOM } from './projection'; + +export type CloudProvider = 'aws' | 'gcp' | 'azure'; + +export interface CloudNode { + nodeType: string; + memoryBytes: number; +} + +export interface CloudService { + provider: CloudProvider; + service: string; + label: string; + maxShards: number; + usableRatio: number; + sourceUrl: string; + asOf: string; + nodes: CloudNode[]; +} + +export interface NodeRecommendation { + provider: CloudProvider; + service: string; + label: string; + nodeType: string; + memoryBytes: number; + shards: number; +} + +const gib = (value: number): number => Math.round(value * 1024 ** 3); +const gb = (value: number): number => Math.round(value * 1e9); +const node = (nodeType: string, memoryBytes: number): CloudNode => ({ nodeType, memoryBytes }); +const AS_OF = '2026-10-12'; + +const R7G: Array<[string, number]> = [ + ['large', 13.07], + ['xlarge', 26.32], + ['2xlarge', 52.82], + ['4xlarge', 105.81], + ['8xlarge', 209.55], + ['12xlarge', 317.77], + ['16xlarge', 419.09], +]; + +const M7G: Array<[string, number]> = [ + ['large', 6.38], + ['xlarge', 12.93], + ['2xlarge', 26.04], + ['4xlarge', 52.26], + ['8xlarge', 103.68], + ['12xlarge', 157.12], + ['16xlarge', 209.55], +]; + +export const CLOUD_SERVICES: CloudService[] = [ + { + provider: 'aws', + service: 'elasticache', + label: 'Amazon ElastiCache', + maxShards: 500, + usableRatio: 0.75, + sourceUrl: 'https://aws.amazon.com/elasticache/pricing/', + asOf: AS_OF, + nodes: [ + node('cache.t4g.micro', gib(0.5)), + node('cache.t4g.small', gib(1.37)), + node('cache.t4g.medium', gib(3.09)), + ...M7G.map(([size, memory]) => node(`cache.m7g.${size}`, gib(memory))), + ...R7G.map(([size, memory]) => node(`cache.r7g.${size}`, gib(memory))), + ], + }, + { + provider: 'aws', + service: 'memorydb', + label: 'Amazon MemoryDB', + maxShards: 500, + usableRatio: 1, + sourceUrl: 'https://aws.amazon.com/memorydb/pricing/', + asOf: AS_OF, + nodes: [ + node('db.t4g.small', gib(1.37)), + node('db.t4g.medium', gib(3.09)), + ...R7G.map(([size, memory]) => node(`db.r7g.${size}`, gib(memory))), + ], + }, + { + provider: 'gcp', + service: 'memorystore-valkey', + label: 'Memorystore for Valkey (1 replica per shard)', + maxShards: 125, + usableRatio: 0.8, + sourceUrl: 'https://cloud.google.com/memorystore/docs/valkey/instance-node-specification', + asOf: AS_OF, + nodes: [ + node('shared-core-nano', gb(1.4)), + node('standard-small', gb(6.5)), + node('highmem-medium', gb(13)), + node('highmem-xlarge', gb(58)), + ], + }, + { + provider: 'azure', + service: 'azure-cache-standard', + label: 'Azure Cache for Redis (Basic/Standard, legacy)', + maxShards: 1, + usableRatio: 0.8, + sourceUrl: 'https://azure.microsoft.com/pricing/details/cache/', + asOf: AS_OF, + nodes: [ + node('C0', gb(0.25)), + node('C1', gb(1)), + node('C2', gb(2.5)), + node('C3', gb(6)), + node('C4', gb(13)), + node('C5', gb(26)), + node('C6', gb(53)), + ], + }, + { + provider: 'azure', + service: 'azure-cache-premium', + label: 'Azure Cache for Redis (Premium, legacy)', + maxShards: 10, + usableRatio: 0.8, + sourceUrl: 'https://azure.microsoft.com/pricing/details/cache/', + asOf: AS_OF, + nodes: [node('P1', gb(6)), node('P2', gb(13)), node('P3', gb(26)), node('P4', gb(53)), node('P5', gb(120))], + }, +]; + +export function smallestNodeFor(bytes: number, provider: CloudProvider, headroom: number = HEADROOM): NodeRecommendation[] { + return CLOUD_SERVICES.filter((service) => service.provider === provider).flatMap((service) => { + const sorted = [...service.nodes].sort((a, b) => a.memoryBytes - b.memoryBytes); + const base = { provider, service: service.service, label: service.label }; + const capacity = (candidate: CloudNode) => candidate.memoryBytes * service.usableRatio * headroom; + const single = sorted.find((candidate) => capacity(candidate) >= bytes); + if (single) return [{ ...base, nodeType: single.nodeType, memoryBytes: single.memoryBytes, shards: 1 }]; + const largest = sorted[sorted.length - 1]; + const shards = Math.ceil(bytes / capacity(largest)); + return shards <= service.maxShards ? [{ ...base, nodeType: largest.nodeType, memoryBytes: largest.memoryBytes, shards }] : []; + }); +} diff --git a/packages/shared/src/utils/cache-simulator/projection.ts b/packages/shared/src/utils/cache-simulator/projection.ts new file mode 100644 index 000000000..3546b8d48 --- /dev/null +++ b/packages/shared/src/utils/cache-simulator/projection.ts @@ -0,0 +1,68 @@ +export const DEFAULT_OVERHEAD_PER_KEY_BYTES = 88; +export const HEADROOM = 0.8; + +export interface ProjectionInputs { + currentKeys: number; + keysPerDay: number; + inPlaceBytesPerDay: number; + ttlDays: number | null; + avgKeyBytes: number; + ci95Bytes: number; + overheadPerKeyBytes: number; + fixedOverheadBytes: number; + maxmemory: number | null; + horizonDays: number; +} + +export interface ProjectionPoint { + day: number; + keys: number; + bytes: number; + low: number; + high: number; +} + +export interface Projection { + points: ProjectionPoint[]; + daysToMaxmemory: number | null; + daysTo80Percent: number | null; +} + +export function fixedOverhead(usedMemory: number, keys: number, avgKeyBytes: number, overheadPerKeyBytes: number): number { + return Math.max(0, usedMemory - keys * (avgKeyBytes + overheadPerKeyBytes)); +} + +export function keysAt(day: number, inputs: Pick): number { + const { currentKeys, keysPerDay, ttlDays } = inputs; + if (ttlDays !== null && ttlDays > 0) { + const steady = Math.max(0, keysPerDay * ttlDays); + return Math.max(0, steady + (currentKeys - steady) * Math.exp(-day / ttlDays)); + } + return Math.max(0, currentKeys + keysPerDay * day); +} + +function firstDayAtOrAbove(points: ProjectionPoint[], limit: number | null): number | null { + if (limit === null || limit <= 0) return null; + return points.find((point) => point.bytes >= limit)?.day ?? null; +} + +export function inPlaceGrowth(growth: { keysPerDay: number; bytesPerDay: number }, perKeyBytes: number): number { + return Math.max(0, growth.bytesPerDay - Math.max(0, growth.keysPerDay) * perKeyBytes); +} + +export function project(inputs: ProjectionInputs): Projection { + const perKey = inputs.avgKeyBytes + inputs.overheadPerKeyBytes; + const margin = inputs.currentKeys > 0 ? inputs.ci95Bytes / inputs.currentKeys : 0; + const points: ProjectionPoint[] = []; + for (let day = 0; day <= inputs.horizonDays; day++) { + const keys = keysAt(day, inputs); + const bytes = keys * perKey + inputs.fixedOverheadBytes + Math.max(0, inputs.inPlaceBytesPerDay) * day; + points.push({ day, keys, bytes, low: Math.max(0, bytes - margin * keys), high: bytes + margin * keys }); + } + const limit = inputs.maxmemory && inputs.maxmemory > 0 ? inputs.maxmemory : null; + return { + points, + daysToMaxmemory: firstDayAtOrAbove(points, limit), + daysTo80Percent: firstDayAtOrAbove(points, limit === null ? null : limit * HEADROOM), + }; +} diff --git a/proprietary/cache-simulator/__tests__/allocation.spec.ts b/proprietary/cache-simulator/__tests__/allocation.spec.ts new file mode 100644 index 000000000..e89f34358 --- /dev/null +++ b/proprietary/cache-simulator/__tests__/allocation.spec.ts @@ -0,0 +1,61 @@ +import { neymanAllocation, splitBudget, STRATUM_FLOOR } from '../allocation'; +import { mean, sampleVariance } from '../sample-stats'; + +describe('sample stats', () => { + it('computes mean and n-1 variance', () => { + expect(mean([90, 110])).toBe(100); + expect(sampleVariance([90, 110])).toBe(200); + }); + + it('returns zero for too few values', () => { + expect(mean([])).toBe(0); + expect(sampleVariance([5])).toBe(0); + }); +}); + +describe('splitBudget', () => { + it('splits in proportion to key counts', () => { + expect(splitBudget(4000, [100, 300, 0])).toEqual([1000, 3000, 0]); + }); + + it('gives every non-empty node at least one sample', () => { + expect(splitBudget(10, [1, 100_000])).toEqual([1, 10]); + }); + + it('returns zeros for an empty keyspace', () => { + expect(splitBudget(4000, [0, 0])).toEqual([0, 0]); + }); +}); + +describe('neymanAllocation', () => { + it('weights by count times spread', () => { + const result = neymanAllocation(300, [ + { type: 'a', estimatedKeys: 1000, std: 10, sampled: 50 }, + { type: 'b', estimatedKeys: 1000, std: 20, sampled: 50 }, + ]); + expect(result.get('a')).toBe(100); + expect(result.get('b')).toBe(200); + }); + + it('falls back to counts when every stratum has zero spread', () => { + const result = neymanAllocation(100, [ + { type: 'a', estimatedKeys: 300, std: 0, sampled: 0 }, + { type: 'b', estimatedKeys: 100, std: 0, sampled: 0 }, + ]); + expect(result.get('a')).toBe(75); + expect(result.get('b')).toBe(30); + }); + + it('applies the floor to zero-spread strata', () => { + const result = neymanAllocation(100, [ + { type: 'a', estimatedKeys: 1000, std: 100, sampled: 10 }, + { type: 'b', estimatedKeys: 50, std: 0, sampled: 5 }, + ]); + expect(result.get('a')).toBe(100); + expect(result.get('b')).toBe(STRATUM_FLOOR - 5); + }); + + it('never asks for more keys than a stratum holds', () => { + expect(neymanAllocation(100, [{ type: 'a', estimatedKeys: 10, std: 5, sampled: 4 }]).get('a')).toBe(6); + }); +}); diff --git a/proprietary/cache-simulator/__tests__/eviction-copy.spec.ts b/proprietary/cache-simulator/__tests__/eviction-copy.spec.ts new file mode 100644 index 000000000..ff814c8ea --- /dev/null +++ b/proprietary/cache-simulator/__tests__/eviction-copy.spec.ts @@ -0,0 +1,36 @@ +import { EVICTION_POLICIES, UNKNOWN_EVICTION_COPY, evictionConsequence } from '@betterdb/shared'; + +describe('eviction copy', () => { + it('describes every known policy', () => { + expect([...EVICTION_POLICIES].sort()).toEqual( + [ + 'allkeys-lfu', + 'allkeys-lrm', + 'allkeys-lru', + 'allkeys-random', + 'noeviction', + 'volatile-lfu', + 'volatile-lrm', + 'volatile-lru', + 'volatile-random', + 'volatile-ttl', + ].sort(), + ); + for (const policy of EVICTION_POLICIES) expect(evictionConsequence(policy)).not.toBe(UNKNOWN_EVICTION_COPY); + }); + + it('explains noeviction as failing writes', () => { + expect(evictionConsequence('noeviction')).toBe('Writes start failing with OOM errors once memory is full; reads keep working.'); + }); + + it('is case-insensitive and falls back for unknown policies', () => { + expect(evictionConsequence('ALLKEYS-LRU')).toBe(evictionConsequence('allkeys-lru')); + expect(evictionConsequence('something-new')).toBe(UNKNOWN_EVICTION_COPY); + }); + + it('falls back for inherited object keys', () => { + expect(evictionConsequence('constructor')).toBe(UNKNOWN_EVICTION_COPY); + expect(evictionConsequence('__proto__')).toBe(UNKNOWN_EVICTION_COPY); + expect(evictionConsequence('toString')).toBe(UNKNOWN_EVICTION_COPY); + }); +}); diff --git a/proprietary/cache-simulator/__tests__/fake-keyspace.ts b/proprietary/cache-simulator/__tests__/fake-keyspace.ts new file mode 100644 index 000000000..2d6389f32 --- /dev/null +++ b/proprietary/cache-simulator/__tests__/fake-keyspace.ts @@ -0,0 +1,63 @@ +import type { SamplerClient } from '../sample-types'; + +export interface FakeKey { + type: string; + size: number | null; +} + +export interface FakeOptions { + memoryError?: string; + scanTypeError?: string; + seed?: number; +} + +function shuffled(names: string[], seed: number): string[] { + let state = seed >>> 0; + const next = () => { + state = (Math.imul(state, 1664525) + 1013904223) >>> 0; + return state / 4294967296; + }; + const result = [...names]; + for (let index = result.length - 1; index > 0; index--) { + const swap = Math.floor(next() * (index + 1)); + [result[index], result[swap]] = [result[swap], result[index]]; + } + return result; +} + +export function fakeClient(keys: Map, options: FakeOptions = {}): SamplerClient { + const names = shuffled([...keys.keys()], options.seed ?? 1); + const call = async (command: string, args: string[]): Promise => { + if (command !== 'SCAN') throw new Error(`unexpected ${command}`); + const [cursor, , count, typeFlag, type] = args; + if (typeFlag === 'TYPE' && options.scanTypeError) throw new Error(options.scanTypeError); + const start = Number(cursor); + const end = Math.min(names.length, start + Number(count)); + const page = names.slice(start, end).filter((name) => !type || keys.get(name)?.type === type); + return [end >= names.length ? '0' : String(end), page]; + }; + const pipeline = () => { + const queued: string[][] = []; + const chain = { + call: (...parts: string[]) => { + queued.push(parts); + return chain; + }, + exec: async () => + queued.map((parts): [Error | null, unknown] => { + const entry = keys.get(parts[parts[0] === 'TYPE' ? 1 : 2]); + if (parts[0] === 'TYPE') return [null, entry?.type ?? 'none']; + if (options.memoryError) return [new Error(options.memoryError), null]; + return [null, entry?.size ?? null]; + }), + }; + return chain; + }; + return { call, getClient: () => ({ pipeline }), getDbSize: async () => keys.size } as unknown as SamplerClient; +} + +export function keyspace(count: number, pick: (index: number) => FakeKey): Map { + const keys = new Map(); + for (let index = 0; index < count; index++) keys.set(`k${index}`, pick(index)); + return keys; +} diff --git a/proprietary/cache-simulator/__tests__/node-catalog.spec.ts b/proprietary/cache-simulator/__tests__/node-catalog.spec.ts new file mode 100644 index 000000000..c1af4fbe9 --- /dev/null +++ b/proprietary/cache-simulator/__tests__/node-catalog.spec.ts @@ -0,0 +1,69 @@ +import { CLOUD_SERVICES, smallestNodeFor, type CloudProvider } from '@betterdb/shared'; + +const GiB = 1024 ** 3; +const providers: CloudProvider[] = ['aws', 'gcp', 'azure']; + +describe('node catalog', () => { + it('covers AWS, GCP and Azure with sources', () => { + expect(new Set(CLOUD_SERVICES.map((service) => service.provider))).toEqual(new Set(providers)); + for (const service of CLOUD_SERVICES) { + expect(service.sourceUrl).toMatch(/^https:\/\//); + expect(service.asOf).toMatch(/^\d{4}-\d{2}-\d{2}$/); + expect(service.nodes.length).toBeGreaterThan(0); + } + }); + + it.each(providers)('picks the smallest node that keeps %s under 80%% of usable memory', (provider) => { + const bytes = 10 * GiB; + for (const pick of smallestNodeFor(bytes, provider).filter((rec) => rec.shards === 1)) { + const service = CLOUD_SERVICES.find((candidate) => candidate.service === pick.service)!; + const usable = service.usableRatio * 0.8; + expect(pick.memoryBytes * usable).toBeGreaterThanOrEqual(bytes); + const smaller = service.nodes.filter((node) => node.memoryBytes < pick.memoryBytes); + expect(smaller.every((node) => node.memoryBytes * usable < bytes)).toBe(true); + } + }); + + it('reserves memory the service keeps back from the dataset', () => { + const byService = Object.fromEntries(CLOUD_SERVICES.map((service) => [service.service, service.usableRatio])); + expect(byService).toEqual({ + elasticache: 0.75, + memorydb: 1, + 'memorystore-valkey': 0.8, + 'azure-cache-standard': 0.8, + 'azure-cache-premium': 0.8, + }); + const [standardSmall] = smallestNodeFor(4.5e9, 'gcp'); + expect(standardSmall.nodeType).toBe('highmem-medium'); + const memorydb = smallestNodeFor(2.4e9, 'aws').find((rec) => rec.service === 'memorydb'); + expect(memorydb?.nodeType).toBe('db.t4g.medium'); + }); + + it('caps GCP shards at one replica per shard and labels Azure as legacy', () => { + const gcp = CLOUD_SERVICES.find((service) => service.provider === 'gcp')!; + expect(gcp.maxShards).toBe(125); + expect(gcp.label).toContain('1 replica per shard'); + for (const service of CLOUD_SERVICES.filter((candidate) => candidate.provider === 'azure')) expect(service.label).toContain('legacy'); + }); + + it('falls back to shards of the largest node', () => { + const bytes = 1024 * GiB; + const gcp = smallestNodeFor(bytes, 'gcp'); + expect(gcp).toHaveLength(1); + const service = CLOUD_SERVICES.find((candidate) => candidate.service === gcp[0].service)!; + const largest = Math.max(...service.nodes.map((node) => node.memoryBytes)); + expect(gcp[0].memoryBytes).toBe(largest); + expect(gcp[0].shards * largest * service.usableRatio * 0.8).toBeGreaterThanOrEqual(bytes); + expect((gcp[0].shards - 1) * largest * service.usableRatio * 0.8).toBeLessThan(bytes); + }); + + it('omits services whose shard limit cannot hold the data', () => { + const services = smallestNodeFor(2048 * GiB, 'azure').map((rec) => rec.service); + expect(services).not.toContain('azure-cache-standard'); + expect(services).not.toContain('azure-cache-premium'); + }); + + it('returns one recommendation per service for small data', () => { + expect(smallestNodeFor(1024, 'aws').map((rec) => rec.service).sort()).toEqual(['elasticache', 'memorydb']); + }); +}); diff --git a/proprietary/cache-simulator/__tests__/projection.spec.ts b/proprietary/cache-simulator/__tests__/projection.spec.ts new file mode 100644 index 000000000..4deac42d0 --- /dev/null +++ b/proprietary/cache-simulator/__tests__/projection.spec.ts @@ -0,0 +1,101 @@ +import { DEFAULT_OVERHEAD_PER_KEY_BYTES, fixedOverhead, inPlaceGrowth, keysAt, project, type ProjectionInputs } from '@betterdb/shared'; + +const base: ProjectionInputs = { + currentKeys: 1000, + keysPerDay: 100, + inPlaceBytesPerDay: 0, + ttlDays: null, + avgKeyBytes: 100, + ci95Bytes: 0, + overheadPerKeyBytes: 0, + fixedOverheadBytes: 0, + maxmemory: null, + horizonDays: 30, +}; + +describe('projection', () => { + it('uses 88 bytes as the default per-key overhead', () => { + expect(DEFAULT_OVERHEAD_PER_KEY_BYTES).toBe(88); + }); + + it('projects linear growth', () => { + const result = project(base); + expect(result.points).toHaveLength(31); + expect(result.points[10]).toMatchObject({ day: 10, keys: 2000, bytes: 200_000 }); + }); + + it('never projects negative keys', () => { + expect(keysAt(10, { ...base, keysPerDay: -500 })).toBe(0); + const result = project({ ...base, keysPerDay: -500, maxmemory: 1_000_000 }); + expect(result.points.every((point) => point.keys >= 0 && point.bytes >= 0)).toBe(true); + expect(result.daysToMaxmemory).toBeNull(); + }); + + it('converges to the TTL steady state', () => { + const inputs = { ...base, currentKeys: 0, keysPerDay: 100, ttlDays: 7 }; + expect(keysAt(7, inputs)).toBeCloseTo(700 * (1 - Math.exp(-1)), 6); + expect(keysAt(70, inputs)).toBeCloseTo(700, 0); + }); + + it('decays toward the steady state from above', () => { + expect(keysAt(1000, { ...base, currentKeys: 10_000, keysPerDay: 100, ttlDays: 7 })).toBeCloseTo(700, 0); + }); + + it('adds per-key and fixed overhead', () => { + const result = project({ ...base, overheadPerKeyBytes: 50, fixedOverheadBytes: 1000 }); + expect(result.points[0].bytes).toBe(1000 * 150 + 1000); + }); + + it('finds the first day at maxmemory and at 80%', () => { + const result = project({ ...base, maxmemory: 150_000 }); + expect(result.daysToMaxmemory).toBe(5); + expect(result.daysTo80Percent).toBe(2); + }); + + it('returns day 0 when already over the limit', () => { + expect(project({ ...base, maxmemory: 50_000 }).daysToMaxmemory).toBe(0); + }); + + it('handles null and zero maxmemory', () => { + expect(project(base).daysToMaxmemory).toBeNull(); + expect(project({ ...base, maxmemory: 0 }).daysToMaxmemory).toBeNull(); + expect(project({ ...base, maxmemory: 0 }).daysTo80Percent).toBeNull(); + }); + + it('anchors the band to the sampled interval and scales it with the key count', () => { + const result = project({ ...base, ci95Bytes: 5000 }); + expect(result.points[0].high - result.points[0].bytes).toBeCloseTo(5000, 6); + expect(result.points[0].bytes - result.points[0].low).toBeCloseTo(5000, 6); + expect(result.points[10].high - result.points[10].bytes).toBeCloseTo(10_000, 6); + }); + + it('has no band without an interval', () => { + const point = project(base).points[0]; + expect(point.low).toBe(point.bytes); + expect(point.high).toBe(point.bytes); + }); + + it('adds in-place byte growth on top of key growth', () => { + const result = project({ ...base, keysPerDay: 0, inPlaceBytesPerDay: 500, maxmemory: 102_000 }); + expect(result.points[10].bytes).toBe(100_000 + 5000); + expect(result.daysToMaxmemory).toBe(4); + }); + + it('derives in-place growth from the byte trend new keys do not explain', () => { + expect(inPlaceGrowth({ keysPerDay: 0, bytesPerDay: 500 }, 100)).toBe(500); + expect(inPlaceGrowth({ keysPerDay: 100, bytesPerDay: 8000 }, 100)).toBe(0); + expect(inPlaceGrowth({ keysPerDay: 10, bytesPerDay: 3000 }, 100)).toBe(2000); + expect(inPlaceGrowth({ keysPerDay: -500, bytesPerDay: 0 }, 100)).toBe(0); + }); + + it('stays finite with zero keys', () => { + const result = project({ ...base, currentKeys: 0, keysPerDay: 0, maxmemory: 100 }); + expect(result.points.every((point) => Number.isFinite(point.bytes))).toBe(true); + expect(result.daysToMaxmemory).toBeNull(); + }); + + it('floors fixed overhead at zero', () => { + expect(fixedOverhead(1000, 10, 100, 88)).toBe(0); + expect(fixedOverhead(10_000, 10, 100, 88)).toBe(10_000 - 10 * 188); + }); +}); diff --git a/proprietary/cache-simulator/__tests__/sample-stats.spec.ts b/proprietary/cache-simulator/__tests__/sample-stats.spec.ts new file mode 100644 index 000000000..51e53997f --- /dev/null +++ b/proprietary/cache-simulator/__tests__/sample-stats.spec.ts @@ -0,0 +1,21 @@ +import { adjustedSkewness } from '../sample-stats'; + +describe('adjustedSkewness', () => { + it('matches a hand-computed value', () => { + expect(adjustedSkewness([1, 2, 6])).toBeCloseTo(Math.sqrt(6) * (6 / (14 / 3) ** 1.5), 10); + expect(adjustedSkewness([1, 2, 6])).toBeCloseTo(1.458, 3); + }); + + it('is zero for symmetric data', () => { + expect(adjustedSkewness([1, 2, 3])).toBeCloseTo(0, 10); + }); + + it('is zero without variance', () => { + expect(adjustedSkewness([5, 5])).toBe(0); + expect(adjustedSkewness([5])).toBe(0); + }); + + it('is untrusted for fewer than three varying values', () => { + expect(adjustedSkewness([1, 2])).toBe(Number.POSITIVE_INFINITY); + }); +}); diff --git a/proprietary/cache-simulator/__tests__/scenarios.spec.ts b/proprietary/cache-simulator/__tests__/scenarios.spec.ts new file mode 100644 index 000000000..de7290f50 --- /dev/null +++ b/proprietary/cache-simulator/__tests__/scenarios.spec.ts @@ -0,0 +1,43 @@ +import { mulberry32, SCENARIOS } from '../../../apps/api/scripts/cache-simulator/scenarios'; + +describe('validation scenarios', () => { + it('defines the research scenarios with their tolerances', () => { + expect(SCENARIOS.map((scenario) => [scenario.name, scenario.tolerance, scenario.runs])).toEqual([ + ['homogeneous-small', 0.05, 10], + ['homogeneous-large', 0.05, 10], + ['mixed', 0.05, 10], + ['pareto', 0.1, 10], + ['ttl-mix', 0.05, 10], + ['rare-giant', null, 10], + ['scale-500k', 0.05, 3], + ]); + }); + + it('is reproducible for a seed', () => { + const first = [...SCENARIOS[2].generate(7)].slice(0, 50); + const second = [...SCENARIOS[2].generate(7)].slice(0, 50); + expect(first).toEqual(second); + expect([...SCENARIOS[2].generate(8)].slice(0, 50)).not.toEqual(first); + }); + + it('gives half the ttl-mix keys a TTL', () => { + const commands = [...SCENARIOS[4].generate(1)]; + const withTtl = commands.filter((command) => command.ttlSeconds !== undefined).length; + expect(withTtl / commands.length).toBeCloseTo(0.5, 1); + }); + + it('produces values in [0, 1) from the PRNG', () => { + const next = mulberry32(42); + for (let index = 0; index < 1000; index++) { + const value = next(); + expect(value).toBeGreaterThanOrEqual(0); + expect(value).toBeLessThan(1); + } + }); + + it('places exactly one dominating key in rare-giant', () => { + const sizes = [...SCENARIOS[5].generate(3)].map((command) => (command.value as string).length); + expect(sizes.filter((size) => size > 100)).toEqual([64 * 1024 * 1024]); + expect(sizes).toHaveLength(50_000); + }); +}); diff --git a/proprietary/cache-simulator/__tests__/size-estimator.spec.ts b/proprietary/cache-simulator/__tests__/size-estimator.spec.ts new file mode 100644 index 000000000..f7898252d --- /dev/null +++ b/proprietary/cache-simulator/__tests__/size-estimator.spec.ts @@ -0,0 +1,161 @@ +import { estimate, keyMemoryCeiling } from '../size-estimator'; +import type { NodeSample } from '../sample-types'; + +const node = (strata: NodeSample['strata'], timedOut = false): NodeSample => ({ + node: 'n', + dbSize: strata.reduce((sum, s) => sum + s.estimatedKeys, 0), + strata, + exhausted: false, + timedOut, +}); + +describe('estimate', () => { + it('is exact when a stratum is fully enumerated', () => { + const result = estimate([node([{ type: 'string', estimatedKeys: 3, exact: true, observed: 0, sizes: [10, 20, 30] }])], 5); + expect(result.totalBytes).toBe(60); + expect(result.ci95Bytes).toBe(0); + expect(result.avgKeyBytes).toBe(20); + expect(result.stdKeyBytes).toBeCloseTo(10, 6); + expect(result.sampleCount).toBe(3); + expect(result.sampledAt).toBe(5); + expect(result.partial).toBe(false); + }); + + it('combines strata with finite population correction', () => { + const result = estimate( + [ + node([ + { type: 'string', estimatedKeys: 1000, exact: true, observed: 0, sizes: [90, 110] }, + { type: 'hash', estimatedKeys: 100, exact: true, observed: 0, sizes: [1000, 1000, 1000] }, + ]), + ], + 0, + ); + expect(result.totalBytes).toBe(200_000); + expect(result.ci95Bytes).toBeCloseTo(1.96 * Math.sqrt((1000 ** 2 / 2) * 200 * (1 - 2 / 1000)) + 165_000, 6); + expect(result.byType[0]).toMatchObject({ type: 'string', keys: 1000, avgBytes: 100, sampled: 2 }); + }); + + it('merges the same type across nodes in the by-type view', () => { + const result = estimate( + [node([{ type: 'string', estimatedKeys: 100, exact: true, observed: 0, sizes: [50, 50] }]), node([{ type: 'string', estimatedKeys: 300, exact: true, observed: 0, sizes: [150, 150] }])], + 0, + ); + expect(result.totalBytes).toBe(50_000); + expect(result.byType).toHaveLength(1); + expect(result.byType[0]).toMatchObject({ type: 'string', keys: 400, avgBytes: 125, sampled: 4 }); + }); + + it('uses the pooled mean for strata without samples', () => { + const result = estimate( + [ + node([ + { type: 'string', estimatedKeys: 100, exact: true, observed: 0, sizes: [10, 30] }, + { type: 'hash', estimatedKeys: 50, exact: true, observed: 0, sizes: [] }, + ]), + ], + 0, + ); + expect(result.totalBytes).toBe(3000); + expect(result.ci95Bytes).toBeCloseTo(1.96 * Math.sqrt((100 ** 2 / 2) * 200 * (1 - 2 / 100) + 50 ** 2 * 200) + 4500, 6); + }); + + it('adds count uncertainty for non-exact strata with different means', () => { + const result = estimate( + [ + node([ + { type: 'string', estimatedKeys: 600, exact: false, observed: 60, sizes: [100, 100] }, + { type: 'hash', estimatedKeys: 400, exact: false, observed: 40, sizes: [300, 300] }, + ]), + ], + 0, + ); + expect(result.totalBytes).toBe(180_000); + expect(result.ci95Bytes).toBeCloseTo(1.96 * Math.sqrt(10_000 * 9600 * 0.9), 6); + }); + + it('adds no count uncertainty when every stratum is exact', () => { + const result = estimate( + [ + node([ + { type: 'string', estimatedKeys: 600, exact: true, observed: 60, sizes: [100, 100] }, + { type: 'hash', estimatedKeys: 400, exact: true, observed: 40, sizes: [300, 300] }, + ]), + ], + 0, + ); + expect(result.ci95Bytes).toBe(0); + }); + + it('uses the pooled variance for a single-sample stratum', () => { + const result = estimate( + [ + node([ + { type: 'string', estimatedKeys: 100, exact: true, observed: 0, sizes: [10, 30] }, + { type: 'hash', estimatedKeys: 50, exact: true, observed: 0, sizes: [20] }, + ]), + ], + 0, + ); + const pooledVariance = 100; + const expected = (100 ** 2 / 2) * 200 * (1 - 2 / 100) + 50 ** 2 * pooledVariance * (1 - 1 / 50); + expect(result.ci95Bytes).toBeCloseTo(1.96 * Math.sqrt(expected) + 4500, 6); + }); + + it('adds the rule-of-three allowance for a heavy-tailed stratum', () => { + const result = estimate([node([{ type: 'string', estimatedKeys: 100, exact: true, observed: 0, sizes: [10, 30] }])], 0); + expect(result.ci95Bytes).toBeCloseTo(1.96 * Math.sqrt((100 ** 2 / 2) * 200 * (1 - 2 / 100)) + 100 * (3 / 2) * 30, 6); + }); + + it('adds no allowance to a non-skewed stratum', () => { + const result = estimate([node([{ type: 'string', estimatedKeys: 1000, exact: true, observed: 0, sizes: [90, 100, 110] }])], 0); + expect(result.ci95Bytes).toBeCloseTo(1.96 * Math.sqrt((1000 ** 2 / 3) * 100 * (1 - 3 / 1000)), 6); + }); + + it('adds no allowance when the stratum is fully sampled', () => { + const result = estimate([node([{ type: 'string', estimatedKeys: 2, exact: true, observed: 0, sizes: [10, 30] }])], 0); + expect(result.ci95Bytes).toBe(0); + }); + + it('widens the interval to the memory the sample cannot explain', () => { + const result = estimate([node([{ type: 'string', estimatedKeys: 100, exact: true, observed: 0, sizes: [10, 10] }])], 0, 5000); + expect(result.totalBytes).toBe(1000); + expect(result.ci95Bytes).toBe(4000); + }); + + it('widens the interval when the estimate exceeds the memory ceiling', () => { + const result = estimate([node([{ type: 'string', estimatedKeys: 100, exact: true, observed: 0, sizes: [1000, 1000] }])], 0, 5000); + expect(result.totalBytes).toBe(100_000); + expect(result.ci95Bytes).toBe(95_000); + }); + + it('treats a nonempty node without measured strata as unsampled', () => { + const result = estimate([{ node: 'n', dbSize: 100, strata: [], exhausted: false, timedOut: true }], 0, 5000); + expect(result.totalBytes).toBe(0); + expect(result.ci95Bytes).toBe(5000); + }); + + it('ignores the memory ceiling when every key was sampled', () => { + const result = estimate([node([{ type: 'string', estimatedKeys: 2, exact: true, observed: 0, sizes: [10, 10] }])], 0, 5000); + expect(result.ci95Bytes).toBe(0); + }); + + it('keeps the statistical interval when it is wider than the unexplained memory', () => { + const withoutCeiling = estimate([node([{ type: 'string', estimatedKeys: 1000, exact: true, observed: 0, sizes: [90, 100, 110] }])], 0); + const withCeiling = estimate([node([{ type: 'string', estimatedKeys: 1000, exact: true, observed: 0, sizes: [90, 100, 110] }])], 0, 100_001); + expect(withCeiling.ci95Bytes).toBe(withoutCeiling.ci95Bytes); + }); + + it('adds the per-key dict entry to the dataset memory for the ceiling', () => { + expect(keyMemoryCeiling(1000, 10)).toBe(1240); + }); + + it('flags partial results', () => { + expect(estimate([node([{ type: 'string', estimatedKeys: 1, exact: true, observed: 0, sizes: [1] }], true)], 0).partial).toBe(true); + }); + + it('handles empty input', () => { + const result = estimate([node([])], 0); + expect(result).toMatchObject({ totalBytes: 0, ci95Bytes: 0, avgKeyBytes: 0, stdKeyBytes: 0, sampleCount: 0, byType: [] }); + }); +}); diff --git a/proprietary/cache-simulator/__tests__/stratified-sampler.spec.ts b/proprietary/cache-simulator/__tests__/stratified-sampler.spec.ts new file mode 100644 index 000000000..dda322a55 --- /dev/null +++ b/proprietary/cache-simulator/__tests__/stratified-sampler.spec.ts @@ -0,0 +1,210 @@ +import { estimate } from '../size-estimator'; +import { sampleNodes, HEAVY_TAIL_MIN_SAMPLES, SamplerUnsupportedError, TARGET_RELATIVE_CI } from '../stratified-sampler'; +import { fakeClient, keyspace } from './fake-keyspace'; + +const string = (size = 100) => ({ type: 'string', size }); +const hash = (index: number) => ({ type: 'hash', size: 1000 + (index % 7) * 150 }); + +describe('sampleNodes', () => { + it('enumerates a small keyspace exactly in the pilot', async () => { + const [result] = await sampleNodes([{ name: 'a', client: fakeClient(keyspace(50, () => string())) }], { budget: 4000, deadlineMs: 60_000 }); + expect(result.exhausted).toBe(true); + expect(result.strata).toEqual([{ type: 'string', estimatedKeys: 50, exact: true, observed: 50, sizes: Array(50).fill(100) }]); + }); + + it('returns no strata for an empty keyspace', async () => { + const [result] = await sampleNodes([{ name: 'a', client: fakeClient(new Map()) }], { budget: 4000, deadlineMs: 60_000 }); + expect(result.strata).toEqual([]); + expect(result.dbSize).toBe(0); + }); + + it('tops up the high-variance type and counts it exactly when the TYPE scan completes', async () => { + const keys = keyspace(1040, (index) => (index % 26 === 0 ? hash(index) : string())); + const [result] = await sampleNodes([{ name: 'a', client: fakeClient(keys) }], { budget: 400, deadlineMs: 60_000 }); + const hashes = result.strata.find((stratum) => stratum.type === 'hash')!; + const strings = result.strata.find((stratum) => stratum.type === 'string')!; + expect(hashes.estimatedKeys).toBe(40); + expect(hashes.sizes).toHaveLength(40); + expect(hashes.exact).toBe(true); + expect(strings.sizes.length).toBeGreaterThanOrEqual(30); + expect(strings.sizes.length + hashes.sizes.length).toBeLessThanOrEqual(400); + expect(strings.estimatedKeys).toBe(1000); + expect(strings.exact).toBe(true); + }); + + it('keeps the total near the budget on a large keyspace', async () => { + const keys = keyspace(5500, (index) => (index % 11 === 0 ? hash(index) : string())); + const [result] = await sampleNodes([{ name: 'a', client: fakeClient(keys) }], { budget: 1000, deadlineMs: 60_000 }); + const sampled = result.strata.reduce((sum, stratum) => sum + stratum.sizes.length, 0); + expect(sampled).toBeLessThanOrEqual(1000); + expect(result.strata.find((stratum) => stratum.type === 'hash')!.sizes.length).toBeGreaterThanOrEqual(30); + expect(result.strata.reduce((sum, stratum) => sum + stratum.estimatedKeys, 0)).toBeGreaterThan(5000); + }); + + it('skips keys deleted while sampling', async () => { + const keys = keyspace(10, (index) => ({ type: 'string', size: index < 3 ? null : 100 })); + const [result] = await sampleNodes([{ name: 'a', client: fakeClient(keys) }], { budget: 4000, deadlineMs: 60_000 }); + expect(result.strata[0].sizes).toHaveLength(7); + }); + + it('splits the budget across nodes by key count', async () => { + const results = await sampleNodes( + [ + { name: 'a', client: fakeClient(keyspace(1000, () => string())) }, + { name: 'b', client: fakeClient(keyspace(3000, () => string())) }, + ], + { budget: 400, deadlineMs: 60_000 }, + ); + expect(results.map((result) => result.strata[0].sizes.length)).toEqual([100, 300]); + }); + + it('reports progress up to the final count', async () => { + const progress: number[] = []; + await sampleNodes([{ name: 'a', client: fakeClient(keyspace(250, () => string())) }], { + budget: 4000, + deadlineMs: 60_000, + onProgress: (sampled) => progress.push(sampled), + }); + expect(progress[progress.length - 1]).toBe(250); + expect([...progress].sort((a, b) => a - b)).toEqual(progress); + }); + + it('stops at the deadline and marks the node timed out', async () => { + const [result] = await sampleNodes([{ name: 'a', client: fakeClient(keyspace(100, () => string())) }], { budget: 4000, deadlineMs: 0 }); + expect(result.timedOut).toBe(true); + }); + + it('fails clearly when MEMORY USAGE is unavailable', async () => { + const client = fakeClient(keyspace(10, () => string()), { memoryError: "ERR unknown command 'MEMORY'" }); + await expect(sampleNodes([{ name: 'a', client }], { budget: 4000, deadlineMs: 60_000 })).rejects.toBeInstanceOf(SamplerUnsupportedError); + }); + + it('fails clearly when SCAN TYPE is unavailable', async () => { + const keys = keyspace(1040, (index) => (index % 26 === 0 ? hash(index) : string())); + const client = fakeClient(keys, { scanTypeError: 'ERR syntax error' }); + await expect(sampleNodes([{ name: 'a', client }], { budget: 400, deadlineMs: 60_000 })).rejects.toBeInstanceOf(SamplerUnsupportedError); + }); + + it('rethrows other MEMORY USAGE errors without marking them unsupported', async () => { + const client = fakeClient(keyspace(10, () => string()), { memoryError: 'LOADING Valkey is loading the dataset in memory' }); + const failure = await sampleNodes([{ name: 'a', client }], { budget: 4000, deadlineMs: 60_000 }).catch((error) => error); + expect(failure).toBeInstanceOf(Error); + expect(failure).not.toBeInstanceOf(SamplerUnsupportedError); + expect(failure.message).toBe('LOADING Valkey is loading the dataset in memory'); + }); + + it('throws when the pipeline returns no replies', async () => { + const client = fakeClient(keyspace(10, () => string())); + (client.getClient as unknown as () => unknown) = () => ({ pipeline: () => ({ call: function () { return this; }, exec: async () => null }) }); + await expect(sampleNodes([{ name: 'a', client }], { budget: 4000, deadlineMs: 60_000 })).rejects.toThrow('Pipeline returned no replies'); + }); + + it('keeps the census estimate when the TYPE scan stops early', async () => { + const keys = keyspace(3000, (index) => ((index < 200 ? index % 10 === 0 : index % 5 === 0) ? hash(index) : string())); + const [result] = await sampleNodes([{ name: 'a', client: fakeClient(keys) }], { budget: 100, deadlineMs: 60_000 }); + const hashes = result.strata.find((stratum) => stratum.type === 'hash')!; + expect(hashes.exact).toBe(false); + expect(hashes.observed).toBeGreaterThan(0); + expect(hashes.estimatedKeys).toBeCloseTo((hashes.observed / 1000) * 3000, 6); + }); + + it('counts types exactly when the census covers the whole keyspace', async () => { + const keys = keyspace(1500, (index) => (index % 5 === 0 ? hash(index) : string())); + const [result] = await sampleNodes([{ name: 'a', client: fakeClient(keys) }], { budget: 200, deadlineMs: 60_000 }); + expect(result.strata.every((stratum) => stratum.exact)).toBe(true); + expect(result.strata.find((stratum) => stratum.type === 'hash')!.estimatedKeys).toBe(300); + expect(result.strata.find((stratum) => stratum.type === 'string')!.estimatedKeys).toBe(1200); + expect(result.strata.reduce((sum, stratum) => sum + stratum.estimatedKeys, 0)).toBe(1500); + }); + + it('renormalises a partial census to the database size', async () => { + const keys = keyspace(6000, (index) => (index % 3 === 0 ? hash(index) : string())); + const [result] = await sampleNodes([{ name: 'a', client: fakeClient(keys) }], { budget: 200, deadlineMs: 60_000 }); + expect(result.strata.some((stratum) => !stratum.exact)).toBe(true); + expect(result.strata.every((stratum) => stratum.observed > 0)).toBe(true); + expect(result.strata.reduce((sum, stratum) => sum + stratum.estimatedKeys, 0)).toBeCloseTo(6000, 6); + }); + + it('never estimates fewer keys than the census saw when DBSIZE is stale', async () => { + const keys = keyspace(6000, (index) => (index % 3 === 0 ? hash(index) : string())); + const client = { ...fakeClient(keys), getDbSize: async () => 1000 }; + const [result] = await sampleNodes([{ name: 'a', client }], { budget: 200, deadlineMs: 60_000 }); + for (const stratum of result.strata) expect(stratum.estimatedKeys).toBeGreaterThanOrEqual(stratum.observed); + expect(result.strata.every((stratum) => stratum.estimatedKeys > 0)).toBe(true); + }); + + describe('adaptive extension', () => { + const heavyTail = () => keyspace(20_000, (index) => ({ type: 'string', size: index % 200 === 0 ? 100_000 : 100 })); + const sampledOf = (strata: { sizes: number[] }[]) => strata.reduce((sum, stratum) => sum + stratum.sizes.length, 0); + + it('samples beyond the budget until the interval is tight', async () => { + const [result] = await sampleNodes([{ name: 'a', client: fakeClient(heavyTail()) }], { budget: 400, hardCap: 20_000, deadlineMs: 60_000 }); + const { totalBytes, ci95Bytes } = estimate([result], 0); + expect(sampledOf(result.strata)).toBeGreaterThan(400); + expect(ci95Bytes / totalBytes <= TARGET_RELATIVE_CI || sampledOf(result.strata) === 20_000).toBe(true); + }); + + it('keeps extending to full enumeration when a narrow interval is untrustworthy', async () => { + const keys = keyspace(1000, (index) => ({ type: 'string', size: index % 50 === 0 ? 160 : 100 })); + const [result] = await sampleNodes([{ name: 'a', client: fakeClient(keys) }], { budget: 100, hardCap: 1000, deadlineMs: 60_000 }); + const final = estimate([result], 0); + expect(sampledOf(result.strata)).toBe(1000); + expect(final.ci95Bytes).toBe(0); + expect(final.totalBytes).toBe(980 * 100 + 20 * 160); + }); + + it('keeps extending a heavy-tailed stratum below the minimum even when the interval is tight', async () => { + const keys = keyspace(30_000, (index) => ({ type: 'string', size: index % 50 === 0 ? 160 : 100 })); + const [early] = await sampleNodes([{ name: 'a', client: fakeClient(keys) }], { budget: 400, deadlineMs: 60_000 }); + const earlyEstimate = estimate([early], 0); + expect(earlyEstimate.ci95Bytes / earlyEstimate.totalBytes).toBeLessThanOrEqual(TARGET_RELATIVE_CI); + const [result] = await sampleNodes([{ name: 'a', client: fakeClient(keys) }], { budget: 400, hardCap: 30_000, deadlineMs: 60_000 }); + expect(sampledOf(result.strata)).toBeGreaterThanOrEqual(HEAVY_TAIL_MIN_SAMPLES); + expect(sampledOf(result.strata)).toBeLessThanOrEqual(30_000); + }); + + it('never exceeds its share of the hard cap', async () => { + const [result] = await sampleNodes([{ name: 'a', client: fakeClient(heavyTail()) }], { budget: 400, hardCap: 1000, deadlineMs: 60_000 }); + expect(sampledOf(result.strata)).toBeLessThanOrEqual(1000); + expect(sampledOf(result.strata)).toBeGreaterThan(400); + }); + + it('stops at the budget when the interval is already tight', async () => { + const [result] = await sampleNodes([{ name: 'a', client: fakeClient(keyspace(20_000, () => string())) }], { + budget: 400, + hardCap: 20_000, + deadlineMs: 60_000, + }); + expect(sampledOf(result.strata)).toBe(400); + }); + + it('never exceeds the budget without a hard cap', async () => { + const [result] = await sampleNodes([{ name: 'a', client: fakeClient(heavyTail()) }], { budget: 400, deadlineMs: 60_000 }); + expect(sampledOf(result.strata)).toBeLessThanOrEqual(400); + }); + }); + + describe('rare dominating key', () => { + const giant = 512 * 1024 * 1024; + const keys = keyspace(50_000, (index) => (index === 31_337 ? string(giant) : string())); + const truth = 49_999 * 100 + giant; + const runs = Array.from({ length: 20 }, (_, index) => index + 1); + + it('covers the true size with the memory ceiling on every scan order', async () => { + for (const seed of runs) { + const [result] = await sampleNodes([{ name: 'a', client: fakeClient(keys, { seed }) }], { budget: 4000, hardCap: 50_000, deadlineMs: 60_000 }); + const { totalBytes, ci95Bytes } = estimate([result], 0, truth + 1024 * 1024); + expect(Math.abs(totalBytes - truth)).toBeLessThanOrEqual(ci95Bytes); + } + }); + + it('misses the key on some scan orders, so the sample alone cannot bound it', async () => { + let missed = 0; + for (const seed of runs) { + const [result] = await sampleNodes([{ name: 'a', client: fakeClient(keys, { seed }) }], { budget: 4000, hardCap: 50_000, deadlineMs: 60_000 }); + if (!result.strata.some((stratum) => stratum.sizes.includes(giant))) missed += 1; + } + expect(missed).toBeGreaterThan(0); + }); + }); +}); diff --git a/proprietary/cache-simulator/allocation.ts b/proprietary/cache-simulator/allocation.ts new file mode 100644 index 000000000..e96fe2a0e --- /dev/null +++ b/proprietary/cache-simulator/allocation.ts @@ -0,0 +1,29 @@ +export const STRATUM_FLOOR = 30; + +export interface AllocationInput { + type: string; + estimatedKeys: number; + std: number; + sampled: number; +} + +export function splitBudget(budget: number, dbSizes: number[]): number[] { + const total = dbSizes.reduce((sum, size) => sum + size, 0); + if (total === 0) return dbSizes.map(() => 0); + return dbSizes.map((size) => (size === 0 ? 0 : Math.max(1, Math.round((budget * size) / total)))); +} + +export function neymanAllocation(remaining: number, strata: AllocationInput[]): Map { + const weights = strata.map((stratum) => stratum.estimatedKeys * stratum.std); + const weightSum = weights.reduce((sum, weight) => sum + weight, 0); + const basis = weightSum > 0 ? weights : strata.map((stratum) => stratum.estimatedKeys); + const basisSum = basis.reduce((sum, value) => sum + value, 0); + const result = new Map(); + strata.forEach((stratum, index) => { + const share = basisSum > 0 ? Math.round((remaining * basis[index]) / basisSum) : 0; + const floor = Math.min(STRATUM_FLOOR, stratum.estimatedKeys); + const desired = Math.min(stratum.estimatedKeys, Math.max(stratum.sampled + share, floor)); + result.set(stratum.type, Math.max(0, Math.round(desired - stratum.sampled))); + }); + return result; +} diff --git a/proprietary/cache-simulator/sample-stats.ts b/proprietary/cache-simulator/sample-stats.ts new file mode 100644 index 000000000..a57625b88 --- /dev/null +++ b/proprietary/cache-simulator/sample-stats.ts @@ -0,0 +1,25 @@ +export function mean(values: number[]): number { + return values.length === 0 ? 0 : values.reduce((sum, value) => sum + value, 0) / values.length; +} + +export function adjustedSkewness(values: number[]): number { + const n = values.length; + const average = mean(values); + const m2 = values.reduce((sum, value) => sum + (value - average) ** 2, 0) / Math.max(1, n); + if (m2 === 0) return 0; + if (n < 3) return Number.POSITIVE_INFINITY; + const m3 = values.reduce((sum, value) => sum + (value - average) ** 3, 0) / n; + return (Math.sqrt(n * (n - 1)) / (n - 2)) * (m3 / m2 ** 1.5); +} + +export function sampleVariance(values: number[]): number { + if (values.length < 2) return 0; + const average = mean(values); + return values.reduce((sum, value) => sum + (value - average) ** 2, 0) / (values.length - 1); +} + +export const HEAVY_TAIL_SKEWNESS = 2; + +export function isHeavyTailed(stratum: { estimatedKeys: number; sizes: number[] }): boolean { + return stratum.estimatedKeys > stratum.sizes.length && adjustedSkewness(stratum.sizes) > HEAVY_TAIL_SKEWNESS; +} diff --git a/proprietary/cache-simulator/sample-types.ts b/proprietary/cache-simulator/sample-types.ts new file mode 100644 index 000000000..9802edc32 --- /dev/null +++ b/proprietary/cache-simulator/sample-types.ts @@ -0,0 +1,19 @@ +import type { DatabasePort } from '@app/common/interfaces/database-port.interface'; + +export type SamplerClient = Pick; + +export interface Stratum { + type: string; + estimatedKeys: number; + exact: boolean; + observed: number; + sizes: number[]; +} + +export interface NodeSample { + node: string; + dbSize: number; + strata: Stratum[]; + exhausted: boolean; + timedOut: boolean; +} diff --git a/proprietary/cache-simulator/size-estimator.ts b/proprietary/cache-simulator/size-estimator.ts new file mode 100644 index 000000000..3b4b4d4e8 --- /dev/null +++ b/proprietary/cache-simulator/size-estimator.ts @@ -0,0 +1,89 @@ +import type { CacheSimulatorEstimate, CacheSimulatorTypeEstimate } from '@betterdb/shared'; +import { isHeavyTailed, mean, sampleVariance } from './sample-stats'; +import type { NodeSample } from './sample-types'; + +const Z95 = 1.96; +export const DICT_ENTRY_BYTES = 24; + +export function keyMemoryCeiling(datasetBytes: number, keys: number): number { + return datasetBytes + keys * DICT_ENTRY_BYTES; +} + +interface TypeTotals { + keys: number; + weighted: number; + second: number; + sampled: number; +} + +function countVariance(strata: NodeSample['strata'], pooledMean: number): number { + const open = strata.filter((stratum) => !stratum.exact && stratum.estimatedKeys > 0); + const observed = open.reduce((sum, stratum) => sum + stratum.observed, 0); + const keys = open.reduce((sum, stratum) => sum + stratum.estimatedKeys, 0); + if (open.length < 2 || observed <= 0 || keys <= observed) return 0; + const means = open.map((stratum) => (stratum.sizes.length > 0 ? mean(stratum.sizes) : pooledMean)); + const weights = open.map((stratum) => stratum.observed / observed); + const center = weights.reduce((sum, weight, index) => sum + weight * means[index], 0); + const spread = weights.reduce((sum, weight, index) => sum + weight * (means[index] - center) ** 2, 0); + return ((keys * keys) / observed) * spread * (1 - observed / keys); +} + +export function estimate(nodes: NodeSample[], sampledAt: number, memoryCeiling?: number): CacheSimulatorEstimate { + const cells = nodes.flatMap((node) => node.strata.filter((stratum) => stratum.estimatedKeys > 0)); + const pooled = cells.flatMap((cell) => cell.sizes); + const pooledMean = mean(pooled); + const pooledVariance = sampleVariance(pooled); + const types = new Map(); + let total = 0; + let variance = nodes.reduce((sum, node) => sum + countVariance(node.strata, pooledMean), 0); + + for (const cell of cells) { + const n = cell.sizes.length; + const N = cell.estimatedKeys; + const cellMean = n > 0 ? mean(cell.sizes) : pooledMean; + const cellVariance = n < 2 && N > n ? pooledVariance : sampleVariance(cell.sizes); + total += N * cellMean; + variance += n === 0 ? N * N * pooledVariance : ((N * N) / n) * cellVariance * Math.max(0, 1 - n / N); + const totals = types.get(cell.type) ?? { keys: 0, weighted: 0, second: 0, sampled: 0 }; + totals.keys += N; + totals.weighted += N * cellMean; + totals.second += N * (cellVariance + cellMean * cellMean); + totals.sampled += n; + types.set(cell.type, totals); + } + + const tailAllowance = cells + .filter(isHeavyTailed) + .reduce((sum, cell) => sum + cell.estimatedKeys * (3 / cell.sizes.length) * cell.sizes.reduce((largest, size) => Math.max(largest, size), 0), 0); + + const byType: CacheSimulatorTypeEstimate[] = [...types].map(([type, totals]) => { + const avgBytes = totals.weighted / totals.keys; + return { + type, + keys: totals.keys, + avgBytes, + stdBytes: Math.sqrt(Math.max(0, totals.second / totals.keys - avgBytes * avgBytes)), + sampled: totals.sampled, + }; + }); + byType.sort((a, b) => b.keys * b.avgBytes - a.keys * a.avgBytes); + + const statistical = Z95 * Math.sqrt(variance) + tailAllowance; + const unsampled = + cells.some((cell) => cell.sizes.length < cell.estimatedKeys) || + nodes.some((node) => node.strata.reduce((sum, stratum) => sum + stratum.sizes.length, 0) < node.dbSize); + const unexplained = memoryCeiling !== undefined && unsampled ? Math.abs(memoryCeiling - total) : 0; + const keys = byType.reduce((sum, entry) => sum + entry.keys, 0); + const second = [...types.values()].reduce((sum, entry) => sum + entry.second, 0); + const avgKeyBytes = keys > 0 ? total / keys : 0; + return { + totalBytes: total, + ci95Bytes: Math.max(statistical, unexplained), + avgKeyBytes, + stdKeyBytes: keys > 0 ? Math.sqrt(Math.max(0, second / keys - avgKeyBytes * avgKeyBytes)) : 0, + sampleCount: byType.reduce((sum, entry) => sum + entry.sampled, 0), + byType, + sampledAt, + partial: nodes.some((node) => node.timedOut), + }; +} diff --git a/proprietary/cache-simulator/stratified-sampler.ts b/proprietary/cache-simulator/stratified-sampler.ts new file mode 100644 index 000000000..b6bb05be4 --- /dev/null +++ b/proprietary/cache-simulator/stratified-sampler.ts @@ -0,0 +1,268 @@ +import { neymanAllocation, splitBudget } from './allocation'; +import type { AllocationInput } from './allocation'; +import { adjustedSkewness, isHeavyTailed, sampleVariance } from './sample-stats'; +import { estimate } from './size-estimator'; +import type { NodeSample, SamplerClient, Stratum } from './sample-types'; + +export const SAMPLE_BUDGET = 4000; +export const SAMPLE_DEADLINE_MS = 60_000; +export const SAMPLE_HARD_CAP = 50_000; +export const TARGET_RELATIVE_CI = 0.05; +export const CENSUS_FACTOR = 10; +export const HEAVY_TAIL_MIN_SAMPLES = 20_000; +const COCHRAN_FACTOR = 25; +export const BATCH_SIZE = 100; +export const PILOT_SHARE = 0.25; +export const PILOT_MIN = 200; +const CLASSIFY_BATCH = 1000; +const SCAN_COUNT = '1000'; +const UNSUPPORTED = /unknown command|unknown subcommand|NOPERM|syntax error/i; + +export class SamplerUnsupportedError extends Error {} + +export interface SamplerNode { + name: string; + client: SamplerClient; +} + +export interface SamplerOptions { + budget: number; + deadlineMs: number; + hardCap?: number; + onProgress?: (sampled: number) => void; +} + +type Reply = [Error | null, unknown]; + +interface NodeState { + name: string; + dbSize: number; + client: SamplerClient; + deadline: number; + seen: Set; + classified: Set; + sizes: Map; + census: Map; + exactCounts: Map; + unmeasured: number; + wrapped: boolean; + counter: { sampled: number }; + onProgress?: (sampled: number) => void; + timedOut: boolean; +} + +export async function sampleNodes(nodes: SamplerNode[], options: SamplerOptions): Promise { + const deadline = Date.now() + options.deadlineMs; + const dbSizes = await Promise.all(nodes.map((node) => node.client.getDbSize())); + const quotas = splitBudget(options.budget, dbSizes); + const hardCap = options.hardCap ?? options.budget; + const caps = hardCap > options.budget ? splitBudget(hardCap, dbSizes) : quotas; + const counter = { sampled: 0 }; + return Promise.all( + nodes.map((node, index) => + sampleNode(quotas[index], caps[index], { + name: node.name, + dbSize: dbSizes[index], + client: node.client, + deadline, + seen: new Set(), + classified: new Set(), + sizes: new Map(), + census: new Map(), + exactCounts: new Map(), + unmeasured: 0, + wrapped: false, + counter, + onProgress: options.onProgress, + timedOut: false, + }), + ), + ); +} + +async function sampleNode(quota: number, cap: number, state: NodeState): Promise { + if (quota === 0) return { node: state.name, dbSize: state.dbSize, strata: [], exhausted: state.dbSize === 0, timedOut: false }; + const pilotQuota = Math.min(quota, Math.max(PILOT_MIN, Math.round(quota * PILOT_SHARE))); + await pilot(state, pilotQuota, quota * CENSUS_FACTOR); + if (state.wrapped) state.census.forEach((count, type) => state.exactCounts.set(type, count)); + if (state.wrapped && state.unmeasured === 0) return nodeResult(state, true); + await grow(state, neymanAllocation(quota - sampledTotal(state), allocationInputs(state))); + if (cap > quota) await extend(state, cap); + return nodeResult(state, false); +} + +async function grow(state: NodeState, extra: Map): Promise { + for (const [type, additional] of extra) { + if (additional === 0 || state.timedOut) continue; + const exact = await topUp(state, type, sizesOf(state, type).length + additional); + if (exact !== null) state.exactCounts.set(type, exact); + } +} + +async function extend(state: NodeState, cap: number): Promise { + while (!expired(state)) { + const measured = sampledTotal(state); + const sample = nodeResult(state, false); + const hasUnmeasured = sample.strata.some((stratum) => stratum.sizes.length < stratum.estimatedKeys); + if (measured >= cap || !hasUnmeasured) return; + const { totalBytes, ci95Bytes } = estimate([sample], 0); + const heavyShort = new Set( + sample.strata.filter((stratum) => isHeavyTailed(stratum) && stratum.sizes.length < Math.min(stratum.estimatedKeys, HEAVY_TAIL_MIN_SAMPLES)).map((stratum) => stratum.type), + ); + if (totalBytes <= 0 || (ci95Bytes / totalBytes <= TARGET_RELATIVE_CI && ciTrusted(sample) && heavyShort.size === 0)) return; + const inputs = allocationInputs(state).filter((input) => heavyShort.size === 0 || heavyShort.has(input.type)); + await grow(state, neymanAllocation(Math.min(measured, cap - measured), inputs)); + if (sampledTotal(state) <= measured) return; + } +} + +function ciTrusted(sample: NodeSample): boolean { + return sample.strata + .filter((stratum) => stratum.estimatedKeys > stratum.sizes.length) + .every((stratum) => stratum.sizes.length >= COCHRAN_FACTOR * adjustedSkewness(stratum.sizes) ** 2); +} + +async function pilot(state: NodeState, pilotQuota: number, censusLimit: number): Promise { + let cursor = '0'; + let dropped = false; + do { + if (expired(state)) return; + const [next, keys] = await scan(state.client, cursor, []); + cursor = next; + const fresh = keys.filter((key) => !state.classified.has(key)); + const page = fresh.slice(0, Math.max(0, censusLimit - state.classified.size)); + dropped ||= page.length < fresh.length; + page.forEach((key) => state.classified.add(key)); + let index = 0; + while (index < page.length && sampledTotal(state) < pilotQuota) { + const take = Math.min(BATCH_SIZE, pilotQuota - sampledTotal(state)); + await measure(state, page.slice(index, index + take), null); + index += take; + } + for (; index < page.length; index += CLASSIFY_BATCH) await classify(state, page.slice(index, index + CLASSIFY_BATCH)); + } while (cursor !== '0' && state.classified.size < censusLimit); + state.wrapped = cursor === '0' && !dropped; +} + +async function topUp(state: NodeState, type: string, target: number): Promise { + const ofType = new Set(); + let cursor = '0'; + do { + if (expired(state)) return null; + const [next, keys] = await scan(state.client, cursor, ['TYPE', type]); + cursor = next; + keys.forEach((key) => ofType.add(key)); + const fresh = keys.filter((key) => !state.seen.has(key)); + let index = 0; + while (index < fresh.length && sizesOf(state, type).length < target) { + const take = Math.min(BATCH_SIZE, target - sizesOf(state, type).length); + await measure(state, fresh.slice(index, index + take), type); + index += take; + } + if (sizesOf(state, type).length >= target) return cursor === '0' ? ofType.size : null; + } while (cursor !== '0'); + return ofType.size; +} + +function raiseOnError(reply: Reply | undefined): void { + if (!reply?.[0]) return; + if (UNSUPPORTED.test(reply[0].message)) throw new SamplerUnsupportedError(reply[0].message); + throw reply[0]; +} + +async function classify(state: NodeState, keys: string[]): Promise { + if (keys.length === 0) return; + const pipeline = state.client.getClient().pipeline(); + for (const key of keys) pipeline.call('TYPE', key); + const executed = await pipeline.exec(); + if (!executed) throw new Error('Pipeline returned no replies'); + (executed as Reply[]).forEach((reply) => { + raiseOnError(reply); + const type = reply ? String(reply[1]) : 'none'; + if (type === 'none') return; + state.census.set(type, (state.census.get(type) ?? 0) + 1); + state.unmeasured += 1; + }); +} + +async function measure(state: NodeState, keys: string[], knownType: string | null): Promise { + if (keys.length === 0) return; + const pipeline = state.client.getClient().pipeline(); + for (const key of keys) { + if (knownType === null) pipeline.call('TYPE', key); + pipeline.call('MEMORY', 'USAGE', key, 'SAMPLES', '0'); + } + const executed = await pipeline.exec(); + if (!executed) throw new Error('Pipeline returned no replies'); + const replies = executed as Reply[]; + const stride = knownType === null ? 2 : 1; + keys.forEach((key, index) => { + state.seen.add(key); + const typeReply: Reply = knownType === null ? replies[index * stride] : [null, knownType]; + const memoryReply = replies[index * stride + stride - 1]; + raiseOnError(typeReply); + raiseOnError(memoryReply); + const type = typeReply && !typeReply[0] ? String(typeReply[1]) : 'none'; + if (type === 'none') return; + if (knownType === null) state.census.set(type, (state.census.get(type) ?? 0) + 1); + const size = memoryReply && !memoryReply[0] && typeof memoryReply[1] === 'number' ? memoryReply[1] : null; + if (size === null) return; + const bucket = state.sizes.get(type); + if (bucket) bucket.push(size); + else state.sizes.set(type, [size]); + state.counter.sampled += 1; + }); + state.onProgress?.(state.counter.sampled); +} + +async function scan(client: SamplerClient, cursor: string, extra: string[]): Promise<[string, string[]]> { + try { + return (await client.call('SCAN', [cursor, 'COUNT', SCAN_COUNT, ...extra])) as [string, string[]]; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + if (UNSUPPORTED.test(message)) throw new SamplerUnsupportedError(message); + throw error; + } +} + +function buildStrata(state: NodeState): Stratum[] { + const exactTotal = [...state.exactCounts.values()].reduce((sum, count) => sum + count, 0); + const openObserved = [...state.census] + .filter(([type]) => !state.exactCounts.has(type)) + .reduce((sum, [, count]) => sum + count, 0); + const remaining = Math.max(openObserved, state.dbSize - exactTotal); + const types = new Set([...state.census.keys(), ...state.exactCounts.keys()]); + return [...types].map((type) => { + const observed = state.census.get(type) ?? 0; + const exactCount = state.exactCounts.get(type); + const estimatedKeys = exactCount ?? (openObserved > 0 ? (remaining * observed) / openObserved : 0); + return { type, estimatedKeys, exact: exactCount !== undefined, observed, sizes: [...sizesOf(state, type)] }; + }); +} + +function allocationInputs(state: NodeState): AllocationInput[] { + return buildStrata(state).map((stratum) => ({ + type: stratum.type, + estimatedKeys: stratum.estimatedKeys, + std: Math.sqrt(sampleVariance(stratum.sizes)), + sampled: stratum.sizes.length, + })); +} + +function nodeResult(state: NodeState, exhausted: boolean): NodeSample { + return { node: state.name, dbSize: state.dbSize, strata: buildStrata(state), exhausted, timedOut: state.timedOut }; +} + +function sizesOf(state: NodeState, type: string): number[] { + return state.sizes.get(type) ?? []; +} + +function sampledTotal(state: NodeState): number { + return [...state.sizes.values()].reduce((sum, sizes) => sum + sizes.length, 0); +} + +function expired(state: NodeState): boolean { + if (Date.now() < state.deadline) return false; + state.timedOut = true; + return true; +}