kimi-code / packages /minidb /bench /message-range.ts

Add files using upload-large-folder tool

4e23b01 verified about 1 month ago

6.2 kB

// bench/message-range.ts
//
// Benchmarks the "show me the user's most recent messages" query:
//   predicate:  role === 'user' AND ts >= since (last 7 days)
//   order:      ts DESC
//   limit:      50   (the typical "recent N" UI pattern)
//
// This is a message-granularity range query with a SMALL LIMIT, which is a
// different shape from the session-level benches:
//   - import-kimi-code / search-baseline index whole sessions (one doc/session).
//   - the metadata lookups (workspace lookup, dt range over sessions) return
//     ALL matches, so at small N a linear filter beats the index.
// Here, an ordered dt index can walk just `limit` entries (O(log N + limit))
// instead of scanning + sorting the whole set (O(N log M)). We measure at the
// real scale (~7.2k user messages) and at scaled-up sizes to find the
// crossover where the index starts winning.
//
// Run:  node --import tsx bench/message-range.ts

import fs from 'node:fs/promises';
import { existsSync, readFileSync } from 'node:fs';
import path from 'node:path';
import os from 'node:os';
import { MiniDb } from '../src/index.js';

const fmt = (n: number) => n.toLocaleString('en-US', { maximumFractionDigits: 0 });
const LIMIT = 50;
const WINDOW_DAYS = 30; // replicated timestamps are spread over the last 30 days
const SINCE_DAYS = 7; // "recent" = last 7 days

interface Msg {
  id: string;
  ts: number;
  role: string;
  text: string;
}

function loadRealMessages(): Msg[] {
  const DATA = path.join(os.homedir(), '.kimi-code');
  const lines = readFileSync(path.join(DATA, 'session_index.jsonl'), 'utf8').trim().split('\n');
  const out: Msg[] = [];
  for (const line of lines) {
    let meta: any;
    try {
      meta = JSON.parse(line);
    } catch {
      continue;
    }
    const wire = path.join(meta.sessionDir, 'agents', 'main', 'wire.jsonl');
    if (!existsSync(wire)) continue;
    let raw: string;
    try {
      raw = readFileSync(wire, 'utf8');
    } catch {
      continue;
    }
    for (const ln of raw.split('\n')) {
      if (!ln) continue;
      let o: any;
      try {
        o = JSON.parse(ln);
      } catch {
        continue;
      }
      if (o.type !== 'context.append_message' || !o.message) continue;
      const m = o.message;
      const ts = typeof o.time === 'number' ? o.time : 0;
      let text = '';
      for (const c of m.content || [])
        if (c && c.type === 'text' && typeof c.text === 'string') text += c.text;
      out.push({ id: `${meta.sessionId}:${out.length}`, ts, role: m.role || 'user', text });
    }
  }
  return out;
}

// Replicate the real messages to N, spreading timestamps uniformly over the
// last WINDOW_DAYS so the "last SINCE_DAYS" match fraction stays ~constant.
function scaleTo(real: Msg[], n: number): Msg[] {
  const now = Date.now();
  const span = WINDOW_DAYS * 864e5;
  const out: Msg[] = Array.from({ length: n }, (_, i) => {
    const r = real[i % real.length]!;
    return {
      id: `m${i}`,
      ts: now - Math.floor((i / n) * span),
      role: r.role,
      text: r.text,
    };
  });
  return out;
}

async function buildMinidb(msgs: Msg[]): Promise<{ db: any; buildMs: number; dir: string }> {
  const dir = await fs.mkdtemp(path.join(os.tmpdir(), 'minidb-msg-'));
  const db = await MiniDb.open({ dir, valueCodec: 'json', fsyncPolicy: 'no', autoCompact: false });
  const t0 = performance.now();
  // concurrent insert in chunks to avoid one giant promise array
  const CHUNK = 5000;
  for (let s = 0; s < msgs.length; s += CHUNK) {
    const p: Promise<unknown>[] = [];
    for (let i = s; i < Math.min(s + CHUNK, msgs.length); i++) {
      const m = msgs[i]!;
      p.push(db.set(m.id, { role: m.role, text: m.text }, { dt: { ts: m.ts } }));
    }
    await Promise.all(p);
  }
  return { db, buildMs: performance.now() - t0, dir };
}

function med(fn: () => void, runs = 9): number {
  const t: number[] = [];
  for (let i = 0; i < runs; i++) {
    const t0 = performance.now();
    fn();
    t.push(performance.now() - t0);
  }
  t.sort((a, b) => a - b);
  return t[(runs / 2) | 0]!;
}

