Skip to content

perf(core): batch the TAINT_FLOW write instead of one Cypher round-trip per flow - #75

Open
pradeepmouli wants to merge 2 commits into
intuit:mainfrom
pradeepmouli:fix-taint-flow-batched-write
Open

perf(core): batch the TAINT_FLOW write instead of one Cypher round-trip per flow#75
pradeepmouli wants to merge 2 commits into
intuit:mainfrom
pradeepmouli:fix-taint-flow-batched-write

Conversation

@pradeepmouli

Copy link
Copy Markdown
Contributor

Problem

taint::write_taint_flows issues the clearing DELETE and one CREATE per flow as separate raw_query calls. Kùzu's raw_query takes a fresh connection each time, so a run with N flows costs N connections and N transactions.

Instrumenting inside detect_taint_flows_with_cache on a 361-file / 4663-function Rust corpus:

analyze-loop   0.022s   (4663 functions, 557 flows)
write          2.512s   (540 active flows, ~4.6ms each)

Writing the flows costs ~110× computing them, and it is the single largest cost of an incremental index — touching one file out of 361 took 4.48s against a 6.51s full reindex, and 2.5s of that was this write.

It is also never atomic, for exactly the reason write_calls_service_edges holds one connection across its batch: the DELETE and the CREATEs run on different connections, so a crash mid-loop deletes every TAINT_FLOW edge and recreates only some of them.

Change

Promotes the write to GraphBackend::replace_taint_flows, following write_calls_service_edges exactly — one held connection with BEGIN/ROLLBACK/COMMIT on Kùzu, one auto-committed statement on Neo4j. No new dependencies.

Clearing is unconditional rather than an early return on an empty slice, so a run that finds no taint still drops the previous run's edges.

before after
one-file incremental index 4.48s 1.51s (2.97×)
post-index analysis block 3.61s 0.61s
taint flow pass 2.96s 0.26s

Detection output is unchanged: 557 flows / 871 inter-procedural / 11 dynamic URLs, before and after. Errors now propagate instead of being dropped by let _ =; the CLI already reports them as a warning.

Also included

The 202 taint patterns are now stored lowercase, and the ten call sites that lowercased them again on every line of every function are gone. Every consumer only ever does a case-insensitive contains, so this removes redundant work rather than adding machinery — it also halves two sibling passes (inter-procedural taint 0.342s → 0.135s, dynamic URLs 0.094s → 0.034s).

That does make an uppercase pattern silently unmatchable, so every_taint_pattern_is_stored_lowercase pins the invariant across all four pattern sets (sources, sinks, sanitizers, HTTP client patterns).

It rides in the same commit because splitting it the other way round would leave an intermediate commit where the call sites match lowercase against patterns that are not yet lowercase.

Tests

New crates/infigraph-core/tests/taint_flow_edges.rs, three cases:

  • every flow lands in one call
  • a second run replaces rather than accumulates (otherwise repeated indexing piles up duplicate edges forever)
  • an empty flow list still clears the previous run's edges

Verification

Run on a worktree built from a clean main, not on my fork:

  • cargo fmt --all -- --check — clean
  • cargo test -p infigraph-core700 passed, 0 failed (37 suites, --test-threads=1)
  • cargo clippy --all-targets -- -D warnings — no findings in any file this PR touches

Two pre-existing issues on main that this PR does not introduce and does not fix, both reproduced on a clean checkout with no changes applied:

  1. clippy -D warnings fails on crates/infigraph-core/src/embed/mod.rs:402 (chunks_exact_to_as_chunks) under rustc 1.98. Worth knowing if CI floats to current stable.
  2. watch::tests::watch_project_detects_changes_through_symlinked_root fails on macOS.

🤖 Generated with Claude Code

https://claude.ai/code/session_01HdmNge4878WaicQ4QPfZme

`taint::write_taint_flows` issued the clearing DELETE and one CREATE per
flow as separate `raw_query` calls. Kùzu's `raw_query` takes a fresh
connection each time, so a run with N flows meant N connections and N
transactions.

Measured on a 361-file / 4663-function Rust corpus, instrumenting inside
`detect_taint_flows_with_cache`:

    analyze-loop   0.022s   (4663 functions, 557 flows)
    write          2.512s   (540 active flows, ~4.6ms each)

That made writing the flows ~110x the cost of computing them, and the
single largest cost of an incremental index: touching one file of 361
took 4.48s against a 6.51s full reindex, and this was 2.5s of it.

