From 66d0fa91cdafd496622557df3bba33c473ebfab9 Mon Sep 17 00:00:00 2001 From: danusha2345 Date: Fri, 30 Oct 2026 17:59:22 -0500 Subject: [PATCH] fix(mcp): idle query workers are recycled, so a burst's memory is given back A session grows its query pool for concurrent reads, and every worker kept its graph caches, SQLite pages and heap until shutdown, so a server that had served one burst held that memory for the rest of its life. After five minutes with no work the pool now retires every worker that has been used and starts one fresh one, which is ready for the next call. A retired worker is asked to close: it closes its projects, runs a full collection and exits by itself, the path worker-teardown.ts measured as safe on Windows. Only one that hasn't exited after 5 seconds is terminated, through the pool's own start-aware path. Calls that arrive while the fresh worker starts are served in process, the same way as before the pool's first worker is up. Changed on top of the PR: - The idle delay is five minutes, not the PR's 60 seconds, so main doesn't get slower than 1.6.2. Agents routinely pause over a minute to edit or build, and the first explore on a fresh worker takes about 200 ms more. The memory still comes back when a session really goes quiet. - The recycle also arms when a burst's last worker reports ready after the last result. Before, that pool stayed at full size until the next call. - A warm worker that has served nothing is never recycled, so that readiness arming can't replace a clean worker over and over. - destroy() waits for retirements underway, so a worker is not left to end with the process. PR #2862 had this; the PR here did not. PR #2862 by @doomsday1906 made the same idle-worker change and is closed in favor of this one. On its own fork it measured idle RSS going from about 673 MiB to 426 MiB. Validation: - Red on main: six pool tests fail, including a real worker that retires with exit code 0 and is never terminated. Each added case was also proven red by undoing its line. - Windows 11 x64, on a busy machine: 480 child processes, each with four real workers on an excalidraw index and retiring right after a burst, all exited 0 (1,920 worker exits). The AGENTS.md crash did not reproduce here even with an owner terminate in place of the close (0 of 480), so this shows the path works, not that the crash can't happen. - excalidraw, four workers after a burst of explores: 611 MB RSS idle on main, 113 MB here. The first explore after a recycle took 299 ms against 93 ms on a warm worker, and the second was back to 92 ms. Original PR #1583 by @danusha2345. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 1 + __tests__/query-pool.test.ts | 161 ++++++++++++++++++++++++++++++++++- src/mcp/query-pool.ts | 108 +++++++++++++++++++++-- src/mcp/query-worker.ts | 22 ++++- 4 files changed, 283 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bb0ab73990..4a86ffa2d3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -113,6 +113,7 @@ and adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). - An MCP tool call whose `projectPath` names a file inside a project, such as the file an agent is looking at, now answers from that project. Before, it failed with "Path is not a directory", even though a path that doesn't exist at all worked. Sensitive folders such as `/etc` are still refused when the path names a file inside them. Thanks @inth3shadows for the report and @maxmilian for the fix. (#2982) - An index left incomplete by an indexing run that was stopped partway, by a crash, an out-of-memory kill or the server's own hang check, no longer gets the background server killed with "Main thread unresponsive" every time it starts: on a large project, restoring the database indexes the run had set aside took longer than the hang check allows, so each restart died at the same point and the index was never repaired. They are now restored in the background, and the project is served once they are back. Thanks @cbeaulieu-gt for the report and the first fix. (#1887) - A new agent session no longer waits up to about 12 minutes for its first answer when the background server is stuck: a server whose process still holds its connection open but never responds made the session wait out a long series of retries. The session now gives up within seconds and answers on its own, and switches back to the background server once that has been restarted. +- The background server now recycles its query workers after five minutes idle, releasing the memory a burst of large queries left behind while keeping one fresh worker ready for the next call. Thanks @danusha2345 and @doomsday1906. - An MCP session whose agent stops reading its answers, while keeping the session open, now shuts down the normal way the next time it answers: it stops its work, frees the project for the next session and exits cleanly, and other sessions sharing the background server carry on. Before, that answer crashed it with "Uncaught exception: Error: EPIPE". Thanks @iiijiashu for the report and the first fix. (#2991) - A tool call whose arguments aren't an object, like `"arguments": "alpha"` from a hand-written MCP client or script, now gets the standard invalid-params error saying the arguments must be an object, and no tool runs. Before, a string or a number made most tools fail and come back as an internal CodeGraph error that told the agent to retry and then to continue without CodeGraph, and an array was run as if it were the arguments. Thanks @inth3shadows for the report and @kjgggggg for the first fix. (#2981) - The bundled Node runtime moves to 24.22.1, Node's October security release, and every release now checks the Node it bundles against the checksums nodejs.org publishes. The new runtime also fixes a crash on Windows as a command like `codegraph sync` exits, after its work is done, printing `Assertion failed: !(handle->flags & UV_HANDLE_CLOSING)`, which made a script or CI run that had succeeded fail. Thanks @Ancienttwo for the report and the first fix, and @IdealBoy709 for the report on the security release. (#2617, #2780) diff --git a/__tests__/query-pool.test.ts b/__tests__/query-pool.test.ts index 7c28f48f04..3fcf73f2ee 100644 --- a/__tests__/query-pool.test.ts +++ b/__tests__/query-pool.test.ts @@ -36,6 +36,10 @@ class FakeWorker implements PoolWorker { private errorCb?: (e: Error) => void; private exitCb?: (code: number) => void; alive = true; + terminated = false; + exitCode: number | undefined; + /** Ignore the pool's 'close', to exercise its terminate fallback. */ + ignoresClose = false; constructor(private behavior: (m: CallMsg) => Action, readyOk: boolean | null = true) { setTimeout(() => { if (this.alive && readyOk !== null) this.emitMessage({ type: 'ready', ok: readyOk }); }, 0); } @@ -51,7 +55,14 @@ class FakeWorker implements PoolWorker { if (this.alive) this.msgCb?.({ type: 'result', id, result }); } postMessage(msg: unknown): void { - const m = msg as CallMsg; + const m = msg as CallMsg | { type: 'close' }; + if (m?.type === 'close') { + // Like the real worker: close up and exit by itself, with code 0. + if (this.ignoresClose) return; + this.alive = false; + setTimeout(() => { this.exitCode = 0; this.exitCb?.(0); }, 0); + return; + } if (!m || m.type !== 'call') return; const action = this.behavior(m); if ('crash' in action) { @@ -63,7 +74,7 @@ class FakeWorker implements PoolWorker { if ('wait' in action) { void action.wait.then((r) => this.reply(m.id, r)); return; } setTimeout(() => this.reply(m.id, action.result), 0); } - terminate(): Promise { this.alive = false; return Promise.resolve(0); } + terminate(): Promise { this.alive = false; this.terminated = true; return Promise.resolve(0); } } const ok = (text: string): ToolResult => ({ content: [{ type: 'text', text }] }); @@ -126,6 +137,127 @@ describe('QueryPool', () => { await pool.destroy(); }); + it('recycles an idle burst back to one fresh worker', async () => { + let release!: () => void; + const gate = new Promise((r) => { release = r; }); + const workers: FakeWorker[] = []; + const pool = new QueryPool({ + root: '/x', size: 4, idleShrinkMs: 20, + createWorker: () => { + const worker = new FakeWorker((m) => ({ + wait: gate.then(() => ok(`r${m.id}`)), + })); + workers.push(worker); + return worker; + }, + }); + + const calls = Promise.all(Array.from({ length: 4 }, (_, i) => pool.run('codegraph_search', { i }))); + await sleep(40); + expect(pool.liveWorkers).toBe(4); + release(); + await calls; + await sleep(40); + + expect(pool.liveWorkers).toBe(1); + expect(workers).toHaveLength(5); // four used isolates replaced by one clean isolate + expect(workers.slice(0, 4).every((w) => !w.alive)).toBe(true); + // Each used worker exited by itself on 'close'; none was terminated. + expect(workers.slice(0, 4).every((w) => w.exitCode === 0 && !w.terminated)).toBe(true); + expect(pool.ready).toBe(true); + const again = await pool.run('codegraph_node', { symbol: 's' }); + expect(again.isError).toBeFalsy(); + await pool.destroy(); + }); + + it('terminates a retired idle worker that does not exit by itself', async () => { + const workers: FakeWorker[] = []; + const pool = new QueryPool({ + root: '/x', size: 1, idleShrinkMs: 20, retireTimeoutMs: 200, + createWorker: () => { + const worker = new FakeWorker(() => ({ result: ok('r') })); + worker.ignoresClose = workers.length === 0; + workers.push(worker); + return worker; + }, + }); + + await pool.run('codegraph_search', {}); + await vi.waitFor(() => expect(workers).toHaveLength(2)); // retired, replaced + expect(workers[0].terminated).toBe(false); // still has time to exit + await vi.waitFor(() => expect(workers[0].terminated).toBe(true)); + expect(workers[1].terminated).toBe(false); + await pool.destroy(); + }); + + it('recycles even when a burst worker reports ready after the last result', async () => { + const workers: FakeWorker[] = []; + let release!: () => void; + const gate = new Promise((r) => { release = r; }); + const pool = new QueryPool({ + root: '/x', size: 2, idleShrinkMs: 20, + createWorker: () => { + // The second worker stays pending until the test lets it report ready. + const worker = new FakeWorker((m) => ({ wait: gate.then(() => ok(`r${m.id}`)) }), workers.length === 0 ? true : null); + workers.push(worker); + return worker; + }, + }); + + await vi.waitFor(() => expect(pool.ready).toBe(true)); + const calls = Promise.all([pool.run('codegraph_search', { i: 0 }), pool.run('codegraph_search', { i: 1 })]); + await vi.waitFor(() => expect(workers).toHaveLength(2)); + release(); + await sleep(40); + expect(workers).toHaveLength(2); // one worker is still starting, so nothing recycled + workers[1]!.emitMessage({ type: 'ready', ok: true }); + await calls; + await vi.waitFor(() => expect(pool.liveWorkers).toBe(1)); + expect(workers).toHaveLength(3); + await pool.destroy(); + }); + + it('never recycles a warm worker that has served nothing', async () => { + const workers: FakeWorker[] = []; + const pool = new QueryPool({ + root: '/x', size: 2, idleShrinkMs: 5, + createWorker: () => { + const worker = new FakeWorker(() => ({ result: ok('r') })); + workers.push(worker); + return worker; + }, + }); + + await vi.waitFor(() => expect(pool.ready).toBe(true)); + await sleep(40); + expect(workers).toHaveLength(1); + + // A recycle leaves one clean worker, which then stays until it is used. + await pool.run('codegraph_search', {}); + await vi.waitFor(() => expect(workers).toHaveLength(2)); + await sleep(40); + expect(workers).toHaveLength(2); + await pool.destroy(); + }); + + it('destroy waits for a retired worker still on its way out', async () => { + const workers: FakeWorker[] = []; + const pool = new QueryPool({ + root: '/x', size: 1, idleShrinkMs: 5, retireTimeoutMs: 60, + createWorker: () => { + const worker = new FakeWorker(() => ({ result: ok('r') })); + worker.ignoresClose = workers.length === 0; + workers.push(worker); + return worker; + }, + }); + + await pool.run('codegraph_search', {}); + await vi.waitFor(() => expect(workers).toHaveLength(2)); + await pool.destroy(); + expect(workers[0]!.terminated).toBe(true); // the timeout ended it, inside destroy() + }); + it('recovers from a worker crash: retries the in-flight call and respawns', async () => { let calls = 0; const pool = new QueryPool({ @@ -465,6 +597,31 @@ describe('MCP query pool with real projects (#1465)', () => { } }, 30000); + it('retires an idle real worker that exits by itself with code 0', async () => { + const alpha = await indexProject('alpha', 'alphaSymbol'); + const terminated: boolean[] = []; + const exits: Promise[] = []; + pool = new QueryPool({ + root: alpha, size: 1, idleShrinkMs: 50, + createWorker: () => { + const worker = new Worker(path.resolve(__dirname, '../dist/mcp/query-worker.js'), { workerData: { root: alpha } }); + const i = terminated.push(false) - 1; + const terminate = worker.terminate.bind(worker); + worker.terminate = () => { + terminated[i] = true; + return terminate(); + }; + exits.push(new Promise((resolve) => worker.once('exit', resolve))); + return worker; + }, + }); + const result = await pool.run('codegraph_search', { query: 'alphaSymbol' }); + expect(result.isError).toBeFalsy(); + expect(await exits[0]).toBe(0); + expect(terminated[0]).toBe(false); + expect(exits).toHaveLength(2); // a fresh worker replaced it + }, 30000); + it('never terminates a real worker before it has started', async () => { const alpha = await indexProject('alpha', 'alphaSymbol'); const startedWhenEnded: boolean[] = []; diff --git a/src/mcp/query-pool.ts b/src/mcp/query-pool.ts index 6b1f473991..653d500391 100644 --- a/src/mcp/query-pool.ts +++ b/src/mcp/query-pool.ts @@ -48,6 +48,15 @@ export interface PoolWorker { /** Default linger before a queued call is answered with busy-guidance. */ const DEFAULT_BUSY_TIMEOUT_MS = 45_000; // < the ~60s MCP client request timeout +/** Recycle a quiescent pool after a burst so worker-local graph/SQLite caches + * do not remain resident for the daemon's whole lifetime. Long enough that an + * agent pausing to edit or build keeps its warm workers: the first explore on + * a fresh one takes about 200 ms more. */ +const DEFAULT_IDLE_SHRINK_MS = 300_000; + +/** How long a retired idle worker gets to exit on 'close' before it is terminated. */ +const RETIRE_TIMEOUT_MS = 5_000; + /** Hard ceiling on pool size regardless of core count / env. */ const MAX_POOL_SIZE = 16; @@ -108,6 +117,10 @@ export interface QueryPoolOptions { softTimeoutMs?: number; /** Retries for an in-flight call whose worker crashed. Default 1. */ maxRetries?: number; + /** Idle delay before recycling back to one fresh worker. Default 5 minutes; 0 disables. */ + idleShrinkMs?: number; + /** How long a retired idle worker gets to exit by itself (tests shorten it). Default 5s. */ + retireTimeoutMs?: number; /** Worker factory (tests inject a fake). Defaults to a real `worker_threads` Worker. */ createWorker?: () => PoolWorker; /** How long destroy() waits on a worker still starting up (tests shorten it). Default 15s. */ @@ -174,14 +187,25 @@ export class QueryPool { private readonly maxSize: number; private readonly softTimeoutMs: number; private readonly maxRetries: number; + private readonly idleShrinkMs: number; + private readonly retireTimeoutMs: number; private readonly createWorker: () => PoolWorker; private readonly startSettleMs: number; + private idleShrinkTimer?: NodeJS.Timeout; + // Retired idle workers asked to exit: the timer that terminates one that + // hasn't by retireTimeoutMs, and what destroy() waits on. + private retiring = new Map; settle: () => void }>(); + // Set by a served call and cleared by a recycle, so a clean warm worker that + // has answered nothing is never recycled for its own sake. + private servedSinceRecycle = false; constructor(opts: QueryPoolOptions) { this.root = opts.root; this.maxSize = Math.max(1, Math.min(opts.size ?? Math.max(1, os.cpus().length - 1), MAX_POOL_SIZE)); this.softTimeoutMs = opts.softTimeoutMs ?? resolveBusyTimeoutMs(); this.maxRetries = opts.maxRetries ?? 1; + this.idleShrinkMs = Math.max(0, opts.idleShrinkMs ?? DEFAULT_IDLE_SHRINK_MS); + this.retireTimeoutMs = opts.retireTimeoutMs ?? RETIRE_TIMEOUT_MS; this.createWorker = opts.createWorker ?? (() => new Worker(WORKER_FILE, { workerData: { root: this.root } })); this.startSettleMs = opts.startSettleMs ?? WORKER_START_SETTLE_MS; this.spawnOne(); // one eager warm worker, ready for the first call @@ -211,8 +235,9 @@ export class QueryPool { * hitting that window was the recurring #662 test flake (and a real * first-call stall for agents). The pool exists for CONCURRENT load, which * by definition arrives after warm-up; the pre-pool in-process path is - * strictly better while nothing is warm. Stays true for the pool's - * lifetime — later crash-respawn gaps are covered by retry + backstop. + * strictly better while nothing is warm, which also holds while an idle + * recycle's fresh worker starts. Stays true across crash-respawn gaps, which + * are covered by retry + backstop. */ get ready(): boolean { return this.everReady && !this.destroyed; @@ -235,7 +260,7 @@ export class QueryPool { this.starts.set(w, { settled, settle }); w.on('message', (m) => this.onMessage(w, (m ?? {}) as WorkerMessage)); w.on('error', () => { this.startSettled(w); this.onWorkerGone(w); }); - w.on('exit', (code) => { this.startSettled(w); if (code !== 0) this.onWorkerGone(w); }); + w.on('exit', (code) => { this.startSettled(w); this.retired(w); if (code !== 0) this.onWorkerGone(w); }); } private onMessage(w: PoolWorker, m: WorkerMessage): void { @@ -252,17 +277,84 @@ export class QueryPool { this.everReady = true; this.idle.push(w); this.drain(); + this.armIdleShrink(); // the last result may have come while this one started return; } if (m.type === 'result') { const job = this.inflight.get(w); this.inflight.delete(w); this.idle.push(w); + this.servedSinceRecycle = true; if (job) this.settle(job, m.result ?? busyGuidance(0)); this.drain(); + this.armIdleShrink(); + } + } + + private clearIdleShrink(): void { + if (this.idleShrinkTimer) clearTimeout(this.idleShrinkTimer); + this.idleShrinkTimer = undefined; + } + + private armIdleShrink(): void { + this.clearIdleShrink(); + if ( + this.idleShrinkMs === 0 || this.destroyed || !this.servedSinceRecycle || + this.queue.length > 0 || this.inflight.size > 0 || this.pendingWorkers.size > 0 || + this.idle.length !== this.workers.size + ) return; + this.idleShrinkTimer = setTimeout(() => this.shrinkIdle(), this.idleShrinkMs); + this.idleShrinkTimer.unref?.(); + } + + private shrinkIdle(): void { + this.idleShrinkTimer = undefined; + if ( + this.destroyed || this.queue.length > 0 || this.inflight.size > 0 || + this.pendingWorkers.size > 0 || this.idle.length !== this.workers.size + ) return; + + // Drop every used isolate, including the last one: a single worker can hold + // hundreds of MB in graph and SQLite caches after a large query. Replace it + // with one clean warm worker so the next burst still avoids a cold queue. + const stale = [...this.workers]; + this.workers.clear(); + this.idle = []; + this.everReady = false; + this.servedSinceRecycle = false; + for (const w of stale) this.retire(w); + this.spawnOne(); + } + + /** + * Ask an idle worker to exit by itself: it closes its connections and + * collects garbage first, so no marking is in flight when its thread ends + * (see worker-teardown.ts). One that hasn't exited by `retireTimeoutMs` is + * terminated. + */ + private retire(w: PoolWorker): void { + const terminate = (): void => { void this.terminateStarted(w).then(() => this.retired(w)); }; + const cap = setTimeout(terminate, this.retireTimeoutMs); + cap.unref?.(); + let settle!: () => void; + const gone = new Promise((resolve) => { settle = resolve; }); + this.retiring.set(w, { cap, gone, settle }); + try { + w.postMessage({ type: 'close' }); + } catch { + clearTimeout(cap); + terminate(); } } + private retired(w: PoolWorker): void { + const retirement = this.retiring.get(w); + if (!retirement) return; + clearTimeout(retirement.cap); + this.retiring.delete(w); + retirement.settle(); + } + // A worker died (crash hook, OOM, segfault, exit≠0). Respawn a replacement and // retry its in-flight job once; a job that keeps crashing workers fails // gracefully so it can't loop the pool forever. @@ -319,6 +411,7 @@ export class QueryPool { /** Run a read tool on the pool. Always resolves (never rejects). */ run(toolName: string, args: Record): Promise { + this.clearIdleShrink(); return new Promise((resolve) => { const job: Job = { id: this.nextId++, toolName, args, resolve, @@ -364,11 +457,13 @@ export class QueryPool { /** * Terminate all workers and answer any outstanding calls gracefully. A worker * still starting up is terminated once it has started (see - * WORKER_START_SETTLE_MS); outstanding calls are answered first. + * WORKER_START_SETTLE_MS); outstanding calls are answered first. A retired + * worker still on its way out is waited for, not left to end with the process. */ async destroy(): Promise { if (this.destroyed) return; this.destroyed = true; + this.clearIdleShrink(); const ws = [...this.workers]; this.workers.clear(); this.pendingWorkers.clear(); @@ -378,6 +473,9 @@ export class QueryPool { } this.inflight.clear(); this.queue = []; - await Promise.all(ws.map((w) => this.terminateStarted(w))); + await Promise.all([ + ...ws.map((w) => this.terminateStarted(w)), + ...[...this.retiring.values()].map((retirement) => retirement.gone), + ]); } } diff --git a/src/mcp/query-worker.ts b/src/mcp/query-worker.ts index 525e9334a6..f03af5eee4 100644 --- a/src/mcp/query-worker.ts +++ b/src/mcp/query-worker.ts @@ -23,6 +23,7 @@ import { parentPort, workerData } from 'worker_threads'; import type { ToolResult } from './tools'; +import { collectBeforeExit } from '../worker-teardown'; interface WorkerInit { root: string | null; @@ -35,6 +36,11 @@ interface CallMessage { args: Record; } +/** The pool retiring this idle worker: close everything and exit. */ +interface CloseMessage { + type: 'close'; +} + // Mirror the engine's lazy-require of the heavy CodeGraph + tools chain. This // module is only ever loaded as a Worker, so the require runs once on spawn. const loadCodeGraph = (): typeof import('../index').default => @@ -50,10 +56,11 @@ if (parentPort) { // opened lazily on first cross-project (projectPath) call by the ToolHandler's // own per-handler cache. openSync does not start a watcher — workers are pure // readers; the single watcher/writer stays on the daemon's main thread. + let cg: import('../index').default | null = null; let handler: InstanceType | null = null; let initError: string | null = null; try { - const cg = root ? loadCodeGraph().openSync(root) : null; + cg = root ? loadCodeGraph().openSync(root) : null; handler = new (loadToolHandler())(cg); } catch (err) { initError = err instanceof Error ? err.message : String(err); @@ -63,7 +70,18 @@ if (parentPort) { // against its crash budget (→ fall back to in-process) without hanging. port.postMessage({ type: 'ready', ok: initError === null, error: initError }); - port.on('message', (msg: CallMessage) => { + port.on('message', (msg: CallMessage | CloseMessage) => { + if (msg?.type === 'close') { + try { + void handler?.closeAll(); // projects opened through projectPath + } catch { /* already closed */ } + try { + cg?.close(); + } catch { /* already closed */ } + // No GC marking in flight when the thread ends (worker-teardown.ts). + collectBeforeExit(); + process.exit(0); + } if (!msg || msg.type !== 'call') return; void serve(msg); });