diff --git a/.changeset/audit-hardening.md b/.changeset/audit-hardening.md new file mode 100644 index 0000000..5413fc7 --- /dev/null +++ b/.changeset/audit-hardening.md @@ -0,0 +1,7 @@ +--- +"webgpu-search": patch +--- + +Harden correctness and safety from a repo audit: shared GPU device pool no longer double-acquires or destroys a live device under concurrent creation; GPU buffers are clamped to `maxStorageBufferBindingSize` (oversized corpora now fall back to CPU instead of silently returning 0 results); batch mutations are atomic when a filter getter throws; high-cardinality string filters use bounded memory; worker INIT/RESTORE keep the previous index on failure and handle messages in order; snapshots reject duplicate/invalid doc IDs and accept typed-array buffers; custom IndexedDB store names are created on existing databases. + +Behaviour changes: boolean filter fields accept only `true`/`false` (or `"true"`/`"false"`); blank number filter bounds throw `InvalidFilterError`; numeric IDs sort before string IDs on score ties; field getters ignore `Object.prototype` properties. diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2e0f28a..6f95761 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -6,6 +6,9 @@ on: pull_request: branches: [ main ] +permissions: + contents: read + jobs: test-and-build: runs-on: ubuntu-latest @@ -13,10 +16,10 @@ jobs: steps: - name: Checkout Code - uses: actions/checkout@v4 + uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4 - name: Setup Bun - uses: oven-sh/setup-bun@v2 + uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2 with: bun-version: 1.4.2 @@ -107,9 +110,12 @@ jobs: - name: Run CPU Baseline Pin (support-matrix reference contract) run: bun run test:cpu-baseline + - name: Run Audit Regression Tests (index integrity, GPU pool, worker/snapshot) + run: bun run test:regressions + - name: Setup Chrome (browser regression gate) id: setup-chrome - uses: browser-actions/setup-chrome@v1 + uses: browser-actions/setup-chrome@c785b87e244131f27c9f19c1a33e2ead956ab7ce # v1 with: chrome-version: stable @@ -127,7 +133,7 @@ jobs: # fall back to PATH lookup for self-hosted / local `act` runs. CHROME_BIN="${{ steps.setup-chrome.outputs.chrome-path }}" if [ -z "$CHROME_BIN" ] || [ ! -x "$CHROME_BIN" ]; then - CHROME_BIN="$(command -v google-chrome-stable || command -v google-chrome || command -v chromium || command -v chrome || echo \"${CHROME_PATH:-}\")" + CHROME_BIN="$(command -v google-chrome-stable || command -v google-chrome || command -v chromium || command -v chrome || echo "${CHROME_PATH:-}")" fi if [ -z "$CHROME_BIN" ] || [ ! -x "$CHROME_BIN" ]; then echo "Chrome/Chromium binary not found or not executable (got '${CHROME_BIN:-}'); failing browser regression gate." >&2 @@ -156,7 +162,7 @@ jobs: - name: Upload Browser Test Logs if: always() - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4 with: name: browser-test-logs path: | diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index e1dbcef..f01f025 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -1,36 +1,45 @@ name: Release +# Runs only after the CI workflow succeeded for a push to main, so an +# untested commit can never be versioned or published. on: - push: - branches: - - main - paths: - - 'packages/**' - - '.changeset/**' + workflow_run: + workflows: [ CI ] + types: [ completed ] + branches: [ main ] concurrency: ${{ github.workflow }}-${{ github.ref }} +permissions: + contents: read + jobs: release: name: Release + if: >- + github.event.workflow_run.conclusion == 'success' && + github.event.workflow_run.event == 'push' && + github.event.workflow_run.head_branch == 'main' runs-on: ubuntu-latest + # No `id-token: write`: npm provenance is not configured. permissions: contents: write pull-requests: write - id-token: write steps: - name: Checkout Repo - uses: actions/checkout@v4 + uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4 with: + # Release exactly the commit CI validated. + ref: ${{ github.event.workflow_run.head_sha }} fetch-depth: 0 - name: Setup Bun - uses: oven-sh/setup-bun@v2 + uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2 with: bun-version: 1.4.2 - name: Setup Node - uses: actions/setup-node@v4 + uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4 with: node-version: 20 registry-url: 'https://registry.npmjs.org' @@ -43,7 +52,7 @@ jobs: - name: Create Release Pull Request or Publish to npm id: changesets - uses: changesets/action@v1 + uses: changesets/action@a45c4d594aa4e2c509dc14a9f2b3b67ba3780d0d # v1.9.0 with: publish: bun run release version: bun run version-packages diff --git a/api/submit-benchmark.ts b/api/submit-benchmark.ts index bf11cca..4997d13 100644 --- a/api/submit-benchmark.ts +++ b/api/submit-benchmark.ts @@ -3,17 +3,26 @@ * * - Write token stays server-side (env `UPSTASH_REDIS_REST_URL` / * `UPSTASH_REDIS_REST_TOKEN`); the browser never sees it. - * - GET ?fingerprint= -> { known: boolean } (dedup probe for the - * "untested hardware" share button). + * - GET ?fingerprint= -> { known: boolean } (has any submission been + * stored for this hardware/workload fingerprint). * - POST { fingerprint, adapter, gpuTier, browser, os, corpusSize, mode, - * query, medians... } -> { ok: true } | { ok: true, duplicate: true }. - * First write wins (SET NX); fingerprints indexed in a Redis set. + * query, medians... } -> 200 { ok: true } | 409 { ok: true, duplicate: true }. + * Submissions are appended to a per-fingerprint capped list + * (`bench:results:`, newest RESULTS_PER_FP kept) so no single + * submitter can pre-claim a fingerprint. Dedup is per (fingerprint, + * hashed IP) via SET NX with a TTL; fingerprints are indexed in the + * `bench:fingerprints` set. + * - Only whitelisted, validated fields are persisted; bodies whose JSON + * encoding exceeds 8 KiB are rejected with 413 (never truncated). * - Best-effort IP rate limit (10/min); fail-open on limiter errors so a - * Redis hiccup never blocks the benchmark itself. - * - Missing env -> 503 { error: 'not-configured' }; the client hides the - * share button in that case. + * Redis hiccup never blocks the benchmark itself. Store errors on the + * read/write path return 502. + * - Missing env -> 503 { error: 'not-configured' }. */ +// Node global (no @types/node in this workspace). +declare const process: { env: Record }; + type VercelReq = { method?: string; query?: Record; @@ -27,9 +36,18 @@ type VercelRes = { setHeader: (name: string, value: string) => void; }; +type Obj = Record; +type RedisCommand = Array; +type StoreReply = { ok: boolean; result: unknown }; + const ALLOWED_SIZES = new Set([100000, 200000, 500000, 1000000, 2000000]); +const GPU_TIERS = ['discrete', 'integrated', 'software', 'unknown']; const FP_RE = /^[a-z0-9-]{8,64}$/; const MAX_BODY_BYTES = 8192; +const RESULTS_PER_FP = 50; +const DEDUP_TTL_S = 60 * 60 * 24 * 30; +const RL_WINDOW_S = 60; +const RL_MAX = 10; function env(name: string): string { return (process.env[name] ?? '').trim().replace(/\/$/, ''); @@ -45,79 +63,165 @@ function clientIp(req: VercelReq): string { return (typeof r === 'string' ? r : 'unknown').slice(0, 64); } -async function upstash(path: string, init?: RequestInit): Promise<{ ok: boolean; data: unknown }> { +async function sha256Hex(input: string): Promise { + const digest = await crypto.subtle.digest('SHA-256', new TextEncoder().encode(input)); + return Array.from(new Uint8Array(digest), b => b.toString(16).padStart(2, '0')).join(''); +} + +function byteLength(s: string): number { + return new TextEncoder().encode(s).length; +} + +/** POSTs a JSON body to Upstash REST; network/HTTP/parse failures -> null. */ +async function upstashPost(path: string, body: unknown): Promise { const url = env('UPSTASH_REDIS_REST_URL'); const token = env('UPSTASH_REDIS_REST_TOKEN'); - if (!url || !token) return { ok: false, data: null }; - const res = await fetch(`${url}${path}`, { - ...init, - headers: { authorization: `Bearer ${token}`, ...(init?.headers ?? {}) } - }); - if (!res.ok) return { ok: false, data: null }; - return { ok: true, data: (await res.json().catch(() => null)) as unknown }; + if (!url || !token) return null; + try { + const res = await fetch(`${url}${path}`, { + method: 'POST', + headers: { authorization: `Bearer ${token}`, 'content-type': 'application/json' }, + body: JSON.stringify(body) + }); + if (!res.ok) return null; + return (await res.json().catch(() => null)) as unknown; + } catch { + return null; + } } -function upstashResult(data: unknown): number | string | null { - if (data && typeof data === 'object' && 'result' in (data as Record)) { - const r = (data as Record).result; - if (typeof r === 'number' || typeof r === 'string') return r; +function toReply(data: unknown): StoreReply { + if (data && typeof data === 'object' && !Array.isArray(data)) { + const d = data as Obj; + if ('error' in d) return { ok: false, result: null }; + if ('result' in d) return { ok: true, result: d.result }; } - return null; + return { ok: false, result: null }; +} + +async function redis(command: RedisCommand): Promise { + return toReply(await upstashPost('', command)); +} + +async function pipeline(commands: RedisCommand[]): Promise { + const data = await upstashPost('/pipeline', commands); + if (!Array.isArray(data) || data.length !== commands.length) return null; + return data.map(toReply); +} + +function isObj(v: unknown): v is Obj { + return typeof v === 'object' && v !== null && !Array.isArray(v); } function isFiniteMs(v: unknown): v is number { return typeof v === 'number' && Number.isFinite(v) && v >= 0 && v <= 60000; } +function isCount(v: unknown): v is number { + return typeof v === 'number' && Number.isInteger(v) && v >= 0 && v <= 2000000; +} + function isShortStr(v: unknown, max = 300): v is string { return typeof v === 'string' && v.length <= max; } -function validatePayload(p: Record): string | null { - if (typeof p.fingerprint !== 'string' || !FP_RE.test(p.fingerprint)) return 'bad fingerprint'; - if (typeof p.adapter !== 'object' || p.adapter === null) return 'bad adapter'; - const a = p.adapter as Record; +/** Validates and returns a record built ONLY from whitelisted fields. */ +function buildRecord(p: Obj): { record: Obj } | { error: string } { + if (typeof p.fingerprint !== 'string' || !FP_RE.test(p.fingerprint)) return { error: 'bad fingerprint' }; + if (!isObj(p.adapter)) return { error: 'bad adapter' }; + const a = p.adapter; for (const k of ['vendor', 'device', 'architecture', 'renderer']) { - if (!isShortStr(a[k])) return `bad adapter.${k}`; - } - if (a['adapterType'] !== undefined && !isShortStr(a['adapterType'], 64)) return 'bad adapter.adapterType'; - const t = p.gpuTier as Record | undefined; - if (!t || !['discrete', 'integrated', 'software', 'unknown'].includes(String(t['tier']))) return 'bad gpuTier'; - if (!isShortStr(p['browser'], 32) || !isShortStr(p['os'], 32)) return 'bad browser/os'; - if (typeof p.corpusSize !== 'number' || !ALLOWED_SIZES.has(p.corpusSize)) return 'bad corpusSize'; - if (p['mode'] !== 'fuzzy' && p['mode'] !== 'substring') return 'bad mode'; - if (!isShortStr(p['query'], 120)) return 'bad query'; + if (!isShortStr(a[k])) return { error: `bad adapter.${k}` }; + } + if (a.adapterType !== undefined && !isShortStr(a.adapterType, 64)) return { error: 'bad adapter.adapterType' }; + if (!isObj(p.gpuTier) || typeof p.gpuTier.tier !== 'string' || !GPU_TIERS.includes(p.gpuTier.tier)) { + return { error: 'bad gpuTier' }; + } + const t = p.gpuTier; + if (t.confidence !== undefined && !isShortStr(t.confidence, 32)) return { error: 'bad gpuTier.confidence' }; + if (t.reason !== undefined && !isShortStr(t.reason, 300)) return { error: 'bad gpuTier.reason' }; + if (!isShortStr(p.browser, 32) || !isShortStr(p.os, 32)) return { error: 'bad browser/os' }; + if (typeof p.corpusSize !== 'number' || !ALLOWED_SIZES.has(p.corpusSize)) return { error: 'bad corpusSize' }; + if (p.mode !== 'fuzzy' && p.mode !== 'substring') return { error: 'bad mode' }; + if (!isShortStr(p.query, 120)) return { error: 'bad query' }; for (const k of ['gpuMedianMs', 'gpuP95Ms', 'cpuMedianMs', 'cpuP95Ms']) { - if (!isFiniteMs(p[k])) return `bad ${k}`; + if (!isFiniteMs(p[k])) return { error: `bad ${k}` }; } for (const k of ['gpuMatches', 'cpuMatches']) { - if (typeof p[k] !== 'number' || !Number.isInteger(p[k]) || (p[k] as number) < 0 || (p[k] as number) > 2000000) { - return `bad ${k}`; - } + if (!isCount(p[k])) return { error: `bad ${k}` }; } - if (p['externals'] !== undefined) { - if (!Array.isArray(p['externals']) || (p['externals'] as unknown[]).length > 8) return 'bad externals'; - for (const e of p['externals'] as Array>) { - if (!isShortStr(e['name'], 32) || !isFiniteMs(e['medianMs']) || !isFiniteMs(e['p95Ms'])) return 'bad external entry'; + const externals: Obj[] = []; + if (p.externals !== undefined) { + if (!Array.isArray(p.externals) || p.externals.length > 8) return { error: 'bad externals' }; + for (const e of p.externals as unknown[]) { + if (!isObj(e) || !isShortStr(e.name, 32) || !isFiniteMs(e.medianMs) || !isFiniteMs(e.p95Ms)) { + return { error: 'bad external entry' }; + } + if (e.matches !== undefined && !isCount(e.matches)) return { error: 'bad external entry' }; + if (e.ran !== undefined && typeof e.ran !== 'boolean') return { error: 'bad external entry' }; + externals.push({ + name: e.name, + medianMs: e.medianMs, + p95Ms: e.p95Ms, + ...(e.matches !== undefined ? { matches: e.matches } : {}), + ...(e.ran !== undefined ? { ran: e.ran } : {}) + }); } } - return null; + return { + record: { + fingerprint: p.fingerprint, + adapter: { + vendor: a.vendor, + device: a.device, + architecture: a.architecture, + renderer: a.renderer, + ...(a.adapterType !== undefined ? { adapterType: a.adapterType } : {}) + }, + gpuTier: { + tier: t.tier, + ...(t.confidence !== undefined ? { confidence: t.confidence } : {}), + ...(t.reason !== undefined ? { reason: t.reason } : {}) + }, + browser: p.browser, + os: p.os, + corpusSize: p.corpusSize, + mode: p.mode, + query: p.query, + gpuMedianMs: p.gpuMedianMs, + gpuP95Ms: p.gpuP95Ms, + gpuMatches: p.gpuMatches, + cpuMedianMs: p.cpuMedianMs, + cpuP95Ms: p.cpuP95Ms, + cpuMatches: p.cpuMatches, + externals, + receivedAt: new Date().toISOString() + } + }; } -async function readJsonBody(req: VercelReq): Promise | null> { - if (req.body && typeof req.body === 'object' && !Array.isArray(req.body)) { - return req.body as Record; - } +type BodyResult = { payload: Obj } | { status: 400 | 413; error: string }; + +function readJsonBody(req: VercelReq): BodyResult { + let parsed: unknown = req.body; if (typeof req.body === 'string') { - if (req.body.length > MAX_BODY_BYTES) return null; + if (byteLength(req.body) > MAX_BODY_BYTES) return { status: 413, error: 'body-too-large' }; try { - const parsed: unknown = JSON.parse(req.body); - if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) return parsed as Record; + parsed = JSON.parse(req.body); } catch { - return null; + return { status: 400, error: 'bad json body' }; } } - return null; + if (!isObj(parsed)) return { status: 400, error: 'bad json body' }; + // Vercel usually hands us an already-parsed object: measure its encoding. + let encoded: string; + try { + encoded = JSON.stringify(parsed); + } catch { + return { status: 400, error: 'bad json body' }; + } + if (byteLength(encoded) > MAX_BODY_BYTES) return { status: 413, error: 'body-too-large' }; + return { payload: parsed }; } export default async function handler(req: VercelReq, res: VercelRes): Promise { @@ -134,12 +238,12 @@ export default async function handler(req: VercelReq, res: VercelRes): Promise 10) { - res.status(429).json({ error: 'rate-limited' }); - return; - } - } catch { - // fail open + // Best-effort rate limit: RL_MAX submits/min/IP. EXPIRE NX runs on every + // request so a key whose TTL was lost still expires. Fail open on errors. + const rlKey = `bench:rl:${ipHash}`; + const rl = await pipeline([ + ['INCR', rlKey], + ['EXPIRE', rlKey, RL_WINDOW_S, 'NX'] + ]); + if (rl && !rl[1]!.ok) { + // Store without EXPIRE NX support: plain EXPIRE so the key can't stick. + await redis(['EXPIRE', rlKey, RL_WINDOW_S]); + } + if (rl && rl[0]!.ok && typeof rl[0]!.result === 'number' && rl[0]!.result > RL_MAX) { + res.status(429).json({ error: 'rate-limited' }); + return; } - const record = JSON.stringify({ ...payload, receivedAt: new Date().toISOString() }).slice(0, MAX_BODY_BYTES); - const setRes = await upstash(`/set/bench:result:${encodeURIComponent(fp)}/${encodeURIComponent(record)}/NX`); - if (!setRes.ok) { + // Per-(fingerprint, submitter) dedup: one submitter can't block others. + const dedupKey = `bench:dedup:${fp}:${ipHash}`; + const claim = await redis(['SET', dedupKey, '1', 'EX', DEDUP_TTL_S, 'NX']); + if (!claim.ok) { res.status(502).json({ error: 'store-unavailable' }); return; } - if (upstashResult(setRes.data) !== 'OK') { + if (claim.result !== 'OK') { res.status(409).json({ ok: true, duplicate: true }); return; } - await upstash(`/sadd/bench:fingerprints/${encodeURIComponent(fp)}`); + + const listKey = `bench:results:${fp}`; + const stored = await pipeline([ + ['RPUSH', listKey, JSON.stringify(built.record)], + ['LTRIM', listKey, -RESULTS_PER_FP, -1], + ['SADD', 'bench:fingerprints', fp] + ]); + if (!stored || !stored.every(r => r.ok)) { + // Release the dedup claim so the client can retry later. + await redis(['DEL', dedupKey]); + res.status(502).json({ error: 'store-unavailable' }); + return; + } res.status(200).json({ ok: true }); } diff --git a/apps/benchmark/src/main.ts b/apps/benchmark/src/main.ts index bf9260e..cdd3a45 100644 --- a/apps/benchmark/src/main.ts +++ b/apps/benchmark/src/main.ts @@ -240,10 +240,14 @@ exportChart?.addEventListener('click', async () => { // `window.__IS_INITIALIZED__` and drives `window.gpuEngine` directly). // The benchmark itself creates per-run engines lazily on click; this // separate eagerly-initialized engine exists only so the CI browser gate -// can exercise real WGSL dispatch without driving the full UI. -const gateGpuEngine = new WebGPUEngine(); -const gateCpuEngine = new CPUEngine(); +// can exercise real WGSL dispatch without driving the full UI. The gate +// always runs against the Vite dev server, so it is dev-only (stripped from +// production builds); `?regression-gate` opts in on a built page. +const regressionGateEnabled = + import.meta.env.DEV || new URLSearchParams(location.search).has('regression-gate'); async function initRegressionGate(): Promise { + const gateGpuEngine = new WebGPUEngine(); + const gateCpuEngine = new CPUEngine(); try { await gateGpuEngine.init(); } catch { @@ -253,4 +257,4 @@ async function initRegressionGate(): Promise { (window as unknown as Record).cpuEngine = gateCpuEngine; (window as unknown as Record).__IS_INITIALIZED__ = true; } -void initRegressionGate(); +if (regressionGateEnabled) void initRegressionGate(); diff --git a/apps/site/src/benchmark-history.ts b/apps/site/src/benchmark-history.ts index 27c2d83..b9f1c54 100644 --- a/apps/site/src/benchmark-history.ts +++ b/apps/site/src/benchmark-history.ts @@ -15,12 +15,25 @@ export interface HistoryRecord { const KEY = 'bench:history:v1'; const MAX_ENTRIES = 20; +const isNum = (v: unknown): v is number => typeof v === 'number' && Number.isFinite(v); +const isStr = (v: unknown): v is string => typeof v === 'string'; + +/** Shape-checks one stored entry; anything malformed is dropped, not rendered. */ +function isHistoryRecord(v: unknown): v is HistoryRecord { + if (typeof v !== 'object' || v === null || Array.isArray(v)) return false; + const r = v as Record; + return isStr(r.at) && !Number.isNaN(Date.parse(r.at)) + && isNum(r.corpusSize) && isStr(r.mode) && isStr(r.query) + && typeof r.gpuRan === 'boolean' && isNum(r.gpuMedianMs) && isNum(r.cpuMedianMs) + && isStr(r.gpuTier) && isStr(r.adapterShort); +} + export function loadHistory(): HistoryRecord[] { try { const raw = localStorage.getItem(KEY); if (!raw) return []; const parsed: unknown = JSON.parse(raw); - return Array.isArray(parsed) ? (parsed as HistoryRecord[]).slice(0, MAX_ENTRIES) : []; + return Array.isArray(parsed) ? parsed.filter(isHistoryRecord).slice(0, MAX_ENTRIES) : []; } catch { return []; } diff --git a/apps/site/src/benchmark-share.ts b/apps/site/src/benchmark-share.ts index 457d4e0..21f3509 100644 --- a/apps/site/src/benchmark-share.ts +++ b/apps/site/src/benchmark-share.ts @@ -77,30 +77,49 @@ export function buildFingerprint( return `webgpu-${fnv1aHex(key)}`; } +/** Server-side limits (api/submit-benchmark.ts); clamp so a valid run is never rejected. */ +const LIMITS = { adapterField: 300, adapterType: 64, tierReason: 300, family: 32, query: 120, externalName: 32, externals: 8, ms: 60000, count: 2000000 } as const; + +const clampStr = (value: string, max: number): string => value.slice(0, max); +const clampMs = (value: number): number => (Number.isFinite(value) ? Math.min(LIMITS.ms, Math.max(0, value)) : 0); +const clampCount = (value: number): number => + (Number.isFinite(value) ? Math.min(LIMITS.count, Math.max(0, Math.floor(value))) : 0); + export function buildSharePayload(result: QuickBenchmarkResult): SharePayload { const ua = typeof navigator !== 'undefined' ? navigator.userAgent : 'headless'; + const field = (value: string | undefined): string => clampStr(value ?? '', LIMITS.adapterField); return { fingerprint: buildFingerprint(result.adapter, result.corpusSize, result.mode, ua), adapter: { - vendor: result.adapter?.vendor ?? '', - device: result.adapter?.device ?? '', - architecture: result.adapter?.architecture ?? '', - renderer: result.adapter?.renderer ?? '', - ...(result.adapter?.adapterType ? { adapterType: result.adapter.adapterType } : {}) + vendor: field(result.adapter?.vendor), + device: field(result.adapter?.device), + architecture: field(result.adapter?.architecture), + renderer: field(result.adapter?.renderer), + ...(result.adapter?.adapterType ? { adapterType: clampStr(result.adapter.adapterType, LIMITS.adapterType) } : {}) + }, + gpuTier: { + tier: result.gpuTier.tier, + confidence: result.gpuTier.confidence, + reason: clampStr(result.gpuTier.reason, LIMITS.tierReason) }, - gpuTier: result.gpuTier, - browser: browserFamily(ua), - os: osFamily(ua), + browser: clampStr(browserFamily(ua), LIMITS.family), + os: clampStr(osFamily(ua), LIMITS.family), corpusSize: result.corpusSize, mode: result.mode, - query: result.query, - gpuMedianMs: result.gpuMedianMs, - gpuP95Ms: result.gpuP95Ms, - gpuMatches: result.gpuMatches, - cpuMedianMs: result.cpuMedianMs, - cpuP95Ms: result.cpuP95Ms, - cpuMatches: result.cpuMatches, - externals: result.externals, + query: clampStr(result.query, LIMITS.query), + gpuMedianMs: clampMs(result.gpuMedianMs), + gpuP95Ms: clampMs(result.gpuP95Ms), + gpuMatches: clampCount(result.gpuMatches), + cpuMedianMs: clampMs(result.cpuMedianMs), + cpuP95Ms: clampMs(result.cpuP95Ms), + cpuMatches: clampCount(result.cpuMatches), + externals: result.externals.slice(0, LIMITS.externals).map(external => ({ + name: clampStr(external.name, LIMITS.externalName), + medianMs: clampMs(external.medianMs), + p95Ms: clampMs(external.p95Ms), + matches: clampCount(external.matches), + ran: external.ran + })), createdAt: new Date().toISOString() }; } @@ -123,7 +142,7 @@ export function markShared(fingerprint: string): void { export async function submitSharedResult( payload: SharePayload -): Promise<{ ok: boolean; duplicate?: boolean; unconfigured?: boolean; error?: string }> { +): Promise<{ ok: boolean; duplicate?: boolean; unconfigured?: boolean; terminal?: boolean; error?: string }> { try { const res = await fetch('/api/submit-benchmark', { method: 'POST', @@ -133,7 +152,10 @@ export async function submitSharedResult( if (res.status === 503) return { ok: false, unconfigured: true }; const data = (await res.json().catch(() => ({}))) as { duplicate?: boolean; error?: string }; if (res.status === 409 || data.duplicate) return { ok: true, duplicate: true }; - if (!res.ok) return { ok: false, error: data.error ?? `HTTP ${res.status}` }; + // 400/413 etc. = this payload will never be accepted: terminal, don't retry. + // 429 (rate limit) and 5xx are transient and may be retried next run. + const terminal = res.status >= 400 && res.status < 500 && res.status !== 429; + if (!res.ok) return { ok: false, terminal, error: data.error ?? `HTTP ${res.status}` }; return { ok: true }; } catch (error) { return { ok: false, error: error instanceof Error ? error.message : String(error) }; diff --git a/apps/site/src/corpus-gen.ts b/apps/site/src/corpus-gen.ts index 900b8c7..eac7c7e 100644 --- a/apps/site/src/corpus-gen.ts +++ b/apps/site/src/corpus-gen.ts @@ -2,9 +2,12 @@ const prefixes = ['src/auth', 'src/search', 'src/components', 'packages/engine', const nouns = ['Auth', 'User', 'Profile', 'Order', 'Invoice', 'Search', 'Query', 'Worker', 'Document', 'Session', 'Token', 'Index']; const suffixes = ['Controller', 'Service', 'Handler', 'Manager', 'Provider', 'Repository', 'View', 'Worker', 'Index']; +/** Deterministic row `i` of the generated corpus. */ +export function corpusRow(i: number): string { + return `${prefixes[i % prefixes.length]}/${nouns[(i * 7 + 3) % nouns.length]}${nouns[(i * 13 + 5) % nouns.length]}${suffixes[(i * 17 + 1) % suffixes.length]}_${i}.ts`; +} + /** Deterministic generated corpus. Same input size always yields the same rows. */ export function generateCorpus(size: number): string[] { - return Array.from({ length: size }, (_, i) => - `${prefixes[i % prefixes.length]}/${nouns[(i * 7 + 3) % nouns.length]}${nouns[(i * 13 + 5) % nouns.length]}${suffixes[(i * 17 + 1) % suffixes.length]}_${i}.ts` - ); + return Array.from({ length: size }, (_, i) => corpusRow(i)); } diff --git a/apps/site/src/corpus.ts b/apps/site/src/corpus.ts index f9c8d1f..a014140 100644 --- a/apps/site/src/corpus.ts +++ b/apps/site/src/corpus.ts @@ -1,10 +1,43 @@ import './fonts.css'; import './style.css'; import { normalizeText, scoreExactMatches, type SearchResultItem } from 'webgpu-search'; -import { generateCorpus } from './corpus-gen'; +import { corpusRow } from './corpus-gen'; const PAGE_SIZE = 100; const MAX_MATCHES_SHOWN = 1000; +/** Sizes the homepage benchmark offers (apps/site/index.html #qb-size). */ +const ALLOWED_SIZES = [100000, 200000, 500000, 1000000, 2000000] as const; +const DEFAULT_SIZE = 100000; +const BUILD_CHUNK = 50000; + +/** Missing/invalid -> default; otherwise the largest offered size <= requested. */ +function parseSize(raw: string | null): number { + if (raw === null || raw.trim() === '') return DEFAULT_SIZE; + const requested = Number(raw); + if (!Number.isFinite(requested) || requested <= 0) return DEFAULT_SIZE; + let size: number = ALLOWED_SIZES[0]; + for (const allowed of ALLOWED_SIZES) if (allowed <= requested) size = allowed; + return size; +} + +const yieldToBrowser = (): Promise => new Promise(resolve => setTimeout(resolve, 0)); + +/** Generates + tokenizes in chunks, yielding between them so the tab stays responsive. */ +async function buildCorpus(size: number, onProgress: (done: number) => void): Promise<{ strings: string[]; tokens: Uint32Array[] }> { + const strings: string[] = new Array(size); + const tokens: Uint32Array[] = new Array(size); + for (let start = 0; start < size; start += BUILD_CHUNK) { + const end = Math.min(size, start + BUILD_CHUNK); + for (let i = start; i < end; i++) { + const row = corpusRow(i); + strings[i] = row; + tokens[i] = normalizeText(row, true).tokens; + } + onProgress(end); + await yieldToBrowser(); + } + return { strings, tokens }; +} const meta = document.querySelector('#c-meta'); const matchesList = document.querySelector('#c-matches'); @@ -55,21 +88,21 @@ function attachPager( render(); } -function main(): void { +async function main(): Promise { const params = new URLSearchParams(location.search); const query = (params.get('query') ?? params.get('q') ?? '').slice(0, 200); - const requested = Number(params.get('size')); - const size = Number.isFinite(requested) - ? Math.min(2000000, Math.max(1, Math.floor(requested))) - : 100000; + const size = parseSize(params.get('size')); if (!query) { if (meta) meta.textContent = 'No query given. Run the benchmark on the homepage first, then follow its link.'; return; } - const strings = generateCorpus(size); - const tokens = strings.map(value => normalizeText(value, true).tokens); + const { strings, tokens } = await buildCorpus(size, done => { + if (meta) meta.textContent = `Generating ${size.toLocaleString()} records… ${Math.round((done / size) * 100)}%`; + }); + if (meta) meta.textContent = `Searching ${size.toLocaleString()} records…`; + await yieldToBrowser(); const queryTokens = normalizeText(query, true).tokens; const response = scoreExactMatches(tokens, queryTokens, 'fuzzy', MAX_MATCHES_SHOWN, strings); const shown: SearchResultItem[] = response.results; @@ -96,4 +129,6 @@ function main(): void { } } -main(); +void main().catch(error => { + if (meta) meta.textContent = `Failed to build corpus: ${error instanceof Error ? error.message : String(error)}`; +}); diff --git a/apps/site/src/main.ts b/apps/site/src/main.ts index 0de33ea..1f7865e 100644 --- a/apps/site/src/main.ts +++ b/apps/site/src/main.ts @@ -114,8 +114,10 @@ async function autoShareResult(result: { const full = share.buildSharePayload(result as never); if (share.alreadyShared(full.fingerprint)) return; // sent before from this browser const res = await share.submitSharedResult(full); - if (!res.ok) return; // backend off / network hiccup: benchmark itself is unaffected - share.markShared(full.fingerprint); + // Mark on success/duplicate and on terminal (4xx validation) failures so a + // payload the server will never accept isn't resent every run. Backend + // off / network hiccup / rate limit: retry next run. + if (res.ok || res.terminal) share.markShared(full.fingerprint); } catch { // share chunk failed to load: benchmark itself is unaffected } diff --git a/bun.lock b/bun.lock index d91cebd..0c32166 100644 --- a/bun.lock +++ b/bun.lock @@ -30,45 +30,6 @@ "vite": "^8.2.2", }, }, - "apps/docs-search": { - "name": "docs-search", - "version": "0.1.0", - "dependencies": { - "fuse.js": "^7.5.0", - "webgpu-search": "workspace:*", - }, - "devDependencies": { - "@webgpu/types": "^0.1.72", - "typescript": "~6.0.2", - "vite": "^8.2.2", - }, - }, - "apps/log-viewer": { - "name": "log-viewer", - "version": "0.1.0", - "dependencies": { - "fuse.js": "^7.5.0", - "webgpu-search": "workspace:*", - }, - "devDependencies": { - "@webgpu/types": "^0.1.72", - "typescript": "~6.0.2", - "vite": "^8.2.2", - }, - }, - "apps/monaco-palette": { - "name": "monaco-palette", - "version": "0.1.0", - "dependencies": { - "fuse.js": "^7.5.0", - "webgpu-search": "workspace:*", - }, - "devDependencies": { - "@webgpu/types": "^0.1.72", - "typescript": "~6.0.2", - "vite": "^8.2.2", - }, - }, "apps/site": { "name": "site", "version": "0.1.0", @@ -402,8 +363,6 @@ "devtools-protocol": ["devtools-protocol@0.0.1666840", "", {}, "sha512-gCcO42XCHKEs7Ag0S7aGYsnJ7hlgrO3qderYqeiY0Eqk+0GFfuvT13IA0hHreJTa2KCdDVyGMeOhdMNmrrTjVg=="], - "docs-search": ["docs-search@workspace:apps/docs-search"], - "emoji-regex": ["emoji-regex@10.6.0", "", {}, "sha512-toUI84YS5YmxW219erniWD0CIVOo46xGKColeNQRgOzDorgBi1v4D71/OFzgD9GO2UGKIv1C3Sp8DAn0+j5w7A=="], "entities": ["entities@7.0.1", "", {}, "sha512-TWrgLOFUQTH994YUyl1yT4uyavY5nNB5muff+RtWaqNVCAK408b5ZnnbNAUEWLTCpum9w6arT70i1XdQ4UeOPA=="], @@ -482,8 +441,6 @@ "locate-character": ["locate-character@3.0.0", "", {}, "sha512-SW13ws7BjaeJ6p7Q6CO2nchbYEc3X3J6WrmTTDto7yMPqVSZTUyY5Tjbid+Ab8gLnATtygYtiDIJGQRRn2ZOiA=="], - "log-viewer": ["log-viewer@workspace:apps/log-viewer"], - "magic-string": ["magic-string@0.30.21", "", { "dependencies": { "@jridgewell/sourcemap-codec": "^1.5.5" } }, "sha512-vd2F4YUyEXKGcLHoq+TEyCjxueSeHnFxyyjNp80yg0XV4vUhnDer/lvvlqM/arB5bXQN5K2/3oinyCRyx8T2CQ=="], "mitt": ["mitt@3.0.1", "", {}, "sha512-vKivATfr97l2/QBCYAkXYDbrIWPM2IIKEl7YPhjCvKlG3kE2gm+uBo6nEXK3M5/Ffh/FLpKExzOQ3JJoJGFKBw=="], @@ -492,8 +449,6 @@ "modern-tar": ["modern-tar@0.8.5", "", {}, "sha512-snEhs+6G5Tjd4I7tLCDOaoln2RgE0bD19RzEKgvgK2hZ5VKy3MpLhLTZ2fWpXSTg4K2cyPwp+VHATFJhxfnOeA=="], - "monaco-palette": ["monaco-palette@workspace:apps/monaco-palette"], - "ms": ["ms@2.1.3", "", {}, "sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA=="], "mz": ["mz@2.7.0", "", { "dependencies": { "any-promise": "^1.0.0", "object-assign": "^4.0.1", "thenify-all": "^1.0.0" } }, "sha512-z81GNO7nnYMEhrGh9LeymoE4+Yr0Wn5McHIZMK5cfQCl+NDX08sCZgUc9/6MHni9IWuFLm1Z3HTCXu2z9fN62Q=="], diff --git a/bunfig.toml b/bunfig.toml new file mode 100644 index 0000000..7ade490 --- /dev/null +++ b/bunfig.toml @@ -0,0 +1,6 @@ +# Load .wgsl shader imports as source text under the Bun runtime (scripts and +# `bun test`), mirroring tsup's `'.wgsl': 'text'` loader used for the build. +preload = ["./scripts/wgsl-loader-plugin.ts"] + +[test] +preload = ["./scripts/wgsl-loader-plugin.ts"] diff --git a/docs/public-api.md b/docs/public-api.md index 1e3b601..64c7a05 100644 --- a/docs/public-api.md +++ b/docs/public-api.md @@ -115,6 +115,13 @@ fail-closed (facets would reflect the unfiltered set). `FacetRequest` / `FacetResult`, `DeterministicRankingOptions`, `AutocompleteOptions` / `SuggestionItem`, `SearchHooks` / `ExtensionHookIds`, `CostBudgetOptions`. + Filter-field value coercion (index and query time alike): default + getters read own properties only (never inherited `Object.prototype` + members such as `constructor`); `number` fields accept finite numbers and + non-blank numeric strings (blank / non-finite values are missing; as range + bounds they throw `InvalidFilterError`); `boolean` fields accept only + `true` / `false` and the strings `"true"` / `"false"` (other indexed values + are missing; other query operands throw `InvalidFilterError`). - Snapshots / IDB: `DocumentIndexSchema`, `SerializeDocumentIndexOptions`, `RestoreDocumentIndexOptions`, `DocumentSnapshotHeader`, `IDBStorageOptions`. @@ -221,9 +228,10 @@ worker, and snapshot-restore paths (`src/ranking.ts:1`): `SearchIndex` weight is uniform so this tier is a no-op there). 3. `exact` DESC — full post-normalization token equality beats partial/fuzzy. 4. `length` ASC — shorter matched-field token span wins. -5. `id` ASC — `compareIdsAsc`: numbers numerically, otherwise `String(id)` - code-unit order (never locale collation). Numeric `2` vs string `"2"` are - id-equal and fall through. +5. `id` ASC — `compareIdsAsc`: a total order over mixed ids. All numeric ids + sort before all string ids; numbers compare numerically, strings by + code-unit order (never locale collation). Numeric `2` sorts before string + `"2"` (they are not id-equal). 6. Implicit final: `docIndex` ASC (insertion order; defensive, ids unique). Properties: wrap-free comparisons (no `|0` subtraction), code-unit string diff --git a/package.json b/package.json index 7693cc7..ba8c3a0 100644 --- a/package.json +++ b/package.json @@ -23,6 +23,7 @@ "test:filtering": "bun scripts/test-filtering.ts", "test:faceting": "bun scripts/test-faceting.ts", "test:diagnostics": "bun test packages/webgpu-search/test/diagnostics.test.ts", + "test:regressions": "bun test packages/webgpu-search/test/document-index-integrity.test.ts packages/webgpu-search/test/engine-pool-regressions.test.ts packages/webgpu-search/test/worker-snapshot-hardening.test.ts", "test:search-modes": "bun test packages/webgpu-search/test/search-modes.test.ts", "test:parity": "bun scripts/test-parity-harness.ts", "test:reliability": "bun scripts/test-reliability.ts", diff --git a/packages/webgpu-search/bunfig.toml b/packages/webgpu-search/bunfig.toml new file mode 100644 index 0000000..82568b8 --- /dev/null +++ b/packages/webgpu-search/bunfig.toml @@ -0,0 +1,6 @@ +# Same .wgsl text loader as the repo root, for Bun runs from this package dir +# (`turbo test`, `bun test` here). Not published (outside package "files"). +preload = ["../../scripts/wgsl-loader-plugin.ts"] + +[test] +preload = ["../../scripts/wgsl-loader-plugin.ts"] diff --git a/packages/webgpu-search/src/cpu-engine.ts b/packages/webgpu-search/src/cpu-engine.ts index 17b4472..e384af1 100644 --- a/packages/webgpu-search/src/cpu-engine.ts +++ b/packages/webgpu-search/src/cpu-engine.ts @@ -2,6 +2,8 @@ import uFuzzy from '@leeoniya/ufuzzy'; import type { SearchResultItem } from './types'; import { nowMs } from './guard'; import { WebGPUSearchError } from './errors'; +import { normalizeText } from './text-normalization'; +import { compareExactResults, scoreSubstringTokens } from './exact-scorer'; export interface CPUSearchResult { query: string; @@ -85,7 +87,12 @@ export class CPUEngine { } /** - * Native JS substring search with ranking symmetry and case-sensitivity support + * Native JS substring scan with GPU scoring symmetry. Records and query go + * through the same `normalizeText` path as the GPU/exact scorer (NFC + + * default case folding when case-insensitive, NFC-only when + * case-sensitive) and are scored in code points via + * `scoreSubstringTokens`, so scores match the WGSL kernel bit-for-bit + * (no UTF-16 unit offsets, no locale-sensitive `toLowerCase`). */ searchNaiveScan( strings: string[], @@ -93,33 +100,31 @@ export class CPUEngine { maxResults: number = 1000, caseSensitive: boolean = false ): CPUSearchResult { - const cleanQuery = query.trim(); - if (!cleanQuery) { + const normalized = !caseSensitive; + const nq = normalizeText(typeof query === 'string' ? query : '', normalized); + if (nq.isEmpty) { return { query: '', totalMatches: 0, results: [], durationMs: 0 }; } const t0 = nowMs(); - const queryTerm = caseSensitive ? cleanQuery : cleanQuery.toLowerCase(); const candidates: SearchResultItem[] = []; let totalMatches = 0; for (let i = 0; i < strings.length; i++) { const s = strings[i] ?? ''; - const target = caseSensitive ? s : s.toLowerCase(); - const matchStart = target.indexOf(queryTerm); - if (matchStart !== -1) { + const rec = normalizeText(typeof s === 'string' ? s : String(s), normalized).tokens; + const r = scoreSubstringTokens(rec, nq.tokens); + if (r.matched) { totalMatches++; - // Formula matching GPU substring scoring: earlier start + shorter string - const score = 1000 - (matchStart * 10) - (s.length - cleanQuery.length); candidates.push({ index: i, - score, + score: r.score, text: s }); } } - candidates.sort((a, b) => b.score - a.score); + candidates.sort(compareExactResults); const results = candidates.slice(0, maxResults); const durationMs = nowMs() - t0; diff --git a/packages/webgpu-search/src/dataset-packing.ts b/packages/webgpu-search/src/dataset-packing.ts index 9705481..997c0f2 100644 --- a/packages/webgpu-search/src/dataset-packing.ts +++ b/packages/webgpu-search/src/dataset-packing.ts @@ -456,11 +456,32 @@ export interface ClampedHeadroomOptions { device?: GPUDevice | null; } +/** + * Effective per-buffer storage limit: `min(maxBufferSize, + * maxStorageBufferBindingSize)`. Forged/partial limits fall back to the + * 128 MB default (fail-closed, mirrors `checkMemoryBudget`). + */ +export function resolveStorageBufferLimit(device?: GPUDevice | null): number { + let maxLimit = 134217728; // 128 MB default safe limit + const l = device?.limits as unknown as { maxBufferSize?: unknown; maxStorageBufferBindingSize?: unknown } | undefined; + if (l !== undefined) { + const a = typeof l.maxBufferSize === 'number' ? l.maxBufferSize : NaN; + const b = typeof l.maxStorageBufferBindingSize === 'number' ? l.maxStorageBufferBindingSize : NaN; + if (Number.isFinite(a) && Number.isFinite(b) && a > 0 && b > 0) { + maxLimit = Math.min(a as number, b as number); + } + } + return maxLimit; +} + /** * Computes buffer capacity with dynamic headroom clamped to adapter limits. - * targetBytes = min(requiredBytes * growthFactor, maxStorageBufferBindingSize). - * If headroom overflows the limit but requiredBytes fits, returns exact requiredBytes - * to avoid triggering an unnecessary CPU fallback. + * targetBytes = min(requiredBytes * growthFactor, limit) where + * limit = min(maxBufferSize, maxStorageBufferBindingSize), 4-byte aligned. + * If headroom overflows the limit but requiredBytes fits, returns exact + * requiredBytes. If requiredBytes itself (a requested *capacity*) exceeds the + * limit, the result is clamped to the limit — never an oversized binding. + * Callers must separately verify the actual data fits (`checkMemoryBudget`). */ export function computeClampedHeadroomBytes( requiredBytes: number, @@ -473,15 +494,8 @@ export function computeClampedHeadroomBytes( ? options.growthFactor : 1.5; - let maxLimit = 134217728; // 128 MB default safe limit - const l = options.device?.limits as unknown as { maxBufferSize?: unknown; maxStorageBufferBindingSize?: unknown } | undefined; - if (l !== undefined) { - const a = typeof l.maxBufferSize === 'number' ? l.maxBufferSize : NaN; - const b = typeof l.maxStorageBufferBindingSize === 'number' ? l.maxStorageBufferBindingSize : NaN; - if (Number.isFinite(a) && Number.isFinite(b) && a > 0 && b > 0) { - maxLimit = Math.min(a as number, b as number); - } - } + const maxLimit = resolveStorageBufferLimit(options.device); + const alignedLimit = Math.max(16, Math.floor(maxLimit / 4) * 4); const alignedReq = Math.max(16, Math.ceil(safeReq / 4) * 4); const desired = Math.max(16, Math.ceil((safeReq * growth) / 4) * 4); @@ -492,7 +506,7 @@ export function computeClampedHeadroomBytes( if (alignedReq <= maxLimit) { return alignedReq; } - return alignedReq; + return alignedLimit; } diff --git a/packages/webgpu-search/src/document-index.ts b/packages/webgpu-search/src/document-index.ts index bd08c36..6cfb517 100644 --- a/packages/webgpu-search/src/document-index.ts +++ b/packages/webgpu-search/src/document-index.ts @@ -85,7 +85,7 @@ import { type TextProfileId } from './text-profile'; import { DocumentBitset } from './filtering/doc-bitset'; -import { ColumnarStore } from './filtering/columnar-store'; +import { ColumnarStore, readOwnField, type PreparedColumnarRow } from './filtering/columnar-store'; import { compileFilter } from './filtering/compile-filter'; import { FacetEngine, @@ -130,6 +130,15 @@ import type { QueryDiagnosticsTimings } from './types'; +/** + * Own-property write that is safe for arbitrary user field names: a plain + * `obj[key] = v` with key `"__proto__"` would replace the prototype instead + * of creating an entry. + */ +function setOwnKey(obj: Record, key: string, value: V): void { + Object.defineProperty(obj, key, { value, enumerable: true, configurable: true, writable: true }); +} + export interface InternalField extends FieldScoreDefinition { name: string; weight: number; @@ -265,7 +274,7 @@ export class DocumentIndex> { throw new TypeError('[webgpu-search] Field name string must not be empty.'); } name = f; - getter = (doc: any) => doc[name]; + getter = (doc: any) => readOwnField(doc, name); } else if (f && typeof f === 'object') { if (typeof f.name !== 'string' || f.name.length === 0) { throw new TypeError('[webgpu-search] Field definition requires a non-empty string name.'); @@ -285,7 +294,7 @@ export class DocumentIndex> { } getter = f.getter; } else { - getter = (doc: any) => doc[name]; + getter = (doc: any) => readOwnField(doc, name); } } else { throw new TypeError('[webgpu-search] Field must be a string or FieldDefinition object.'); @@ -328,7 +337,7 @@ export class DocumentIndex> { throw new TypeError('[webgpu-search] Filter field name string must not be empty.'); } name = ff; - getter = (doc: any) => doc[name]; + getter = (doc: any) => readOwnField(doc, name); hasGetter = false; } else if (ff && typeof ff === 'object') { if (typeof ff.name !== 'string' || ff.name.trim().length === 0) { @@ -343,7 +352,7 @@ export class DocumentIndex> { getter = ff.getter; hasGetter = true; } else { - getter = (doc: any) => doc[name]; + getter = (doc: any) => readOwnField(doc, name); hasGetter = false; } } else { @@ -2013,9 +2022,9 @@ export class DocumentIndex> { const rawStr = this.rawFieldStrings[dIdx][f] ?? ''; const ranges = alignHighlights(rawStr, query, alignOpts); if (ranges.length > 0) { - highlights[fieldDef.name] = ranges; + setOwnKey(highlights, fieldDef.name, ranges); if (tag) { - highlightedText[fieldDef.name] = renderHighlightedText(rawStr, ranges, tag, escapeHtml); + setOwnKey(highlightedText, fieldDef.name, renderHighlightedText(rawStr, ranges, tag, escapeHtml)); } } } @@ -2027,9 +2036,9 @@ export class DocumentIndex> { const rawStr = this.rawFieldStrings[dIdx][fIdx] ?? ''; const ranges = alignHighlights(rawStr, query, alignOpts); if (ranges.length > 0) { - highlights[fName] = ranges; + setOwnKey(highlights, fName, ranges); if (tag) { - highlightedText[fName] = renderHighlightedText(rawStr, ranges, tag, escapeHtml); + setOwnKey(highlightedText, fName, renderHighlightedText(rawStr, ranges, tag, escapeHtml)); } } } @@ -2040,9 +2049,9 @@ export class DocumentIndex> { if (primaryFIdx !== undefined) { const primaryRaw = this.rawFieldStrings[dIdx][primaryFIdx] ?? ''; const ranges = alignHighlights(primaryRaw, query, alignOpts); - highlights[item.matchedField] = ranges; + setOwnKey(highlights, item.matchedField, ranges); if (tag) { - highlightedText[item.matchedField] = renderHighlightedText(primaryRaw, ranges, tag, escapeHtml); + setOwnKey(highlightedText, item.matchedField, renderHighlightedText(primaryRaw, ranges, tag, escapeHtml)); } } @@ -2054,9 +2063,9 @@ export class DocumentIndex> { const auxRaw = this.rawFieldStrings[dIdx][auxFIdx] ?? ''; const ranges = alignHighlights(auxRaw, query, alignOpts); aux.highlights = ranges; - highlights[aux.field] = ranges; + setOwnKey(highlights, aux.field, ranges); if (tag) { - highlightedText[aux.field] = renderHighlightedText(auxRaw, ranges, tag, escapeHtml); + setOwnKey(highlightedText, aux.field, renderHighlightedText(auxRaw, ranges, tag, escapeHtml)); } } } @@ -2067,7 +2076,7 @@ export class DocumentIndex> { if (item.matches) { for (let m = 0; m < item.matches.length; m++) { const aux = item.matches[m]; - if (highlights[aux.field] !== undefined) { + if (Object.prototype.hasOwnProperty.call(highlights, aux.field)) { aux.highlights = highlights[aux.field]; } } @@ -2188,6 +2197,7 @@ export class DocumentIndex> { rawStrings: string[]; tokens: Uint32Array[]; tokenCount: number; + columnar: PreparedColumnarRow; } const preparedUpdates: PreparedDocUpdate[] = []; const updatedIdsInBatch = new Set(); @@ -2209,8 +2219,10 @@ export class DocumentIndex> { throw new DuplicateIdError(id); } updatedIdsInBatch.add(id); + // Filter getters run here (before any mutation) so a throwing + // getter rejects the whole batch with the index untouched. const prepared = this.prepareDocFields(doc, id); - preparedUpdates.push(prepared); + preparedUpdates.push({ ...prepared, columnar: this.columnarStore.prepare(doc) }); } } @@ -2220,6 +2232,7 @@ export class DocumentIndex> { rawStrings: string[]; tokens: Uint32Array[]; tokenCount: number; + columnar: PreparedColumnarRow; } const preparedAdds: PreparedDocAdd[] = []; const addedIdsInBatch = new Set(); @@ -2247,7 +2260,7 @@ export class DocumentIndex> { addedIdsInBatch.add(id); const prepared = this.prepareDocFields(doc, id); - preparedAdds.push(prepared); + preparedAdds.push({ ...prepared, columnar: this.columnarStore.prepare(doc) }); } } @@ -2304,7 +2317,7 @@ export class DocumentIndex> { this.docIds.push(upd.id); this.idToDocIndex.set(upd.id, newD); this.rawFieldStrings.push(upd.rawStrings); - this.columnarStore.add(newD, upd.doc); + this.columnarStore.addPrepared(newD, upd.columnar); const docRows: number[] = new Array(this.sortedFields.length); for (let f = 0; f < this.sortedFields.length; f++) { @@ -2345,7 +2358,7 @@ export class DocumentIndex> { this.docIds.push(ad.id); this.idToDocIndex.set(ad.id, newD); this.rawFieldStrings.push(ad.rawStrings); - this.columnarStore.add(newD, ad.doc); + this.columnarStore.addPrepared(newD, ad.columnar); const docRows: number[] = new Array(this.sortedFields.length); for (let f = 0; f < this.sortedFields.length; f++) { @@ -2772,7 +2785,7 @@ export class DocumentIndex> { if (uf && typeof uf === 'object' && typeof uf.getter === 'function') { getter = uf.getter; } else { - getter = (doc: any) => doc[name]; + getter = (doc: any) => readOwnField(doc, name); } return { name, weight, getter, originalIndex }; }); diff --git a/packages/webgpu-search/src/filtering/columnar-store.ts b/packages/webgpu-search/src/filtering/columnar-store.ts index fc30d81..52b9b6a 100644 --- a/packages/webgpu-search/src/filtering/columnar-store.ts +++ b/packages/webgpu-search/src/filtering/columnar-store.ts @@ -1,10 +1,24 @@ /** * Columnar attribute storage for structured pre-filtering and facet aggregations. - * Backed by typed arrays (Float64Array, Uint32Array) and DocumentBitsets with geometric growth. + * Backed by typed arrays (Float64Array, Uint32Array), DocumentBitsets and + * adaptive per-value posting lists (sparse sets promoted to bitsets only when + * dense), so memory stays proportional to the stored (row, value) entries + * rather than `distinctValues x capacity`. + * + * Value coercion (applied identically at index and query time): + * - `number`: finite numbers, or non-blank numeric strings; anything else + * (blank strings, NaN, +/-Infinity, booleans, objects) is treated as missing. + * - `boolean`: only `true`/`false` and the exact strings `"true"`/`"false"`; + * any other value is treated as missing at index time and rejected with + * `InvalidFilterError` at query time. + * - Default getters read own properties only (plus non-`Object.prototype` + * inherited ones such as class accessors), never `constructor`/`toString`. + * * 100% portable across browser main thread, Web Workers, Node.js, and SSR (zero DOM references). */ import { DocumentBitset } from './doc-bitset'; +import { PostingList } from './posting-list'; import type { FilterFieldDefinition, FilterFieldType, @@ -26,7 +40,8 @@ interface StringColumn extends BaseColumn { stringTable: string[]; stringToCode: Map; codes: Uint32Array; - invertedIndex: Map; + /** code -> rows holding it; entries with zero rows are removed eagerly. */ + invertedIndex: Map; } interface NumberColumn extends BaseColumn { @@ -41,12 +56,63 @@ interface BooleanColumn extends BaseColumn { interface StringArrayColumn extends BaseColumn { type: 'string[]'; - tagInverted: Map; + /** tag -> rows holding it; entries with zero rows are removed eagerly. */ + tagInverted: Map; docTags: Map; } type Column = StringColumn | NumberColumn | BooleanColumn | StringArrayColumn; +/** + * Default field accessor: reads `doc[name]` only when it is an own property + * or an inherited one not provided by `Object.prototype` (so class accessors + * still work, but `constructor`, `toString`, `__proto__`, ... never leak in as + * field values). + */ +export function readOwnField(doc: any, name: string): any { + if (doc === null || doc === undefined) return undefined; + if (Object.prototype.hasOwnProperty.call(doc, name)) return doc[name]; + if (name in Object.prototype) return undefined; + return doc[name]; +} + +/** + * Strict numeric coercion shared by index and query paths: finite numbers + * and non-blank numeric strings only. Returns NaN for anything else + * (`Number("")`/`Number(" ")` would otherwise silently yield 0). + */ +function coerceNumber(raw: unknown): number { + if (typeof raw === 'number') return Number.isFinite(raw) ? raw : NaN; + if (typeof raw === 'string') { + if (raw.trim().length === 0) return NaN; + const n = Number(raw); + return Number.isFinite(n) ? n : NaN; + } + return NaN; +} + +/** + * Strict boolean coercion shared by index and query paths: `true`/`false` + * and the exact strings `"true"`/`"false"`; `undefined` for anything else. + */ +function coerceBoolean(raw: unknown): boolean | undefined { + if (raw === true || raw === 'true') return true; + if (raw === false || raw === 'false') return false; + return undefined; +} + +/** Query-time boolean operand: rejects non-boolean values fail-closed. */ +function toBooleanOperand(raw: unknown, field: string): boolean { + const b = coerceBoolean(raw); + if (b === undefined) { + throw new InvalidFilterError( + `Boolean field "${field}" expects true/false (or "true"/"false"), got ${JSON.stringify(raw)}.`, + field + ); + } + return b; +} + /** Validates a numeric range bound, throwing InvalidFilterError on non-numeric input. */ function toNumberBound(raw: unknown, op: string, field: string): number { if (typeof raw !== 'number' && typeof raw !== 'string') { @@ -55,7 +121,7 @@ function toNumberBound(raw: unknown, op: string, field: string): number { field ); } - const n = Number(raw); + const n = coerceNumber(raw); if (Number.isNaN(n)) { throw new InvalidFilterError( `Range operator "${op}" on field "${field}" expects a numeric bound, got ${JSON.stringify(String(raw))}.`, @@ -65,6 +131,53 @@ function toNumberBound(raw: unknown, op: string, field: string): number { return n; } +/** + * Normalized column value produced by `ColumnarStore.prepare`: + * `undefined` = missing; otherwise string / finite number / boolean / + * non-empty deduped string[] matching the column type. + */ +export type PreparedColumnValue = string | number | boolean | string[] | undefined; + +/** One normalized value per registered column, in column order. */ +export type PreparedColumnarRow = PreparedColumnValue[]; + +function normalizeColumnValue(type: FilterFieldType, rawVal: unknown): PreparedColumnValue { + if (rawVal === null || rawVal === undefined) return undefined; + switch (type) { + case 'string': + return String(rawVal); + case 'number': { + const n = coerceNumber(rawVal); + return Number.isFinite(n) ? n : undefined; + } + case 'boolean': + return coerceBoolean(rawVal); + case 'string[]': { + if (!Array.isArray(rawVal) || rawVal.length === 0) return undefined; + const seen = new Set(); + const tags: string[] = []; + for (let i = 0; i < rawVal.length; i++) { + const item = rawVal[i]; + if (item === null || item === undefined) continue; + const t = String(item); + if (seen.has(t)) continue; + seen.add(t); + tags.push(t); + } + return tags.length > 0 ? tags : undefined; + } + } + return undefined; +} + +function assertDocIndex(docIndex: number, op: string): void { + if (!Number.isInteger(docIndex) || docIndex < 0) { + throw new RangeError( + `[webgpu-search] ColumnarStore.${op} requires a non-negative integer docIndex, got ${String(docIndex)}.` + ); + } +} + export interface ColumnarStoreOptions { initialCapacity?: number; growthFactor?: number; @@ -100,7 +213,7 @@ export class ColumnarStore> { private registerField(fieldDef: FilterFieldDefinition): void { const name = fieldDef.name; - const getter = fieldDef.getter ?? ((doc: any) => doc[name]); + const getter = fieldDef.getter ?? ((doc: any) => readOwnField(doc, name)); const type: FilterFieldType = fieldDef.type ?? 'string'; const cap = this.capacity; @@ -193,8 +306,8 @@ export class ColumnarStore> { const newCodes = new Uint32Array(newCap); newCodes.set(col.codes); col.codes = newCodes; - for (const bs of col.invertedIndex.values()) { - bs.ensureCapacity(newCap); + for (const pl of col.invertedIndex.values()) { + pl.ensureCapacity(newCap); } break; } @@ -210,8 +323,8 @@ export class ColumnarStore> { break; } case 'string[]': { - for (const bs of col.tagInverted.values()) { - bs.ensureCapacity(newCap); + for (const pl of col.tagInverted.values()) { + pl.ensureCapacity(newCap); } break; } @@ -261,10 +374,35 @@ export class ColumnarStore> { } } + /** + * Runs every column getter for `doc` and normalizes the results into a + * plain per-column value vector. This is the only step that executes user + * code (getters, `String()` conversions) and may throw; it performs no + * mutation, so callers needing all-or-nothing semantics (e.g. + * `DocumentIndex.applyBatch`) can prepare every doc before mutating + * anything and then apply with `addPrepared`, which cannot throw on value + * conversion. + */ + prepare(doc: TDoc): PreparedColumnarRow { + const out: PreparedColumnValue[] = new Array(this.columns.size); + let i = 0; + for (const col of this.columns.values()) { + out[i++] = normalizeColumnValue(col.type, col.getter(doc)); + } + return out; + } + add(docIndex: number, doc: TDoc): void { - if (!Number.isInteger(docIndex) || docIndex < 0) { - throw new RangeError( - `[webgpu-search] ColumnarStore.add requires a non-negative integer docIndex, got ${String(docIndex)}.` + assertDocIndex(docIndex, 'add'); + this.addPrepared(docIndex, this.prepare(doc)); + } + + /** Applies a row produced by `prepare` (no user code runs here). */ + addPrepared(docIndex: number, row: PreparedColumnarRow): void { + assertDocIndex(docIndex, 'add'); + if (!Array.isArray(row) || row.length !== this.columns.size) { + throw new TypeError( + `[webgpu-search] ColumnarStore.addPrepared expects a prepared row of ${this.columns.size} values.` ); } if (docIndex >= this.capacity) { @@ -272,31 +410,18 @@ export class ColumnarStore> { } this.activeDocs.set(docIndex); + let i = 0; for (const col of this.columns.values()) { - const rawVal = col.getter(doc); // Clear-then-set so re-add on an occupied slot never leaks stale - // string codes or tag bits (mirrors update()). + // string codes or tag bits. this.clearDocColumnValue(col, docIndex); - this.setDocColumnValue(col, docIndex, rawVal); + this.setDocColumnValue(col, docIndex, row[i++]); } } update(docIndex: number, doc: TDoc): void { - if (!Number.isInteger(docIndex) || docIndex < 0) { - throw new RangeError( - `[webgpu-search] ColumnarStore.update requires a non-negative integer docIndex, got ${String(docIndex)}.` - ); - } - if (docIndex >= this.capacity) { - this.ensureCapacity(docIndex + 1); - } - this.activeDocs.set(docIndex); - - for (const col of this.columns.values()) { - const rawVal = col.getter(doc); - this.clearDocColumnValue(col, docIndex); - this.setDocColumnValue(col, docIndex, rawVal); - } + assertDocIndex(docIndex, 'update'); + this.addPrepared(docIndex, this.prepare(doc)); } remove(docIndex: number): void { @@ -308,42 +433,39 @@ export class ColumnarStore> { } } - private setDocColumnValue(col: Column, docIndex: number, rawVal: any): void { - if (rawVal === null || rawVal === undefined) { + private setDocColumnValue(col: Column, docIndex: number, val: PreparedColumnValue): void { + if (val === undefined) { col.presence.clear(docIndex); return; } switch (col.type) { case 'string': { - const str = String(rawVal); + const str = val as string; col.presence.set(docIndex); let code = col.stringToCode.get(str); if (code === undefined) { code = col.stringTable.length; col.stringTable.push(str); col.stringToCode.set(str, code); - const bs = new DocumentBitset(this.capacity); - col.invertedIndex.set(code, bs); + } + let pl = col.invertedIndex.get(code); + if (!pl) { + pl = new PostingList(); + col.invertedIndex.set(code, pl); } col.codes[docIndex] = code; - col.invertedIndex.get(code)!.set(docIndex); + pl.add(docIndex, this.capacity); break; } case 'number': { - const num = typeof rawVal === 'number' ? rawVal : Number(rawVal); - if (Number.isFinite(num)) { - col.presence.set(docIndex); - col.values[docIndex] = num; - } else { - col.presence.clear(docIndex); - col.values[docIndex] = NaN; - } + col.presence.set(docIndex); + col.values[docIndex] = val as number; break; } case 'boolean': { col.presence.set(docIndex); - if (Boolean(rawVal)) { + if (val === true) { col.trueBitset.set(docIndex); } else { col.trueBitset.clear(docIndex); @@ -351,35 +473,19 @@ export class ColumnarStore> { break; } case 'string[]': { - if (!Array.isArray(rawVal) || rawVal.length === 0) { - col.presence.clear(docIndex); - return; - } - // Per-doc dedupe: facet counts are per-doc (not per-occurrence) to - // match string/number/boolean branches and tagInverted bitsets. - const seen = new Set(); - const tags: string[] = []; - for (let i = 0; i < rawVal.length; i++) { - const item = rawVal[i]; - if (item !== null && item !== undefined) { - const t = String(item); - if (seen.has(t)) continue; - seen.add(t); - tags.push(t); - let bs = col.tagInverted.get(t); - if (!bs) { - bs = new DocumentBitset(this.capacity); - col.tagInverted.set(t, bs); - } - bs.set(docIndex); + // Tags are already per-doc deduped by normalizeColumnValue: facet + // counts are per-doc (not per-occurrence). + const tags = val as string[]; + for (let i = 0; i < tags.length; i++) { + let pl = col.tagInverted.get(tags[i]); + if (!pl) { + pl = new PostingList(); + col.tagInverted.set(tags[i], pl); } + pl.add(docIndex, this.capacity); } - if (tags.length > 0) { - col.presence.set(docIndex); - col.docTags.set(docIndex, tags); - } else { - col.presence.clear(docIndex); - } + col.presence.set(docIndex); + col.docTags.set(docIndex, tags); break; } } @@ -392,7 +498,11 @@ export class ColumnarStore> { switch (col.type) { case 'string': { const code = col.codes[docIndex]; - col.invertedIndex.get(code)?.clear(docIndex); + const pl = col.invertedIndex.get(code); + if (pl) { + pl.delete(docIndex, this.capacity); + if (pl.size === 0) col.invertedIndex.delete(code); + } col.codes[docIndex] = 0; break; } @@ -408,7 +518,11 @@ export class ColumnarStore> { const prevTags = col.docTags.get(docIndex); if (prevTags) { for (let i = 0; i < prevTags.length; i++) { - col.tagInverted.get(prevTags[i])?.clear(docIndex); + const pl = col.tagInverted.get(prevTags[i]); + if (pl) { + pl.delete(docIndex, this.capacity); + if (pl.size === 0) col.tagInverted.delete(prevTags[i]); + } } col.docTags.delete(docIndex); } @@ -428,25 +542,45 @@ export class ColumnarStore> { for (const [name, col] of this.columns.entries()) { switch (col.type) { case 'string': { + // Re-dictionary live values only: codes with zero surviving rows + // are dropped. Codes are assigned in old-code order so the table + // stays deterministic. const newCol: StringColumn = { name, type: 'string', getter: col.getter, presence: new DocumentBitset(newCap), - stringTable: col.stringTable.slice(), - stringToCode: new Map(col.stringToCode), + stringTable: [], + stringToCode: new Map(), codes: new Uint32Array(newCap), invertedIndex: new Map() }; - for (const [code] of col.invertedIndex.entries()) { - newCol.invertedIndex.set(code, new DocumentBitset(newCap)); + const liveOld: number[] = []; + for (const [oldD] of oldToNewDocIndexMap.entries()) { + if (col.presence.has(oldD)) liveOld.push(col.codes[oldD]); + } + liveOld.sort((a, b) => a - b); + const oldToNewCode = new Map(); + for (let i = 0; i < liveOld.length; i++) { + const oc = liveOld[i]; + if (oldToNewCode.has(oc)) continue; + const nc = newCol.stringTable.length; + const str = col.stringTable[oc]; + oldToNewCode.set(oc, nc); + newCol.stringTable.push(str); + newCol.stringToCode.set(str, nc); } for (const [oldD, newD] of oldToNewDocIndexMap.entries()) { if (col.presence.has(oldD)) { newCol.presence.set(newD); - const code = col.codes[oldD]; + const code = oldToNewCode.get(col.codes[oldD])!; newCol.codes[newD] = code; - newCol.invertedIndex.get(code)?.set(newD); + let pl = newCol.invertedIndex.get(code); + if (!pl) { + pl = new PostingList(); + newCol.invertedIndex.set(code, pl); + } + pl.add(newD, newCap); } } this.columns.set(name, newCol); @@ -499,9 +633,6 @@ export class ColumnarStore> { tagInverted: new Map(), docTags: new Map() }; - for (const [tag] of col.tagInverted.entries()) { - newCol.tagInverted.set(tag, new DocumentBitset(newCap)); - } for (const [oldD, newD] of oldToNewDocIndexMap.entries()) { if (col.presence.has(oldD)) { newCol.presence.set(newD); @@ -509,7 +640,12 @@ export class ColumnarStore> { if (tags) { newCol.docTags.set(newD, tags.slice()); for (let i = 0; i < tags.length; i++) { - newCol.tagInverted.get(tags[i])?.set(newD); + let pl = newCol.tagInverted.get(tags[i]); + if (!pl) { + pl = new PostingList(); + newCol.tagInverted.set(tags[i], pl); + } + pl.add(newD, newCap); } } } @@ -540,11 +676,11 @@ export class ColumnarStore> { if (code === undefined) { return DocumentBitset.none(this.capacity); } - const bs = col.invertedIndex.get(code); - return bs ? bs.and(active) : DocumentBitset.none(this.capacity); + const pl = col.invertedIndex.get(code); + return pl ? pl.toBitset(this.capacity).andInPlace(active) : DocumentBitset.none(this.capacity); } case 'number': { - const target = typeof value === 'number' ? value : Number(value); + const target = coerceNumber(value); if (!Number.isFinite(target)) { return DocumentBitset.none(this.capacity); } @@ -573,7 +709,7 @@ export class ColumnarStore> { return result; } case 'boolean': { - const target = Boolean(value); + const target = toBooleanOperand(value, name); if (target) { return col.trueBitset.and(active); } else { @@ -582,8 +718,8 @@ export class ColumnarStore> { } case 'string[]': { const tag = String(value); - const bs = col.tagInverted.get(tag); - return bs ? bs.and(active) : DocumentBitset.none(this.capacity); + const pl = col.tagInverted.get(tag); + return pl ? pl.toBitset(this.capacity).andInPlace(active) : DocumentBitset.none(this.capacity); } } } @@ -610,8 +746,7 @@ export class ColumnarStore> { } else { const code = col.stringToCode.get(String(v)); if (code !== undefined) { - const bs = col.invertedIndex.get(code); - if (bs) result.orInPlace(bs); + col.invertedIndex.get(code)?.orInto(result); } } } @@ -624,8 +759,7 @@ export class ColumnarStore> { if (v === null) { includeNull = true; } else { - const bs = col.tagInverted.get(String(v)); - if (bs) result.orInPlace(bs); + col.tagInverted.get(String(v))?.orInto(result); } } if (includeNull) { @@ -641,7 +775,7 @@ export class ColumnarStore> { if (v === null) { includeNull = true; } else { - const n = typeof v === 'number' ? v : Number(v); + const n = coerceNumber(v); if (Number.isFinite(n)) numSet.add(n); } } @@ -676,7 +810,7 @@ export class ColumnarStore> { const v = values[i]; if (v === null) { result.orInPlace(active.andNot(col.presence)); - } else if (Boolean(v)) { + } else if (toBooleanOperand(v, name)) { result.orInPlace(col.trueBitset); } else { result.orInPlace(col.presence.andNot(col.trueBitset)); @@ -897,8 +1031,7 @@ export class ColumnarStore> { if (gteStr !== undefined && !(str >= gteStr)) continue; if (ltStr !== undefined && !(str < ltStr)) continue; if (lteStr !== undefined && !(str <= lteStr)) continue; - const bs = col.invertedIndex.get(code); - if (bs) result.orInPlace(bs); + col.invertedIndex.get(code)?.orInto(result); } return result.andInPlace(active); } diff --git a/packages/webgpu-search/src/filtering/posting-list.ts b/packages/webgpu-search/src/filtering/posting-list.ts new file mode 100644 index 0000000..fd44af7 --- /dev/null +++ b/packages/webgpu-search/src/filtering/posting-list.ts @@ -0,0 +1,117 @@ +/** + * Adaptive per-value posting list for columnar filter indexes. + * + * Low-frequency values are kept as a sparse `Set` of row indices + * (memory proportional to the value's document count); a value is promoted + * to a capacity-sized `DocumentBitset` only once it is dense enough that the + * bitset is the cheaper representation. Dense lists are demoted back to + * sparse when they become sparse relative to the store capacity (after + * removals or capacity growth), so total memory stays proportional to the + * number of (row, value) entries rather than `distinctValues x capacity`. + * 100% portable (zero DOM references). + */ + +import { DocumentBitset } from './doc-bitset'; + +/** Minimum doc count before a value may be promoted to a dense bitset. */ +const MIN_DENSE_COUNT = 64; + +/** Promote when count exceeds capacity / PROMOTE_DIVISOR (a Set costs ~16+ B/entry vs 1 bit/row). */ +const PROMOTE_DIVISOR = 128; + +/** Demote when count drops below capacity / DEMOTE_DIVISOR (hysteresis). */ +const DEMOTE_DIVISOR = 256; + +function promoteThreshold(capacity: number): number { + return Math.max(MIN_DENSE_COUNT, Math.floor(capacity / PROMOTE_DIVISOR)); +} + +function demoteThreshold(capacity: number): number { + return Math.max(MIN_DENSE_COUNT >>> 1, Math.floor(capacity / DEMOTE_DIVISOR)); +} + +export class PostingList { + private sparse: Set | null = new Set(); + private dense: DocumentBitset | null = null; + private count = 0; + + get size(): number { + return this.count; + } + + /** True when backed by a capacity-sized bitset (exposed for tests/diagnostics). */ + get isDense(): boolean { + return this.dense !== null; + } + + has(docIndex: number): boolean { + return this.dense ? this.dense.has(docIndex) : this.sparse!.has(docIndex); + } + + add(docIndex: number, capacity: number): void { + if (this.dense) { + if (this.dense.has(docIndex)) return; + this.dense.set(docIndex); + this.count++; + return; + } + const sparse = this.sparse!; + if (sparse.has(docIndex)) return; + sparse.add(docIndex); + this.count++; + if (this.count > promoteThreshold(capacity)) { + const bs = new DocumentBitset(capacity); + for (const d of sparse) bs.set(d); + this.dense = bs; + this.sparse = null; + } + } + + delete(docIndex: number, capacity: number): void { + if (this.dense) { + if (!this.dense.has(docIndex)) return; + this.dense.clear(docIndex); + this.count--; + if (this.count < demoteThreshold(capacity)) this.demote(); + return; + } + if (this.sparse!.delete(docIndex)) this.count--; + } + + /** Called on store capacity growth: grow dense bitsets or demote them. */ + ensureCapacity(capacity: number): void { + if (!this.dense) return; + if (this.count < demoteThreshold(capacity)) { + this.demote(); + } else { + this.dense.ensureCapacity(capacity); + } + } + + /** ORs this posting list into `target` (sized to the store capacity). */ + orInto(target: DocumentBitset): void { + if (this.dense) { + target.orInPlace(this.dense); + return; + } + for (const d of this.sparse!) target.set(d); + } + + /** Materializes this list as a fresh bitset of `capacity`. */ + toBitset(capacity: number): DocumentBitset { + if (this.dense) { + const out = new DocumentBitset(capacity); + out.orInPlace(this.dense); + return out; + } + const out = new DocumentBitset(capacity); + for (const d of this.sparse!) out.set(d); + return out; + } + + private demote(): void { + const dense = this.dense!; + this.sparse = new Set(dense.getMatchingIndices()); + this.dense = null; + } +} diff --git a/packages/webgpu-search/src/gpu-device-pool.ts b/packages/webgpu-search/src/gpu-device-pool.ts index 21003cb..7096cb4 100644 --- a/packages/webgpu-search/src/gpu-device-pool.ts +++ b/packages/webgpu-search/src/gpu-device-pool.ts @@ -25,9 +25,9 @@ export function assertValidPowerPreference(value: unknown): asserts value is GPU export class GpuDevicePool { private static sharedDevice: GPUDevice | null = null; - private static sharedAdapter: GPUAdapter | null = null; private static sharedAdapterInfo: AdapterInfo | null = null; private static refCount: number = 0; + private static inFlight: Promise | null = null; private static deviceLostListeners: Set<(reason: string) => void> = new Set(); /** @@ -118,22 +118,57 @@ export class GpuDevicePool { return null; } + // 4. Single in-flight acquisition: concurrent callers share one + // requestAdapter/requestDevice so the pool never creates (and leaks) a + // second device or resets the refcount under a live holder. Each caller + // takes its own reference once the shared acquisition settles. + if (!this.inFlight) { + const p = this.createSharedDevice(options?.powerPreference); + this.inFlight = p; + const clear = () => { + if (this.inFlight === p) this.inFlight = null; + }; + p.then(clear, clear); + } + const device = await this.inFlight; + // Device may have been lost between creation and this continuation. + if (!device || this.sharedDevice !== device) { + // Superseded while in flight (e.g. simulated loss): nobody holds a + // reference to this stale device, so release it instead of leaking. + if (device && this.sharedDevice !== device) device.destroy(); + return null; + } + this.refCount++; + return { + device, + adapterInfo: this.sharedAdapterInfo, + isShared: true + }; + } + + /** + * Request adapter + device and install it as the pooled device with + * refCount 0 (callers increment). Resolves null on failure; never throws. + */ + private static async createSharedDevice( + powerPreference?: GPUPowerPreference + ): Promise { try { - const powerPref = options?.powerPreference || 'high-performance'; - this.sharedAdapter = await navigator.gpu.requestAdapter({ powerPreference: powerPref }); - if (!this.sharedAdapter) { - this.sharedAdapter = await navigator.gpu.requestAdapter(); + const powerPref = powerPreference || 'high-performance'; + let adapter = await navigator.gpu.requestAdapter({ powerPreference: powerPref }); + if (!adapter) { + adapter = await navigator.gpu.requestAdapter(); } - if (!this.sharedAdapter) { + if (!adapter) { return null; } // Inspect adapter info - let info: any = (this.sharedAdapter as any).info || {}; - if (!info.vendor && !info.device && 'requestAdapterInfo' in this.sharedAdapter) { + let info: any = (adapter as any).info || {}; + if (!info.vendor && !info.device && 'requestAdapterInfo' in adapter) { try { - info = await (this.sharedAdapter as any).requestAdapterInfo(); + info = await (adapter as any).requestAdapterInfo(); } catch { info = {}; } @@ -141,35 +176,35 @@ export class GpuDevicePool { const unmaskedRenderer = this.getUnmaskedRenderer(); let vendor = info.vendor || ''; - let device = info.device || ''; + let deviceName = info.device || ''; let architecture = info.architecture || ''; const rawAdapterType = typeof info.type === 'string' && info.type ? String(info.type) : undefined; - if (!device && unmaskedRenderer) device = unmaskedRenderer; + if (!deviceName && unmaskedRenderer) deviceName = unmaskedRenderer; if (!vendor) { - if (/nvidia/i.test(unmaskedRenderer) || /nvidia/i.test(device)) vendor = 'NVIDIA'; - else if (/intel/i.test(unmaskedRenderer) || /intel/i.test(device)) vendor = 'Intel'; - else if (/amd|radeon/i.test(unmaskedRenderer) || /amd|radeon/i.test(device)) vendor = 'AMD'; - else if (/apple/i.test(unmaskedRenderer) || /apple/i.test(device)) vendor = 'Apple'; + if (/nvidia/i.test(unmaskedRenderer) || /nvidia/i.test(deviceName)) vendor = 'NVIDIA'; + else if (/intel/i.test(unmaskedRenderer) || /intel/i.test(deviceName)) vendor = 'Intel'; + else if (/amd|radeon/i.test(unmaskedRenderer) || /amd|radeon/i.test(deviceName)) vendor = 'AMD'; + else if (/apple/i.test(unmaskedRenderer) || /apple/i.test(deviceName)) vendor = 'Apple'; else vendor = 'Unknown GPU Vendor'; } - if (!device) device = 'WebGPU Generic Device'; + if (!deviceName) deviceName = 'WebGPU Generic Device'; if (!architecture) architecture = 'Default'; - const limits = this.sharedAdapter.limits; + const limits = adapter.limits; const requiredFeatures: GPUFeatureName[] = []; - const hasTimestamp = this.sharedAdapter.features.has('timestamp-query'); + const hasTimestamp = adapter.features.has('timestamp-query'); if (hasTimestamp) { requiredFeatures.push('timestamp-query'); } - this.sharedAdapterInfo = { + const adapterInfo: AdapterInfo = { vendor, architecture, - device, + device: deviceName, description: info.description || unmaskedRenderer || (typeof navigator !== 'undefined' ? navigator.userAgent : 'WebGPU Device'), - renderer: unmaskedRenderer || device, + renderer: unmaskedRenderer || deviceName, ...(rawAdapterType ? { adapterType: rawAdapterType } : {}), maxBufferSizeMB: Math.round(limits.maxBufferSize / (1024 * 1024)), maxStorageBindingSizeMB: Math.round(limits.maxStorageBufferBindingSize / (1024 * 1024)), @@ -179,8 +214,9 @@ export class GpuDevicePool { }; // Request device with graceful limit fallback + let device: GPUDevice; try { - this.sharedDevice = await this.sharedAdapter.requestDevice({ + device = await adapter.requestDevice({ requiredFeatures, requiredLimits: { maxBufferSize: limits.maxBufferSize, @@ -190,22 +226,28 @@ export class GpuDevicePool { }); } catch { try { - this.sharedDevice = await this.sharedAdapter.requestDevice({ requiredFeatures }); + device = await adapter.requestDevice({ requiredFeatures }); } catch { - this.sharedDevice = await this.sharedAdapter.requestDevice(); + device = await adapter.requestDevice(); } } - this.refCount = 1; + this.sharedAdapterInfo = adapterInfo; + this.sharedDevice = device; + this.refCount = 0; - // Handle device loss gracefully - this.sharedDevice.lost.then(lostInfo => { - const reason = lostInfo.message || 'WebGPU device lost'; - console.warn('WebGPU device lost:', reason); + // Handle device loss gracefully. Identity-checked: a stale handler for + // a previously pooled device must not tear down the current one, and + // intentional destroy() (pool release → reason 'destroyed') is not a + // loss event — listeners would otherwise push healthy indexes to CPU. + device.lost?.then?.(lostInfo => { + if (this.sharedDevice !== device) return; this.sharedDevice = null; - this.sharedAdapter = null; this.refCount = 0; - for (const listener of this.deviceLostListeners) { + if ((lostInfo as { reason?: unknown } | undefined)?.reason === 'destroyed') return; + const reason = lostInfo?.message || 'WebGPU device lost'; + console.warn('WebGPU device lost:', reason); + for (const listener of [...this.deviceLostListeners]) { try { listener(reason); } catch (e) { @@ -214,11 +256,7 @@ export class GpuDevicePool { } }, () => {}); - return { - device: this.sharedDevice, - adapterInfo: this.sharedAdapterInfo, - isShared: true - }; + return device; } catch (err) { console.warn('Failed to acquire WebGPU device:', err); return null; @@ -226,19 +264,26 @@ export class GpuDevicePool { } /** - * Release device reference. When reference count drops to 0, shared device is destroyed. + * Release a device reference. Only the currently pooled device is + * refcounted: releasing a stale device (from before a loss/re-acquire) or a + * non-shared injected device is a no-op, so it can never decrement — and + * destroy — another holder's device. When the count drops to 0 the shared + * device is destroyed. */ - static releaseDevice(_device: GPUDevice, isShared: boolean = true): void { + static releaseDevice(device: GPUDevice, isShared: boolean = true): void { if (!isShared) return; + if (!device || device !== this.sharedDevice) return; this.refCount = Math.max(0, this.refCount - 1); - if (this.refCount === 0 && this.sharedDevice) { + if (this.refCount === 0) { + const dev = this.sharedDevice; + // Null first so the lost handler (reason 'destroyed') sees a stale + // identity and stays silent. + this.sharedDevice = null; try { - this.sharedDevice.destroy(); + dev.destroy(); } catch { // ignore } - this.sharedDevice = null; - this.sharedAdapter = null; } } @@ -289,7 +334,6 @@ export class GpuDevicePool { */ static simulateDeviceLoss(reason: string = 'Simulated device loss'): void { this.sharedDevice = null; - this.sharedAdapter = null; this.refCount = 0; for (const listener of [...this.deviceLostListeners]) { try { diff --git a/packages/webgpu-search/src/ranking.ts b/packages/webgpu-search/src/ranking.ts index 3fa05d6..8a20d73 100644 --- a/packages/webgpu-search/src/ranking.ts +++ b/packages/webgpu-search/src/ranking.ts @@ -93,29 +93,28 @@ function compareStringsAsc(a: string, b: string): number { } /** - * Deterministic ID comparison: + * Deterministic ID comparison (a strict total order over mixed id types): * - number vs number: numeric ascending (both must be finite; NaN/Infinity * throw fail-closed to preserve the total-order contract). - * - otherwise: String(id) code-unit ascending (covers string/string and - * mixed string/number pairs deterministically; note numeric `2` and string - * `"2"` compare id-equal and fall through to `docIndex`). + * - number vs string: every number sorts before every string (type-ranked, + * so mixed corpora stay transitive; numeric `2` and string `"2"` are NOT + * id-equal: `2` sorts first). + * - string vs string: code-unit ascending (never locale collation). */ export function compareIdsAsc(a: string | number, b: string | number): number { - if (typeof a === 'number' && typeof b === 'number') { - if (!Number.isFinite(a) || !Number.isFinite(b)) { - throw new TypeError( - `[webgpu-search] compareIdsAsc requires finite numeric ids, got ${String(a)} vs ${String(b)}.` - ); - } - if (a !== b) return a > b ? 1 : -1; - return 0; - } - if (typeof a === 'number' && !Number.isFinite(a)) { + const aNum = typeof a === 'number'; + const bNum = typeof b === 'number'; + if (aNum && !Number.isFinite(a)) { throw new TypeError(`[webgpu-search] compareIdsAsc requires finite numeric id, got ${String(a)}.`); } - if (typeof b === 'number' && !Number.isFinite(b)) { + if (bNum && !Number.isFinite(b)) { throw new TypeError(`[webgpu-search] compareIdsAsc requires finite numeric id, got ${String(b)}.`); } + if (aNum && bNum) { + if (a !== b) return (a as number) > (b as number) ? 1 : -1; + return 0; + } + if (aNum !== bNum) return aNum ? -1 : 1; return compareStringsAsc(String(a), String(b)); } diff --git a/packages/webgpu-search/src/search-index.ts b/packages/webgpu-search/src/search-index.ts index 3afd219..1064033 100644 --- a/packages/webgpu-search/src/search-index.ts +++ b/packages/webgpu-search/src/search-index.ts @@ -191,8 +191,12 @@ export class SearchIndex { return index; } + // Hoisted so the catch can release the engine: a throw after init() + // succeeds (pipeline creation / loadDataset budget/OOM) would + // otherwise leak the pool refcount and GPU buffers. + let gpu: WebGPUEngine | null = null; try { - const gpu = new WebGPUEngine(); + gpu = new WebGPUEngine(); const initialized = await gpu.init( options.device !== undefined || options.powerPreference !== undefined ? { device: options.device, powerPreference: options.powerPreference } @@ -218,6 +222,7 @@ export class SearchIndex { return index; } else { + try { gpu.destroy(); } catch {} index.engineState = transitionEngineState( index.engineState, 'cpu', @@ -227,6 +232,7 @@ export class SearchIndex { ); } } catch (gpuErr) { + try { gpu?.destroy(); } catch {} console.warn('[webgpu-search] WebGPU initialization failed, falling back to CPU:', gpuErr); index.engineState = transitionEngineState( index.engineState, @@ -741,8 +747,11 @@ export class SearchIndex { this.gpuEngine = null; } this.vramAllocatedBytes = 0; + // Hoisted so the catch can release a post-init failure (pipeline / + // loadDataset throw) instead of leaking the pool refcount + buffers. + let gpu: WebGPUEngine | null = null; try { - const gpu = new WebGPUEngine(); + gpu = new WebGPUEngine(); const initialized = await gpu.init( deviceToUse !== undefined || powerPreference !== undefined ? { device: deviceToUse, powerPreference } @@ -773,6 +782,7 @@ export class SearchIndex { this.subscribeDeviceLost(); return true; } catch { + try { gpu?.destroy(); } catch {} this.engineState = transitionEngineState( this.engineState, 'cpu', diff --git a/packages/webgpu-search/src/snapshot-codec.ts b/packages/webgpu-search/src/snapshot-codec.ts index 945698c..cc1aa35 100644 --- a/packages/webgpu-search/src/snapshot-codec.ts +++ b/packages/webgpu-search/src/snapshot-codec.ts @@ -205,6 +205,56 @@ export function validateColumnarPayload( } } +function isArrayBufferLike(value: unknown): value is ArrayBuffer { + if (typeof ArrayBuffer !== 'undefined' && value instanceof ArrayBuffer) return true; + // Cross-realm ArrayBuffer (e.g. from another iframe / vm context). + return Object.prototype.toString.call(value) === '[object ArrayBuffer]'; +} + +function isSharedArrayBufferLike(value: unknown): boolean { + return Object.prototype.toString.call(value) === '[object SharedArrayBuffer]'; +} + +/** + * Normalizes snapshot input to a standalone `ArrayBuffer`. + * + * - `ArrayBuffer` is returned as-is (zero-copy). + * - `ArrayBufferView` (typed arrays, `DataView`, Node `Buffer`) is narrowed + * to its exact `[byteOffset, byteOffset + byteLength)` window; a view that + * spans its whole non-shared backing buffer is returned zero-copy, + * otherwise the window is copied. + * - `SharedArrayBuffer` (or views over one) is copied into a private + * `ArrayBuffer` so decoding never observes concurrent writes. + * + * Returns `null` for anything else so callers can raise their own + * fail-closed `IncompatibleIndexError`. + */ +export function toSnapshotArrayBuffer(input: unknown): ArrayBuffer | null { + if (input === null || input === undefined) return null; + if (isArrayBufferLike(input)) return input; + if (isSharedArrayBufferLike(input)) { + const src = new Uint8Array(input as ArrayBufferLike); + const copy = new Uint8Array(src.byteLength); + copy.set(src); + return copy.buffer; + } + if (typeof ArrayBuffer !== 'undefined' && ArrayBuffer.isView(input)) { + const view = input as ArrayBufferView; + const backing = view.buffer as ArrayBufferLike; + if ( + isArrayBufferLike(backing) && + view.byteOffset === 0 && + view.byteLength === backing.byteLength + ) { + return backing; + } + const copy = new Uint8Array(view.byteLength); + copy.set(new Uint8Array(backing, view.byteOffset, view.byteLength)); + return copy.buffer; + } + return null; +} + /** * Parses a snapshot header, accepting both legacy (48 bytes, * magic `0x55324433`, version 3) and canonical (56 bytes, magic @@ -213,11 +263,10 @@ export function validateColumnarPayload( * * Legacy headers report `columnarByteLength: 0` (absent segment). */ -export function decodeSnapshotHeader(buffer: ArrayBuffer): DocumentSnapshotHeader { - const buf = buffer as unknown as { byteLength?: unknown; slice?: unknown }; - const byteLen = typeof buf?.byteLength === 'number' ? (buf.byteLength as number) : NaN; - const canSlice = typeof (buf as any)?.slice === 'function'; - if (!Number.isFinite(byteLen) || byteLen < LEGACY_SNAPSHOT_HEADER_BYTES || !canSlice) { +export function decodeSnapshotHeader(input: ArrayBuffer | ArrayBufferView): DocumentSnapshotHeader { + const buffer = toSnapshotArrayBuffer(input); + const byteLen = buffer ? buffer.byteLength : NaN; + if (!buffer || !Number.isFinite(byteLen) || byteLen < LEGACY_SNAPSHOT_HEADER_BYTES) { throw new IncompatibleIndexError(LEGACY_SNAPSHOT_MAGIC, 'neutered/short'); } @@ -557,11 +606,17 @@ export interface RestoredDocumentSnapshot> { * offsets, and document records. */ export function decodeSnapshot>( - buffer: ArrayBuffer, + input: ArrayBuffer | ArrayBufferView, options?: RestoreDocumentIndexOptions ): RestoredDocumentSnapshot { - if (!buffer || typeof (buffer as ArrayBuffer).byteLength !== 'number') { - throw new IncompatibleIndexError('arraybuffer', typeof buffer); + const inputLen = (input as { byteLength?: unknown } | null | undefined)?.byteLength; + if (typeof inputLen === 'number' && inputLen > MAX_SNAPSHOT_BYTES) { + // Cap before normalization so oversized views are never copied. + throw new IncompatibleIndexError(`snapshot-bytes<=${MAX_SNAPSHOT_BYTES}`, inputLen); + } + const buffer = toSnapshotArrayBuffer(input); + if (!buffer) { + throw new IncompatibleIndexError('arraybuffer', typeof input); } if ((buffer as ArrayBuffer).byteLength > MAX_SNAPSHOT_BYTES) { throw new IncompatibleIndexError(`snapshot-bytes<=${MAX_SNAPSHOT_BYTES}`, (buffer as ArrayBuffer).byteLength); @@ -712,6 +767,26 @@ export function decodeSnapshot>( Array.isArray(schema.docIds) ? schema.docIds.length : typeof schema.docIds ); } + // Fail-closed id validation mirroring DocumentIndex id rules: every id + // must be a non-empty string or finite number, and unique (Map/Set key + // semantics, so `1` and `"1"` are distinct ids). + const seenIds = new Set(); + for (let i = 0; i < schema.docIds.length; i++) { + const id = schema.docIds[i] as unknown; + const validScalar = + (typeof id === 'string' && id.length > 0) || + (typeof id === 'number' && Number.isFinite(id)); + if (!validScalar) { + throw new IncompatibleIndexError( + `schema.docIds[${i}] non-empty string or finite number`, + id === null ? 'null' : Array.isArray(id) ? 'array' : typeof id + ); + } + if (seenIds.has(id as DocumentId)) { + throw new IncompatibleIndexError(`schema.docIds unique`, `duplicate ${typeof id} id at [${i}]`); + } + seenIds.add(id as DocumentId); + } docIds = schema.docIds; } @@ -727,6 +802,15 @@ export function decodeSnapshot>( if (!Array.isArray(records) || records.length !== header.docCount) { throw new IncompatibleIndexError(`array of ${header.docCount} records`, records?.length); } + for (let i = 0; i < records.length; i++) { + const doc = records[i] as unknown; + if (doc === null || typeof doc !== 'object' || Array.isArray(doc)) { + throw new IncompatibleIndexError( + `docs[${i}] object`, + doc === null ? 'null' : Array.isArray(doc) ? 'array' : typeof doc + ); + } + } } else { // Decoupled document storage if (options?.documents && Array.isArray(options.documents)) { @@ -758,7 +842,7 @@ export function decodeSnapshot>( * (canonical + legacy migration path). */ export async function restoreSnapshot>( - buffer: ArrayBuffer, + buffer: ArrayBuffer | ArrayBufferView, options?: RestoreDocumentIndexOptions ): Promise> { const snapshot = decodeSnapshot(buffer, options); diff --git a/packages/webgpu-search/src/snapshot-idb.ts b/packages/webgpu-search/src/snapshot-idb.ts index 0e9bb46..77c621f 100644 --- a/packages/webgpu-search/src/snapshot-idb.ts +++ b/packages/webgpu-search/src/snapshot-idb.ts @@ -1,4 +1,4 @@ -import { restoreSnapshot } from './snapshot-codec'; +import { restoreSnapshot, toSnapshotArrayBuffer } from './snapshot-codec'; import type { DocumentIndex } from './document-index'; import type { SearchWorkerClient } from './worker/worker-client'; import type { @@ -26,31 +26,32 @@ function resolveIDBFactory(options?: IDBStorageOptions): IDBFactory { return idb as IDBFactory; } -/** - * Opens the IndexedDB database, provisioning required object stores cleanly. - */ -export function openSearchDatabase(options?: IDBStorageOptions): Promise { - const idb = resolveIDBFactory(options); - const dbName = options?.dbName ?? DEFAULT_IDB_DATABASE_NAME; - const snapshotStoreName = options?.snapshotStoreName ?? DEFAULT_SNAPSHOT_STORE_NAME; - const docStoreName = options?.docStoreName ?? DEFAULT_DOCUMENT_STORE_NAME; - +function openRequest( + idb: IDBFactory, + dbName: string, + version: number | undefined, + requiredStores: string[] +): Promise { return new Promise((resolve, reject) => { - // Open with version 1 (or allow auto-increment) - const req = idb.open(dbName, 1); + const req = version === undefined ? idb.open(dbName) : idb.open(dbName, version); req.onupgradeneeded = () => { const db = req.result; - if (!db.objectStoreNames.contains(snapshotStoreName)) { - db.createObjectStore(snapshotStoreName); - } - if (!db.objectStoreNames.contains(docStoreName)) { - db.createObjectStore(docStoreName); + for (const name of requiredStores) { + if (!db.objectStoreNames.contains(name)) { + db.createObjectStore(name); + } } }; req.onsuccess = () => { - resolve(req.result); + const db = req.result; + // Yield to future version upgrades (e.g. another caller provisioning a + // custom store) instead of blocking them for the connection lifetime. + db.onversionchange = () => { + db.close(); + }; + resolve(db); }; req.onerror = () => { @@ -58,11 +59,58 @@ export function openSearchDatabase(options?: IDBStorageOptions): Promise { + // The upgrade proceeds once other connections close (ours close on + // `versionchange`); the request stays pending until then. console.warn(`[webgpu-search] IndexedDB database "${dbName}" open request blocked by other open tabs/connections.`); }; }); } +const MAX_STORE_PROVISION_ATTEMPTS = 3; + +/** + * Opens the IndexedDB database, provisioning required object stores cleanly. + * + * Opens at the current version first; if any required store (including a + * custom `snapshotStoreName` / `docStoreName`) is missing on an existing + * database, the connection is closed and the database is reopened at + * `version + 1` so `onupgradeneeded` can create the missing stores. + */ +export async function openSearchDatabase(options?: IDBStorageOptions): Promise { + const idb = resolveIDBFactory(options); + const dbName = options?.dbName ?? DEFAULT_IDB_DATABASE_NAME; + const snapshotStoreName = options?.snapshotStoreName ?? DEFAULT_SNAPSHOT_STORE_NAME; + const docStoreName = options?.docStoreName ?? DEFAULT_DOCUMENT_STORE_NAME; + const requiredStores = snapshotStoreName === docStoreName ? [snapshotStoreName] : [snapshotStoreName, docStoreName]; + + let db = await openRequest(idb, dbName, undefined, requiredStores); + for (let attempt = 0; attempt < MAX_STORE_PROVISION_ATTEMPTS; attempt++) { + const missing = requiredStores.some((name) => !db.objectStoreNames.contains(name)); + if (!missing) { + return db; + } + const nextVersion = db.version + 1; + db.close(); + try { + db = await openRequest(idb, dbName, nextVersion, requiredStores); + } catch (err) { + // A concurrent opener may have upgraded past `nextVersion` first + // (VersionError); reopen at the current version and re-check. + if ((err as { name?: string } | null)?.name !== 'VersionError') { + throw err; + } + db = await openRequest(idb, dbName, undefined, requiredStores); + } + } + if (requiredStores.some((name) => !db.objectStoreNames.contains(name))) { + db.close(); + throw new Error( + `[webgpu-search] Failed to provision IndexedDB object stores (${requiredStores.join(', ')}) in database "${dbName}".` + ); + } + return db; +} + /** * Transaction-safe helper to save a DocumentIndex or pre-serialized ArrayBuffer snapshot to IndexedDB. * @@ -72,7 +120,7 @@ export function openSearchDatabase(options?: IDBStorageOptions): Promise | SearchWorkerClient | ArrayBuffer, + indexOrBuffer: DocumentIndex | SearchWorkerClient | ArrayBuffer | ArrayBufferView, options?: SaveIDBOptions ): Promise { if (!indexOrBuffer) { @@ -84,14 +132,13 @@ export async function saveIndexToIDB( let docsToStore: any[] | undefined = options?.documents; // 1. Serialization phase (OUTSIDE transaction to avoid auto-commit) - const isBuffer = - indexOrBuffer instanceof ArrayBuffer || - (typeof ArrayBuffer !== 'undefined' && - typeof (indexOrBuffer as any)?.byteLength === 'number' && - typeof (indexOrBuffer as any)?.slice === 'function'); - - if (isBuffer) { - snapshot = indexOrBuffer as ArrayBuffer; + // ArrayBuffer is stored as-is; ArrayBufferViews (typed arrays, DataView, + // Buffer) are narrowed to their exact byte window so the stored value is + // always a plain ArrayBuffer that decodes correctly on load. + const normalizedBuffer = toSnapshotArrayBuffer(indexOrBuffer); + + if (normalizedBuffer) { + snapshot = normalizedBuffer; } else if (typeof (indexOrBuffer as any)?.serialize === 'function') { snapshot = await (indexOrBuffer as any).serialize({ decoupled }); if (decoupled && !docsToStore && typeof (indexOrBuffer as any)?.getRecords === 'function') { @@ -187,7 +234,12 @@ export async function loadIndexFromIDB>( const snapStore = tx.objectStore(snapshotStoreName); const snapReq = snapStore.get(key); snapReq.onsuccess = () => { - snapshotResult = (snapReq.result as ArrayBuffer) ?? null; + const raw = snapReq.result as unknown; + // Normalize legacy typed-array records to a plain ArrayBuffer; any + // other non-empty value is passed through so restore fails closed. + snapshotResult = raw === undefined || raw === null + ? null + : (toSnapshotArrayBuffer(raw) ?? (raw as ArrayBuffer)); }; if (shouldLoadDocs && db!.objectStoreNames.contains(docStoreName)) { diff --git a/packages/webgpu-search/src/webgpu-engine.ts b/packages/webgpu-search/src/webgpu-engine.ts index 937fbe6..fbf9565 100644 --- a/packages/webgpu-search/src/webgpu-engine.ts +++ b/packages/webgpu-search/src/webgpu-engine.ts @@ -5,6 +5,7 @@ import { checkMemoryBudget, computeClampedHeadroomBytes, deserializeDataset, + resolveStorageBufferLimit, packDataset, validatePackedOffsets, type PackedDataset, @@ -161,10 +162,13 @@ export class WebGPUEngine { async init( customDevice?: GPUDevice | { device?: GPUDevice; powerPreference?: GPUPowerPreference } ): Promise { - // Guard re-entry: dispose existing GPU buffers before re-creating so - // init() twice does not leak the first set (benchmark/power users). + // Guard re-entry: dispose existing GPU buffers AND release the previous + // device reference before re-acquiring, so init() twice leaks neither the + // first buffer set nor a pool refcount (benchmark/power users). if (this.device) { + this.generation++; this.disposeGpuBuffers(); + this.releaseCurrentDevice(); } // Normalize overloads: bare GPUDevice (legacy) or options bag. // powerPreference is validated fail-closed (unknown throws) and forwarded @@ -203,7 +207,7 @@ export class WebGPUEngine { hasTimestampQuery: injected.features ? injected.features.has('timestamp-query') : false }; this.attachLostHandler(); - return await this.setupPipelinesAndBuffers(); + return await this.setupOrTeardown(); } const acquired = await GpuDevicePool.acquireDevice( @@ -218,7 +222,32 @@ export class WebGPUEngine { this.isSharedDevice = acquired.isShared; this.deviceReleased = false; this.attachLostHandler(); - return await this.setupPipelinesAndBuffers(); + return await this.setupOrTeardown(); + } + + /** + * Pipeline/buffer setup with fail-closed teardown: a throw (shader + * compile, pipeline creation, OOM) destroys partial buffers and releases + * the device reference so a failed init never leaks a pool refcount. + */ + private async setupOrTeardown(): Promise { + try { + return await this.setupPipelinesAndBuffers(); + } catch (err) { + this.destroy(); + throw err; + } + } + + /** Release the held device reference exactly once (deviceReleased guard). */ + private releaseCurrentDevice(): void { + if (this.device && !this.deviceReleased) { + this.deviceReleased = true; + try { + GpuDevicePool.releaseDevice(this.device, this.isSharedDevice); + } catch {} + } + this.device = null; } private attachLostHandler(): void { @@ -235,13 +264,7 @@ export class WebGPUEngine { this.offsetsBuffer = this.recordsBuffer = this.uniformBuffer = this.queryBuffer = null; this.outputBuffer = this.stagingBuffer = this.queryResolveBuffer = this.queryStagingBuffer = null; this.querySet = null; - if (this.device !== null && !this.deviceReleased) { - this.deviceReleased = true; - try { - GpuDevicePool.releaseDevice(this.device, this.isSharedDevice); - } catch {} - } - this.device = null; + this.releaseCurrentDevice(); this.generation++; }, () => {}); } @@ -561,6 +584,26 @@ export class WebGPUEngine { const allocOffsets = Math.max(reqOffsets, targetOffsets); const allocRecords = Math.max(reqRecords, targetRecords); + // Per-buffer binding limit: headroom is clamped to the limit, so only an + // actual corpus above the limit can land here. Throw the budget error + // (callers fall back to CPU) instead of creating an oversized binding + // that fails validation and silently yields 0 results. + const bindingLimit = resolveStorageBufferLimit(this.device); + if (allocOffsets > bindingLimit || allocRecords > bindingLimit) { + Object.assign(this, { + currentDatasetSize: prev.size, + currentStrings: prev.strings, + currentTokens: prev.tokens, + currentOffsets: prev.offsets, + normalized: prev.normalized, + profileId: prev.profileId, + unicodeVersion: prev.unicodeVersion, + scoringVersion: prev.scoringVersion, + }); + const over = allocRecords > bindingLimit ? 'records' : 'offsets'; + throw new Error(`[webgpu-search] Dataset ${over} buffer too large. Falling back to CPU.`); + } + const t0 = nowMs(); let newOffsetsBuffer: GPUBuffer | null = null; let newRecordsBuffer: GPUBuffer | null = null; @@ -899,53 +942,63 @@ export class WebGPUEngine { const tReadbackStart = nowMs(); let gpuExecutionMs: number | null = null; - if (this.queryStagingBuffer && !skipTimestamps) { + // Generation-epoch abort for in-flight mapAsync (takes no signal, so + // abort is cooperative discard). Captured BEFORE the first await (the + // timestamp readback): destroy()/ensureCandidateCapacity()/loadDataset() + // during that await bump the epoch and null/replace buffers, which must + // surface as AbortError — never a TypeError or a read of a fresh + // (zeroed) staging buffer. Buffers are pinned locally for the same reason. + // Rejected: bare Promise.race — abandons the mapping, later + // destroy-while-mapped error. mapAsync rejection itself (e.g. destroy + // raced it via unmap) also maps to AbortError when the epoch moved, so + // hybrid doesn't misclassify it as fallback-eligible GPU failure. + const gen = this.generation; + const stagingBuffer = this.stagingBuffer; + const tsStaging = this.queryStagingBuffer; + + if (tsStaging && !skipTimestamps) { try { - await this.queryStagingBuffer.mapAsync(MapMode.READ); - const tu = new BigUint64Array(this.queryStagingBuffer.getMappedRange()); - if (tu[1] >= (tu[0]) && (tu[0]) > 0n) { - gpuExecutionMs = Number((tu[1]) - (tu[0])) / 1_000_000; + await tsStaging.mapAsync(MapMode.READ); + if (gen === this.generation) { + const tu = new BigUint64Array(tsStaging.getMappedRange()); + if (tu[1] >= (tu[0]) && (tu[0]) > 0n) { + gpuExecutionMs = Number((tu[1]) - (tu[0])) / 1_000_000; + } } } catch (tsErr) { - console.warn('timestamp read failed:', tsErr); + if (gen === this.generation) console.warn('timestamp read failed:', tsErr); } finally { - try { this.queryStagingBuffer.unmap(); } catch {} + try { tsStaging.unmap(); } catch {} } } + // Torn down / reallocated during the timestamp await. + if (gen !== this.generation || !this.device || this.stagingBuffer !== stagingBuffer) { + throw abortError(); + } + let totalMatches = 0; let candidateCount = 0; let hasOverflow = false; let candidates: SearchResultItem[] = []; - // Generation-epoch abort for in-flight mapAsync (takes no signal, so - // abort is cooperative discard). Capture before await; on mismatch unmap - // and discard (stale dataset / torn-down engine). Rejected: bare - // Promise.race — abandons the mapping, later destroy-while-mapped error. - // mapAsync rejection itself (e.g. destroy raced it via unmap) also maps - // to AbortError when the epoch moved, so hybrid doesn't misclassify it - // as fallback-eligible GPU failure. - const gen = this.generation; try { - await this.stagingBuffer.mapAsync(MapMode.READ); + await stagingBuffer.mapAsync(MapMode.READ); } catch (mapErr) { if (gen !== this.generation) { - try { this.stagingBuffer.unmap(); } catch {} + try { stagingBuffer.unmap(); } catch {} throw abortError(); } throw mapErr; } try { - if (gen !== this.generation) { - try { this.stagingBuffer.unmap(); } catch {} - throw abortError(); - } - // Re-check liveness after the await: destroy() nulls device/buffers. - if (!this.device || !this.stagingBuffer) { + // Re-check epoch + liveness after the await: destroy() nulls + // device/buffers, reallocation swaps them. + if (gen !== this.generation || !this.device || this.stagingBuffer !== stagingBuffer) { throw abortError(); } throwIfAborted(options.signal); - const arrayBuffer = this.stagingBuffer.getMappedRange(); + const arrayBuffer = stagingBuffer.getMappedRange(); const u32Read = new Uint32Array(arrayBuffer); const i32Read = new Int32Array(arrayBuffer); @@ -966,7 +1019,7 @@ export class WebGPUEngine { }; } } finally { - try { this.stagingBuffer.unmap(); } catch {} + try { stagingBuffer.unmap(); } catch {} } throwIfAborted(options.signal); @@ -1033,13 +1086,7 @@ export class WebGPUEngine { this.currentDatasetSize = 0; this.allocatedOffsetsByteLength = 0; this.allocatedRecordsByteLength = 0; - if (this.device && !this.deviceReleased) { - this.deviceReleased = true; - try { - GpuDevicePool.releaseDevice(this.device, this.isSharedDevice); - } catch {} - } - this.device = null; + this.releaseCurrentDevice(); const p = this.searchMutex.then(() => {}, () => {}); this.searchMutex = p.catch(() => {}); } diff --git a/packages/webgpu-search/src/worker/protocol.ts b/packages/webgpu-search/src/worker/protocol.ts index 83b1767..296b8ac 100644 --- a/packages/webgpu-search/src/worker/protocol.ts +++ b/packages/webgpu-search/src/worker/protocol.ts @@ -67,6 +67,20 @@ export interface WorkerRestorePayload { options?: RestoreDocumentIndexOptions; } +/** + * RESTORE success result: the validated schema of the restored index, so the + * client commits field/id definitions only after the worker accepted the + * snapshot (checksum + schema validation happen once, worker-side). + * `records` is present only when the snapshot embedded documents and the + * caller did not supply `options.documents`. + */ +export interface WorkerRestoreResult { + fields: Array<{ name: string; weight: number }>; + filterFields: Array<{ name: string; type?: string }>; + idField?: string; + records?: Array>; +} + /** * Serialize an Error object (including custom library errors) into a structured-clone-safe format. */ diff --git a/packages/webgpu-search/src/worker/search-worker.ts b/packages/webgpu-search/src/worker/search-worker.ts index fc32532..b887220 100644 --- a/packages/webgpu-search/src/worker/search-worker.ts +++ b/packages/webgpu-search/src/worker/search-worker.ts @@ -1,5 +1,5 @@ import { DocumentIndex } from '../document-index'; -import { restoreDocumentIndex } from '../snapshot-codec'; +import { decodeSnapshotHeader, restoreDocumentIndex } from '../snapshot-codec'; import type { WorkerMessageType } from '../types'; import { serializeError, @@ -9,6 +9,7 @@ import { type WorkerRequest, type WorkerResponse, type WorkerRestorePayload, + type WorkerRestoreResult, type WorkerSearchPayload, type WorkerSerializePayload } from './protocol'; @@ -44,7 +45,56 @@ export function startSearchWorker(customScope?: any): void { let index: DocumentIndex | null = null; let activeAbortController: AbortController | null = null; + let activeQueryId: number = 0; + // Highest SEARCH queryId seen (updated on arrival, before queueing). let latestQueryId: number = 0; + // Every queryId <= abortedUpTo was cancelled by the client. + let abortedUpTo: number = -1; + + /** + * Message serialization (reader/writer queue). State-mutating requests + * (INIT / RESTORE / MUTATE / DESTROY) run exclusively, after every earlier + * request has settled. Read requests (SEARCH / SERIALIZE / STATS) wait only + * for earlier exclusive requests, so concurrent searches keep their + * supersede/abort semantics. ABORT is handled immediately on arrival. + */ + let allTail: Promise = Promise.resolve(); + let exclusiveTail: Promise = Promise.resolve(); + const EXCLUSIVE_TYPES: ReadonlySet = new Set(['INIT', 'RESTORE', 'MUTATE', 'DESTROY']); + + /** + * Session epoch, bumped when DESTROY arrives. Requests are tagged with the + * epoch they arrived in; replies (and not-yet-started work) from an older + * epoch are dropped so a reused worker never answers a new client's + * request ids with a previous session's results. + */ + let sessionEpoch = 0; + const requestEpoch = new WeakMap(); + const reply = (req: WorkerRequest, msg: WorkerResponse, transfer?: Transferable[]): void => { + if (requestEpoch.get(req) !== sessionEpoch) { + return; + } + if (transfer) { + scope.postMessage(msg, transfer); + } else { + scope.postMessage(msg); + } + }; + + const swapIndex = (next: DocumentIndex | null): void => { + const prev = index; + index = next; + if (prev && prev !== next) { + prev.destroy(); + } + }; + + const abortActiveSearch = (): void => { + if (activeAbortController) { + activeAbortController.abort(); + activeAbortController = null; + } + }; type WorkerHandler = (req: WorkerRequest) => Promise | void; @@ -57,16 +107,14 @@ export function startSearchWorker(customScope?: any): void { 'INIT', async (req) => { try { - if (index) { - index.destroy(); - index = null; - } const payload = req.payload as WorkerInitPayload; const records = Array.isArray(payload?.records) ? payload.records : []; - index = await DocumentIndex.create(records, payload?.options ?? { fields: [] }); - scope.postMessage({ id: req.id, success: true } satisfies WorkerResponse); + // Build-then-swap: a failed INIT leaves the previous index live. + const next = await DocumentIndex.create(records, payload?.options ?? { fields: [] }); + swapIndex(next); + reply(req, { id: req.id, success: true } satisfies WorkerResponse); } catch (err) { - scope.postMessage({ + reply(req, { id: req.id, success: false, error: serializeError(err) @@ -81,7 +129,7 @@ export function startSearchWorker(customScope?: any): void { const queryId = payload?.queryId ?? 0; if (!index) { - scope.postMessage({ + reply(req, { id: req.id, success: false, error: serializeError( @@ -91,17 +139,16 @@ export function startSearchWorker(customScope?: any): void { return; } - if (queryId < latestQueryId) { - // Superseded by newer query + if (queryId < latestQueryId || queryId <= abortedUpTo) { + // Superseded by a newer query or cancelled while queued. return; } - latestQueryId = queryId; - if (activeAbortController) { - activeAbortController.abort(); - } - activeAbortController = new AbortController(); - const signal = activeAbortController.signal; + abortActiveSearch(); + const controller = new AbortController(); + activeAbortController = controller; + activeQueryId = queryId; + const signal = controller.signal; try { const resp = await index.search(payload.query, { @@ -109,6 +156,9 @@ export function startSearchWorker(customScope?: any): void { signal }); + if (activeAbortController === controller) { + activeAbortController = null; + } if (signal.aborted || queryId < latestQueryId) { return; } @@ -118,16 +168,19 @@ export function startSearchWorker(customScope?: any): void { resp.results = resp.results.map((item) => ({ ...item, doc: undefined })); } - scope.postMessage({ + reply(req, { id: req.id, success: true, result: resp } satisfies WorkerResponse); } catch (err: any) { + if (activeAbortController === controller) { + activeAbortController = null; + } if (signal.aborted || queryId < latestQueryId) { return; } - scope.postMessage({ + reply(req, { id: req.id, success: false, error: serializeError(err) @@ -139,11 +192,15 @@ export function startSearchWorker(customScope?: any): void { 'ABORT', (req) => { const payload = req.payload as WorkerAbortPayload; - if (payload && payload.queryId >= latestQueryId) { - if (activeAbortController) { - activeAbortController.abort(); - activeAbortController = null; - } + const queryId = payload?.queryId; + if (typeof queryId !== 'number' || !Number.isFinite(queryId)) { + return; + } + if (queryId > abortedUpTo) { + abortedUpTo = queryId; + } + if (activeAbortController && activeQueryId <= queryId) { + abortActiveSearch(); } } ], @@ -151,7 +208,7 @@ export function startSearchWorker(customScope?: any): void { 'MUTATE', async (req) => { if (!index) { - scope.postMessage({ + reply(req, { id: req.id, success: false, error: serializeError( @@ -163,13 +220,13 @@ export function startSearchWorker(customScope?: any): void { try { const payload = req.payload as WorkerMutatePayload; const res = await index.applyBatch(payload.batch, payload.options); - scope.postMessage({ + reply(req, { id: req.id, success: true, result: res } satisfies WorkerResponse); } catch (err) { - scope.postMessage({ + reply(req, { id: req.id, success: false, error: serializeError(err) @@ -181,7 +238,7 @@ export function startSearchWorker(customScope?: any): void { 'SERIALIZE', (req) => { if (!index) { - scope.postMessage({ + reply(req, { id: req.id, success: false, error: serializeError( @@ -193,12 +250,9 @@ export function startSearchWorker(customScope?: any): void { try { const payload = req.payload as WorkerSerializePayload | undefined; const buf = index.serialize(payload?.options); - scope.postMessage( - { id: req.id, success: true, result: buf } satisfies WorkerResponse, - [buf] - ); + reply(req, { id: req.id, success: true, result: buf } satisfies WorkerResponse, [buf]); } catch (err) { - scope.postMessage({ + reply(req, { id: req.id, success: false, error: serializeError(err) @@ -211,14 +265,29 @@ export function startSearchWorker(customScope?: any): void { async (req) => { try { const payload = req.payload as WorkerRestorePayload; - if (index) { - index.destroy(); - index = null; + const hasCallerDocs = Array.isArray(payload?.options?.documents); + // Validate-then-swap: a corrupt snapshot leaves the previous index live. + const header = decodeSnapshotHeader(payload?.buffer); + const next = await restoreDocumentIndex(payload?.buffer, payload?.options); + let result: WorkerRestoreResult; + try { + result = { + fields: next.getSortedFields().map((f) => ({ name: f.name, weight: f.weight })), + filterFields: next.getFilterFieldDefinitions().map((ff) => ({ + name: ff.name as string, + ...(ff.type !== undefined ? { type: ff.type } : {}) + })), + ...(typeof next.options.idField === 'string' ? { idField: next.options.idField } : {}), + ...(header.docsByteLength > 0 && !hasCallerDocs ? { records: next.getRecords() } : {}) + }; + } catch (err) { + next.destroy(); + throw err; } - index = await restoreDocumentIndex(payload.buffer, payload.options); - scope.postMessage({ id: req.id, success: true } satisfies WorkerResponse); + swapIndex(next); + reply(req, { id: req.id, success: true, result } satisfies WorkerResponse); } catch (err) { - scope.postMessage({ + reply(req, { id: req.id, success: false, error: serializeError(err) @@ -230,7 +299,7 @@ export function startSearchWorker(customScope?: any): void { 'STATS', (req) => { if (!index) { - scope.postMessage({ + reply(req, { id: req.id, success: false, error: serializeError( @@ -241,13 +310,13 @@ export function startSearchWorker(customScope?: any): void { } try { const stats = index.getStats(); - scope.postMessage({ + reply(req, { id: req.id, success: true, result: stats } satisfies WorkerResponse); } catch (err) { - scope.postMessage({ + reply(req, { id: req.id, success: false, error: serializeError(err) @@ -258,23 +327,32 @@ export function startSearchWorker(customScope?: any): void { [ 'DESTROY', (req) => { - if (index) { - index.destroy(); - index = null; - } - if (activeAbortController) { - activeAbortController.abort(); - activeAbortController = null; - } - scope.postMessage({ id: req.id, success: true } satisfies WorkerResponse); + swapIndex(null); + abortActiveSearch(); + reply(req, { id: req.id, success: true } satisfies WorkerResponse); } ] ]); - const onMessage = async (event: MessageEvent) => { + const runGuarded = (handler: WorkerHandler, req: WorkerRequest) => async (): Promise => { + if (requestEpoch.get(req) !== sessionEpoch) { + return; // Superseded by a later DESTROY before it started. + } + try { + await handler(req); + } catch (err) { + reply(req, { + id: req.id, + success: false, + error: serializeError(err) + } satisfies WorkerResponse); + } + }; + + const onMessage = (event: MessageEvent): Promise => { const req = event?.data as WorkerRequest; if (!req || typeof req.id !== 'number' || typeof req.type !== 'string') { - return; + return Promise.resolve(); } const handler = handlers.get(req.type as WorkerMessageType); @@ -286,9 +364,50 @@ export function startSearchWorker(customScope?: any): void { new Error(`[webgpu-search] Unknown worker request type: ${(req as any).type}`) ) } satisfies WorkerResponse); - return; + return Promise.resolve(); + } + + // Never let a rejected run (e.g. reply() throwing DataCloneError) poison + // the queue tails and silently skip every later request. + const guarded = runGuarded(handler, req); + const run = (): Promise => guarded().catch(() => undefined); + + if (req.type === 'SEARCH') { + // Supersede on arrival so a newer query cancels the in-flight one even + // while it waits behind a queued state mutation. + const queryId = (req.payload as WorkerSearchPayload | undefined)?.queryId ?? 0; + if (queryId > latestQueryId) { + latestQueryId = queryId; + if (activeAbortController && activeQueryId < queryId) { + abortActiveSearch(); + } + } + } else if (req.type === 'DESTROY') { + // Cancel in-flight work early; the index teardown itself is queued + // behind any pending INIT/RESTORE so it cannot be undone by them. + // Query sequencing restarts so a reused (non-owned) worker accepts a + // fresh client's queryIds. + abortActiveSearch(); + latestQueryId = 0; + abortedUpTo = -1; + sessionEpoch++; + } + + requestEpoch.set(req, sessionEpoch); + + if (req.type === 'ABORT') { + return run(); + } + + if (EXCLUSIVE_TYPES.has(req.type)) { + const p = allTail.then(run); + allTail = p; + exclusiveTail = p; + return p; } - await handler(req); + const p = exclusiveTail.then(run); + allTail = Promise.all([allTail, p]); + return p; }; if (typeof scope.addEventListener === 'function') { diff --git a/packages/webgpu-search/src/worker/worker-client.ts b/packages/webgpu-search/src/worker/worker-client.ts index 1aeb09c..f03fc67 100644 --- a/packages/webgpu-search/src/worker/worker-client.ts +++ b/packages/webgpu-search/src/worker/worker-client.ts @@ -17,11 +17,12 @@ import { type WorkerMessageType, type WorkerRequest, type WorkerResponse, + type WorkerRestoreResult, type WorkerSearchPayload } from './protocol'; import { abortError, throwIfAborted } from '../guard'; -import { decodeSnapshotHeader, MAX_SNAPSHOT_BYTES } from '../snapshot-codec'; -import { LEGACY_SNAPSHOT_HEADER_BYTES, SNAPSHOT_HEADER_BYTES, SNAPSHOT_MAGIC, IncompatibleIndexError } from '../text-profile'; +import { MAX_SNAPSHOT_BYTES, toSnapshotArrayBuffer } from '../snapshot-codec'; +import { IncompatibleIndexError } from '../text-profile'; import { IncompatibleHookError, IncompatibleOptionError } from '../errors'; import { hasAnyHook, normalizeSearchHooks } from '../hooks'; @@ -52,6 +53,16 @@ interface InternalFieldDef { getter: (doc: TDoc) => string | string[] | undefined | null; } +/** + * Doc-extraction config: how the client derives ids and serializable + * field/filter values for records it sends to the worker. + */ +interface ClientIndexConfig { + getId: (doc: TDoc) => DocumentId; + fieldDefinitions: InternalFieldDef[]; + filterDefinitions: Array<{ name: string; getter: (doc: any) => any }>; +} + interface PendingQuery { queryId: number; resolve: (res: DocumentSearchResponse) => void; @@ -78,9 +89,39 @@ export class SearchWorkerClient> { private nextQueryId: number = 1; private readonly stringIsolated: boolean; private readonly docMap: Map = new Map(); + /** + * Outgoing config: matches the worker state implied by every request + * posted so far (assuming success). Used to serialize MUTATE payloads. + */ private getId: (doc: TDoc) => DocumentId = (doc: any) => doc?.id; private fieldDefinitions: InternalFieldDef[] = []; private filterDefinitions: Array<{ name: string; getter: (doc: any) => any }> = []; + /** Config of the most recent INIT/RESTORE the worker confirmed. */ + private committedConfig: ClientIndexConfig = { + getId: this.getId, + fieldDefinitions: this.fieldDefinitions, + filterDefinitions: this.filterDefinitions + }; + /** + * Bumped whenever an INIT/RESTORE is posted. The worker processes + * requests in order and replies in order, so commits happen in completion + * order; a failed request only rolls the outgoing config back when no + * later INIT/RESTORE was posted after it. + */ + private stateGeneration: number = 0; + /** + * Settles once every RESTORE posted so far has settled. The restored + * schema is only known after the worker validates the snapshot, so + * mutations issued meanwhile wait before serializing their documents. + */ + private restoreBarrier: Promise | null = null; + /** + * Settles once every deferred request has posted. Requests normally post + * synchronously (worker processes them in arrival order); while a request + * is deferred (a mutation awaiting a pending RESTORE), later requests queue + * behind it so the worker still sees them in call order. + */ + private deferredPosts: Promise | null = null; private messageListener: ((event: MessageEvent) => void) | null = null; private errorListener: ((err: any) => void) | null = null; @@ -296,15 +337,83 @@ export class SearchWorkerClient> { }); } - private extractSerializableDoc(doc: TDoc, id: DocumentId): Record { + private currentConfig(): ClientIndexConfig { + return { + getId: this.getId, + fieldDefinitions: this.fieldDefinitions, + filterDefinitions: this.filterDefinitions + }; + } + + private setOutgoingConfig(config: ClientIndexConfig): void { + this.getId = config.getId; + this.fieldDefinitions = config.fieldDefinitions; + this.filterDefinitions = config.filterDefinitions; + } + + /** + * Runs `run` (which must post its request synchronously before its first + * `await`) immediately when nothing is deferred, otherwise after `waitFor` + * and every earlier deferred request have settled/posted. + */ + private runInPostOrder(waitFor: Promise | null, run: () => Promise): Promise { + const prior = this.deferredPosts; + if (!waitFor && !prior) { + return run(); + } + let release!: () => void; + const posted = new Promise((resolve) => { + release = resolve; + }); + this.deferredPosts = posted; + return Promise.all([waitFor, prior]).then(() => { + try { + return run(); + } finally { + release(); + if (this.deferredPosts === posted) { + this.deferredPosts = null; + } + } + }); + } + + /** Posts an INIT/RESTORE and commits/rolls back config in completion order. */ + private async sendStateRequest( + type: 'INIT' | 'RESTORE', + payload: unknown, + outgoing: ClientIndexConfig | null, + transfer?: Transferable[] + ): Promise<{ result: T; isLatest: boolean }> { + const generation = ++this.stateGeneration; + if (outgoing) { + this.setOutgoingConfig(outgoing); + } + try { + const result = await this.sendRequest(type, payload, transfer); + return { result, isLatest: this.stateGeneration === generation }; + } catch (err) { + if (this.stateGeneration === generation) { + // Worker kept its previous index; mirror it. + this.setOutgoingConfig(this.committedConfig); + } + throw err; + } + } + + private extractSerializableDoc( + doc: TDoc, + id: DocumentId, + config: ClientIndexConfig = this.currentConfig() + ): Record { const out: Record = { [INTERNAL_WORKER_ID_KEY]: id }; - for (let i = 0; i < this.fieldDefinitions.length; i++) { - const f = this.fieldDefinitions[i]; + for (let i = 0; i < config.fieldDefinitions.length; i++) { + const f = config.fieldDefinitions[i]; const val = f.getter(doc); out[f.name] = val !== undefined && val !== null ? val : ''; } - for (let i = 0; i < this.filterDefinitions.length; i++) { - const ff = this.filterDefinitions[i]; + for (let i = 0; i < config.filterDefinitions.length; i++) { + const ff = config.filterDefinitions[i]; const val = ff.getter(doc); if (val !== undefined) { out[ff.name] = val; @@ -313,7 +422,14 @@ export class SearchWorkerClient> { return out; } - async init( + init( + optionsOrRecords?: DocumentIndexOptions | TDoc[], + maybeOptions?: DocumentIndexOptions + ): Promise { + return this.runInPostOrder(null, () => this.initNow(optionsOrRecords, maybeOptions)); + } + + private async initNow( optionsOrRecords?: DocumentIndexOptions | TDoc[], maybeOptions?: DocumentIndexOptions ): Promise { @@ -343,20 +459,22 @@ export class SearchWorkerClient> { throw new TypeError('[webgpu-search] idField must be a string property name or function.'); } + // Stage config locally; it is only committed once the worker confirms. + let stagedGetId: (doc: TDoc) => DocumentId; if (typeof options.idField === 'function') { - this.getId = options.idField; + stagedGetId = options.idField; } else if (typeof options.idField === 'string') { const prop = options.idField; - this.getId = (doc: any) => doc[prop]; + stagedGetId = (doc: any) => doc[prop]; } else { - this.getId = (doc: any) => doc.id; + stagedGetId = (doc: any) => doc.id; } if (!Array.isArray(options.fields) || options.fields.length === 0) { throw new TypeError('[webgpu-search] DocumentIndex expects options.fields to be a non-empty array.'); } - this.fieldDefinitions = options.fields.map((f) => { + const stagedFieldDefs: InternalFieldDef[] = options.fields.map((f) => { if (typeof f === 'string') { if (!f) { throw new TypeError('[webgpu-search] Field name cannot be an empty string.'); @@ -406,16 +524,20 @@ export class SearchWorkerClient> { } } } - this.filterDefinitions = filterDefs; + const staged: ClientIndexConfig = { + getId: stagedGetId, + fieldDefinitions: stagedFieldDefs, + filterDefinitions: filterDefs + }; const nextDocMap = new Map(); const serializableRecords: Record[] = []; for (let i = 0; i < records.length; i++) { const doc = records[i]; - const id = this.getId(doc); + const id = stagedGetId(doc); nextDocMap.set(id, doc); if (this.stringIsolated) { - serializableRecords.push(this.extractSerializableDoc(doc, id)); + serializableRecords.push(this.extractSerializableDoc(doc, id, staged)); } else { serializableRecords.push(doc as any); } @@ -431,7 +553,7 @@ export class SearchWorkerClient> { assertNoWorkerExtensions(options.hooks, 'init'); } - const workerFields = this.fieldDefinitions.map((f) => ({ + const workerFields = stagedFieldDefs.map((f) => ({ name: f.name, weight: f.weight })); @@ -455,18 +577,27 @@ export class SearchWorkerClient> { fields: workerFields }; - await this.sendRequest('INIT', { + await this.sendStateRequest('INIT', { options: workerOptions, records: serializableRecords - }); + }, staged); + // Commit in completion order (worker replies in request order). + this.committedConfig = staged; this.docMap.clear(); for (const [id, doc] of nextDocMap) { this.docMap.set(id, doc); } } - async search( + search( + query: string, + options?: DocumentSearchOptions + ): Promise> { + return this.runInPostOrder(null, () => this.searchNow(query, options)); + } + + private async searchNow( query: string, options?: DocumentSearchOptions ): Promise> { @@ -596,7 +727,13 @@ export class SearchWorkerClient> { }); } - async applyBatch(batch: MutationBatch, options?: AddOptions): Promise { + applyBatch(batch: MutationBatch, options?: AddOptions): Promise { + // Field/id definitions are unknown until a pending RESTORE settles, so + // the mutation (and every request issued after it) waits for it. + return this.runInPostOrder(this.restoreBarrier, () => this.applyBatchNow(batch, options)); + } + + private async applyBatchNow(batch: MutationBatch, options?: AddOptions): Promise { if (this.isDestroyed) { throw new Error('[webgpu-search] SearchWorkerClient has been destroyed.'); } @@ -678,22 +815,32 @@ export class SearchWorkerClient> { return Array.from(this.docMap.values()); } - async serialize(options?: SerializeDocumentIndexOptions): Promise { - if (this.isDestroyed) { - throw new Error('[webgpu-search] SearchWorkerClient has been destroyed.'); - } - return this.sendRequest('SERIALIZE', { options }); + serialize(options?: SerializeDocumentIndexOptions): Promise { + return this.runInPostOrder(null, async () => { + if (this.isDestroyed) { + throw new Error('[webgpu-search] SearchWorkerClient has been destroyed.'); + } + return this.sendRequest('SERIALIZE', { options }); + }); + } + + restore(buffer: ArrayBuffer | ArrayBufferView, options?: RestoreDocumentIndexOptions): Promise { + return this.runInPostOrder(null, () => this.restoreNow(buffer, options)); } - async restore(buffer: ArrayBuffer, options?: RestoreDocumentIndexOptions): Promise { + private async restoreNow(buffer: ArrayBuffer | ArrayBufferView, options?: RestoreDocumentIndexOptions): Promise { if (this.isDestroyed) { throw new Error('[webgpu-search] SearchWorkerClient has been destroyed.'); } if (!buffer || typeof (buffer as any).byteLength !== 'number') { throw new TypeError('[webgpu-search] restore expects an ArrayBuffer.'); } - if ((buffer as ArrayBuffer).byteLength > MAX_SNAPSHOT_BYTES) { - throw new IncompatibleIndexError(`snapshot-bytes<=${MAX_SNAPSHOT_BYTES}`, (buffer as ArrayBuffer).byteLength); + if (buffer.byteLength > MAX_SNAPSHOT_BYTES) { + throw new IncompatibleIndexError(`snapshot-bytes<=${MAX_SNAPSHOT_BYTES}`, buffer.byteLength); + } + const snapshotBuffer = toSnapshotArrayBuffer(buffer); + if (!snapshotBuffer) { + throw new TypeError('[webgpu-search] restore expects an ArrayBuffer.'); } if ((options?.options as unknown as Record | undefined)?.extensions !== undefined) { throw new IncompatibleHookError( @@ -704,118 +851,9 @@ export class SearchWorkerClient> { if (options?.options?.hooks) { assertNoWorkerExtensions(options?.options?.hooks, 'restore'); } - - let stagedFieldDefs = this.fieldDefinitions; - let stagedFilterDefs = this.filterDefinitions; - let stagedGetId = this.getId; - const stagedDocMap = new Map(); - - try { - const header = decodeSnapshotHeader(buffer); - // Version-aware header width: snapshot is 56 B, legacy legacy snapshot is 48 B. - // Derive before the length guard so truncated snapshot buffers cannot - // slip through the legacy 48 B threshold. - const headerBytesForGuard = header.magic === SNAPSHOT_MAGIC ? SNAPSHOT_HEADER_BYTES : LEGACY_SNAPSHOT_HEADER_BYTES; - const userFields = options?.options?.fields; - const userFieldMap = new Map(); - if (Array.isArray(userFields)) { - for (const uf of userFields) { - if (typeof uf === 'string') { - userFieldMap.set(uf, uf); - } else if (uf && typeof uf === 'object' && typeof uf.name === 'string') { - userFieldMap.set(uf.name, uf); - } - } - } - - if (buffer.byteLength >= headerBytesForGuard + header.schemaByteLength) { - // snapshot inserts a columnar segment between offsets and docs; derive - // header width and columnar length from the parsed header so legacy - // legacy snapshot snapshots (no columnar segment) still stage correctly. - const headerBytes = header.magic === SNAPSHOT_MAGIC ? SNAPSHOT_HEADER_BYTES : LEGACY_SNAPSHOT_HEADER_BYTES; - const schemaBytes = new Uint8Array(buffer, headerBytes, header.schemaByteLength); - const schemaStr = new TextDecoder().decode(schemaBytes); - const schema = JSON.parse(schemaStr); - if (schema && Array.isArray(schema.fields)) { - stagedFieldDefs = schema.fields.map((f: any) => { - const uf = userFieldMap.get(f.name); - const getter = typeof uf === 'object' && uf !== null && typeof uf.getter === 'function' - ? uf.getter - : (doc: any) => doc[f.name]; - return { - name: f.name, - weight: f.weight ?? 1.0, - getter - }; - }); - } - if (schema && Array.isArray(schema.filterFields)) { - const userFilterFields = options?.options?.filterFields; - const userFilterMap = new Map(); - if (Array.isArray(userFilterFields)) { - for (const uff of userFilterFields) { - if (typeof uff === 'string') { - userFilterMap.set(uff, uff); - } else if (uff && typeof uff === 'object' && typeof uff.name === 'string') { - userFilterMap.set(uff.name, uff); - } - } - } - stagedFilterDefs = schema.filterFields.map((ff: any) => { - const uf = userFilterMap.get(ff.name); - const getter = typeof uf === 'object' && uf !== null && typeof uf.getter === 'function' - ? uf.getter - : (doc: any) => doc[ff.name]; - return { - name: ff.name, - getter - }; - }); - } - if (typeof options?.options?.idField === 'function') { - stagedGetId = options.options.idField; - } else if (typeof options?.options?.idField === 'string') { - const prop = options.options.idField; - stagedGetId = (doc: any) => doc[prop] ?? doc[INTERNAL_WORKER_ID_KEY] ?? doc.id; - } else if (schema.idField && typeof schema.idField === 'string') { - const prop = schema.idField; - stagedGetId = (doc: any) => doc[prop] ?? doc[INTERNAL_WORKER_ID_KEY] ?? doc.id; - } else { - stagedGetId = (doc: any) => doc[INTERNAL_WORKER_ID_KEY] ?? doc.id; - } - } - - // Repopulate stagedDocMap if documents are provided or embedded - if (options?.documents && Array.isArray(options.documents)) { - for (const doc of options.documents) { - const id = stagedGetId(doc); - stagedDocMap.set(id, doc); - } - } else if (header.docsByteLength > 0) { - const headerBytes = header.magic === SNAPSHOT_MAGIC ? SNAPSHOT_HEADER_BYTES : LEGACY_SNAPSHOT_HEADER_BYTES; - const columnarLen = header.columnarByteLength ?? 0; - const docsOffset = - headerBytes + - header.schemaByteLength + - header.tokenCount * 4 + - (header.rowCount + 1) * 4 + - columnarLen; - const docsBytes = new Uint8Array(buffer, docsOffset, header.docsByteLength); - const docsStr = new TextDecoder().decode(docsBytes); - const docs = JSON.parse(docsStr); - if (Array.isArray(docs)) { - for (const doc of docs) { - const id = stagedGetId(doc); - if (doc && doc[INTERNAL_WORKER_ID_KEY] !== undefined && doc.id === undefined) { - doc.id = doc[INTERNAL_WORKER_ID_KEY]; - } - stagedDocMap.set(id, doc); - } - } - } - } catch { - // Allow worker to handle fail-closed validation and throw IncompatibleIndexError - } + // Header, checksum, schema, docIds and documents are validated once, + // worker-side; client state commits only after the worker succeeds and + // worker errors (IncompatibleIndexError, ...) propagate rehydrated. // Sanitize options to avoid DataCloneError over postMessage. // Hooks are rejected fail-closed above; strip them defensively so @@ -841,25 +879,122 @@ export class SearchWorkerClient> { } : undefined } : undefined; + // Transfer only on opt-in; partial views were already copied by + // toSnapshotArrayBuffer, so a view's wider backing store is never detached. const shouldTransfer = options?.transfer === true; - const toSend = shouldTransfer ? buffer : buffer.slice(0); - await this.sendRequest('RESTORE', { buffer: toSend, options: sanitizedOptions }, [toSend]); + const toSend = shouldTransfer ? snapshotBuffer : snapshotBuffer.slice(0); - // Commit only after successful restore response from worker - this.fieldDefinitions = stagedFieldDefs; - this.filterDefinitions = stagedFilterDefs; - this.getId = stagedGetId; + const request = this.sendStateRequest( + 'RESTORE', + { buffer: toSend, options: sanitizedOptions }, + null, + [toSend] + ); + const barrier: Promise = Promise.all([this.restoreBarrier, request]).then( + () => undefined, + () => undefined + ); + this.restoreBarrier = barrier; + let outcome: { result: WorkerRestoreResult; isLatest: boolean }; + try { + outcome = await request; + } finally { + // Clear once the latest pending restore settles (no-op if superseded). + void barrier.then(() => { + if (this.restoreBarrier === barrier) this.restoreBarrier = null; + }); + } + + // Commit only after the worker validated and swapped in the snapshot. + const staged = this.buildRestoredConfig(outcome.result, options); + const stagedDocMap = new Map(); + if (options?.documents && Array.isArray(options.documents)) { + for (const doc of options.documents) { + stagedDocMap.set(staged.getId(doc), doc); + } + } else if (Array.isArray(outcome.result?.records)) { + for (const doc of outcome.result.records as any[]) { + if (doc && typeof doc === 'object') { + if (doc[INTERNAL_WORKER_ID_KEY] !== undefined && doc.id === undefined) { + doc.id = doc[INTERNAL_WORKER_ID_KEY]; + } + stagedDocMap.set(staged.getId(doc), doc as TDoc); + } + } + } + + this.committedConfig = staged; + if (outcome.isLatest) { + this.setOutgoingConfig(staged); + } this.docMap.clear(); for (const [k, v] of stagedDocMap) { this.docMap.set(k, v); } } - async getStats(): Promise { - if (this.isDestroyed) { - throw new Error('[webgpu-search] SearchWorkerClient has been destroyed.'); + private buildRestoredConfig( + result: WorkerRestoreResult | undefined, + options?: RestoreDocumentIndexOptions + ): ClientIndexConfig { + if (!result || !Array.isArray(result.fields) || !Array.isArray(result.filterFields)) { + throw new IncompatibleIndexError('worker restore schema', typeof result); + } + const userFieldMap = new Map(); + for (const uf of options?.options?.fields ?? []) { + if (typeof uf === 'string') { + userFieldMap.set(uf, uf); + } else if (uf && typeof uf === 'object' && typeof uf.name === 'string') { + userFieldMap.set(uf.name, uf); + } } - return this.sendRequest('STATS'); + const fieldDefinitions: InternalFieldDef[] = result.fields.map((f) => { + const uf = userFieldMap.get(f.name); + const getter = typeof uf === 'object' && uf !== null && typeof uf.getter === 'function' + ? uf.getter + : (doc: any) => doc[f.name]; + return { name: f.name, weight: f.weight ?? 1.0, getter }; + }); + + const userFilterMap = new Map(); + for (const uff of options?.options?.filterFields ?? []) { + if (typeof uff === 'string') { + userFilterMap.set(uff, uff); + } else if (uff && typeof uff === 'object' && typeof uff.name === 'string') { + userFilterMap.set(uff.name, uff); + } + } + const filterDefinitions = result.filterFields.map((ff) => { + const uf = userFilterMap.get(ff.name); + const getter = typeof uf === 'object' && uf !== null && typeof uf.getter === 'function' + ? uf.getter + : (doc: any) => doc[ff.name]; + return { name: ff.name, getter }; + }); + + let getId: (doc: TDoc) => DocumentId; + const userIdField = options?.options?.idField; + if (typeof userIdField === 'function') { + getId = userIdField; + } else if (typeof userIdField === 'string') { + const prop = userIdField; + getId = (doc: any) => doc[prop] ?? doc[INTERNAL_WORKER_ID_KEY] ?? doc.id; + } else if (typeof result.idField === 'string' && result.idField) { + const prop = result.idField; + getId = (doc: any) => doc[prop] ?? doc[INTERNAL_WORKER_ID_KEY] ?? doc.id; + } else { + getId = (doc: any) => doc[INTERNAL_WORKER_ID_KEY] ?? doc.id; + } + return { getId, fieldDefinitions, filterDefinitions }; + } + + getStats(): Promise { + return this.runInPostOrder(null, async () => { + if (this.isDestroyed) { + throw new Error('[webgpu-search] SearchWorkerClient has been destroyed.'); + } + return this.sendRequest('STATS'); + }); } async destroy(): Promise { diff --git a/packages/webgpu-search/test/document-index-integrity.test.ts b/packages/webgpu-search/test/document-index-integrity.test.ts new file mode 100644 index 0000000..8dbc87f --- /dev/null +++ b/packages/webgpu-search/test/document-index-integrity.test.ts @@ -0,0 +1,245 @@ +import { describe, test, expect } from 'bun:test'; +import { + ColumnarStore, + compileFilter, + compareIdsAsc, + DocumentIndex, + InvalidFilterError +} from '../src/index'; + +const cpu = { preferGpu: false } as const; + +describe('applyBatch atomicity with throwing filter getters', () => { + test('a throwing filter getter rejects the whole batch and leaves the index intact', async () => { + const idx = await DocumentIndex.create([{ id: 1, title: 'apple', cat: 'a' }] as any[], { + ...cpu, + fields: ['title'], + filterFields: [ + { + name: 'cat', + type: 'string', + getter: (d: any) => { + if (d.bad) throw new Error('boom'); + return d.cat; + } + } + ] + } as any); + + await expect( + idx.add([{ id: 2, title: 'banana', cat: 'b' }, { id: 3, title: 'cherry', bad: true }] as any[]) + ).rejects.toThrow('boom'); + await expect(idx.update({ id: 1, title: 'apricot', bad: true } as any)).rejects.toThrow('boom'); + + expect(idx.getStats().docCount).toBe(1); + expect((await idx.search('banana')).results).toHaveLength(0); + expect((await idx.search('apple', { filter: { cat: 'a' } })).results.map((r) => r.id)).toEqual([1]); + + // Subsequent mutations still work (no stale idToDocIndex / misaligned rows). + await idx.add([{ id: 2, title: 'banana', cat: 'b' }, { id: 4, title: 'durian', cat: 'd' }] as any[]); + await idx.update({ id: 4, title: 'eggplant', cat: 'e' } as any); + expect((await idx.search('banana', { filter: { cat: 'b' } })).results.map((r) => r.id)).toEqual([2]); + expect((await idx.search('eggplant', { filter: { cat: 'e' } })).results.map((r) => r.id)).toEqual([4]); + expect((await idx.search('durian')).results).toHaveLength(0); + expect(idx.getStats().docCount).toBe(3); + }); +}); + +describe('default getters ignore inherited Object.prototype members', () => { + test('search field and filter field named "constructor"', async () => { + const idx = await DocumentIndex.create( + [{ id: 1, title: 'apple' }, { id: 2, title: 'pear', constructor: 'acme' }] as any[], + { ...cpu, fields: ['title', 'constructor'], filterFields: [{ name: 'constructor', type: 'string' }] } as any + ); + // Without own-property access, `Object` would be stringified as "function Object() { [native code] }". + expect((await idx.search('native code')).results).toHaveLength(0); + expect((await idx.search('acme')).results.map((r) => r.id)).toEqual([2]); + expect( + (await idx.search('apple', { filter: { constructor: { exists: false } } } as any)).results.map((r) => r.id) + ).toEqual([1]); + }); + + test('ColumnarStore default getter reads own props and class accessors only', () => { + class Item { + constructor(private readonly c: string) {} + get cat(): string { + return this.c; + } + } + const s = new ColumnarStore([ + { name: 'toString', type: 'string' }, + { name: 'cat', type: 'string' } + ]); + s.init([{}, new Item('x')]); + expect(compileFilter({ toString: { exists: true } } as any, s).getMatchingIndices()).toEqual([]); + expect(compileFilter({ cat: 'x' }, s).getMatchingIndices()).toEqual([1]); + }); +}); + +describe('highlights with prototype-sensitive field names', () => { + test('field named __proto__ keeps highlights as an own key', async () => { + const idx = await DocumentIndex.create([{ id: 1, a: 'hello' }] as any[], { + ...cpu, + fields: [{ name: '__proto__', getter: (d: any) => d.a }] + } as any); + const r = (await idx.search('hello', { highlight: { tag: 'mark' } } as any)).results[0]; + expect(r).toBeDefined(); + expect(Object.getPrototypeOf(r.highlights)).toBe(Object.prototype); + expect(Object.prototype.hasOwnProperty.call(r.highlights, '__proto__')).toBe(true); + expect(Object.keys(r.highlights!)).toEqual(['__proto__']); + expect(Array.isArray((r.highlights as any)['__proto__'])).toBe(true); + }); +}); + +describe('compareIdsAsc total order', () => { + test('mixed number/string ids are transitive (numbers first)', () => { + expect(compareIdsAsc(9, 10)).toBe(-1); + expect(compareIdsAsc(10, '10a')).toBe(-1); + expect(compareIdsAsc(9, '10a')).toBe(-1); + expect(compareIdsAsc('10a', 9)).toBe(1); + expect(compareIdsAsc(2, '2')).toBe(-1); + expect(compareIdsAsc('10', '9')).toBe(-1); + expect(compareIdsAsc('a', 'a')).toBe(0); + expect(() => compareIdsAsc(NaN, 'a')).toThrow(TypeError); + const ids: Array = ['b', 10, '10a', 9, '9', 2, 'a']; + for (const perm of [ids, [...ids].reverse(), [...ids].sort()]) { + expect([...perm].sort(compareIdsAsc)).toEqual([2, 9, 10, '10a', '9', 'a', 'b']); + } + }); + + test('search ordering of tied mixed ids is input-order independent', async () => { + const run = async (docs: any[]) => + (await (await DocumentIndex.create(docs, { ...cpu, fields: ['t'] } as any)).search('x')).results.map( + (r) => r.id + ); + const a = await run([{ id: '10a', t: 'x' }, { id: 10, t: 'x' }, { id: 9, t: 'x' }]); + const b = await run([{ id: 10, t: 'x' }, { id: 9, t: 'x' }, { id: '10a', t: 'x' }]); + expect(a).toEqual([9, 10, '10a']); + expect(b).toEqual(a); + }); +}); + +describe('ColumnarStore strict value coercion', () => { + test('blank / non-finite numeric range bounds are rejected', () => { + const s = new ColumnarStore([{ name: 'n', type: 'number' }]); + s.init([{ n: 0 }, { n: 5 }]); + for (const bad of ['', ' ', 'abc', 'Infinity', NaN, Infinity]) { + expect(() => compileFilter({ n: { gte: bad as any } }, s)).toThrow(InvalidFilterError); + } + expect(compileFilter({ n: { gte: '1' } }, s).getMatchingIndices()).toEqual([1]); + }); + + test('blank strings are missing for number fields at index and query time', () => { + const s = new ColumnarStore([{ name: 'n', type: 'number' }]); + s.init([{ n: '' }, { n: ' 3 ' }, { n: 0 }]); + expect(compileFilter({ n: { exists: true } }, s).getMatchingIndices()).toEqual([1, 2]); + expect(compileFilter({ n: '' }, s).getMatchingIndices()).toEqual([]); + expect(compileFilter({ n: 3 }, s).getMatchingIndices()).toEqual([1]); + }); + + test('boolean fields accept only booleans and "true"/"false"', () => { + const s = new ColumnarStore([{ name: 'b', type: 'boolean' }]); + s.init([{ b: 'false' }, { b: false }, { b: true }, { b: 'true' }, { b: 1 }, { b: 'yes' }, {}]); + expect(compileFilter({ b: false }, s).getMatchingIndices()).toEqual([0, 1]); + expect(compileFilter({ b: true }, s).getMatchingIndices()).toEqual([2, 3]); + expect(compileFilter({ b: 'false' }, s).getMatchingIndices()).toEqual([0, 1]); + expect(compileFilter({ b: null }, s).getMatchingIndices()).toEqual([4, 5, 6]); + expect(compileFilter({ b: { in: [true, null] } }, s).getMatchingIndices()).toEqual([2, 3, 4, 5, 6]); + expect(() => compileFilter({ b: 'yes' }, s)).toThrow(InvalidFilterError); + expect(() => compileFilter({ b: 0 } as any, s)).toThrow(InvalidFilterError); + expect(() => compileFilter({ b: { in: [1] } } as any, s)).toThrow(InvalidFilterError); + }); +}); + +describe('ColumnarStore memory-proportional value postings', () => { + const distinctBitsetBytes = (s: ColumnarStore, name: string): number => { + const col = s.getColumn(name) as any; + const map: Map = col.type === 'string' ? col.invertedIndex : col.tagInverted; + let dense = 0; + for (const pl of map.values()) if (pl.isDense) dense++; + return dense; + }; + + test('high-cardinality values stay sparse; low-cardinality values go dense', () => { + const N = 40_000; + const s = new ColumnarStore([ + { name: 'sku', type: 'string' }, + { name: 'cat', type: 'string' }, + { name: 'tags', type: 'string[]' } + ]); + s.init(Array.from({ length: N }, (_, i) => ({ sku: 's' + i, cat: 'c' + (i % 4), tags: ['t' + i, 'common'] }))); + expect(distinctBitsetBytes(s, 'sku')).toBe(0); + expect(distinctBitsetBytes(s, 'cat')).toBe(4); + expect(distinctBitsetBytes(s, 'tags')).toBe(1); + + expect(compileFilter({ sku: 's123' }, s).getMatchingIndices()).toEqual([123]); + expect(compileFilter({ sku: ['s1', 's39999'] }, s).getMatchingIndices()).toEqual([1, 39999]); + expect(compileFilter({ cat: 'c1', sku: { gte: 's39990', lt: 's4' } }, s).getMatchingIndices()).toEqual([ + 39993, 39997 + ]); + expect(compileFilter({ tags: 't7' }, s).getMatchingIndices()).toEqual([7]); + expect(compileFilter({ tags: 'common' }, s).popcount()).toBe(N); + expect(compileFilter({ tags: { nin: ['common'] } }, s).popcount()).toBe(0); + }); + + test('remove + compact drops dead codes and keeps filters/facets correct', () => { + const s = new ColumnarStore([ + { name: 'sku', type: 'string' }, + { name: 'tags', type: 'string[]' } + ]); + const N = 200; + s.init(Array.from({ length: N }, (_, i) => ({ sku: 's' + i, tags: ['t' + (i % 3), 'u' + i] }))); + const map = new Map(); + for (let i = 0; i < N; i++) { + if (i % 2 === 0) s.remove(i); + else map.set(i, map.size); + } + // Removed values are released from postings eagerly. + expect((s.getColumn('sku') as any).invertedIndex.size).toBe(N / 2); + expect((s.getColumn('tags') as any).tagInverted.has('u0')).toBe(false); + + s.compact(map, map.size); + const col = s.getColumn('sku') as any; + expect(col.stringTable).toHaveLength(N / 2); + expect(col.stringToCode.has('s0')).toBe(false); + expect(compileFilter({ sku: 's0' }, s).popcount()).toBe(0); + expect(compileFilter({ sku: 's3' }, s).getMatchingIndices()).toEqual([1]); + expect(compileFilter({ tags: 't0' }, s).popcount()).toBe(33); // odd i with i % 3 === 0 + const facet = s.termsDistribution('sku', [0, 1])!; + expect(facet).toEqual([ + { value: 's1', count: 1 }, + { value: 's3', count: 1 } + ]); + // Re-adding a dropped value after compaction works. + s.add(map.size, { sku: 's0', tags: ['u0'] }); + expect(compileFilter({ sku: 's0' }, s).getMatchingIndices()).toEqual([map.size]); + }); + + test('dense postings demote when they become sparse', () => { + const s = new ColumnarStore([{ name: 'cat', type: 'string' }]); + const N = 4096; + s.init(Array.from({ length: N }, (_, i) => ({ cat: i < 2048 ? 'hot' : 'x' + i }))); + const pl = () => (s.getColumn('cat') as any).invertedIndex.get(0); + expect(pl().isDense).toBe(true); + for (let i = 0; i < 2040; i++) s.remove(i); + expect(pl().isDense).toBe(false); + expect(compileFilter({ cat: 'hot' }, s).getMatchingIndices()).toEqual([2040, 2041, 2042, 2043, 2044, 2045, 2046, 2047]); + }); +}); + +describe('ColumnarStore.prepare / addPrepared', () => { + test('prepare runs getters without mutating; addPrepared applies', () => { + let calls = 0; + const s = new ColumnarStore([ + { name: 'a', type: 'string', getter: (d: any) => { calls++; return d.a; } }, + { name: 'n', type: 'number' } + ]); + const row = s.prepare({ a: 'x', n: '7' }); + expect(calls).toBe(1); + expect(s.getActiveDocs().popcount()).toBe(0); + s.addPrepared(0, row); + expect(calls).toBe(1); + expect(compileFilter({ a: 'x', n: 7 }, s).getMatchingIndices()).toEqual([0]); + expect(() => s.addPrepared(1, [])).toThrow(TypeError); + }); +}); diff --git a/packages/webgpu-search/test/engine-pool-regressions.test.ts b/packages/webgpu-search/test/engine-pool-regressions.test.ts new file mode 100644 index 0000000..afedafe --- /dev/null +++ b/packages/webgpu-search/test/engine-pool-regressions.test.ts @@ -0,0 +1,287 @@ +/** + * Regressions for GPU device pool / engine lifecycle audit findings: + * 1. Concurrent GpuDevicePool.acquireDevice shares one in-flight acquisition + * (no double requestDevice, refcount counts every holder); releaseDevice + * honours device identity; intentional destroy() does not fire + * onDeviceLost listeners. + * 2. Headroom capacity is clamped to min(maxBufferSize, + * maxStorageBufferBindingSize); actual data above the limit throws the + * budget error (CPU fallback) instead of binding an oversized buffer. + * 3. Teardown/reallocation during the timestamp readback await surfaces as + * AbortError, never a TypeError or a silent empty result. + * 4. Failed / repeated init() never leaks a pool reference. + * 5. Native CPU baseline scores in normalized code points (GPU symmetry). + * + * Run: bun test packages/webgpu-search/test/engine-pool-regressions.test.ts + */ +import { describe, test, expect, afterEach } from 'bun:test'; +import { createMockAdapter } from 'vgpu/mock'; +import { + WebGPUEngine, + CPUEngine, + GpuDevicePool, + computeClampedHeadroomBytes, + normalizeText, + scoreSubstringTokens, +} from '../src/index'; + +const NAV_DESC = Object.getOwnPropertyDescriptor(globalThis, 'navigator'); + +function restoreNavigator(): void { + if (NAV_DESC) Object.defineProperty(globalThis, 'navigator', NAV_DESC); + else delete (globalThis as { navigator?: unknown }).navigator; +} + +interface FakeDevice { + id: number; + lost: Promise<{ reason: string; message: string }>; + resolveLost: (info: { reason: string; message: string }) => void; + destroyed: boolean; + destroy(): void; +} + +/** Install a fake navigator.gpu whose devices resolve `lost` like real WebGPU. */ +function installFakeGpu(makeDevice?: () => Promise) { + const created: FakeDevice[] = []; + let requestDeviceCalls = 0; + const limits = { + maxBufferSize: 1 << 28, + maxStorageBufferBindingSize: 1 << 27, + maxComputeWorkgroupsPerDimension: 65535, + maxComputeInvocationsPerWorkgroup: 256, + }; + const adapter = { + features: new Set(), + limits, + info: { vendor: 'test', device: 'fake' }, + requestDevice: async () => { + requestDeviceCalls++; + if (makeDevice) return makeDevice(); + await new Promise((r) => setTimeout(r, 5)); + let resolveLost!: FakeDevice['resolveLost']; + const lost = new Promise<{ reason: string; message: string }>((r) => (resolveLost = r)); + const d: FakeDevice = { + id: created.length + 1, + lost, + resolveLost, + destroyed: false, + destroy() { + if (this.destroyed) return; + this.destroyed = true; + resolveLost({ reason: 'destroyed', message: 'Device destroyed' }); + }, + }; + created.push(d); + return d; + }, + }; + Object.defineProperty(globalThis, 'navigator', { + value: { gpu: { requestAdapter: async () => adapter } }, + configurable: true, + writable: true, + }); + return { created, calls: () => requestDeviceCalls }; +} + +const tick = () => new Promise((r) => setTimeout(r, 0)); + +afterEach(() => { + restoreNavigator(); +}); + +describe('GpuDevicePool concurrency + identity (finding 1)', () => { + test('concurrent acquires share one device and refcount both holders', async () => { + const fake = installFakeGpu(); + const [a, b] = await Promise.all([GpuDevicePool.acquireDevice(), GpuDevicePool.acquireDevice()]); + expect(a).not.toBeNull(); + expect(b).not.toBeNull(); + expect(a!.device).toBe(b!.device); + expect(fake.calls()).toBe(1); + expect(GpuDevicePool.getRefCount()).toBe(2); + + GpuDevicePool.releaseDevice(a!.device, true); + expect((b!.device as unknown as FakeDevice).destroyed).toBe(false); + expect(GpuDevicePool.getRefCount()).toBe(1); + expect(GpuDevicePool.hasSharedDevice()).toBe(true); + + GpuDevicePool.releaseDevice(b!.device, true); + expect((b!.device as unknown as FakeDevice).destroyed).toBe(true); + expect(GpuDevicePool.getRefCount()).toBe(0); + expect(GpuDevicePool.hasSharedDevice()).toBe(false); + }); + + test('releasing a stale / foreign device never decrements the pooled device', async () => { + installFakeGpu(); + const first = await GpuDevicePool.acquireDevice(); + GpuDevicePool.simulateDeviceLoss('stale-test'); + const second = await GpuDevicePool.acquireDevice(); + expect(second!.device).not.toBe(first!.device); + expect(GpuDevicePool.getRefCount()).toBe(1); + + // Stale holder releases its (lost) device: the live device must survive. + GpuDevicePool.releaseDevice(first!.device, true); + GpuDevicePool.releaseDevice({} as GPUDevice, true); + expect(GpuDevicePool.getRefCount()).toBe(1); + expect((second!.device as unknown as FakeDevice).destroyed).toBe(false); + + GpuDevicePool.releaseDevice(second!.device, true); + expect(GpuDevicePool.hasSharedDevice()).toBe(false); + }); + + test('intentional destroy does not fire onDeviceLost; real loss fires once', async () => { + const fake = installFakeGpu(); + const reasons: string[] = []; + const unsub = GpuDevicePool.onDeviceLost((r) => reasons.push(r)); + try { + const a = await GpuDevicePool.acquireDevice(); + GpuDevicePool.releaseDevice(a!.device, true); // refcount 0 -> destroy() + await tick(); + expect(reasons).toEqual([]); + + const b = await GpuDevicePool.acquireDevice(); + (b!.device as unknown as FakeDevice).resolveLost({ reason: 'unknown', message: 'GPU hung' }); + await tick(); + expect(reasons).toEqual(['GPU hung']); + expect(GpuDevicePool.hasSharedDevice()).toBe(false); + expect(GpuDevicePool.getRefCount()).toBe(0); + + // Late loss of an already-replaced device is ignored (identity check). + const c = await GpuDevicePool.acquireDevice(); + fake.created[0]!.resolveLost({ reason: 'unknown', message: 'stale' }); + await tick(); + expect(reasons).toEqual(['GPU hung']); + expect(GpuDevicePool.hasSharedDevice()).toBe(true); + GpuDevicePool.releaseDevice(c!.device, true); + } finally { + unsub(); + } + }); +}); + +async function mockDevice(limits?: Record): Promise { + const dev = (await createMockAdapter({ features: ['timestamp-query'] as any }).requestDevice()).gpu as any; + return (limits ? { ...dev, limits: { ...dev.limits, ...limits } } : dev) as GPUDevice; +} + +describe('headroom clamp vs binding limit (finding 2)', () => { + test('computeClampedHeadroomBytes never exceeds min(maxBufferSize, maxStorageBufferBindingSize)', () => { + const device = { limits: { maxBufferSize: 256 * 2 ** 20, maxStorageBufferBindingSize: 128 * 2 ** 20 } } as any; + const bytes = computeClampedHeadroomBytes(160_000_000, { growthFactor: 1.0, device }); + expect(bytes).toBe(128 * 2 ** 20); + // Unaligned limit floors to a 4-byte multiple. + const odd = { limits: { maxBufferSize: 1001, maxStorageBufferBindingSize: 1001 } } as any; + expect(computeClampedHeadroomBytes(5000, { device: odd })).toBe(1000); + }); + + test('oversized tokenCapacity request is clamped at allocation', async () => { + const limit = 1 << 20; + const dev = await mockDevice({ maxBufferSize: limit, maxStorageBufferBindingSize: limit }); + const sizes: number[] = []; + const createBuffer = dev.createBuffer.bind(dev); + (dev as any).createBuffer = (d: GPUBufferDescriptor) => { + sizes.push(d.size); + return createBuffer(d); + }; + const engine = new WebGPUEngine(); + await engine.init(dev); + await engine.loadDataset(['alpha', 'beta'], { tokenCapacity: 10_000_000, rowCapacity: 10_000_000 }); + expect(engine.allocatedRecordsBytes).toBeLessThanOrEqual(limit); + expect(engine.allocatedOffsetsBytes).toBeLessThanOrEqual(limit); + expect(Math.max(...sizes)).toBeLessThanOrEqual(limit); + engine.destroy(); + }); + + test('actual data above the binding limit throws the budget error', async () => { + const limit = 70_000; // output buffer (8 + 8192*8) fits; records do not + const dev = await mockDevice({ maxBufferSize: limit, maxStorageBufferBindingSize: limit }); + const engine = new WebGPUEngine(); + await engine.init(dev); + const rows = Array.from({ length: 200 }, (_, i) => `row-${i}-`.padEnd(120, 'x')); + await expect(engine.loadDataset(rows)).rejects.toThrow(/too large/); + expect(engine.currentSize).toBe(0); + engine.destroy(); + }); +}); + +describe('teardown during timestamp readback (finding 3)', () => { + async function readyEngine() { + const engine = new WebGPUEngine(); + await engine.init(await mockDevice()); + await engine.loadDataset(['hello world', 'help']); + return engine; + } + + test('destroy() during timestamp mapAsync rejects with AbortError', async () => { + const engine = await readyEngine(); + (engine as any).queryStagingBuffer = { + mapAsync: async () => { engine.destroy(); }, + getMappedRange: () => new ArrayBuffer(16), + unmap() {}, + destroy() {}, + }; + const err = await engine.search('hel', { mode: 'fuzzy' }).then(() => null, (e) => e); + expect(err).not.toBeNull(); + expect((err as Error).name).toBe('AbortError'); + }); + + test('output reallocation during timestamp mapAsync rejects with AbortError', async () => { + const engine = await readyEngine(); + (engine as any).queryStagingBuffer = { + mapAsync: async () => { engine.ensureCandidateCapacity(20_000); }, + getMappedRange: () => new ArrayBuffer(16), + unmap() {}, + destroy() {}, + }; + const err = await engine.search('hel', { mode: 'fuzzy' }).then(() => null, (e) => e); + expect((err as Error | null)?.name).toBe('AbortError'); + engine.destroy(); + }); +}); + +describe('init() lifecycle never leaks pool references (finding 4)', () => { + test('init() twice on the pooled device keeps refcount at 1', async () => { + const mock = await mockDevice(); + installFakeGpu(async () => mock); + const engine = new WebGPUEngine(); + expect(await engine.init()).toBe(true); + expect(GpuDevicePool.getRefCount()).toBe(1); + expect(await engine.init()).toBe(true); + expect(GpuDevicePool.getRefCount()).toBe(1); + engine.destroy(); + expect(GpuDevicePool.getRefCount()).toBe(0); + expect(GpuDevicePool.hasSharedDevice()).toBe(false); + }); + + test('pipeline creation failure after acquire releases the device', async () => { + const mock = await mockDevice(); + const broken = { + ...(mock as any), + createComputePipelineAsync: async () => { throw new Error('pipeline boom'); }, + }; + installFakeGpu(async () => broken); + const engine = new WebGPUEngine(); + await expect(engine.init()).rejects.toThrow(/pipeline boom/); + expect(engine.isReady).toBe(false); + expect(GpuDevicePool.getRefCount()).toBe(0); + expect(GpuDevicePool.hasSharedDevice()).toBe(false); + }); +}); + +describe('native CPU baseline scoring symmetry (finding 5)', () => { + test('astral prefix scores in code points like the exact/GPU scorer', () => { + const rec = '\u{1F600}abc'; + const exact = scoreSubstringTokens(normalizeText(rec, true).tokens, normalizeText('abc', true).tokens); + const native = new CPUEngine().searchNaiveScan([rec], 'abc', 10, false); + expect(native.totalMatches).toBe(1); + expect(native.results[0]!.score).toBe(exact.score); + expect(exact.score).toBe(989); + }); + + test('NFC-equivalent inputs match and score identically', () => { + const composed = 'café'; + const decomposed = 'café'; + const res = new CPUEngine().searchNaiveScan([decomposed], composed, 10, false); + expect(res.totalMatches).toBe(1); + expect(res.results[0]!.score).toBe(1000); + }); +}); diff --git a/packages/webgpu-search/test/worker-snapshot-hardening.test.ts b/packages/webgpu-search/test/worker-snapshot-hardening.test.ts new file mode 100644 index 0000000..bfa3664 --- /dev/null +++ b/packages/webgpu-search/test/worker-snapshot-hardening.test.ts @@ -0,0 +1,338 @@ +/** + * Regression suite: worker message serialization + build-then-swap, client + * restore/init commit ordering, snapshot docIds/docs validation, + * ArrayBufferView snapshot inputs, and IndexedDB custom store provisioning. + * + * Run: bun test packages/webgpu-search/test/worker-snapshot-hardening.test.ts + */ +import { describe, test, expect } from 'bun:test'; +import { DocumentIndex } from '../src/document-index'; +import { SearchWorkerClient } from '../src/worker/worker-client'; +import { startSearchWorker } from '../src/worker/search-worker'; +import { + decodeSnapshot, + decodeSnapshotHeader, + restoreSnapshot, + toSnapshotArrayBuffer +} from '../src/snapshot-codec'; +import { loadIndexFromIDB, openSearchDatabase, restoreIndexFromIDB, saveIndexToIDB } from '../src/snapshot-idb'; +import { crc32Parts } from '../src/dataset-packing'; +import { IncompatibleIndexError } from '../src/text-profile'; +import { SNAPSHOT_HEADER_BYTES } from '../src/text-profile'; + +const OPTS = { fields: ['t'], preferGpu: false } as const; + +async function makeSnapshot(records: Array>, decoupled = false): Promise { + const idx = await DocumentIndex.create(records, OPTS as any); + try { + return idx.serialize({ decoupled }); + } finally { + idx.destroy(); + } +} + +/** Rewrites schema and/or docs JSON of a canonical snapshot with a valid CRC. */ +function forgeSnapshot( + buf: ArrayBuffer, + mutate: { schema?: (s: any) => any; docs?: (d: any) => any } +): ArrayBuffer { + const h = decodeSnapshotHeader(buf); + const u8 = new Uint8Array(buf); + let cur = SNAPSHOT_HEADER_BYTES; + const schemaBytes = u8.slice(cur, cur + h.schemaByteLength); cur += h.schemaByteLength; + const tokens = u8.slice(cur, cur + h.tokenCount * 4); cur += h.tokenCount * 4; + const offsets = u8.slice(cur, cur + (h.rowCount + 1) * 4); cur += (h.rowCount + 1) * 4; + const columnar = u8.slice(cur, cur + h.columnarByteLength); cur += h.columnarByteLength; + const docsBytes = u8.slice(cur, cur + h.docsByteLength); + const enc = new TextEncoder(); + const dec = new TextDecoder(); + const newSchema = mutate.schema ? enc.encode(JSON.stringify(mutate.schema(JSON.parse(dec.decode(schemaBytes))))) : schemaBytes; + const newDocs = mutate.docs && docsBytes.length > 0 ? enc.encode(JSON.stringify(mutate.docs(JSON.parse(dec.decode(docsBytes))))) : docsBytes; + const total = SNAPSHOT_HEADER_BYTES + newSchema.length + tokens.length + offsets.length + columnar.length + newDocs.length; + const out = new Uint8Array(total); + out.set(u8.subarray(0, SNAPSHOT_HEADER_BYTES)); + let o = SNAPSHOT_HEADER_BYTES; + for (const part of [newSchema, tokens, offsets, columnar, newDocs]) { out.set(part, o); o += part.length; } + const dv = new DataView(out.buffer); + dv.setUint32(36, newSchema.length, true); + dv.setUint32(40, newDocs.length, true); + dv.setUint32(52, crc32Parts([out.subarray(0, 52), newSchema, tokens, offsets, columnar, newDocs]), true); + return out.buffer; +} + +function createWorkerHarness() { + const out: any[] = []; + let listener: (e: any) => Promise = async () => {}; + const scope = { postMessage: (m: any) => out.push(m), addEventListener: (_: string, l: any) => (listener = l) }; + startSearchWorker(scope); + return { out, send: (d: any) => listener({ data: d }), last: (id: number) => out.filter((m) => m.id === id).pop() }; +} + +function createMockWorkerPair() { + const clientListeners: Array<(e: any) => void> = []; + const workerListeners: Array<(e: any) => void> = []; + const clientWorker: any = { + postMessage(data: any) { queueMicrotask(() => workerListeners.forEach((l) => l({ data }))); }, + addEventListener(ev: string, l: any) { if (ev === 'message') clientListeners.push(l); }, + removeEventListener(ev: string, l: any) { const i = clientListeners.indexOf(l); if (ev === 'message' && i >= 0) clientListeners.splice(i, 1); }, + terminate() {} + }; + const workerScope = { + postMessage(data: any) { queueMicrotask(() => clientListeners.forEach((l) => l({ data }))); }, + addEventListener(ev: string, l: any) { if (ev === 'message') workerListeners.push(l); } + }; + startSearchWorker(workerScope); + return clientWorker; +} + +describe('worker build-then-swap + serialization', () => { + test('failed RESTORE / INIT keep the previous index live', async () => { + const w = createWorkerHarness(); + await w.send({ id: 1, type: 'INIT', payload: { options: OPTS, records: [{ id: 1, t: 'hello' }] } }); + await w.send({ id: 2, type: 'RESTORE', payload: { buffer: new ArrayBuffer(10) } }); + expect(w.last(2).success).toBe(false); + expect(w.last(2).error.name).toBe('IncompatibleIndexError'); + await w.send({ id: 3, type: 'INIT', payload: { options: OPTS, records: [{ id: 1, t: 'a' }, { id: 1, t: 'b' }] } }); + expect(w.last(3).success).toBe(false); + await w.send({ id: 4, type: 'SEARCH', payload: { queryId: 1, query: 'hello' } }); + expect(w.last(4).success).toBe(true); + expect(w.last(4).result.results[0].id).toBe(1); + }); + + test('overlapping INITs + DESTROY end with no index and no leaked instance', async () => { + const w = createWorkerHarness(); + const created: any[] = []; + const destroyed = new Set(); + const origCreate = DocumentIndex.create; + const origDestroy = DocumentIndex.prototype.destroy; + (DocumentIndex as any).create = async function (...args: any[]) { + const i = await (origCreate as any).apply(this, args); created.push(i); return i; + }; + DocumentIndex.prototype.destroy = function () { destroyed.add(this); return origDestroy.call(this); }; + try { + const pA = w.send({ id: 10, type: 'INIT', payload: { options: OPTS, records: [{ id: 1, t: 'a' }] } }); + const pB = w.send({ id: 11, type: 'INIT', payload: { options: OPTS, records: [{ id: 2, t: 'b' }] } }); + await Promise.resolve(); // let INIT A start building + await Promise.resolve(); + const pD = w.send({ id: 12, type: 'DESTROY' }); + await Promise.all([pA, pB, pD]); + } finally { + (DocumentIndex as any).create = origCreate; + DocumentIndex.prototype.destroy = origDestroy; + } + // Replies superseded by DESTROY are dropped (session epoch); INIT B never starts. + expect(w.out.filter((m) => m.id >= 10 && m.id <= 12).map((m) => [m.id, m.success])).toEqual([[12, true]]); + expect(created.length).toBe(1); + for (const i of created) expect(destroyed.has(i)).toBe(true); + await w.send({ id: 13, type: 'STATS' }); + expect(w.last(13).success).toBe(false); + }); + + test('SEARCH issued during RESTORE waits for it', async () => { + const w = createWorkerHarness(); + const good = await makeSnapshot([{ id: 1, t: 'hello' }]); + await Promise.all([ + w.send({ id: 20, type: 'RESTORE', payload: { buffer: good } }), + w.send({ id: 21, type: 'SEARCH', payload: { queryId: 5, query: 'hello' } }) + ]); + expect(w.last(20).success).toBe(true); + expect(w.last(20).result.fields.map((f: any) => f.name)).toEqual(['t']); + expect(w.last(21).success).toBe(true); + expect(w.last(21).result.results.length).toBe(1); + }); +}); + +describe('worker session epoch', () => { + test('requests queued before DESTROY never answer after it', async () => { + const w = createWorkerHarness(); + const pInit = w.send({ id: 1, type: 'INIT', payload: { options: OPTS, records: [{ id: 1, t: 'hello' }] } }); + const pSearch = w.send({ id: 2, type: 'SEARCH', payload: { queryId: 1, query: 'hello' } }); + const pDestroy = w.send({ id: 3, type: 'DESTROY' }); + await Promise.all([pInit, pSearch, pDestroy]); + expect(w.out.map((m) => m.id)).toEqual([3]); + // A fresh session on the same worker restarts query ids. + await w.send({ id: 1, type: 'INIT', payload: { options: OPTS, records: [{ id: 9, t: 'again' }] } }); + await w.send({ id: 2, type: 'SEARCH', payload: { queryId: 1, query: 'again' } }); + expect(w.last(2).success).toBe(true); + expect(w.last(2).result.results[0].id).toBe(9); + }); +}); + +describe('client restore/init commit ordering', () => { + test('restore; add; search (not awaited) preserve call order', async () => { + const source = new SearchWorkerClient({ worker: createMockWorkerPair() }); + await source.init([{ id: 'x', t: 'other' }], { fields: ['t'], preferGpu: false } as any); + const snap = await source.serialize(); + await source.destroy(); + const client = new SearchWorkerClient({ worker: createMockWorkerPair() }); + await client.init([{ id: 'a', title: 'alpha' }], { fields: ['title'], preferGpu: false } as any); + const pRestore = client.restore(snap); + const pAdd = client.add({ id: 'y', t: 'yonder' }); + const pSearch = client.search('yonder'); + await pRestore; + await pAdd; + const res = await pSearch; + expect(res.results.map((r: any) => r.id)).toEqual(['y']); + expect(res.results[0].doc).toEqual({ id: 'y', t: 'yonder' }); + await client.destroy(); + }); + + test('corrupt restore rejects with IncompatibleIndexError and keeps client state', async () => { + const client = new SearchWorkerClient({ worker: createMockWorkerPair() }); + await client.init([{ id: 'a', title: 'alpha' }], { fields: ['title'], preferGpu: false } as any); + const good = await makeSnapshot([{ id: 'x', t: 'other' }]); + const corrupt = good.slice(0); + new Uint8Array(corrupt)[corrupt.byteLength - 2] ^= 0xff; // flip a docs byte → CRC mismatch + await expect(client.restore(corrupt)).rejects.toBeInstanceOf(IncompatibleIndexError); + expect(client.getRecords()).toEqual([{ id: 'a', title: 'alpha' }]); + const res = await client.search('alpha'); + expect(res.results[0].doc).toEqual({ id: 'a', title: 'alpha' }); + // Outgoing field defs still match the live worker index. + await client.add({ id: 'b', title: 'beta' }); + expect((await client.search('beta')).results[0].id).toBe('b'); + await client.destroy(); + }); + + test('failed init (worker-side) does not install new field defs', async () => { + const client = new SearchWorkerClient({ worker: createMockWorkerPair() }); + await client.init([{ id: 'a', title: 'alpha' }], { fields: ['title'], preferGpu: false } as any); + await expect( + client.init([{ id: 'd', body: 'x' }, { id: 'd', body: 'y' }], { fields: ['body'], preferGpu: false } as any) + ).rejects.toThrow(); + await client.add({ id: 'b', title: 'beta' }); + expect((await client.search('beta')).results[0].id).toBe('b'); + await client.destroy(); + }); + + test('successful restore commits worker-validated schema and records', async () => { + const client = new SearchWorkerClient({ worker: createMockWorkerPair() }); + await client.init([{ id: 'a', title: 'alpha' }], { fields: ['title'], preferGpu: false } as any); + const source = new SearchWorkerClient({ worker: createMockWorkerPair() }); + await source.init([{ id: 'x', t: 'other' }], { fields: ['t'], preferGpu: false } as any); + const snap = await source.serialize(); + await source.destroy(); + await client.restore(new Uint8Array(snap)); + expect(client.getRecords().map((d: any) => d.id)).toEqual(['x']); + await client.add({ id: 'y', t: 'yonder' }); + expect((await client.search('yonder')).results[0].id).toBe('y'); + await client.destroy(); + }); +}); + +describe('snapshot codec validation', () => { + test('rejects duplicate / non-scalar / empty docIds, accepts 1 and "1"', async () => { + const base = await makeSnapshot([{ id: 'p', t: 'a' }, { id: 'q', t: 'b' }], true); + for (const ids of [['p', 'p'], [{}, 'q'], [[], 'q'], ['', 'q'], [null, 'q'], [Number.NaN, 'q']]) { + const forged = forgeSnapshot(base, { schema: (s) => ({ ...s, docIds: ids }) }); + expect(() => decodeSnapshot(forged)).toThrow(IncompatibleIndexError); + } + const mixed = forgeSnapshot(base, { schema: (s) => ({ ...s, docIds: [1, '1'] }) }); + expect(decodeSnapshot(mixed).docIds).toEqual([1, '1']); + }); + + test('rejects non-object embedded docs with IncompatibleIndexError', async () => { + const base = await makeSnapshot([{ id: 'p', t: 'a' }, { id: 'q', t: 'b' }]); + for (const bad of [null, 5, 'str', []]) { + const forged = forgeSnapshot(base, { docs: (d) => [d[0], bad] }); + expect(() => decodeSnapshot(forged)).toThrow(IncompatibleIndexError); + } + }); + + test('accepts ArrayBufferViews, honouring byteOffset', async () => { + const buf = await makeSnapshot([{ id: 'p', t: 'hello' }]); + const padded = new Uint8Array(buf.byteLength + 16); + padded.set(new Uint8Array(buf), 8); + const view = padded.subarray(8, 8 + buf.byteLength); + expect(decodeSnapshotHeader(view).docCount).toBe(1); + expect(decodeSnapshot(view).docIds).toEqual(['p']); + expect(decodeSnapshot(new DataView(padded.buffer, 8, buf.byteLength)).docIds).toEqual(['p']); + const restored = await restoreSnapshot(new Uint8Array(buf)); + expect((await restored.search('hello')).results[0].id).toBe('p'); + restored.destroy(); + expect(toSnapshotArrayBuffer(new Uint8Array(buf))).toBe(buf); + expect(toSnapshotArrayBuffer({ byteLength: 10, slice() {} })).toBeNull(); + expect(() => decodeSnapshot({ byteLength: 100, slice() {} } as any)).toThrow(IncompatibleIndexError); + }); +}); + +/** Strict versioned IDB mock: stores only created during upgrades. */ +function createStrictIDB() { + const dbs = new Map> }>(); + const mkReq = () => ({ result: undefined as any, error: null as any, onsuccess: null as any, onerror: null as any, onupgradeneeded: null as any, onblocked: null as any }); + const connection = (rec: { version: number; stores: Map> }, upgrading: boolean) => ({ + get version() { return rec.version; }, + objectStoreNames: { contains: (n: string) => rec.stores.has(n) }, + createObjectStore(n: string) { + if (!upgrading) throw Object.assign(new Error('not in versionchange'), { name: 'InvalidStateError' }); + rec.stores.set(n, new Map()); + }, + onversionchange: null as any, + close() {}, + transaction(names: string[], _mode: string) { + for (const n of names) if (!rec.stores.has(n)) throw Object.assign(new Error(`missing store ${n}`), { name: 'NotFoundError' }); + const tx: any = { oncomplete: null, onerror: null, onabort: null, error: null }; + let pending = 0; + const done = () => queueMicrotask(() => { if (pending === 0) tx.oncomplete?.(); }); + tx.objectStore = (n: string) => { + const m = rec.stores.get(n)!; + const op = (fn: (r: any) => void) => { const r = mkReq(); pending++; queueMicrotask(() => { fn(r); r.onsuccess?.(); pending--; done(); }); return r; }; + return { put: (v: any, k: any) => op(() => m.set(k, v)), get: (k: any) => op((r) => { r.result = m.get(k); }), delete: (k: any) => op(() => m.delete(k)) }; + }; + queueMicrotask(done); + return tx; + } + }); + const factory: any = { + dbs, + open(name: string, version?: number) { + const req = mkReq(); + queueMicrotask(() => { + let rec = dbs.get(name); + if (!rec) { rec = { version: 0, stores: new Map() }; dbs.set(name, rec); } + const target = version ?? Math.max(rec.version, 1); + if (target < rec.version) { req.error = Object.assign(new Error('VersionError'), { name: 'VersionError' }); req.onerror?.(); return; } + if (target > rec.version) { + rec.version = target; + req.result = connection(rec, true); + req.onupgradeneeded?.(); + } + req.result = connection(rec, false); + req.onsuccess?.(); + }); + return req; + } + }; + return factory; +} + +describe('IndexedDB store provisioning + buffer normalization', () => { + test('custom store names on an existing DB are provisioned via version bump', async () => { + const idb = createStrictIDB(); + const buf = await makeSnapshot([{ id: 'p', t: 'hello' }]); + await saveIndexToIDB(buf, { indexedDB: idb }); + expect(idb.dbs.get('webgpu_search_db').version).toBe(1); + const custom = { indexedDB: idb, snapshotStoreName: 'snaps2', docStoreName: 'docs2' }; + await saveIndexToIDB(buf, custom); + expect(idb.dbs.get('webgpu_search_db').version).toBe(2); + const loaded = await loadIndexFromIDB(custom); + expect(loaded?.snapshot.byteLength).toBe(buf.byteLength); + const db = await openSearchDatabase(custom); // no further bump needed + expect(db.version).toBe(2); + }); + + test('typed-array snapshots are stored and loaded as exact ArrayBuffers', async () => { + const idb = createStrictIDB(); + const buf = await makeSnapshot([{ id: 'p', t: 'hello' }]); + const padded = new Uint8Array(buf.byteLength + 4); + padded.set(new Uint8Array(buf), 4); + await saveIndexToIDB(padded.subarray(4), { indexedDB: idb }); + const stored = idb.dbs.get('webgpu_search_db').stores.get('index_snapshots').get('default_index'); + expect(stored instanceof ArrayBuffer).toBe(true); + expect(stored.byteLength).toBe(buf.byteLength); + // Legacy record saved as a typed array is normalized on load. + idb.dbs.get('webgpu_search_db').stores.get('index_snapshots').set('default_index', new Uint8Array(buf)); + const restored = await restoreIndexFromIDB({ indexedDB: idb }); + expect((await restored!.search('hello')).results[0].id).toBe('p'); + restored!.destroy(); + }); +}); diff --git a/scripts/check-bundle-size.ts b/scripts/check-bundle-size.ts index 0241e41..c3aa691 100644 --- a/scripts/check-bundle-size.ts +++ b/scripts/check-bundle-size.ts @@ -127,7 +127,13 @@ * pre-contract baseline guarantees false failures on the first commit. * Total index bumped 82 KB -> 90 KB and delta cap 34 KB -> 40 KB; worker entry * gets its own 76 KB total budget (no historical baseline: gate on total only). - * Budget: 90 KB gzip per file (dist/index.js + dist/index.cjs); + * (audit hardening) batch-atomic filter prepare, posting-list filter storage + * (bounded memory for high-cardinality string fields), serialized worker queue + + * commit-after-success client, snapshot docId/buffer validation, GPU pool + * in-flight acquisition + binding-limit clamping: measured index.js 93,294 B + * gzip (+4.8 KB). Decision: bump total 90 KB -> 96 KB, delta cap 40 KB -> 46 KB. + * Rationale: these are correctness/memory fixes with no lazy-loadable seam. + * Budget: 96 KB gzip per file (dist/index.js + dist/index.cjs); * 76 KB gzip per file (dist/worker.js + dist/worker.cjs). * * Fail-closed: missing dist or dist older than library sources fails @@ -141,8 +147,8 @@ import { stat, readFile } from 'node:fs/promises'; const BASELINE_RAW = 234907; const BASELINE_GZIP = 49246; -const DELTA_CAP_GZIP = 40 * 1024; -const TOTAL_BUDGET_INDEX_GZIP = 90 * 1024; +const DELTA_CAP_GZIP = 46 * 1024; +const TOTAL_BUDGET_INDEX_GZIP = 96 * 1024; const TOTAL_BUDGET_WORKER_GZIP = 76 * 1024; diff --git a/scripts/wgsl-loader-plugin.ts b/scripts/wgsl-loader-plugin.ts new file mode 100644 index 0000000..c43edb8 --- /dev/null +++ b/scripts/wgsl-loader-plugin.ts @@ -0,0 +1,18 @@ +/** + * Bun runtime plugin: `import x from './shaders/*.wgsl'` yields the shader + * SOURCE TEXT (default export), matching tsup (`loader: { '.wgsl': 'text' }`) + * and the benchmark Vite plugin. Without it Bun resolves unknown extensions + * to the file path string, so headless tests would compile a path instead of + * the real WGSL. Registered via `preload` in the root bunfig.toml. + */ +import { plugin } from 'bun'; + +plugin({ + name: 'wgsl-text-loader', + setup(build) { + build.onLoad({ filter: /\.wgsl$/ }, async ({ path }) => ({ + exports: { default: await Bun.file(path).text() }, + loader: 'object', + })); + }, +});