It was also never atomic, for the same reason `write_calls_service_edges`
holds one connection across its batch: the DELETE and the CREATEs ran on
different connections, so a crash mid-loop deleted every TAINT_FLOW edge
and recreated only some of them.

Promotes the write to `GraphBackend::replace_taint_flows`, following
`write_calls_service_edges` exactly -- one held connection with
BEGIN/ROLLBACK/COMMIT on Kùzu, one auto-committed statement on Neo4j.
Clearing is unconditional rather than an early return on an empty slice,
so a run that finds no taint still drops the previous run's edges.

    one-file incremental index   4.48s -> 1.51s   (2.97x)
    post-index analysis block    3.61s -> 0.61s
    taint flow pass              2.96s -> 0.26s

Detection output is unchanged (557 flows / 871 inter-procedural / 11
dynamic URLs). Errors now propagate instead of being dropped by `let _ =`;
the CLI already reports them as a warning.

Also stores the 202 taint patterns lowercase and drops the ten call sites
that lowercased them again on every line of every function. Every consumer
only ever does a case-insensitive `contains`, so this removes redundant
work rather than adding machinery -- it also halves two sibling passes
(inter-procedural taint 0.342s -> 0.135s, dynamic URLs 0.094s -> 0.034s).
It does make an uppercase pattern silently unmatchable, so a new test
pins the invariant across all four pattern sets.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HdmNge4878WaicQ4QPfZme
`cargo clippy --all-targets -- -D warnings` -- the command this repo's CI
runs -- fails on `main` under rustc 1.98, before any change is applied:

    error: using `chunks_exact` with a constant chunk size
      --> crates/infigraph-core/src/embed/mod.rs:402:14

`chunks_exact_to_as_chunks` was added to clippy after this code was
written, and CI resolves `stable` at run time, so `main` went red without
anything in the repo changing.

Uses a targeted `#[allow]` rather than rewriting to `as_chunks::<4>()`,
which is the mechanical fix the lint suggests but would raise the effective
MSRV to 1.88 for a decode loop that is already correct.

Unrelated to the rest of this PR, and separable -- it is here only because
otherwise the PR sits red on a failure it did not cause.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HdmNge4878WaicQ4QPfZme
@pradeepmouli

Copy link
Copy Markdown
Contributor Author

Pushed b124b96, which adds a targeted #[allow(clippy::chunks_exact_to_as_chunks)] in embed/mod.rs.

That is unrelated to this PR's change and is separable — it's here only so this PR doesn't sit red on a failure it didn't cause. cargo clippy --all-targets -- -D warnings, the exact command CI runs, fails on main before any change is applied:

error: using `chunks_exact` with a constant chunk size
  --> crates/infigraph-core/src/embed/mod.rs:402:14

The lint postdates the code and CI resolves stable at run time, so main went red without anything in the repo changing. Used an #[allow] rather than the as_chunks::<4>() rewrite the lint suggests, since that would raise the effective MSRV to 1.88 for a decode loop that is already correct — happy to drop this commit if you'd rather fix it separately on main.

The second pre-existing failure noted in the description is untouched: watch::tests::watch_project_detects_changes_through_symlinked_root fails on macOS on a clean checkout.

@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 replaces the per-flow raw_query loop for TAINT_FLOW edges with a batched replace_taint_flows backend method (Kuzu single-transaction, Neo4j single-statement), and pre-lowercases all ~200 taint/HTTP-client patterns to eliminate a per-line to_lowercase() allocation, pinned by a new invariant test. The Kuzu path, the lowercasing (behavior-preserving — comparisons were already against lowercased patterns), the pinning test, and the three new integration tests are solid work; however the Neo4j implementation has a Cypher cardinality bug that makes it never write any flows on a fresh graph, the new write path takes no advisory lock (explicit repo invariant), and the documented/tested empty-slice-clears-stale-edges contract is unreachable because both callers still guard with !all_flows.is_empty().


[BLOCKER] correctness — crates/infigraph-core/src/graph/neo4j_backend.rs:1876

The single-statement Cypher 'MATCH ()-[r:TAINT_FLOW]->() DELETE r WITH 1 AS _cleared UNWIND $flows ... CREATE ...' is cardinality-broken: MATCH with zero existing TAINT_FLOW edges yields zero rows, so the UNWIND/CREATE tail never executes and nothing is ever written; with N>1 existing edges, 'WITH 1 AS _cleared' preserves N rows so each flow is created N times. The in-code comment ('Correct for an empty slice too') covers empty $flows but misses the empty-MATCH case. Fix: 'OPTIONAL MATCH ()-[r:TAINT_FLOW]->() DELETE r WITH count(r) AS _cleared UNWIND ...' (aggregation collapses to exactly one row for both N=0 and N>1), or wrap the delete in a CALL {} subquery. Neo4j is feature-gated so CI's test suite will not catch this — it needs a test alongside the fix.

Failure scenario: Fresh Neo4j-backed project, first taint run with 500 detected flows: MATCH finds 0 existing TAINT_FLOW edges, query completes 'successfully', zero edges written. Every subsequent run also starts from zero edges, so the backend is permanently unable to record any taint flow. Conversely, a graph that somehow had 3 pre-existing edges would get every new flow written 3 times.

[MAJOR] concurrency — crates/infigraph-core/src/graph/kuzu_backend.rs:506

replace_taint_flows is a new write path (DELETE + CREATEs in an explicit transaction) but takes no advisory write lock, violating the repo invariant 'the graph DB is single-writer; write paths must take an advisory lock — a new write path needs one too'. GraphStore::connection() acquires no lock, and neither caller holds one: the CLI index path (infigraph-cli/src/index.rs:151 — no write_lock anywhere in that file) and the MCP tool path (infigraph-mcp/src/tools/analysis/taint.rs:11) both reach this method unlocked. Sibling write paths (bulk upsert at kuzu_backend.rs:433, GraphStore::remove_file, store_write.rs) all take self.store.write_lock()?. Since no caller holds the flock, adding 'let _lock = self.store.write_lock()?;' inside the method is safe and matches the upsert convention. Note: write_calls_service_edges (line 478), which this method is explicitly modeled on, has the identical pre-existing gap and should get the same treatment.

Failure scenario: An MCP detect_taint_flows call runs concurrently with a CLI reindex in another process. The reindex holds the flock for its upsert while the unlocked replace_taint_flows opens a write transaction on the same Kuzu DB — the writes are unserialized, so one side fails with a Kuzu writer-conflict error (surfaced to the user as a spurious taint/index failure) instead of being serialized by the lock.

[MAJOR] correctness — crates/infigraph-core/src/taint/mod.rs:129

The trait doc and the new test 'replace_taint_flows_with_no_flows_clears_previous_ones' pin the contract 'a run that finds no taint must still drop the previous run's edges', but both production callers (taint/mod.rs:129-131 and 158-160) still guard with 'if !all_flows.is_empty()' and the PR does not remove the guards — so the zero-flows-found case never reaches replace_taint_flows and stale TAINT_FLOW edges persist. (The all-sanitized case does reach it, since all_flows is non-empty and the sanitized filter empties the edge list — the gap is specifically zero flows found.) This is pre-existing behavior, but the PR writes the opposite contract into the trait docs and tests a path production cannot reach. Fix is trivial: drop both is_empty guards, and ideally add an end-to-end test (index tainted code, fix it, reindex, assert no TAINT_FLOW edges remain).

Failure scenario: A repo has one SQL-injection taint flow recorded. The developer fixes the code, reindexes: analyze_function finds zero flows, all_flows.is_empty() short-circuits, write_taint_flows never runs, and the stale TAINT_FLOW edge stays in the graph — detect_taint_flows queries and security reports keep flagging a vulnerability that was fixed.

[NIT] style — crates/infigraph-core/src/taint/sinks.rs:540

Lowercasing introduced duplicate patterns: the Redirect sink now lists "response.redirect(" twice (pre-existing lowercase entry plus lowered "Response.Redirect("), and the EnvVar source in sources.rs lists "os.getenv(" twice (plus lowered "os.Getenv("). Harmless dead weight — each line is scanned against the same substring twice — but worth deduplicating while touching these tables.

Failure scenario: No incorrect behavior; only redundant contains() checks on every scanned line.

[NIT] scope — crates/infigraph-core/src/embed/mod.rs:401

The #[allow(clippy::chunks_exact_to_as_chunks)] in load_embeddings is an unrelated drive-by (presumably keeping clippy -D warnings green on a newer toolchain); fine to keep, but it belongs in its own commit or at least a mention in the PR description.

Failure scenario: None — cosmetic/scope hygiene only.

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