1// Code search, measured and queried the way the browser does it: the built2// www served by demo/serve.ts, and site/assets/codesearch.ts (sqlite-wasm,3// the range VFS, the query) reading code.db over HTTP. Node has no4// synchronous XHR, so a worker thread fetches each range while this one5// waits on an Atomics flag, the way a browser worker waits on its XHR.6//7// node test/code-search-bench.ts bench <state> [query...] sizes, requests, bytes8// node test/code-search-bench.ts http <state> <owner/repo> <query> [limit]9// node test/code-search-bench.ts sql <code.db> <query> node:sqlite, no HTTP10//11// <state> is a forge's .forge directory (www/ inside). `http` and `sql` print12// "path:line: text" per hit; `http` then prints what the query fetched and13// how the server answered, for code.db and for sqlite-wasm's .wasm.14import fs from 'node:fs';15import net from 'node:net';16import path from 'node:path';17import zlib from 'node:zlib';18import { spawn, type ChildProcess } from 'node:child_process';19import { Worker } from 'node:worker_threads';20import { fileURLToPath } from 'node:url';21import { CodeIndex, indexed, installVfs, lineMatch, matchAll, MORE, words, type GetRange, type Sqlite3 } from '../site/assets/codesearch.ts';2223const here = path.dirname(fileURLToPath(import.meta.url));24const [mode, ...args] = process.argv.slice(2);2526// Ranges are at most 32 blocks of 16K; a whole small file is under 256K.27const MAX = 4 << 20;28const FETCHER = `29const { parentPort } = require('node:worker_threads');30parentPort.on('message', async ({ url, from, to, sab, auth }) => {31 const ctl = new Int32Array(sab, 0, 4);32 const body = new Uint8Array(sab, 1024);33 const meta = new Uint8Array(sab, 16, 1008);34 try {35 const r = await fetch(url, { headers: { Range: from < 0 ? 'bytes=' + from : 'bytes=' + from + '-' + (to - 1), ...auth } });36 const data = new Uint8Array(await r.arrayBuffer());37 if (data.length > body.length) throw new Error('response too large: ' + data.length);38 body.set(data);39 const info = JSON.stringify({ status: r.status, range: r.headers.get('content-range'), type: r.headers.get('content-type'), enc: r.headers.get('content-encoding'), tag: r.headers.get('etag') ?? r.headers.get('last-modified') ?? '' });40 meta.set(new TextEncoder().encode(info).subarray(0, 1000));41 ctl[1] = data.length;42 ctl[2] = Math.min(info.length, 1000);43 } catch (e) {44 const info = JSON.stringify({ error: String(e) });45 meta.set(new TextEncoder().encode(info).subarray(0, 1000));46 ctl[1] = -1;47 ctl[2] = Math.min(info.length, 1000);48 }49 Atomics.store(ctl, 0, 1);50 Atomics.notify(ctl, 0);51});`;5253// A blocking range fetch.54function syncRange(): { get: GetRange; stop: () => void } {55 const w = new Worker(FETCHER, { eval: true });56 const sab = new SharedArrayBuffer(1024 + MAX);57 const ctl = new Int32Array(sab, 0, 4);58 const get: GetRange = (url, from, to) => {59 Atomics.store(ctl, 0, 0);60 w.postMessage({ url, from, to, sab, auth });61 Atomics.wait(ctl, 0, 0);62 const info = JSON.parse(new TextDecoder().decode(new Uint8Array(sab, 16, ctl[2]).slice()));63 if (ctl[1] < 0) throw new Error(info.error);64 if (info.status !== 200 && info.status !== 206) throw new Error(`${url}: HTTP ${info.status}`);65 if (info.enc && info.enc !== 'identity') throw new Error(`${url} is served with Content-Encoding: ${info.enc}`);66 const data = new Uint8Array(sab, 1024, ctl[1]).slice();67 const whole = info.status === 200;68 const m = /(\d+)-\d+\/(\d+)$/.exec(info.range ?? '');69 return { whole, data, start: whole ? 0 : Number(m?.[1]), size: whole ? data.length : Number(m?.[2]), tag: info.tag };70 };71 return { get, stop: () => void w.terminate() };72}7374// FORGE_BENCH_TOKEN: a token for every request, to read a private repo's75// code.db through the private server (which demo/serve.ts hands misses to).76const auth: Record<string, string> = process.env.FORGE_BENCH_TOKEN ? { Authorization: `Bearer ${process.env.FORGE_BENCH_TOKEN}` } : {};77// The worker's prefetches fetch, as a browser's would with its cookie.78const plainFetch = globalThis.fetch;79globalThis.fetch = (url, init) => plainFetch(url, { ...init, headers: { ...(init?.headers as Record<string, string>), ...auth } });8081// demo/serve.ts on a free port (it also takes the next one, for events).82async function serve(state: string): Promise<{ base: string; stop: () => void }> {83 const free = (): Promise<number> => new Promise((ok, no) => {84 const s = net.createServer().once('error', no).listen(0, '127.0.0.1', () => {85 const p = (s.address() as net.AddressInfo).port;86 s.close(() => ok(p));87 });88 });89 for (let tries = 0; ; tries++) {90 const port = await free();91 const child: ChildProcess = spawn(process.execPath, [path.join(here, '..', 'demo', 'serve.ts'), String(port)], { env: { ...process.env, FORGE_STATE: state }, stdio: ['ignore', 'pipe', 'pipe'] });92 const up = await new Promise<boolean>((ok) => {93 child.stdout!.on('data', (d: Buffer) => d.toString().includes('demo forge on') && ok(true));94 child.on('exit', () => ok(false));95 });96 if (up) return { base: `http://127.0.0.1:${port}`, stop: () => void child.kill() };97 if (tries > 5) throw new Error('demo/serve.ts did not start');98 }99}100101async function wasm(): Promise<Sqlite3> {102 const init = (await import('@sqlite.org/sqlite-wasm')).default as unknown as (o: object) => Promise<Sqlite3>;103 return init({ print: () => {}, printErr: () => {} });104}105106const kb = (n: number): string => (n < 1024 ? `${n}B` : n < 1 << 20 ? `${(n / 1024).toFixed(1)}K` : `${(n / (1 << 20)).toFixed(2)}M`);107108// node:sqlite, without its experimental warning.109function open(db: string): import('node:sqlite').DatabaseSync {110 const emit = process.emitWarning;111 process.emitWarning = (() => {}) as typeof emit;112 const { DatabaseSync } = process.getBuiltinModule('node:sqlite') as typeof import('node:sqlite');113 process.emitWarning = emit;114 return new DatabaseSync(db, { readOnly: true });115}116117function sqlHits(db: string, q: string): string[] {118 const ws = words(q);119 if (!indexed(ws).length) return [];120 const d = open(db);121 const rows = d.prepare('select f.path, c.line, c.text from code join chunk c on c.id = code.rowid join file f on f.id = c.file where code match ?').all(matchAll(ws)) as { path: string; line: number; text: string }[];122 const out: string[] = [];123 for (const r of rows) r.text.split('\n').forEach((l, i) => lineMatch(l, ws) >= 0 && out.push(`${r.path}:${r.line + i}: ${l.trim()}`));124 return out.sort();125}126127if (mode === 'sql') {128 for (const h of sqlHits(args[0], args[1])) console.log(h);129} else if (mode === 'http') {130 const [state, slug, q, limit] = args;131 const srv = await serve(state);132 const f = syncRange();133 try {134 const sqlite3 = await wasm();135 const ix = new CodeIndex(sqlite3, installVfs(sqlite3), `${srv.base}/${slug}/code.db`, f.get);136 const r = await ix.search(q, Number(limit || MORE));137 for (const h of r.hits) console.log(`${h.path}:${h.line}: ${h.text}`);138 console.log(`requests ${r.requests} rounds ${r.rounds} bytes ${r.bytes}`);139 const head = await fetch(`${srv.base}/${slug}/code.db`, { headers: { Range: 'bytes=0-99', ...auth } });140 console.log(`range ${head.status} ${head.headers.get('content-range')} type ${head.headers.get('content-type')} encoding ${head.headers.get('content-encoding') ?? 'none'} length ${(await head.arrayBuffer()).byteLength}`);141 console.log(`cache ${head.headers.get('cache-control')}`);142 // What the browser instantiates sqlite-wasm from.143 const wasmFile = fs.readdirSync(path.join(state, 'www', 'assets')).find((a) => a.endsWith('.wasm'));144 const w = await fetch(`${srv.base}/assets/${wasmFile}`, { method: 'HEAD' });145 console.log(`wasm ${w.status} type ${w.headers.get('content-type')} cache ${w.headers.get('cache-control')}`);146 } finally {147 f.stop();148 srv.stop();149 }150} else if (mode === 'bench') {151 const [state, ...qs] = args;152 const queries = qs.length ? qs : ['function', 'escapeHtml', 'createHash', 'readFileSync(', 'TODO', 'import fs from', 'Cache-Control', 'zzqqxx'];153 const www = path.join(state, 'www');154 const dbs: string[] = [];155 for (const owner of fs.readdirSync(www)) {156 const d = path.join(www, owner);157 if (!fs.statSync(d).isDirectory()) continue;158 for (const repo of fs.readdirSync(d)) if (fs.existsSync(path.join(d, repo, 'code.db'))) dbs.push(`${owner}/${repo}`);159 }160 const assets = fs.readdirSync(path.join(www, 'assets'));161 for (const a of assets.filter((f) => /^(sqlite3|codesearch)\./.test(f))) {162 const data = fs.readFileSync(path.join(www, 'assets', a));163 console.log(`asset ${a}: ${kb(data.length)}, ${kb(zlib.gzipSync(data).length)} gzipped (once, then cached for good)`);164 }165 const srv = await serve(state);166 const f = syncRange();167 try {168 const sqlite3 = await wasm();169 const vfs = installVfs(sqlite3);170 for (const slug of dbs) {171 const file = path.join(www, slug, 'code.db');172 const d = open(file);173 const { files, text } = d.prepare('select (select count(*) from file) files, (select sum(length(cast(text as blob))) from chunk) text').get() as { files: number; text: number };174 d.close();175 const size = fs.statSync(file).size;176 console.log(`\n${slug}: code.db ${kb(size)} for ${files} files, ${kb(text)} of text (${(size / text).toFixed(1)}x)`);177 console.log('query'.padEnd(16) + 'hits'.padStart(5) + ' cold: rounds requests bytes again after the others: rounds requests bytes');178 const url = `${srv.base}/${slug}/code.db`;179 for (const q of queries) {180 // Cold: a fresh index, as on a page that never searched code.181 const cold = new CodeIndex(sqlite3, vfs, url, f.get);182 const c = await cold.search(q);183 const again = await cold.search(q);184 cold.close();185 // After the others: the queries before it warmed some pages.186 const later = new CodeIndex(sqlite3, vfs, url, f.get);187 for (const o of queries) if (o !== q) await later.search(o);188 const l = await later.search(q);189 later.close();190 const col = (n: number | string, w: number): string => String(n).padStart(w);191 console.log(`${q.padEnd(16)}${col(c.hits.length + (c.more ? '+' : ''), 5)} ${col(c.rounds, 13)} ${col(c.requests, 8)} ${col(kb(c.bytes), 6)} ${col(again.requests, 6)} ${col(l.rounds, 25)} ${col(l.requests, 8)} ${col(kb(l.bytes), 6)}`);192 }193 }194 } finally {195 f.stop();196 srv.stop();197 }198} else {199 console.error('usage: node test/code-search-bench.ts bench <state> [query...] | http <state> <owner/repo> <query> [limit] | sql <code.db> <query>');200 process.exit(2);201}