Skip to content

Commit fbda999

Browse files
authored
perf(synth): shard near-duplicate signing across the resolver pool (#609)
* probe: shard near-duplicate signing across the resolver pool * probe: store signed shingle arrays and build Sets only for pair bodies * fix(graph): fall back to in-process signing when the nearDup signer fails Review of #609: a rejected pool/worker signing task left the stale bodies unsigned and the pairs stale. Catch the rejection and run the same in-process loop the no-pool path uses, logging under SYNTH_TIMINGS. The new test forces a rejecting signer and asserts the rows and pairs match the in-process result. * fix(graph): reset the near-duplicate ran flag for each full index Review of #609: the post-synthesis hook set the instance field true and a later indexAll on the same CodeGraph instance skipped the fallback even when the hook did not run (no refs, no pool, resolver skipped). Reset it at the start of each full index.
1 parent 723df83 commit fbda999

7 files changed

Lines changed: 204 additions & 21 deletions

File tree

‎__tests__/near-duplicates.test.ts‎

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import * as os from 'os';
1111
import CodeGraph from '../src/index';
1212
import { execFileSync } from 'child_process';
1313
import { ToolHandler } from '../src/mcp/tools';
14-
import { forEachBandBucket, refreshNearDuplicates, signature, similarity } from '../src/graph/near-duplicates';
14+
import { forEachBandBucket, refreshNearDuplicates, shingles, signature, signatureFromSet, similarity } from '../src/graph/near-duplicates';
1515

1616
const urlOf = (name: string, label: string) => [
1717
`export function ${name}(input: string | URL): string {`,
@@ -173,15 +173,39 @@ describe('near-duplicates in the index', () => {
173173

174174
// dsUrlOf changes: its two pairs are re-scored, the points pair is not read.
175175
db.prepare("UPDATE nodes SET updated_at = updated_at + 1 WHERE name = 'dsUrlOf'").run();
176-
const r = refreshNearDuplicates(deps);
176+
const r = await refreshNearDuplicates(deps);
177177
expect(r.signed).toBe(1);
178178
expect(r.rescored).toBe(2);
179179
expect(reads.some((p) => p.includes('points'))).toBe(false);
180180
expect(snapshot()).toEqual(before);
181181

182182
// Every body re-signed: a full recompute stores the same pairs.
183183
db.prepare("UPDATE nodes SET updated_at = updated_at + 1 WHERE kind = 'function'").run();
184-
refreshNearDuplicates(deps);
184+
await refreshNearDuplicates(deps);
185+
expect(snapshot()).toEqual(before);
186+
187+
// The off-thread signer path stores the same signatures and pairs as the
188+
// in-process loop (the pool worker computes shingles/signatureFromSet).
189+
db.prepare("UPDATE nodes SET updated_at = updated_at + 1 WHERE kind = 'function'").run();
190+
const sign = async (bodies: Array<{ id: string; name: string; filePath: string; startLine: number; endLine: number; updatedAt: number }>) =>
191+
bodies.map((b) => {
192+
const set = shingles(deps.readFile(b.filePath)!.split(/\r?\n/).slice(b.startLine - 1, b.endLine).join('\n'), b.name);
193+
const sig = set ? signatureFromSet(set) : null;
194+
return {
195+
id: b.id,
196+
sig: sig ? new Uint8Array(sig.buffer, sig.byteOffset, sig.byteLength).slice() : new Uint8Array(0),
197+
set: set ? Uint32Array.from(set) : new Uint32Array(0),
198+
};
199+
});
200+
const viaSigner = await refreshNearDuplicates({ ...deps, sign });
201+
expect(viaSigner.signed).toBeGreaterThan(0);
202+
expect(snapshot()).toEqual(before);
203+
204+
// A worker/pool rejection falls back to the in-process loop; the rows and
205+
// pairs must be the same, not stale or empty.
206+
db.prepare("UPDATE nodes SET updated_at = updated_at + 1 WHERE kind = 'function'").run();
207+
const viaFallback = await refreshNearDuplicates({ ...deps, sign: async () => { throw new Error('pool down'); } });
208+
expect(viaFallback.signed).toBeGreaterThan(0);
185209
expect(snapshot()).toEqual(before);
186210
});
187211

‎src/codegraph.ts‎

Lines changed: 27 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ import type { FileWatcher, WatchOptions, PendingFile } from './sync';
5858
import { EXTRACTION_VERSION } from './extraction/extraction-version';
5959
import { isGeneratedFile } from './extraction/generated-detection';
6060
import { hasGrammarLoadFailure, isFileLevelOnlyLanguage } from './extraction/grammars';
61-
import { refreshNearDuplicates } from './graph/near-duplicates';
61+
import { refreshNearDuplicates, type NearDupSigner } from './graph/near-duplicates';
6262
import { getCodeGraphDir } from './directory';
6363
import { deriveProjectNameTokens } from './search/query-utils';
6464
import { findSeedSource, copySeedIndex, type SeedSource } from './sync/worktree-seed';
@@ -189,6 +189,8 @@ export class CodeGraph {
189189
// method call away from the constructor body.
190190
private lazyOrchestrator?: ExtractionOrchestrator;
191191
private lazyResolver?: ReferenceResolver;
192+
/** Set by the resolver's post-synthesis near-dup hook so indexAll does not repeat it. */
193+
private nearDuplicatesRan = false;
192194
private get orchestrator(): ExtractionOrchestrator {
193195
return (this.lazyOrchestrator ??= new (extraction().ExtractionOrchestrator)(this.projectRoot, this.queries));
194196
}
@@ -655,6 +657,10 @@ export class CodeGraph {
655657
return { success: false, filesIndexed: 0, filesSkipped: 0, filesErrored: 0, nodesCreated: 0, edgesCreated: 0, errors: [{ message: 'Could not acquire file lock - another process may be indexing', severity: 'error' as const }], durationMs: 0 };
656658
}
657659
const startedAt = Date.now();
660+
// Each full index owns its near-duplicate pass: the post-synthesis hook
661+
// sets this, and a later index on the same instance must run it again
662+
// even when that hook did not (no refs, no pool, resolver skipped).
663+
this.nearDuplicatesRan = false;
658664
// Defer WAL auto-checkpointing for the whole bulk run (#1231): the
659665
// default 1000-page interval re-writes hot pages into the main DB file
660666
// over and over — ~95% of all disk I/O during a bulk index, and a
@@ -827,7 +833,8 @@ export class CodeGraph {
827833
total: totalPasses,
828834
});
829835
},
830-
walValve
836+
walValve,
837+
true
831838
);
832839
if (process.env.CODEGRAPH_SYNTH_TIMINGS) console.error(`[phase-timing] resolution: ${Date.now() - tResolve}ms`);
833840

