diff --git a/gitnexus/src/core/lbug/lbug-adapter.ts b/gitnexus/src/core/lbug/lbug-adapter.ts index 621d26da6..5f5a076d4 100644 --- a/gitnexus/src/core/lbug/lbug-adapter.ts +++ b/gitnexus/src/core/lbug/lbug-adapter.ts @@ -34,7 +34,16 @@ import type { GraphEmitManifest } from './graph-emit-sink.js'; import type { PdgEmitManifest } from './pdg-emit-sink.js'; import { PDG_EDGE_TYPES } from './pdg-emit-sink.js'; import { getNodeLabel as deriveNodeLabel, type WriteStreamFactory } from './rel-pair-routing.js'; -import { EMBEDDABLE_LABELS, type CachedEmbedding } from '../embeddings/types.js'; +import { EMBEDDABLE_LABELS } from '../embeddings/types.js'; +import { + abortCachedEmbeddingsBuilder, + createCachedEmbeddingsBuilder, + emptyCachedEmbeddingsSnapshot, + finalizeCachedEmbeddingsSnapshot, + ingestCachedEmbeddingRow, + type CachedEmbeddingsSnapshot, + type LoadCachedEmbeddingsOptions, +} from '../embeddings/embedding-restore-spill.js'; import { extensionManager, getFtsCapability, @@ -2069,28 +2078,31 @@ export const getLbugStats = async (): Promise<{ /** * Load cached embeddings from LadybugDB before a rebuild. - * Returns all embedding vectors so they can be re-inserted after the graph is reloaded, - * avoiding expensive re-embedding of unchanged nodes. + * + * Streams `CodeEmbedding` rows with `hasNext`/`getNext` under `withConnLock` + * (#2264, #3306). Vectors are spilled to a temp Float32 file once the table + * exceeds the in-memory row limit so incremental analyze cannot OOM the V8 + * heap by materializing every `number[]` up front. Small tables still return + * in-RAM `embeddings` for existing callers/tests. * * Detects old schema (no chunkIndex column) and returns empty cache to trigger rebuild. */ -export const loadCachedEmbeddings = async (): Promise<{ - embeddingNodeIds: Set; - embeddings: CachedEmbedding[]; -}> => { +export const loadCachedEmbeddings = async ( + options?: LoadCachedEmbeddingsOptions, +): Promise => { const c = conn; if (!c) { - return { embeddingNodeIds: new Set(), embeddings: [] }; + return emptyCachedEmbeddingsSnapshot(); } // The whole read runs inside the connection lock (#2264 review P2). It's safe // today only by call-ordering (loadCachedEmbeddings runs before the WAL driver // starts), but the lock makes it robust to future reordering — a concurrent // CHECKPOINT on the singleton connection is the documented corruption trigger. - // Leaf read: no nested withConnLock-wrapped helpers inside. + // Leaf read: no nested withConnLock-wrapped helpers inside. Do NOT call + // `streamQuery` here — that path is unlocked and would race a CHECKPOINT. return withConnLock(async () => { - const embeddingNodeIds = new Set(); - const embeddings: CachedEmbedding[] = []; + const builder = createCachedEmbeddingsBuilder(options); try { // Schema migration detection: query with new columns to verify schema version. // Old schema only had (nodeId, embedding); new schema adds (id, chunkIndex, startLine, endLine, contentHash). @@ -2104,51 +2116,46 @@ export const loadCachedEmbeddings = async (): Promise<{ ); await readQueryRows(check); } catch { - return { embeddingNodeIds: new Set(), embeddings: [] }; + abortCachedEmbeddingsBuilder(builder); + return emptyCachedEmbeddingsSnapshot(); } - // Try to read contentHash alongside chunk columns - let rows: any; + let queryResult: lbug.QueryResult | lbug.QueryResult[] | undefined; let hasContentHash = true; try { - rows = await c.query( - `MATCH (e:${EMBEDDING_TABLE_NAME}) RETURN e.nodeId AS nodeId, e.chunkIndex AS chunkIndex, e.startLine AS startLine, e.endLine AS endLine, e.embedding AS embedding, e.contentHash AS contentHash`, - ); - } catch (err: any) { - // Fallback for legacy DBs without contentHash column - const msg = err?.message ?? ''; - if (isMissingColumnOrTableError(msg)) { - hasContentHash = false; - rows = await c.query( - `MATCH (e:${EMBEDDING_TABLE_NAME}) RETURN e.nodeId AS nodeId, e.chunkIndex AS chunkIndex, e.startLine AS startLine, e.endLine AS endLine, e.embedding AS embedding`, + try { + queryResult = await c.query( + `MATCH (e:${EMBEDDING_TABLE_NAME}) RETURN e.nodeId AS nodeId, e.chunkIndex AS chunkIndex, e.startLine AS startLine, e.endLine AS endLine, e.embedding AS embedding, e.contentHash AS contentHash`, ); - } else { - throw err; + } catch (err: any) { + // Fallback for legacy DBs without contentHash column + const msg = err?.message ?? ''; + if (isMissingColumnOrTableError(msg)) { + hasContentHash = false; + queryResult = await c.query( + `MATCH (e:${EMBEDDING_TABLE_NAME}) RETURN e.nodeId AS nodeId, e.chunkIndex AS chunkIndex, e.startLine AS startLine, e.endLine AS endLine, e.embedding AS embedding`, + ); + } else { + throw err; + } } - } - for (const row of await readQueryRows(rows)) { - const nodeId = String(row.nodeId ?? row[0] ?? ''); - if (!nodeId) continue; - embeddingNodeIds.add(nodeId); - const embedding = row.embedding ?? row[4]; - if (embedding) { - embeddings.push({ - nodeId, - chunkIndex: Number(row.chunkIndex ?? row[1] ?? 0), - startLine: Number(row.startLine ?? row[2] ?? 0), - endLine: Number(row.endLine ?? row[3] ?? 0), - embedding: Array.isArray(embedding) - ? embedding.map(Number) - : Array.from(embedding as any).map(Number), - contentHash: hasContentHash ? (row.contentHash ?? row[5] ?? undefined) : undefined, - }); + const results = Array.isArray(queryResult) ? queryResult : [queryResult]; + const result = results[0]; + while (await result.hasNext()) { + const row = await result.getNext(); + ingestCachedEmbeddingRow(builder, row, hasContentHash); } + return finalizeCachedEmbeddingsSnapshot(builder); + } catch (err) { + abortCachedEmbeddingsBuilder(builder); + throw err; + } finally { + if (queryResult) await closeQueryResults(queryResult); } - } catch { - /* embedding table may not exist */ + } catch (err) { + abortCachedEmbeddingsBuilder(builder); + throw err; } - - return { embeddingNodeIds, embeddings }; }); }; diff --git a/gitnexus/test/integration/load-cached-embeddings-spill.test.ts b/gitnexus/test/integration/load-cached-embeddings-spill.test.ts new file mode 100644 index 000000000..2d98da77e --- /dev/null +++ b/gitnexus/test/integration/load-cached-embeddings-spill.test.ts @@ -0,0 +1,106 @@ +/** + * Real-DB coverage for #3306: loadCachedEmbeddings must stream CodeEmbedding + * rows instead of getAll()+map(Number) of the whole table, and must be able + * to spill vectors so incremental analyze does not keep every embedding in + * the V8 heap. + */ +import fs from 'node:fs'; +import path from 'node:path'; +import { afterEach, describe, expect, it } from 'vitest'; +import { createTempDir, type TestDBHandle } from '../helpers/test-db.js'; +import { EMBEDDING_DIMS } from '../../src/core/lbug/schema.js'; +import { batchInsertEmbeddings } from '../../src/core/embeddings/embedding-pipeline.js'; +import { + disposeEmbeddingSpill, + materializeCachedEmbeddings, +} from '../../src/core/embeddings/embedding-restore-spill.js'; + +describe('loadCachedEmbeddings streaming (#3306)', () => { + let tmp: TestDBHandle | undefined; + + afterEach(async () => { + const adapter = await import('../../src/core/lbug/lbug-adapter.js'); + try { + await adapter.closeLbug(); + } catch { + /* already closed */ + } + await tmp?.cleanup(); + tmp = undefined; + }); + + async function seedDb(rowCount: number) { + tmp = await createTempDir('gitnexus-lbug-'); + const dbPath = path.join(tmp.dbPath, 'lbug'); + const adapter = await import('../../src/core/lbug/lbug-adapter.js'); + await adapter.initLbug(dbPath); + const rows = Array.from({ length: rowCount }, (_, i) => ({ + nodeId: `Function:src/f${i}.ts:fn${i}:1`, + chunkIndex: 0, + startLine: 1, + endLine: 3, + embedding: Array.from({ length: EMBEDDING_DIMS }, (__, d) => (d === 0 ? i + 1 : 0)), + contentHash: `hash-${i}`, + })); + await batchInsertEmbeddings(adapter.executeWithReusedStatement, rows); + return { adapter, rows }; + } + + it('materializes a small table in RAM (skip-fts / mock-compatible shape)', async () => { + const { adapter, rows } = await seedDb(3); + const cached = await adapter.loadCachedEmbeddings(); + expect(cached.spill).toBeUndefined(); + expect(cached.embeddings).toHaveLength(3); + expect(cached.rows).toHaveLength(3); + expect(cached.embeddingNodeIds.size).toBe(3); + expect(cached.embeddings.map((e) => e.nodeId).sort()).toEqual(rows.map((r) => r.nodeId).sort()); + expect(cached.embeddings.find((e) => e.nodeId === rows[1]!.nodeId)?.embedding[0]).toBe(2); + }); + + it('streams into a spill file when the in-memory limit is 0 and restores a subset', async () => { + const { adapter, rows } = await seedDb(12); + const cached = await adapter.loadCachedEmbeddings({ inMemoryRowLimit: 0 }); + try { + expect(cached.embeddings).toEqual([]); + expect(cached.spill?.rowCount).toBe(12); + expect(cached.rows).toHaveLength(12); + const subset = materializeCachedEmbeddings( + cached, + cached.rows.filter((_, i) => i % 5 === 0), + ); + expect(subset).toHaveLength(3); + expect(subset[0]?.embedding[0]).toBe(1); + expect(subset[1]?.embedding[0]).toBe(6); + expect(subset[2]?.embedding[0]).toBe(11); + expect(subset[0]?.contentHash).toBe(rows[0]!.contentHash); + } finally { + disposeEmbeddingSpill(cached.spill); + } + }); + + it('flips from RAM to spill once a non-zero in-memory limit is crossed', async () => { + const { adapter } = await seedDb(8); + const cached = await adapter.loadCachedEmbeddings({ inMemoryRowLimit: 4 }); + try { + expect(cached.embeddings).toEqual([]); + expect(cached.spill?.rowCount).toBe(8); + expect(cached.rows).toHaveLength(8); + expect(fs.statSync(cached.spill!.path).size).toBe(12 + 8 * EMBEDDING_DIMS * 4); + const subset = materializeCachedEmbeddings(cached, cached.rows.slice(2, 4)); + expect(subset).toHaveLength(2); + expect(subset[0]?.embedding[0]).toBe(3); + expect(subset[1]?.embedding[0]).toBe(4); + } finally { + disposeEmbeddingSpill(cached.spill); + } + }); + + it('surfaces a spill write failure instead of adopting an empty snapshot', async () => { + const { adapter } = await seedDb(3); + const spillDir = path.join(tmp!.dbPath, 'not-a-directory'); + fs.writeFileSync(spillDir, 'x'); + await expect(adapter.loadCachedEmbeddings({ inMemoryRowLimit: 0, spillDir })).rejects.toThrow( + /ENOTDIR|not a directory|ENOSPC|EACCES/i, + ); + }); +}); diff --git a/gitnexus/vitest.config.ts b/gitnexus/vitest.config.ts index 4b2625805..1239862a3 100644 --- a/gitnexus/vitest.config.ts +++ b/gitnexus/vitest.config.ts @@ -100,6 +100,7 @@ export default defineConfig({ 'test/integration/analyze-wal-checkpoint-failure.test.ts', 'test/integration/lbug-non-ascii-path.test.ts', 'test/integration/lbug-conn-serialization.test.ts', + 'test/integration/load-cached-embeddings-spill.test.ts', 'test/integration/group/manifest-resolve-symbol-2325.test.ts', 'test/integration/group/manifest-synthetic-impact-lbug.test.ts', 'test/integration/group/http-route-resolve-symbol.test.ts', @@ -179,6 +180,7 @@ export default defineConfig({ 'test/integration/analyze-wal-checkpoint-failure.test.ts', 'test/integration/lbug-non-ascii-path.test.ts', 'test/integration/lbug-conn-serialization.test.ts', + 'test/integration/load-cached-embeddings-spill.test.ts', 'test/integration/group/manifest-resolve-symbol-2325.test.ts', 'test/integration/group/manifest-synthetic-impact-lbug.test.ts', 'test/integration/group/http-route-resolve-symbol.test.ts',