diff --git a/src/store/memoryIndex.ts b/src/store/memoryIndex.ts index fb2b1f2..daefb2c 100644 --- a/src/store/memoryIndex.ts +++ b/src/store/memoryIndex.ts @@ -45,6 +45,8 @@ export interface MemoryIndexHit { score: number; } +import { withOpenTimeout } from "./pgOpenGuard.js"; + let db: PGliteInstance | undefined; let initPromise: Promise | undefined; let disabled = false; @@ -135,12 +137,19 @@ async function openPgLite( if (!mod) return undefined; const dir = indexDir(); mkdirSync(dir, { recursive: true }); - const pg = await new mod.PGlite({ - dataDir: dir, - extensions: { vector: mod.vector }, - }); - await pg.exec("CREATE EXTENSION IF NOT EXISTS vector;"); - await pg.exec(` + // Bounded open: PGlite is single-writer over a shared dataDir, so a second + // pi process on the same dir can block here forever. Without the ceiling the + // never-settling promise gets cached in initPromise and every later caller + // awaits it — which is how a stalled index wedged a whole pi turn. + let openTimedOut = false; + const pg = await withOpenTimeout( + (async () => { + const inst = await new mod.PGlite({ + dataDir: dir, + extensions: { vector: mod.vector }, + }); + await inst.exec("CREATE EXTENSION IF NOT EXISTS vector;"); + await inst.exec(` CREATE TABLE IF NOT EXISTS memory_index ( repo_id TEXT NOT NULL, memory_id INTEGER NOT NULL, @@ -149,9 +158,26 @@ async function openPgLite( PRIMARY KEY (repo_id, memory_id) ); `); - await pg.exec( - "CREATE INDEX IF NOT EXISTS memory_index_hnsw ON memory_index USING hnsw (embedding vector_cosine_ops);", + await inst.exec( + "CREATE INDEX IF NOT EXISTS memory_index_hnsw ON memory_index USING hnsw (embedding vector_cosine_ops);", + ); + return inst; + })(), + (reason) => { + openTimedOut = true; + logWarn(`init ${reason}`); + }, ); + if (!pg) { + if (openTimedOut) { + // Don't leave the dead open cached, and don't retry on the next call — + // a contended dataDir would just burn another full timeout per caller. + // Same terminal state as any other init failure: fall back to the scan. + initPromise = undefined; + disabled = true; + } + return undefined; + } db = pg; return pg; } catch (err) { diff --git a/src/store/pgOpenGuard.test.ts b/src/store/pgOpenGuard.test.ts new file mode 100644 index 0000000..aae663b --- /dev/null +++ b/src/store/pgOpenGuard.test.ts @@ -0,0 +1,89 @@ +/** + * pgOpenGuard.test.ts — the PGlite open must never hang a turn. + * + * Regression cover for the wedge: a stalled `await new PGlite(...)` was cached + * in initPromise, so every later caller awaited a promise that could not settle + * and the pi turn awaiting it never ended. + */ + +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { withOpenTimeout, pgOpenTimeoutMs, DEFAULT_PG_OPEN_TIMEOUT_MS } from "./pgOpenGuard.js"; + +test("a never-settling open resolves to undefined instead of hanging", async () => { + const never = new Promise(() => { + /* deliberately never settles — the wedge */ + }); + const reasons: string[] = []; + const t0 = Date.now(); + + const result = await withOpenTimeout(never, (r) => reasons.push(r), 50); + + assert.equal(result, undefined, "caller gets undefined and can fall back"); + assert.ok(Date.now() - t0 < 5_000, "returned promptly rather than hanging"); + assert.equal(reasons.length, 1, "onTimeout fired exactly once"); + assert.match(reasons[0], /timed out after 50ms/); +}); + +test("a successful open passes its value through untouched", async () => { + const reasons: string[] = []; + const result = await withOpenTimeout(Promise.resolve("pg"), (r) => reasons.push(r), 5_000); + assert.equal(result, "pg"); + assert.deepEqual(reasons, [], "no timeout reported on the happy path"); +}); + +test("a rejected open propagates so the corrupt-dir retry still runs", async () => { + const reasons: string[] = []; + await assert.rejects( + () => withOpenTimeout(Promise.reject(new Error("Aborted()")), (r) => reasons.push(r), 5_000), + /Aborted/, + "rejection reaches the caller's catch, which owns the wipe-and-retry path", + ); + assert.deepEqual(reasons, [], "a rejection is not reported as a timeout"); +}); + +test("an abandoned open is closed if it settles after the timeout", async () => { + let closed = false; + let release: (v: { close: () => void }) => void = () => {}; + const late = new Promise<{ close: () => void }>((r) => { + release = r; + }); + + const result = await withOpenTimeout(late, () => {}, 25); + assert.equal(result, undefined, "timed out first"); + + // The open finally completes, long after we stopped waiting for it. + release({ + close: () => { + closed = true; + }, + }); + await late; + await new Promise((r) => setImmediate(r)); + + assert.ok(closed, "the orphaned instance was closed, not left holding the dataDir"); +}); + +test("timeout of 0 disables the guard (unbounded, original behavior)", async () => { + const result = await withOpenTimeout(Promise.resolve("pg"), () => {}, 0); + assert.equal(result, "pg"); +}); + +test("pgOpenTimeoutMs honors the env override and rejects junk", async () => { + const prev = process.env.MEGACOMPACT_PGLITE_OPEN_TIMEOUT_MS; + try { + process.env.MEGACOMPACT_PGLITE_OPEN_TIMEOUT_MS = "1234"; + assert.equal(pgOpenTimeoutMs(), 1234); + + process.env.MEGACOMPACT_PGLITE_OPEN_TIMEOUT_MS = "0"; + assert.equal(pgOpenTimeoutMs(), 0, "0 is a valid opt-out, not junk"); + + for (const junk of ["", " ", "abc", "-5"]) { + process.env.MEGACOMPACT_PGLITE_OPEN_TIMEOUT_MS = junk; + assert.equal(pgOpenTimeoutMs(), DEFAULT_PG_OPEN_TIMEOUT_MS, `junk "${junk}" falls back`); + } + } finally { + if (prev === undefined) delete process.env.MEGACOMPACT_PGLITE_OPEN_TIMEOUT_MS; + else process.env.MEGACOMPACT_PGLITE_OPEN_TIMEOUT_MS = prev; + } +}); diff --git a/src/store/pgOpenGuard.ts b/src/store/pgOpenGuard.ts new file mode 100644 index 0000000..ba544b7 --- /dev/null +++ b/src/store/pgOpenGuard.ts @@ -0,0 +1,93 @@ +/** + * pgOpenGuard.ts — bound the PGlite open so a stalled WASM init can never wedge + * a pi turn. + * + * Both index modules (vectorIndex / memoryIndex) cache their in-flight open in a + * module-level `initPromise`. That cache is what turns a single stalled open + * into a permanent hang: PGlite is a single-writer WASM Postgres over a shared + * dataDir (~/.pi/mega-compact-vector), so a second pi process opening the same + * dir can block indefinitely. `await new PGlite(...)` then never settles, the + * never-settling promise is cached, and every later caller awaits that same dead + * promise — with no timers and no sockets left, node reports + * "Promise resolution is still pending but the event loop has already resolved" + * and the pi turn that awaited it never ends. + * + * withOpenTimeout() puts a ceiling on that wait. On timeout the caller gets + * undefined (both modules already degrade to a synchronous scan), and the + * abandoned open is disowned: if it does eventually settle, the instance is + * closed so a stray PGlite can't keep the loop alive or hold the dataDir lock. + * + * A rejected open is NOT swallowed — it propagates so the callers' existing + * corrupt-dir detection (Aborted / RuntimeError → wipe + one retry) still runs. + * Only the timeout resolves to undefined. + */ + +/** Sentinel so a legitimately-undefined open is distinguishable from a timeout. */ +const TIMED_OUT = Symbol("pglite-open-timeout"); + +/** Default ceiling for a PGlite open. Generous — a cold WASM + HNSW init is slow. */ +export const DEFAULT_PG_OPEN_TIMEOUT_MS = 30_000; + +/** Resolve the open timeout. 0 (or negative) disables the guard entirely. */ +export function pgOpenTimeoutMs(): number { + const raw = process.env.MEGACOMPACT_PGLITE_OPEN_TIMEOUT_MS; + if (raw === undefined || raw.trim() === "") return DEFAULT_PG_OPEN_TIMEOUT_MS; + const n = Number(raw); + if (!Number.isFinite(n) || n < 0) return DEFAULT_PG_OPEN_TIMEOUT_MS; + return n; +} + +/** A PGlite-ish handle we may need to dispose of after abandoning it. */ +interface Closable { + close?: () => Promise | unknown; +} + +/** + * Race `open` against the configured timeout. + * + * Resolves to the opened value, or to undefined when the open outruns the + * timeout (`onTimeout` fires first so the caller can log and flip its own + * disabled state). Rejections propagate to the caller unchanged. + */ +export async function withOpenTimeout( + open: Promise, + onTimeout: (reason: string) => void, + timeoutMs: number = pgOpenTimeoutMs(), +): Promise { + // Guard disabled — preserve the original unbounded behavior verbatim. + if (timeoutMs <= 0) return open; + + let timer: ReturnType | undefined; + const expiry = new Promise((resolve) => { + timer = setTimeout(() => resolve(TIMED_OUT), timeoutMs); + // Never hold the process open on account of the guard itself. + timer.unref?.(); + }); + + try { + // `open` is raced as-is so a rejection rejects the race — and therefore + // this function — leaving the caller's corrupt-retry path intact. + const winner = await Promise.race([open, expiry]); + + if (winner === TIMED_OUT) { + // Disown the open. If it ever settles, close the instance so an orphaned + // PGlite cannot keep the event loop alive or hold the dataDir lock. + void open + .then((late) => { + try { + void (late as Closable | undefined)?.close?.(); + } catch { + /* ignore */ + } + }) + .catch(() => { + /* the abandoned open failed on its own — nothing left to release */ + }); + onTimeout(`timed out after ${timeoutMs}ms`); + return undefined; + } + return winner as T; + } finally { + if (timer) clearTimeout(timer); + } +} diff --git a/src/store/vectorIndex.ts b/src/store/vectorIndex.ts index dd9dee6..bb94299 100644 --- a/src/store/vectorIndex.ts +++ b/src/store/vectorIndex.ts @@ -40,6 +40,8 @@ export interface VectorIndexHit { score: number; } +import { withOpenTimeout } from "./pgOpenGuard.js"; + let db: PGliteInstance | undefined; let initPromise: Promise | undefined; let disabled = false; @@ -131,12 +133,19 @@ async function openPgLite( if (!mod) return undefined; const dir = indexDir(); mkdirSync(dir, { recursive: true }); - const pg = await new mod.PGlite({ - dataDir: dir, - extensions: { vector: mod.vector }, - }); - await pg.exec("CREATE EXTENSION IF NOT EXISTS vector;"); - await pg.exec(` + // Bounded open: PGlite is single-writer over a shared dataDir, so a second + // pi process on the same dir can block here forever. Without the ceiling the + // never-settling promise gets cached in initPromise and every later caller + // awaits it — which is how a stalled index wedged a whole pi turn. + let openTimedOut = false; + const pg = await withOpenTimeout( + (async () => { + const inst = await new mod.PGlite({ + dataDir: dir, + extensions: { vector: mod.vector }, + }); + await inst.exec("CREATE EXTENSION IF NOT EXISTS vector;"); + await inst.exec(` CREATE TABLE IF NOT EXISTS vector_index ( repo_id TEXT NOT NULL, session_id TEXT NOT NULL, @@ -145,10 +154,27 @@ async function openPgLite( PRIMARY KEY (repo_id, session_id, checkpoint_id) ); `); - // HNSW index over cosine distance for fast NN. Created idempotently. - await pg.exec( - "CREATE INDEX IF NOT EXISTS vector_index_hnsw ON vector_index USING hnsw (embedding vector_cosine_ops);", + // HNSW index over cosine distance for fast NN. Created idempotently. + await inst.exec( + "CREATE INDEX IF NOT EXISTS vector_index_hnsw ON vector_index USING hnsw (embedding vector_cosine_ops);", + ); + return inst; + })(), + (reason) => { + openTimedOut = true; + logWarn(`init ${reason}`); + }, ); + if (!pg) { + if (openTimedOut) { + // Don't leave the dead open cached, and don't retry on the next call — + // a contended dataDir would just burn another full timeout per caller. + // Same terminal state as any other init failure: fall back to the scan. + initPromise = undefined; + disabled = true; + } + return undefined; + } db = pg; return pg; } catch (err) {