Missing parse-time concurrency

[iacore]

Related: #320

Below are what my agent created. If this is good, I'll implement it.

---

Parse Concurrency Plan for CodeGraph

1. Current Bottleneck Summary

ExtractionOrchestrator.indexAll() in src/extraction/index.ts already reads files in parallel batches of FILE_IO_BATCH_SIZE = 10 (lines 887-913), but parsing is strictly serial. The loop inside the batch reader waits for await requestParse(filePath, content) for each file before sending the next one (lines 909-956). The single requestParse function routes every file through one parse-worker.js thread, so at most one CPU core is doing tree-sitter work at any time.

Key serial points:

• src/extraction/index.ts:909 — result = await requestParse(filePath, content) inside the per-file loop.
• src/extraction/index.ts:775-824 — requestParse queues one message at a time to a single parseWorker.
• src/extraction/index.ts:672-680 — only one Worker is spawned per indexAll() invocation.
• src/extraction/index.ts:989-1056 — retry passes also reuse the same serial requestParse path.

This is a CPU-bound bottleneck. Parsing in extractFromSource() is synchronous tree-sitter work, so adding worker threads that each own an independent WASM heap/parser cache is the correct way to scale.

2. Proposed Architecture

Worker Pool Size Heuristic

Introduce a constant PARSE_WORKER_POOL_SIZE computed at runtime:

const PARSE_WORKER_POOL_SIZE = Math.max(2, Math.min(8, os.cpus().length - 1));

• Cap at 8 to avoid spawning more WASM heaps than memory can tolerate.
• Reserve one CPU for the main thread and any UI/shimmer progress work.
• Allow override via environment variable CODEGRAPH_PARSE_WORKERS for benchmarking and constrained CI runners.

File → Worker Dispatch Strategy

Use a simple round-robin dispatcher over the worker pool. Each worker maintains its own parser cache and parse-count state, so no main-thread state is needed to assign work. The dispatcher sends parse messages immediately up to a per-worker in-flight limit; if every worker is at capacity, it waits on a Promise-based backpressure slot before sending the next file.

Pseudo-code:

const workerIndex = nextWorker++ % poolSize;
await pool[workerIndex].acquireSlot(); // backpressure
pool[workerIndex].post({ type: 'parse', id, filePath, content, frameworkNames });

Reasoning for round-robin vs. work-stealing:

• Keeps per-worker state local (parse counts, loaded grammars).
• Avoids central contention on a work queue; each worker's MessagePort is the queue.
• Easy to reason about timeouts and crash recovery: a crashed worker only loses its own in-flight items.

Result Ordering

Results are returned by request id and correlated with the original file in a Map<number, PendingParse>. Ordering does not need to match submission order. The main thread increments processed and reports progress as each result arrives, then stores the result in SQLite on the main thread. Because DB writes are serialized, preserving input order has no correctness benefit.

Main-Thread DB Serialization

storeExtractionResult() stays on the main thread and stays unchanged. Each worker result is handed back to the main thread as a { type: 'parse-result', id, result } message; the main thread resolves the pending promise, increments counters, and calls storeExtractionResult() synchronously. SQLite writes remain single-threaded and safe.

3. Required File Changes and New Files

New file: src/extraction/parse-worker-pool.ts

Responsibilities:

• Spawn N Worker instances from parse-worker.js.
• Load grammars in each worker once via load-grammars message.
• Expose requestParse(filePath, content, frameworkNames): Promise<ExtractionResult>.
• Track per-worker in-flight count and enforce a MAX_IN_FLIGHT_PER_WORKER cap.
• Round-robin dispatch across workers.
• Handle worker error / exit events:
  • Reject all pending parses for that worker with the crash reason.
  • Restart the worker with a fresh WASM heap.
  • Mark restart so that requestParse retries can be retried by the caller.
• Implement terminate() to shut down all workers and clear timers.
• Implement inProcessFallback: boolean mode that skips workers and calls extractFromSource directly.

Type shape (illustrative):

export interface ParseWorkerPool {
  requestParse(filePath: string, content: string, frameworkNames?: string[]): Promise<ExtractionResult>;
  terminate(): Promise<void>;
}

export function createParseWorkerPool(options: {
  workerScriptPath: string;
  languages: Language[];
  poolSize: number;
  maxInFlightPerWorker: number;
  parseTimeoutMs: number;
  workerRecycleInterval: number;
  log?: (msg: string) => void;
  fallback?: boolean;
}): ParseWorkerPool;

Modify: src/extraction/index.ts

