Skip to content

perf(docs): persist the doc-search BM25 index across sessions - #55

Open
pradeepmouli wants to merge 4 commits into
intuit:mainfrom
pradeepmouli:perf/doc-bm25-cache
Open

perf(docs): persist the doc-search BM25 index across sessions#55
pradeepmouli wants to merge 4 commits into
intuit:mainfrom
pradeepmouli:perf/doc-bm25-cache

Conversation

@pradeepmouli

Copy link
Copy Markdown
Contributor

Summary

Doc search (hybrid_doc_search_in_dir) currently rebuilds its entire DocBM25Index in memory on every query — loading all chunks from the store and re-tokenizing the whole corpus — while code search already persists its BM25 index to .infigraph/bm25_cache.bin. This PR gives doc search the same treatment:

  • DocBM25Index gains save/load/docs(). The on-disk format mirrors code search's BM25Index::save idiom (leading version byte + length-prefixed LE binary).
  • hybrid_doc_search_in_dir loads docs_bm25_cache.bin when it is at least as new as docs_embeddings.bin — the same mtime freshness anchor code search uses, so a doc reindex invalidates the cache automatically. On a cache hit the store.get_all_chunks() DB read is skipped too (the cache carries the chunk texts).
  • Miss/stale/corrupt cache → rebuild in memory exactly as before, then a best-effort atomic save (pid-unique temp file + rename). A cache failure can never fail or change a search — results are identical with and without the cache.

Two deliberate hardenings over the idiom being mirrored:

  • load is fully bounds-checked and returns Err (never panics) on truncated/corrupt input, since a bad cache must degrade to a rebuild.
  • The temp-file name is per-process unique, so concurrent searches can't race on the same temp path.