@@ -850,7 +857,7 @@ export class CodeGraph {
850857
await this.resolver.resolveDeferredThisMemberRefs();
851858
if (process.env.CODEGRAPH_SYNTH_TIMINGS) console.error(`[synth-timing] deferredThisMember: ${Date.now() - tDeferred}ms`);
852859
}
853-
if (result.success) this.refreshNearDuplicates();
860+
if (result.success && !this.nearDuplicatesRan) await this.refreshNearDuplicates();
854861

855862
// Refresh planner stats + checkpoint the WAL after bulk writes.
856863
// Off-thread (worker connection): on a multi-GB index this is minutes
@@ -1859,7 +1866,7 @@ export class CodeGraph {
18591866
throw new Error('Reconciliation did not complete synthesized-edge refresh.');
18601867
}
18611868
if (filesChanged || result.filesRemoved > 0) {
1862-
this.refreshNearDuplicates(result.changedFilePaths, result.filesRemoved > 0);
1869+
await this.refreshNearDuplicates(result.changedFilePaths, result.filesRemoved > 0);
18631870
}
18641871

18651872
// Refresh planner stats + checkpoint the WAL after bulk writes.
@@ -2255,20 +2262,21 @@ export class CodeGraph {
22552262
* Only changed bodies are re-signed, so a sync that touched no function costs
22562263
* one query. Never fails an index: the pairs are an annotation.
22572264
*/
2258-
private refreshNearDuplicates(changedFiles?: readonly string[], force = false): void {
2265+
private async refreshNearDuplicates(changedFiles?: readonly string[], force = false, sign?: NearDupSigner): Promise<void> {
22592266
const t = Date.now();
22602267
try {
22612268
const deprioritized = this.queries.getDeprioritizedPathMatcher();
22622269
// The index's content-based generated flag, not just the path check.
22632270
const generated = this.queries.generatedPredicateFor(this.queries.getAllFiles().map((f) => f.path));
2264-
const r = refreshNearDuplicates({
2271+
const r = await refreshNearDuplicates({
22652272
queries: this.queries,
22662273
readFile: (p) => {
22672274
try { return fs.readFileSync(path.join(this.projectRoot, p), 'utf-8'); } catch { return null; }
22682275
},
22692276
isExcluded: (p) => generated(p) || isGeneratedFile(p) || (deprioritized?.(p) ?? false),
22702277
changedFiles,
22712278
force,
2279+
sign,
22722280
});
22732281
if (process.env.CODEGRAPH_SYNTH_TIMINGS) console.error(`[synth-timing] nearDuplicates: ${Date.now() - t}ms (${r.signed} signed, ${r.rescored} pairs re-scored, ${r.pairs ?? 'unchanged'} pairs)`);
22742282
} catch (error) {
@@ -2288,10 +2296,22 @@ export class CodeGraph {
22882296
// kernel conn is live on the db (the rusqlite build's wal-index locks
22892297
// can't see node:sqlite's — an off-thread checkpoint could rewrite the
22902298
// -shm header under a writer commit). Structural so tests can fake it.
2291-
walValve?: { stop(): void; start(): void; drain(): Promise<void>; backpressure(): Promise<void> | null } | null
2299+
walValve?: { stop(): void; start(): void; drain(): Promise<void>; backpressure(): Promise<void> | null } | null,
2300+
/** Run the near-duplicate refresh inside the resolver's pool lifetime. */
2301+
withNearDuplicates = false
22922302
): Promise<ResolutionResult> {
22932303
return this.resolver.resolveAndPersistBatched(onProgress, undefined, onSynthesisProgress, {
22942304
dbPath: this.db.getPath(),
2305+
nearDuplicates: withNearDuplicates
2306+
? (pool) => {
2307+
const sign: NearDupSigner | undefined = pool
2308+
? (bodies) => pool.signNearDuplicates(bodies, this.projectRoot)
2309+
: undefined;
2310+
return this.refreshNearDuplicates(undefined, false, sign).finally(() => {
2311+
this.nearDuplicatesRan = true;
2312+
});
2313+
}
2314+
: undefined,
22952315
// Bulk-edge-load hooks: on big runs the resolver drops the non-unique
22962316
// edge indexes for the batch loop and recreates them before synthesis
22972317
// (which reads kind-keyed). Concurrent readers (a daemon serving this

‎src/graph/near-duplicates.ts‎

Lines changed: 67 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -128,10 +128,32 @@ function asSignature(bytes: Uint8Array): Uint32Array {
128128
return new Uint32Array(bytes.buffer.slice(bytes.byteOffset, bytes.byteOffset + bytes.byteLength));
129129
}
130130

131+
/** One body a sharded signer computes a signature for. */
132+
export interface NearDupBody {
133+
id: string;
134+
name: string;
135+
filePath: string;
136+
startLine: number;
137+
endLine: number;
138+
updatedAt: number;
139+
}
140+
141+
/** A signed body: the MinHash bytes plus its shingle set (empty when unreadable). */
142+
export interface NearDupSignature {
143+
id: string;
144+
sig: Uint8Array;
145+
set: Uint32Array;
146+
}
147+
148+
/** Off-thread signer for a batch of bodies (the resolver pool in a full index). */
149+
export type NearDupSigner = (bodies: NearDupBody[]) => Promise<NearDupSignature[]>;
150+
131151
export interface NearDuplicateDeps {
132152
queries: QueryBuilder;
133153
/** The file's text, or null when it cannot be read. */
134154
readFile(filePath: string): string | null;
155+
/** When set, compute the signatures off-thread instead of reading files here. */
156+
sign?: NearDupSigner;
135157
/** Generated or configured-away files, left out like tests and vendored code. */
136158
isExcluded?(filePath: string): boolean;
137159
/**
@@ -152,7 +174,7 @@ export interface NearDuplicateDeps {
152174
* the pairs. Returns how many signatures were written and pairs stored
153175
* (`pairs` is null when nothing changed and the stored pairs stand).
154176
*/
155-
export function refreshNearDuplicates(deps: NearDuplicateDeps): { signed: number; pairs: number | null; rescored: number } {
177+
export async function refreshNearDuplicates(deps: NearDuplicateDeps): Promise<{ signed: number; pairs: number | null; rescored: number }> {
156178
const { queries } = deps;
157179
const excluded = (p: string) => isTestFile(p) || VENDORED.test(p) || (deps.isExcluded?.(p) ?? false);
158180
const tSign = Date.now();
@@ -172,14 +194,47 @@ export function refreshNearDuplicates(deps: NearDuplicateDeps): { signed: number
172194
const rows: Array<{ nodeId: string; nodeUpdatedAt: number; sig: Uint8Array }> = [];
173195
// The shingle sets just computed, keyed by node id: scoring reuses them
174196
// instead of re-tokenizing every body that lands in a candidate pair.
175-
const freshSets = new Map<string, Set<number>>();
176-
for (const [filePath, bodies] of stale) {
177-
const lines = deps.readFile(filePath)?.split(/\r?\n/);
178-
for (const b of bodies) {
179-
const set = lines ? shingles(lines.slice(b.startLine - 1, b.endLine).join('\n'), b.name) : null;
180-
const sig = set ? signatureFromSet(set) : null;
181-
if (set) freshSets.set(b.id, set);
182-
rows.push({ nodeId: b.id, nodeUpdatedAt: b.updatedAt, sig: sig ? new Uint8Array(sig.buffer) : new Uint8Array(0) });
197+
// The in-process path stores Sets; the off-thread signer stores its
198+
// transferred Uint32Arrays and the pair pass builds a Set only for a body
199+
// that actually lands in a candidate pair (36k Sets was most of the handoff).
200+
const freshSets = new Map<string, Set<number> | Uint32Array>();
201+
let signedOffThread = false;
202+
if (deps.sign) {
203+
// The bodies are grouped by file for the in-process path; the signer takes
204+
// them as one batch and the worker/pool groups what it can.
205+
const bodies: NearDupBody[] = [];
206+
for (const [filePath, list] of stale) {
207+
for (const b of list) {
208+
bodies.push({ id: b.id, name: b.name, filePath, startLine: b.startLine, endLine: b.endLine, updatedAt: b.updatedAt });
209+
}
210+
}
211+
try {
212+
const signed = bodies.length > 0 ? await deps.sign(bodies) : [];
213+
const byId = new Map(bodies.map((b) => [b.id, b]));
214+
for (const s of signed) {
215+
const b = byId.get(s.id);
216+
if (!b) continue;
217+
if (s.set.length > 0) freshSets.set(b.id, s.set);
218+
rows.push({ nodeId: b.id, nodeUpdatedAt: b.updatedAt, sig: s.sig.length > 0 ? s.sig : new Uint8Array(0) });
219+
}
220+
signedOffThread = true;
221+
} catch (err) {
222+
// A worker/pool failure must not leave the signatures stale: sign the
223+
// same bodies in-process below, exactly as a run without a pool does.
224+
if (process.env.CODEGRAPH_SYNTH_TIMINGS) {
225+
console.error(`[synth-timing] nearDup signer failed, signing in-process: ${err instanceof Error ? err.message : String(err)}`);
226+
}
227+
}
228+
}
229+
if (!signedOffThread) {
230+
for (const [filePath, bodies] of stale) {
231+
const lines = deps.readFile(filePath)?.split(/\r?\n/);
232+
for (const b of bodies) {
233+
const set = lines ? shingles(lines.slice(b.startLine - 1, b.endLine).join('\n'), b.name) : null;
234+
const sig = set ? signatureFromSet(set) : null;
235+
if (set) freshSets.set(b.id, set);
236+
rows.push({ nodeId: b.id, nodeUpdatedAt: b.updatedAt, sig: sig ? new Uint8Array(sig.buffer) : new Uint8Array(0) });
237+
}
183238
}
184239
}
185240
// A re-signed body's cached pair scores are stale; a dropped body's are moot.
@@ -229,7 +284,7 @@ export function forEachBandBucket(
229284
}
230285

231286
/** LSH over the stored signatures, one band at a time so memory stays linear in the bodies. */
232-
function pairNearDuplicates(deps: NearDuplicateDeps, freshSets?: ReadonlyMap<string, Set<number>>): { kept: number; rescored: number } {
287+
function pairNearDuplicates(deps: NearDuplicateDeps, freshSets?: ReadonlyMap<string, Set<number> | Uint32Array>): { kept: number; rescored: number } {
233288
const { queries } = deps;
234289
// Every signature in one scan (rowid order), then each band's buckets from
235290
// the bytes in memory: sixteen SQL scans with a per-row substr and base64
@@ -271,7 +326,8 @@ function pairNearDuplicates(deps: NearDuplicateDeps, freshSets?: ReadonlyMap<str
271326
const node = queries.getNodeById(row.nodeId);
272327
// A set signed this refresh is already in memory; only unchanged bodies
273328
// pay the tokenize pass.
274-
const cached = node ? freshSets?.get(node.id) : undefined;
329+
const cachedRaw = node ? freshSets?.get(node.id) : undefined;
330+
const cached = cachedRaw instanceof Set ? cachedRaw : cachedRaw ? new Set(cachedRaw) : undefined;
275331
const text = !cached && node ? lines(node.filePath)?.slice(node.startLine - 1, node.endLine).join('\n') : undefined;
276332
hit = { nodeId: row.nodeId, sig: asSignature(row.sig), set: cached ?? (text ? shingles(text, node!.name) : null) };
277333
sigs.set(rowid, hit);

‎src/resolution/index.ts‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2280,6 +2280,9 @@ export class ReferenceResolver {
22802280
// a 22GB WAL on a 4.6GB DB (migration plan §7a.1).
22812281
parallel?: {
22822282
dbPath: string;
2283+
/** Full index: sign near-duplicate bodies after synthesis, while the
2284+
* pool is still alive; null pool means the in-process path. */
2285+
nearDuplicates?: (pool: ResolverPool | null) => Promise<void> | void;
22832286
bulkEdgeLoad?: { begin: () => void; end: () => void | Promise<void> };
22842287
/** unresolved_refs index window for the batched loop — the loop only
22852288
* reads the status index + PK; dropping the sync-path ref indexes cuts
@@ -3051,6 +3054,18 @@ export class ReferenceResolver {
30513054
// synthesis is additive and optional; ignore failures
30523055
}
30533056
if (process.env.CODEGRAPH_SYNTH_TIMINGS) console.error(`[phase-timing] callback-synthesis: ${Date.now() - tSynth}ms`);
3057+
// Full index: sign near-duplicate bodies now, while the pool is still
3058+
// open and idle. The callback falls back to the in-process path when no
3059+
// pool engaged; it never fails the index.
3060+
if (parallel?.nearDuplicates) {
3061+
try {
3062+
await parallel.nearDuplicates(pool);
3063+
} catch (err) {
3064+
logDebug('Near-duplicate refresh failed; the index stands', {
3065+
error: err instanceof Error ? err.message : String(err),
3066+
});
3067+
}
3068+
}
30543069
} finally {
30553070
if (pool) await pool.destroy().catch(() => undefined);
30563071
// The workers' kernel-reader snapshot is per-run scratch — remove it.

0 commit comments

Comments
 (0)