async function runCase(label: string, msgs: Msg[]) {
  const since = Date.now() - SINCE_DAYS * 864e5;

  // ---- minidb dt index ----
  const { db, buildMs, dir } = await buildMinidb(msgs);
  let indexHits = 0;
  const indexMs = med(() => {
    const res = db.dtRange('ts', { gte: since, reverse: true, limit: LIMIT });
    indexHits = res.length;
  });
  // variant: return ALL matches (no limit), to mirror the earlier metadata case
  let allHits = 0;
  const indexAllMs = med(() => {
    const res = db.dtRange('ts', { gte: since });
    allHits = res.length;
  });
  if (global.gc) global.gc();
  const heap = process.memoryUsage().heapUsed;
  await db.close();
  await fs.rm(dir, { recursive: true, force: true });

  // ---- naive baseline ----
  let naiveHits = 0;
  const naiveMs = med(() => {
    const res = msgs.filter((m) => m.ts >= since).sort((a, b) => b.ts - a.ts).slice(0, LIMIT);
    naiveHits = res.length;
  });
  let naiveAllHits = 0;
  const naiveAllMs = med(() => {
    const res = msgs.filter((m) => m.ts >= since);
    naiveAllHits = res.length;
  });

  console.log(`\n--- ${label}: N=${fmt(msgs.length)} messages ---`);
  console.log(`  build: minidb ${(buildMs / 1000).toFixed(2)}s  |  heap ${(heap / 1024 / 1024).toFixed(0)} MiB`);
  console.log(
    `  top-${LIMIT} recent:  minidb dtRange ${indexMs.toFixed(3)} ms (${indexHits})  vs  naive filter+sort ${naiveMs.toFixed(3)} ms (${naiveHits})  ->  ${(naiveMs / indexMs).toFixed(1)}x`,
  );
  console.log(
    `  ALL matches:        minidb dtRange ${indexAllMs.toFixed(3)} ms (${fmt(allHits)})  vs  naive filter ${naiveAllMs.toFixed(3)} ms (${fmt(naiveAllHits)})  ->  ${(naiveAllMs / indexAllMs).toFixed(1)}x`,
  );
}

async function main() {
  const real = loadRealMessages();
  console.log(`loaded ${fmt(real.length)} real user messages`);
  const sizes = [real.length, 100_000, 1_000_000];
  for (const n of sizes) {
    const msgs = n === real.length ? real : scaleTo(real, n);
    await runCase(n === real.length ? 'real data' : `scaled`, msgs);
  }
  console.log('\ndone.');
}