Group/combined doc stores get their own cache automatically (keyed off the store's artifact dir; combined stores publish per-generation directories, so no cross-generation staleness is possible).

Testing

  • 3 new unit tests: save/load round-trip preserves scores; empty/wrong-version/truncated inputs rejected without panic; atomic save leaves no temp file.
  • 3 new integration tests against a real DocStore: first search writes the cache and repeat results are identical; corrupt cache falls back silently and is rewritten; stale cache (newer embeddings) is rebuilt.
  • Full infigraph-docs suite: 93/93 passing on this branch (rebased on current main); cargo fmt/clippy -D warnings/cargo check --workspace clean.

🤖 Generated with Claude Code

https://claude.ai/code/session_01SjmvwHuwV5r7ZeZpJLp5oR

@murari316 murari316 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Automated review (Fable, adversarially verified). Recommendation: request-changes.

The PR persists the doc-search BM25 index to .infigraph/docs_bm25_cache.bin with a versioned, bounds-checked binary format, mtime freshness anchored on docs_embeddings.bin, and graceful fallback-to-rebuild on any load failure. The core design is solid and well-tested (roundtrip, corrupt-cache, stale-cache tests), but it misses cleanup integration (DocIndex::clean leaves the full-text cache on disk) and has a query-time/reindex race that can pin a stale cache as fresh until the next reindex.


[MAJOR] data-retention — crates/infigraph-docs/src/lib.rs:127-129

DocIndex::clean() removes docs.kuzu, docs_embeddings.bin, and the HNSW files but the PR does not add removal of docs_bm25_cache.bin, and the cache file contains the verbatim full text of every indexed chunk. clean_docs is supposed to wipe doc-index artifacts; the new artifact survives it.

Failure scenario: User indexes internal Confluence pages, runs one doc search (cache written), then runs clean_docs to purge the doc index. docs.kuzu and embeddings are deleted, but .infigraph/docs_bm25_cache.bin still holds the complete text of every chunk on disk indefinitely. It is never served (anchor missing -> not fresh), so this is orphaned sensitive data plus disk bloat, not wrong results. One-line fix: remove_file the cache alongside lib.rs:127-129.

[MAJOR] concurrency — crates/infigraph-docs/src/search.rs:233-249 (load_or_build_doc_index, PR hunk)

TOCTOU between reading chunks and saving the cache: if a doc reindex rewrites docs_embeddings.bin while a search is between store.get_all_chunks() and idx.save(), the cache built from pre-reindex chunks lands with an mtime newer than the new anchor, so doc_bm25_cache_fresh (c >= a) treats the stale cache as fresh from then on.

Failure scenario: Doc watcher fires a reindex (update_doc_embeddings rewrites the anchor) while a search_docs call is mid-flight rebuilding the cache. The stale cache is saved after the anchor and wins the mtime comparison. Every subsequent search silently drops chunks deleted in the reindex (hydration lookup returns None -> filter_map discards) and never BM25-matches newly added docs — persisting until the NEXT reindex bumps the anchor, which on a quiet repo can be days. Fix: capture the anchor mtime before get_all_chunks and skip the save (or set the cache mtime to the captured anchor mtime) if the anchor changed.

[MINOR] concurrency — crates/infigraph-docs/src/search.rs:128-131 (save, PR hunk)

The temp-file name is per-process (path.with_extension("{pid}.tmp")), but the MCP server handles concurrent tool calls on multiple threads in ONE process — two simultaneous cache misses share the same tmp path, so the code comment 'concurrent writers can't interleave' is wrong for the in-process case.

Failure scenario: Two search_docs requests hit the MCP server concurrently after a reindex; both miss the cache and call save(). Thread A's fs::write truncates/overwrites the tmp file while thread B is mid-write, and the rename publishes a torn file. Self-healing because load() validates and falls back to rebuild, so impact is wasted rebuilds — but the comment claims a guarantee the code doesn't provide. Use tempfile::NamedTempFile or add a per-call unique suffix.

[MINOR] error-handling — crates/infigraph-docs/src/search.rs:129-131 (save, PR hunk)

save() leaks the .tmp file when the write succeeds but the rename fails — the error is returned without cleaning up, and per-PID names mean leaked files accumulate rather than being overwritten.

Failure scenario: On Windows (CI runs there per repo invariants), std::fs::rename over a destination another process currently has open can fail with permission denied. Each failed save from a fresh process leaves a distinct docs_bm25_cache..tmp in .infigraph, accumulating over time. Add a remove_file(&tmp) on the rename error path (and arguably on the write error path).

[MINOR] correctness — crates/infigraph-docs/src/search.rs:read_str closure in load (PR hunk)

read_str uses String::from_utf8_lossy, silently accepting invalid UTF-8, which contradicts the function's own doc ('Any structural problem is an Err ... callers treat a bad cache as a miss and rebuild').

Failure scenario: Bit-rot flips bytes inside a chunk-id string without changing lengths: all bounds checks pass, load() returns Ok, but the mangled id (with U+FFFD) fails the get_chunk_details hydration lookup, so matching results are silently dropped from search output instead of the corrupt cache triggering a rebuild. Use String::from_utf8(...).map_err(...) so corruption becomes an Err.

[MINOR] test-coverage — crates/infigraph-docs/tests/doc_bm25_cache.rs:125-155 (stale_cache_is_rebuilt_when_embeddings_are_newer)

No test verifies end-to-end invalidation semantics: the stale-cache test only asserts the cache file's mtime advanced, not that search results actually reflect content indexed after the cache was created.

Failure scenario: A future refactor could rewrite the cache file (mtime advances, test passes) while still serving the old in-memory/cached doc set — e.g., saving the loaded stale index back out. Extend the test to index_doc a new file with a distinctive term after the first search, then assert a search for that term returns the new chunk.

[NIT] correctness — crates/infigraph-docs/src/search.rs:load (PR hunk)

load() never checks that pos == data.len() after parsing, so files with trailing garbage (e.g., the torn-write outcome from the tmp-path collision where a shorter payload was renamed over a longer partial one) load as valid.

Failure scenario: A structurally-complete prefix followed by junk bytes parses successfully and the junk is ignored; combined with the lossy-UTF-8 issue this weakens the 'validated load' safety net. Add anyhow::ensure!(pos == data.len()).

[NIT] correctness — crates/infigraph-docs/src/search.rs:doc_bm25_cache_fresh (PR hunk)

Freshness uses c >= a on mtimes; on filesystems with coarse mtime granularity a reindex that rewrites docs_embeddings.bin within the same timestamp tick as a cache save leaves the stale cache 'fresh'. Consistent with the existing code-search bm25_cache pattern, so noting rather than blocking; the relatedly mtime-based test's 50ms sleep could be flaky on coarse-granularity filesystems.

Generated with Claude Code (Fable). Findings verified by an adversarial refute pass; false positives removed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants