From a41dfa5ff6f2d09bcc151a7a81cbd57688c00258 Mon Sep 17 00:00:00 2001 From: Garry Tan Date: Wed, 3 Jun 2026 01:29:34 -0700 Subject: [PATCH] fix(core): drain fire-and-forget sinks before disconnect via a background-work registry (#1762) New src/core/background-work.ts registry (Map, ordered drain, awaited abort). facts-queue (order 0, abort=shutdown), last-retrieved (1), and eval-capture (3, now self-tracked) register as sinks. Both CLI exit paths (op-dispatch finally + handleCliOnly finally) drain the registry before engine.disconnect() so a PGLite db.close() can't race in-flight work into the re-pump busy-loop that pinned the single-writer lock. Op-dispatch error path converts process.exit(1) to exitCode+return so the finally still drains. --- src/cli.ts | 92 +++++++++++++++--------------- src/core/background-work.ts | 111 ++++++++++++++++++++++++++++++++++++ src/core/eval-capture.ts | 55 ++++++++++++++++++ src/core/facts/queue.ts | 16 ++++++ src/core/last-retrieved.ts | 9 +++ 5 files changed, 238 insertions(+), 45 deletions(-) create mode 100644 src/core/background-work.ts diff --git a/src/cli.ts b/src/cli.ts index 4ecb5e4e0..7a4dfa9fe 100755 --- a/src/cli.ts +++ b/src/cli.ts @@ -16,7 +16,7 @@ import type { AIGatewayConfig } from './core/ai/types.ts'; import type { BrainEngine } from './core/engine.ts'; import { operations, OperationError } from './core/operations.ts'; import type { Operation, OperationContext } from './core/operations.ts'; -import { awaitPendingLastRetrievedWrites, type DrainOutcome } from './core/last-retrieved.ts'; +import { drainAllBackgroundWorkForCliExit } from './core/background-work.ts'; import { shouldForceExitAfterMain } from './core/cli-force-exit.ts'; import { serializeMarkdown } from './core/markdown.ts'; import { parseGlobalFlags, setCliOptions, getCliOptions } from './core/cli-options.ts'; @@ -253,7 +253,10 @@ async function main() { console.warn( `[cli] engine.disconnect() did not return within ${DISCONNECT_HARD_DEADLINE_MS}ms — force-exiting`, ); - process.exit(0); + // v0.42.11.0 (codex): honor an exit code an errored op already set — + // a bare process.exit(0) here would mask a failed op as success if the + // drain/disconnect then hangs. + process.exit(process.exitCode ?? 0); }, DISCONNECT_HARD_DEADLINE_MS); // unref so the timer itself doesn't keep the event loop alive — only // the actual pending work (PGLite WASM handle) does. Without unref, @@ -261,7 +264,6 @@ async function main() { forceExitTimer.unref?.(); } - let drainResult: DrainOutcome = { outcome: 'drained', pending: 0 }; try { const ctx = await makeContext(engine, params); const rawResult = await op.handler(ctx, params); @@ -272,55 +274,32 @@ async function main() { const result = JSON.parse(JSON.stringify(rawResult)); const output = formatResult(op.name, result); if (output) process.stdout.write(output); - if (op.name === 'query') { - const { awaitPendingSearchCacheWrites } = await import('./core/search/hybrid.ts'); - await awaitPendingSearchCacheWrites(); - } - // Drain unconditionally for every op — empty-set fast-path is a - // few microseconds. Not per-op-name gated: that was the original - // PR #1259 mistake that left search and get_page exposed. - drainResult = await awaitPendingLastRetrievedWrites(); } catch (e: unknown) { - // C9 fix: drain BEFORE process.exit so a successful op that throws - // during stdout/format still gets its bumpLastRetrievedAt UPDATE - // a chance to commit. Bounded by the drain's own 5s timeout; the - // outer hard-exit timer above bounds the disconnect path. - try { await awaitPendingLastRetrievedWrites(); } catch { /* best-effort */ } + // v0.42.11.0 (codex D4): on error, set exitCode + return so the `finally` + // STILL runs (drains every background-work sink + disconnects). A bare + // process.exit(1) here would skip the finally → skip the drain + disconnect + // (leaves facts/cache/eval-capture writes racing teardown). The finally's + // drain bounds teardown; the outer hard-deadline timer bounds a hung one. if (e instanceof OperationError) { console.error(`Error [${e.code}]: ${e.message}`); if (e.suggestion) console.error(` Fix: ${e.suggestion}`); - process.exit(1); + } else { + console.error(e instanceof Error ? e.message : String(e)); } - console.error(e instanceof Error ? e.message : String(e)); - process.exit(1); + process.exitCode = 1; } finally { - // v0.41.25.0 (#1570) — drain the facts:absorb queue BEFORE disconnect - // so the fire-and-forget queue worker has a live engine to write its - // log against. Closes the bug class that absorb-log.ts:87-100 names: - // facts subsystem holds an engine reference past CLI exit, fires its - // post-completion log against a dead singleton, surfaces as a 'No - // database connection' stderr line on every `gbrain capture`. - // - // 1s timeout is per codex finding 10 from the v0.41.25 plan review: - // ops that don't enqueue facts (most read paths) pay only the - // 0-pending fast-path cost (~microseconds). Capture / import / sync - // that DO enqueue pay up to 1s while in-flight Haiku calls finish. - // Lazy-import keeps this off the hot path for ops that never touch - // the facts queue at all. - try { - const { getFactsQueue } = await import('./core/facts/queue.ts'); - await getFactsQueue().drainPending({ timeout: 1000 }); - } catch { /* best-effort; never block disconnect on drain failure */ } + // v0.42.11.0 — drain ALL fire-and-forget sinks (facts, last-retrieved, + // search-cache, eval-capture) via the background-work registry BEFORE + // disconnect, so a PGLite db.close() can't race in-flight work into the + // re-pump busy-loop (#1762). facts drains first (order 0) so its abort-path + // DB logIngest gets the freshest live-engine window. 1s per-sink timeout: + // read paths with no pending work pay the ~0ms fast path; capture/import + // that DO enqueue pay up to 1s (+ facts shutdown grace) while in-flight + // Haiku finishes. The unref'd hard-deadline timer above is the backstop if + // disconnect or a lingering socket keeps Bun's loop alive. + await drainAllBackgroundWorkForCliExit({ timeoutMs: 1000 }); await engine.disconnect(); if (forceExitTimer) clearTimeout(forceExitTimer); - // Narrow force-exit: only when the drain timed out AND we are NOT - // running a daemon. The drain helper already stderr-warned with the - // pending count, so the diagnostic signal is preserved. Without - // this guard a hung underlying promise can still keep Bun's loop - // alive past disconnect — Codex outside-voice finding #1. - if (drainResult.outcome === 'timeout' && shouldForceExitAfterMain()) { - process.exit(0); - } } } @@ -1772,7 +1751,30 @@ async function handleCliOnly(command: string, args: string[]) { } } } finally { - if (command !== 'serve') await engine.disconnect(); + // v0.42.11.0 (#1762) — the CLI_ONLY path (which owns `gbrain capture`) + // lacked the op-dispatch drain-before-disconnect contract. `put_page` fires + // a fire-and-forget facts:absorb job AFTER printing the receipt; on a + // multi-chunk page that job is in flight when this finally tears the engine + // down, and `engine.disconnect()` nulling PGLite's _db mid-job spins + // db.close() into a 100%-CPU busy-loop that pins the single-writer lock. + // Drain every background-work sink first (facts shutdown() abort cancels a + // hung Haiku), THEN disconnect. The drain-before-disconnect is the causal + // fix; the force-exit defense below is secondary (it CANNOT preempt a WASM + // busy-loop on a pinned JS thread — that's exactly why the drain matters). + if (command !== 'serve') { + const forceExit = shouldForceExitAfterMain(); + let hardExitTimer: ReturnType | undefined; + if (forceExit) { + hardExitTimer = setTimeout(() => { + console.warn('[cli] engine.disconnect() did not return within 10000ms — force-exiting'); + process.exit(process.exitCode ?? 0); + }, 10_000); + hardExitTimer.unref?.(); + } + await drainAllBackgroundWorkForCliExit(); + await engine.disconnect(); + if (hardExitTimer) clearTimeout(hardExitTimer); + } } } diff --git a/src/core/background-work.ts b/src/core/background-work.ts new file mode 100644 index 000000000..2d6900232 --- /dev/null +++ b/src/core/background-work.ts @@ -0,0 +1,111 @@ +/** + * v0.42.11.0 (#1762 / #1745 / #1775 reliability wave) — process background-work + * registry. Single source of truth for "drain every fire-and-forget sink before + * the CLI exits / disconnects." + * + * WHY THIS EXISTS (rule-of-four): four independent fire-and-forget sinks each + * write to the DB after an op returns its response — + * - `last-retrieved.ts` UPDATE pages.last_retrieved_at (#1247/#1269/#1290) + * - `facts/queue.ts` facts:absorb Haiku job + logIngest (#1762) + * - `search/hybrid.ts` query_cache write + * - `eval-capture.ts` eval_candidates INSERT + * On PGLite, if `engine.disconnect()` nulls `_db` while one of these is in + * flight, the sink's "not connected" error path re-pumps via queueMicrotask and + * spins `db.close()` into a 100%-CPU busy-loop that pins the single-writer lock + * (the #1762 incident). The fix is to DRAIN every sink before disconnect. A + * registry (not a hand-written N-call helper) makes that structural: a future + * 5th sink that registers is auto-drained, and the drain is invoked from THREE + * exit points (op-dispatch success finally, op-dispatch error catch, CLI_ONLY + * finally) without repeating the sink list at each. + * + * register (at module import) ─┐ + * last-retrieved (order 1) │ + * facts (order 0) ├─► Map + * search-cache (order 2) │ + * eval-capture (order 3) ┘ + * │ CLI exit + * ▼ + * drainAllBackgroundWorkForCliExit ──► sort by (order, name) + * for each: await drain(timeoutMs) + * if unfinished>0 && abort: + * await abort() ◄─ facts shutdown() + * ▼ + * engine.disconnect() (caller) + * + * Registration MUST live in the enqueue-owning module (so "module not imported + * ⇒ no work enqueued ⇒ nothing to drain" holds). The Map is keyed by name so a + * re-import / test mock REPLACES rather than duplicating (an array would + * double-register). + */ + +export interface BackgroundWorkDrainer { + /** Stable identity; also the Map key (idempotent registration). */ + name: string; + /** + * Explicit drain order — lower runs first. Facts is 0 so its abort-path DB + * `logIngest` gets the freshest live-engine window before the fast + * last-retrieved / search-cache drains. Ties break by name for determinism. + */ + order: number; + /** Resolve when in-flight work settles OR the bound elapses; report leftovers. */ + drain(timeoutMs: number): Promise<{ unfinished: number }>; + /** + * Optional hard-stop for stragglers (facts-queue: `shutdown()`). AWAITED by + * the registry so the aborted job's DB write settles against a live engine + * BEFORE the caller disconnects. Only invoked when `drain` reports unfinished. + */ + abort?(): Promise; +} + +const drainers = new Map(); + +/** Register (or replace, by name) a fire-and-forget sink drainer. */ +export function registerBackgroundWorkDrainer(d: BackgroundWorkDrainer): void { + drainers.set(d.name, d); +} + +/** + * Test seam — registers a drainer and returns an unregister handle. Preferred + * over a blunt reset: real sink modules register at import time and won't re-run + * that top-level side effect on a second import, so a global clear would + * silently drop the production drainers for the rest of the test process. + */ +export function __registerDrainerForTest(d: BackgroundWorkDrainer): () => void { + drainers.set(d.name, d); + return () => { drainers.delete(d.name); }; +} + +/** Test seam — snapshot of registered drainer names (sorted), for assertions. */ +export function __listDrainerNamesForTest(): string[] { + return [...drainers.keys()].sort(); +} + +/** + * CLI-EXIT-ONLY. `abort()` is a permanent process-level state change on a sink + * (the facts queue's `shutdown()` sets `shuttingDown=true` for the process + * lifetime). NEVER call this in a long-lived process (`gbrain serve`). Drains + * every registered sink before `engine.disconnect()` so a PGLite `db.close()` + * can't race in-flight work into the re-pump busy-loop (#1762). + * + * Best-effort and non-throwing: one sink's failure never blocks the others or + * the subsequent disconnect. + */ +export async function drainAllBackgroundWorkForCliExit(opts?: { timeoutMs?: number }): Promise { + const timeoutMs = opts?.timeoutMs ?? 2000; + const ordered = [...drainers.values()].sort( + (a, b) => a.order - b.order || a.name.localeCompare(b.name), + ); + for (const d of ordered) { + try { + const { unfinished } = await d.drain(timeoutMs); + if (unfinished > 0 && d.abort) { + // codex #9: AWAIT — the facts:absorb job writes its absorb-log to the + // DB on settle; the abort must finish against a live engine before the + // caller disconnects. + await d.abort(); + } + } catch { + /* best-effort; never block disconnect on one sink's failure */ + } + } +} diff --git a/src/core/eval-capture.ts b/src/core/eval-capture.ts index 9983e892b..59cb0489f 100644 --- a/src/core/eval-capture.ts +++ b/src/core/eval-capture.ts @@ -43,6 +43,7 @@ import type { } from './types.ts'; import type { GBrainConfig } from './config.ts'; import { scrubPii } from './eval-capture-scrub.ts'; +import { registerBackgroundWorkDrainer } from './background-work.ts'; // HybridSearchMeta is canonical in src/core/types.ts and exported via the // public `gbrain/types` subpath. Surfaced from hybridSearch via the @@ -157,6 +158,21 @@ export async function captureEvalCandidate( engine: BrainEngine, ctx: CaptureContext, opts: { scrub_pii?: boolean } = {}, +): Promise { + // v0.42.11.0 — track the fire-and-forget promise so the background-work + // registry can drain it before CLI disconnect. Callers still `void` this; the + // returned promise is back-compat for any awaiter. The async DB write + // (logEvalCandidate) is the same lock-pin / disconnect-race class as the other + // sinks — on PGLite an undrained capture racing db.close() can wedge. + const p = doCaptureEvalCandidate(engine, ctx, opts); + trackEvalCapture(p); + return p; +} + +async function doCaptureEvalCandidate( + engine: BrainEngine, + ctx: CaptureContext, + opts: { scrub_pii?: boolean } = {}, ): Promise { try { const input = buildEvalCandidateInput(ctx, opts); @@ -174,6 +190,45 @@ export async function captureEvalCandidate( } } +// v0.42.11.0 — bounded drain + registration (mirrors last-retrieved). The 4th +// fire-and-forget DB-write sink (codex caught this one missing from the +// registry). order 3, no abort (bare INSERT, nothing to hard-stop). +const pendingEvalCaptures = new Set>(); + +function trackEvalCapture(promise: Promise): void { + pendingEvalCaptures.add(promise); + promise.finally(() => pendingEvalCaptures.delete(promise)).catch(() => { /* swallow */ }); +} + +export async function awaitPendingEvalCaptures(timeoutMs = 5_000): Promise<{ unfinished: number }> { + if (pendingEvalCaptures.size === 0) return { unfinished: 0 }; + const snapshot = [...pendingEvalCaptures]; + let timer: ReturnType | undefined; + const timeout = new Promise<'timeout'>((resolve) => { + timer = setTimeout(() => resolve('timeout'), timeoutMs); + }); + const drain = Promise.allSettled(snapshot).then(() => 'drained' as const); + const outcome = await Promise.race([drain, timeout]); + if (timer) clearTimeout(timer); + if (outcome === 'timeout') { + const unfinished = pendingEvalCaptures.size; + for (const pr of snapshot) pendingEvalCaptures.delete(pr); + return { unfinished }; + } + return { unfinished: 0 }; +} + +/** Test seam — clears the pending eval-capture set. */ +export function _resetPendingEvalCapturesForTests(): void { + pendingEvalCaptures.clear(); +} + +registerBackgroundWorkDrainer({ + name: 'eval-capture', + order: 3, + drain: (ms) => awaitPendingEvalCaptures(ms), +}); + /** * Check whether capture is enabled for this process. * diff --git a/src/core/facts/queue.ts b/src/core/facts/queue.ts index 475863c7a..805e40f59 100644 --- a/src/core/facts/queue.ts +++ b/src/core/facts/queue.ts @@ -18,6 +18,8 @@ * concurrency + dropping under load. */ +import { registerBackgroundWorkDrainer } from '../background-work.ts'; + export interface FactsQueueCounters { enqueued: number; completed: number; @@ -253,3 +255,17 @@ export function getFactsQueue(opts?: FactsQueueOpts): FactsQueue { export function __resetFactsQueueForTests(): void { _singleton = null; } + +// v0.42.11.0 — register as a background-work sink (order 0 — drained FIRST so +// its abort-path DB logIngest gets the freshest live-engine window). `abort` = +// shutdown(): sets shuttingDown=true (pump short-circuits) + fires internalAbort +// (the facts:absorb job forwards it to gateway.chat, cancelling a hung Haiku the +// drain-only fix can't). Registry AWAITS the abort so logIngest settles against +// a live engine before disconnect (#1762). `drainPending` itself stays +// non-aborting — the abort is the registry's separate post-drain step. +registerBackgroundWorkDrainer({ + name: 'facts', + order: 0, + drain: (ms) => getFactsQueue().drainPending({ timeout: ms }).then((r) => ({ unfinished: r.unfinished })), + abort: () => getFactsQueue().shutdown(), +}); diff --git a/src/core/last-retrieved.ts b/src/core/last-retrieved.ts index 34ceaa473..f5560992a 100644 --- a/src/core/last-retrieved.ts +++ b/src/core/last-retrieved.ts @@ -35,6 +35,7 @@ import type { BrainEngine } from './engine.ts'; import { isUndefinedColumnError } from './utils.ts'; +import { registerBackgroundWorkDrainer } from './background-work.ts'; let _trackRetrievalCache: { ts: number; enabled: boolean } | null = null; const TRACK_RETRIEVAL_CACHE_TTL_MS = 30_000; @@ -125,6 +126,14 @@ export function _peekPendingLastRetrievedWritesForTests(): number { return pendingLastRetrievedWrites.size; } +// v0.42.11.0 — register as a background-work sink (order 1; no abort — bare +// UPDATEs, nothing to hard-stop). Drained before CLI disconnect. +registerBackgroundWorkDrainer({ + name: 'last-retrieved', + order: 1, + drain: (ms) => awaitPendingLastRetrievedWrites(ms).then((r) => ({ unfinished: r.pending })), +}); + /** * Resolve `search.track_retrieval` config with a 30s in-process cache so * hot-path callers don't pay a SELECT per search. Default-on: missing