main().catch((e) => {
  console.error(e);
  process.exit(1);
});

	// bench/message-range.ts
	//
	// Benchmarks the "show me the user's most recent messages" query:
	// predicate: role === 'user' AND ts >= since (last 7 days)
	// order: ts DESC
	// limit: 50 (the typical "recent N" UI pattern)
	//
	// This is a message-granularity range query with a SMALL LIMIT, which is a
	// different shape from the session-level benches:
	// - import-kimi-code / search-baseline index whole sessions (one doc/session).
	// - the metadata lookups (workspace lookup, dt range over sessions) return
	// ALL matches, so at small N a linear filter beats the index.
	// Here, an ordered dt index can walk just `limit` entries (O(log N + limit))
	// instead of scanning + sorting the whole set (O(N log M)). We measure at the
	// real scale (~7.2k user messages) and at scaled-up sizes to find the
	// crossover where the index starts winning.
	//
	// Run: node --import tsx bench/message-range.ts

	import fs from 'node:fs/promises';
	import { existsSync, readFileSync } from 'node:fs';
	import path from 'node:path';
	import os from 'node:os';
	import { MiniDb } from '../src/index.js';

	const fmt = (n: number) => n.toLocaleString('en-US', { maximumFractionDigits: 0 });
	const LIMIT = 50;
	const WINDOW_DAYS = 30; // replicated timestamps are spread over the last 30 days
	const SINCE_DAYS = 7; // "recent" = last 7 days

	interface Msg {
	id: string;
	ts: number;
	role: string;
	text: string;
	}

	function loadRealMessages(): Msg[] {
	const DATA = path.join(os.homedir(), '.kimi-code');
	const lines = readFileSync(path.join(DATA, 'session_index.jsonl'), 'utf8').trim().split('\n');
	const out: Msg[] = [];
	for (const line of lines) {
	let meta: any;
	try {
	meta = JSON.parse(line);
	} catch {
	continue;
	}
	const wire = path.join(meta.sessionDir, 'agents', 'main', 'wire.jsonl');
	if (!existsSync(wire)) continue;
	let raw: string;
	try {
	raw = readFileSync(wire, 'utf8');
	} catch {
	continue;
	}
	for (const ln of raw.split('\n')) {
	if (!ln) continue;
	let o: any;
	try {
	o = JSON.parse(ln);
	} catch {
	continue;
	}
	if (o.type !== 'context.append_message' \|\| !o.message) continue;
	const m = o.message;
	const ts = typeof o.time === 'number' ? o.time : 0;
	let text = '';
	for (const c of m.content \|\| [])
	if (c && c.type === 'text' && typeof c.text === 'string') text += c.text;
	out.push({ id: `${meta.sessionId}:${out.length}`, ts, role: m.role \|\| 'user', text });
	}
	}
	return out;
	}

	// Replicate the real messages to N, spreading timestamps uniformly over the
	// last WINDOW_DAYS so the "last SINCE_DAYS" match fraction stays ~constant.
	function scaleTo(real: Msg[], n: number): Msg[] {
	const now = Date.now();
	const span = WINDOW_DAYS * 864e5;
	const out: Msg[] = Array.from({ length: n }, (_, i) => {
	const r = real[i % real.length]!;
	return {
	id: `m${i}`,
	ts: now - Math.floor((i / n) * span),
	role: r.role,
	text: r.text,
	};
	});
	return out;
	}

	async function buildMinidb(msgs: Msg[]): Promise<{ db: any; buildMs: number; dir: string }> {
	const dir = await fs.mkdtemp(path.join(os.tmpdir(), 'minidb-msg-'));
	const db = await MiniDb.open({ dir, valueCodec: 'json', fsyncPolicy: 'no', autoCompact: false });
	const t0 = performance.now();
	// concurrent insert in chunks to avoid one giant promise array
	const CHUNK = 5000;
	for (let s = 0; s < msgs.length; s += CHUNK) {
	const p: Promise<unknown>[] = [];
	for (let i = s; i < Math.min(s + CHUNK, msgs.length); i++) {
	const m = msgs[i]!;
	p.push(db.set(m.id, { role: m.role, text: m.text }, { dt: { ts: m.ts } }));
	}
	await Promise.all(p);
	}
	return { db, buildMs: performance.now() - t0, dir };
	}

	function med(fn: () => void, runs = 9): number {
	const t: number[] = [];
	for (let i = 0; i < runs; i++) {
	const t0 = performance.now();
	fn();
	t.push(performance.now() - t0);
	}
	t.sort((a, b) => a - b);
	return t[(runs / 2) \| 0]!;
	}

	async function runCase(label: string, msgs: Msg[]) {
	const since = Date.now() - SINCE_DAYS * 864e5;

	// ---- minidb dt index ----
	const { db, buildMs, dir } = await buildMinidb(msgs);
	let indexHits = 0;
	const indexMs = med(() => {
	const res = db.dtRange('ts', { gte: since, reverse: true, limit: LIMIT });
	indexHits = res.length;
	});
	// variant: return ALL matches (no limit), to mirror the earlier metadata case
	let allHits = 0;
	const indexAllMs = med(() => {
	const res = db.dtRange('ts', { gte: since });
	allHits = res.length;
	});
	if (global.gc) global.gc();
	const heap = process.memoryUsage().heapUsed;
	await db.close();
	await fs.rm(dir, { recursive: true, force: true });

	// ---- naive baseline ----
	let naiveHits = 0;
	const naiveMs = med(() => {
	const res = msgs.filter((m) => m.ts >= since).sort((a, b) => b.ts - a.ts).slice(0, LIMIT);
	naiveHits = res.length;
	});
	let naiveAllHits = 0;
	const naiveAllMs = med(() => {
	const res = msgs.filter((m) => m.ts >= since);
	naiveAllHits = res.length;
	});

	console.log(`\n--- ${label}: N=${fmt(msgs.length)} messages ---`);
	console.log(` build: minidb ${(buildMs / 1000).toFixed(2)}s \| heap ${(heap / 1024 / 1024).toFixed(0)} MiB`);
	console.log(
	` top-${LIMIT} recent: minidb dtRange ${indexMs.toFixed(3)} ms (${indexHits}) vs naive filter+sort ${naiveMs.toFixed(3)} ms (${naiveHits}) -> ${(naiveMs / indexMs).toFixed(1)}x`,
	);
	console.log(
	` ALL matches: minidb dtRange ${indexAllMs.toFixed(3)} ms (${fmt(allHits)}) vs naive filter ${naiveAllMs.toFixed(3)} ms (${fmt(naiveAllHits)}) -> ${(naiveAllMs / indexAllMs).toFixed(1)}x`,
	);
	}

	async function main() {
	const real = loadRealMessages();
	console.log(`loaded ${fmt(real.length)} real user messages`);
	const sizes = [real.length, 100_000, 1_000_000];
	for (const n of sizes) {
	const msgs = n === real.length ? real : scaleTo(real, n);
	await runCase(n === real.length ? 'real data' : `scaled`, msgs);
	}
	console.log('\ndone.');
	}

	main().catch((e) => {
	console.error(e);
	process.exit(1);
	});