• Replace the single-worker lifecycle block (lines 666-824) with a createParseWorkerPool(...) instance created after needed languages are known.
• Keep the in-process fallback path when parse-worker.js is missing.
• Change the inner batch loop (lines 885-956) to dispatch all parsed files in the batch concurrently via pool.requestParse() and await the array of results.
• Keep retry passes (lines 978-1073) serial or allow limited concurrency (e.g., one retry at a time per worker). Because retries are rare and already require a fresh worker, serial retries are acceptable; the plan leaves this decision to the implementer but recommends serial retries for simplicity.
• On abort, call await pool.terminate() and return the aborted result.
• At the end of indexAll, call await pool.terminate().

Modify: src/extraction/parse-worker.ts (minor)

No structural change required. The existing load-grammars, parse, and shutdown message handlers already support one-file-at-a-time work. The pool sends the same messages. Optionally expose parse-count tracking to the pool so it can recycle workers from the main thread, but the worker already recycles parsers internally every PARSER_RESET_INTERVAL = 5000. Recommendation: keep parser reset inside the worker; add per-worker recycle from the main thread only if benchmarking shows benefit.

Modify: src/extraction/grammars.ts (none expected)

The per-module parser cache is correct because each worker is a separate module isolate. The main thread should not load grammars when workers are used, matching current behavior.

4. Memory and Stability Rules

WASM Heap Per Worker

• Each worker process has its own web-tree-sitter WASM linear memory and its own parserCache / languageCache in grammars.ts.
• No shared parser instances between workers.
• Main thread must not call loadGrammarsForLanguages() when the worker pool is active.

Recycle Policy

• Keep the existing PARSER_RESET_INTERVAL = 5000 resets inside the worker (line 44 in parse-worker.ts).
• Add a per-worker workerParseCount in the pool and recycle the whole worker after WORKER_RECYCLE_INTERVAL = 250 total parses per worker, same as today. Recycle should:
  1. Stop assigning new work to the worker.
  2. Wait for in-flight parses to finish.
  3. terminate() the worker.
  4. Spawn a replacement and re-send load-grammars.

Max In-Flight Parses / Backpressure

• Cap each worker at MAX_IN_FLIGHT_PER_WORKER = 4. With 8 workers this allows up to 32 parses in flight.
• If all workers are at capacity, requestParse awaits a slot-release promise rather than queuing unbounded work.
• This prevents memory blow-up when the main thread reads files much faster than workers parse them.

Large File Handling

• Keep the existing MAX_FILE_SIZE check and timeout scaling (PARSE_TIMEOUT_MS + 10s per 100KB) exactly as today.
• Keep FILE_IO_BATCH_SIZE = 10 so file reads do not run far ahead of the parser pool.

5. Error Handling

Timeouts

• Each requestParse sets its own setTimeout keyed by request id.
• On timeout, reject the caller, mark the worker as "needs recycle", and terminate it.
• Other in-flight parses on the same worker are rejected with Worker error: terminated due to timeout on another request.

Worker Crashes

• On worker.on('error') or non-zero exit, reject every pending parse for that worker with a crash reason.
• Restart the worker automatically with load-grammars.
• A retried parse should be routed to any healthy worker; the caller decides whether to retry.

Retry-After-Crash

• Preserve the existing retry pass in indexAll() (lines 978-1073).
• Before each retry, ensure the worker pool has at least one healthy worker; the retry itself calls pool.requestParse() and will use round-robin assignment.
• The crash reason propagated from the pool must include the original file path so the retry filter (e.message.includes('Worker exited') || e.message.includes('memory access out of bounds')) still matches.

Graceful Fallback to In-Process Serial Parsing

• If parse-worker.js is missing (e.g. in unbuilt tests), keep the current behavior: set fallback = true, load grammars in-process, and call extractFromSource() synchronously.
• If all workers fail to spawn (e.g. worker_threads unavailable), throw a clear error unless fallback is enabled.

6. Progress Reporting and Abort-Signal Behavior

Progress Reporting

• Progress is reported when a parse result returns, not when it is dispatched.
• The main thread updates processed++ inside the per-result handler and calls onProgress({ phase: 'parsing', current: processed, total, currentFile }).
• Because results arrive out of order, currentFile shows the most recently completed file, not the file currently parsing. This is acceptable for a progress bar.

Abort Signal

• Check signal?.aborted before each batch and inside result handlers.
• On abort:
  1. Stop reading new batches.
  2. Call await pool.terminate() to kill all workers.
  3. Reject any still-pending parses.
  4. Return the aborted IndexResult exactly as today.
• The existing abort check before requestParse (lines 885-897) moves to the result await point.

7. Verification Strategy

Unit / Integration Tests

1. Pool unit test (__tests__/parse-worker-pool.test.ts):

   • Spawn a pool with 2 workers and dispatch 10 tiny files.
   • Assert all resolve, results have correct nodes/errors, and no result is lost.
   • Assert that with a worker that crashes on a specific input, the pool restarts the worker and rejects only that request.

2. Concurrency correctness test (__tests__/extraction-concurrency.test.ts):

   • Create a temp project with 20 small TypeScript files.
   • Run ExtractionOrchestrator.indexAll() with CODEGRAPH_PARSE_WORKERS=4.
   • Assert total files indexed equals 20, total nodes/edges match the serial run, and the DB contains all files.

3. Abort-signal test:

   • Start indexing a large temp project, abort mid-run.
   • Assert the pool terminates and the returned result has success: false with an Aborted error.

4. Fallback test:

   • Run indexing without building parse-worker.js (or mock fs.existsSync to return false).
   • Assert it completes using in-process parsing.

5. Memory-stability regression test:

   • Parse a file that triggers the worker to crash with memory access out of bounds.
   • Assert the pool recycles the worker, the retry pass succeeds, and no cascading failures occur.

Benchmark Command

Add a lightweight benchmark script to compare serial vs. parallel parse throughput:

# Build the worker so the pool path is exercised
npm run build

# Serial baseline (1 worker)
CODEGRAPH_PARSE_WORKERS=1 npx tsx scripts/benchmark-parse.ts /path/to/project

# Parallel (heuristic pool size)
CODEGRAPH_PARSE_WORKERS=0 npx tsx scripts/benchmark-parse.ts /path/to/project

# Explicit 4 workers
CODEGRAPH_PARSE_WORKERS=4 npx tsx scripts/benchmark-parse.ts /path/to/project

scripts/benchmark-parse.ts should:

1. Scan the target project.
2. Create a fresh .codegraph DB in a temp directory.
3. Call indexAll() with a no-op onProgress.
4. Print durationMs, filesIndexed, filesErrored, and parses/second.
5. Run each configuration at least 3 times and report the median.

8. Rollback / Kill-Switch Plan

Environment Variable Kill-Switch

Add CODEGRAPH_PARSE_WORKERS:

• 0 or unset single-worker path — force the existing serial single-worker behavior.
• 1 — single worker (useful for debugging).
• 2-8 — pool size.

If CODEGRAPH_PARSE_WORKERS is missing, default to the heuristic pool size after the feature is proven stable; during the rollout phase default to 1 so maintainers can opt-in.

Code Rollback

• Keep requestParse and the single-worker lifecycle in src/extraction/index.ts behind a usePool flag during the rollout. Once the pool is stable, remove the old block in a follow-up cleanup PR.
• parse-worker.ts message protocol remains backward-compatible, so reverting only index.ts and deleting parse-worker-pool.ts restores the old behavior.

Verification Before Full Rollout

1. Run the full test suite: npm test.
2. Run the benchmark on at least three real projects (small, medium, large).
3. Confirm filesIndexed, filesErrored, and nodesCreated are identical between CODEGRAPH_PARSE_WORKERS=1 and CODEGRAPH_PARSE_WORKERS=4 for the same project.
4. Confirm abort signal terminates cleanly and leaves no orphan worker processes.
5. Confirm the fallback path still works when dist/extraction/parse-worker.js is absent.

If any discrepancy is found, set CODEGRAPH_PARSE_WORKERS=1 as the default and file a bug before re-enabling the pool.

[colbymchenry]

Implemented and merged in #1025 — ships in the next release. Thanks for the detailed plan; the implementation follows it closely (worker pool with idle-worker dispatch, per-worker grammar load, recycle-to-reclaim-WASM-heap, per-file timeout, crash→reject→respawn, graceful in-process fallback when the compiled worker is absent).

A few notes on where it landed:

• One deliberate divergence from the plan. The plan says "Ordering does not need to match submission order … preserving input order has no correctness benefit." That turned out not to hold here: indexAll is followed by a resolution phase that resolves an ambiguous reference (several same-named symbols) to a candidate by node DB-insertion order. Committing parse results in completion order made that order timing-dependent, so the graph drifted run-to-run (edge counts varied ~0.04% while node counts stayed identical — the giveaway). So the pool parses in parallel but commits to SQLite in file order (a small bounded reorder buffer, with backpressure on the dispatched-but-uncommitted count so a slow file can't grow it without bound). The resulting graph is byte-identical to the old single-worker output.
• Config: pool size defaults to clamp(cores-1, 1, 8); set CODEGRAPH_PARSE_WORKERS to pick a count, or CODEGRAPH_PARSE_WORKERS=1 to get the exact previous single-worker behavior.
• On the OOM you hit in Support large-scale repository indexing #320: validated peak memory is flat-to-lower with the pool, not multiplied by core count — workers recycle their WASM heaps independently and spread large-file parses across smaller heaps (on redis, peak RSS actually dropped vs. the single worker). So defaulting to multi-worker is safe.
• Honest expectation on speed: the win is real but modest and depends on how parse-heavy your repo is — parsing is often a minority of total index time (file storage + resolution are serial and frequently dominate). In my testing: ~+9% on django (Python), ~+17% on redis (C, parse-heavier). Not an N× speedup.

Closing as done.